#!/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()