Add AAAI depth experiments and diagnostic figures

1 files changed, 174 insertions, 0 deletions

diff --git a/scripts/parameter_matched_depth_control.py b/scripts/parameter_matched_depth_control.py
new file mode 100644
index 0000000..24f1f2b
--- /dev/null
+++ b/scripts/parameter_matched_depth_control.py
@@ -0,0 +1,174 @@
+#!/usr/bin/env python3
+"""Parameter-count-matched supplement for the FA finite-time depth sweep."""
+
+from __future__ import annotations
+
+import argparse
+import csv
+import json
+import math
+import sys
+from dataclasses import asdict
+from pathlib import Path
+from types import SimpleNamespace
+
+import matplotlib
+
+matplotlib.use("Agg")
+import matplotlib.pyplot as plt
+import numpy as np
+import torch
+
+SCRIPT_DIR = Path(__file__).resolve().parent
+if str(SCRIPT_DIR) not in sys.path:
+    sys.path.insert(0, str(SCRIPT_DIR))
+
+import aaai_depth_experiments as depth_exp  # noqa: E402
+
+
+def parse_args() -> argparse.Namespace:
+    parser = argparse.ArgumentParser(description=__doc__)
+    parser.add_argument("--depths", type=int, nargs="+", default=[1, 2, 3, 4, 6])
+    parser.add_argument("--reference-depth", type=int, default=2)
+    parser.add_argument("--reference-width", type=int, default=64)
+    parser.add_argument("--input-dim", type=int, default=16)
+    parser.add_argument("--output-dim", type=int, default=4)
+    parser.add_argument("--train-samples", type=int, default=64)
+    parser.add_argument("--data-seed", type=int, default=2027)
+    parser.add_argument("--init-seeds", type=int, default=4)
+    parser.add_argument("--feedback-seeds", type=int, default=8)
+    parser.add_argument("--lr", type=float, default=1e-3)
+    parser.add_argument("--horizon", type=int, default=50)
+    parser.add_argument("--early-step", type=int, default=20)
+    parser.add_argument("--feedback-scale", choices=["relu", "fan-in", "unit"], default="relu")
+    parser.add_argument("--torch-threads", type=int, default=16)
+    parser.add_argument(
+        "--outdir", type=Path, default=Path("outputs/parameter_matched_depth_control")
+    )
+    return parser.parse_args()
+
+
+def parameter_count(input_dim: int, output_dim: int, depth: int, width: int) -> int:
+    return input_dim * width + (depth - 1) * width * width + width * output_dim
+
+
+def closest_width(input_dim: int, output_dim: int, depth: int, target: int) -> int:
+    upper = max(2, math.ceil(target / (input_dim + output_dim)) + 2)
+    return min(
+        range(1, upper + 1),
+        key=lambda width: abs(parameter_count(input_dim, output_dim, depth, width) - target),
+    )
+
+
+def write_rows(path: Path, rows: list[depth_exp.FiniteRow]) -> None:
+    with path.open("w", newline="") as handle:
+        writer = csv.DictWriter(handle, fieldnames=list(depth_exp.FiniteRow.__annotations__))
+        writer.writeheader()
+        for row in rows:
+            writer.writerow(asdict(row))
+
+
+def main() -> None:
+    args = parse_args()
+    torch.set_num_threads(args.torch_threads)
+    args.outdir.mkdir(parents=True, exist_ok=True)
+    target = parameter_count(
+        args.input_dim, args.output_dim, args.reference_depth, args.reference_width
+    )
+    widths = {
+        depth: closest_width(args.input_dim, args.output_dim, depth, target)
+        for depth in args.depths
+    }
+    generator = torch.Generator().manual_seed(args.data_seed)
+    x = torch.randn(args.train_samples, args.input_dim, generator=generator, dtype=torch.float64)
+    y = torch.randn(args.train_samples, args.output_dim, generator=generator, dtype=torch.float64)
+
+    rows: list[depth_exp.FiniteRow] = []
+    for depth in args.depths:
+        run_args = SimpleNamespace(
+            depths=[depth],
+            finite_widths=[widths[depth]],
+            input_dim=args.input_dim,
+            output_dim=args.output_dim,
+            finite_init_seeds=args.init_seeds,
+            finite_feedback_seeds=args.feedback_seeds,
+            lr=args.lr,
+            horizon=args.horizon,
+            early_step=args.early_step,
+            feedback_scale=args.feedback_scale,
+        )
+        rows.extend(depth_exp.run_finite_time(run_args, x, y))
+
+    write_rows(args.outdir / "parameter_matched_rows.csv", rows)
+    metrics = depth_exp.summarize_finite(rows)
+    depth_exp.write_dict_csv(args.outdir / "parameter_matched_metrics.csv", metrics)  # type: ignore[arg-type]
+
+    summaries: list[dict[str, object]] = []
+    for depth in args.depths:
+        subset = [row for row in rows if row.depth == depth]
+        values = np.asarray([row.empirical_gap for row in subset])
+        init_means = np.asarray(
+            [
+                np.mean([row.empirical_gap for row in subset if row.init_seed == init_seed])
+                for init_seed in sorted({row.init_seed for row in subset})
+            ]
+        )
+        summaries.append(
+            {
+                "depth": depth,
+                "width": widths[depth],
+                "parameter_count": parameter_count(
+                    args.input_dim, args.output_dim, depth, widths[depth]
+                ),
+                "parameter_count_error": parameter_count(
+                    args.input_dim, args.output_dim, depth, widths[depth]
+                )
+                - target,
+                "rows": len(subset),
+                "empirical_gap_mean": float(values.mean()),
+                "empirical_gap_sem": float(
+                    init_means.std(ddof=1) / math.sqrt(len(init_means))
+                ),
+            }
+        )
+    depth_exp.write_dict_csv(args.outdir / "parameter_matched_summary.csv", summaries)
+
+    fig, ax = plt.subplots(figsize=(6.6, 4.7), dpi=180)
+    ax.errorbar(
+        [int(row["depth"]) for row in summaries],
+        [float(row["empirical_gap_mean"]) for row in summaries],
+        yerr=[float(row["empirical_gap_sem"]) for row in summaries],
+        marker="o",
+        capsize=3,
+        color="#6f4c9b",
+    )
+    for row in summaries:
+        ax.annotate(
+            f"w={row['width']}",
+            (int(row["depth"]), float(row["empirical_gap_mean"])),
+            xytext=(0, 9),
+            textcoords="offset points",
+            ha="center",
+            fontsize=8,
+        )
+    ax.set_xlabel("hidden-layer depth")
+    ax.set_ylabel("FA loss - BP loss at fixed horizon")
+    ax.set_title(f"Parameter-matched depth control (target P={target})")
+    ax.grid(alpha=0.16)
+    fig.tight_layout()
+    fig.savefig(args.outdir / "parameter_matched_gap_by_depth.png", bbox_inches="tight")
+    plt.close(fig)
+
+    payload = {
+        "config": {key: str(value) if isinstance(value, Path) else value for key, value in vars(args).items()},
+        "target_parameter_count": target,
+        "widths": {str(depth): width for depth, width in widths.items()},
+        "rows": len(rows),
+        "summary": summaries,
+    }
+    (args.outdir / "summary.json").write_text(json.dumps(payload, indent=2) + "\n")
+    print(f"summary: {args.outdir / 'summary.json'}")
+
+
+if __name__ == "__main__":
+    main()