From 84fa70be406248e660488ddc4c56ca5ea0034189 Mon Sep 17 00:00:00 2001 From: YurenHao0426 Date: Tue, 9 Jun 2026 15:30:20 -0500 Subject: Add teacher-task test-gap sweep and figure MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit SGD teacher sweep (widths 8-96, 24 FA runs/width): train gap soft and monotone (0.249 -> 0.0069) while the test gap changes sign — FA generalizes better than BP at intermediate width (-0.120 +/- 0.013 at w=32) and settles slightly positive once both fit. Scopes the paper's optimization-gap claims and preempts the generalization question. Co-Authored-By: Claude Fable 5 --- scripts/plot_teacher_test_gap.py | 74 ++++++++++++++++++++++++++++++++++++++++ 1 file changed, 74 insertions(+) create mode 100644 scripts/plot_teacher_test_gap.py (limited to 'scripts/plot_teacher_test_gap.py') diff --git a/scripts/plot_teacher_test_gap.py b/scripts/plot_teacher_test_gap.py new file mode 100644 index 0000000..8c4e0de --- /dev/null +++ b/scripts/plot_teacher_test_gap.py @@ -0,0 +1,74 @@ +#!/usr/bin/env python3 +"""Two-panel teacher-task figure: train gap (soft ramp) vs test gap (sign flip). + +Reads width_summary.csv from a teacher-task sweep produced by +downstream_capacity_sweep.py and plots, side by side: + left - FA/BP train gap vs width (log y): the optimization cost, soft ramp; + right - FA/BP test gap vs width (linear y, error bars): task-dependent, + crosses zero (FA can generalize better at intermediate width). +""" + +from __future__ import annotations + +import argparse +import csv +from pathlib import Path + +import matplotlib + +matplotlib.use("Agg") +import matplotlib.pyplot as plt +import numpy as np + + +def parse_args() -> argparse.Namespace: + p = argparse.ArgumentParser(description="Teacher-task train vs test gap figure.") + p.add_argument("--run-dir", type=Path, + default=Path("outputs/teacher_test_gap_sgd_T8000")) + return p.parse_args() + + +def main() -> None: + args = parse_args() + rows = [] + with (args.run_dir / "width_summary.csv").open() as fh: + for row in csv.DictReader(fh): + rows.append({k: float(v) for k, v in row.items()}) + rows.sort(key=lambda r: r["width"]) + + widths = np.array([r["width"] for r in rows]) + n_fa = np.array([r["runs_fa"] for r in rows]) + train_gap = np.array([r["train_gap_mean"] for r in rows]) + train_se = np.array([r["train_gap_std"] for r in rows]) / np.sqrt(n_fa) + test_gap = np.array([r["gap_mean"] for r in rows]) + test_se = np.array([r["gap_std"] for r in rows]) / np.sqrt(n_fa) + + fig, axes = plt.subplots(1, 2, figsize=(11.2, 4.4), dpi=180) + axes[0].errorbar(widths, train_gap, yerr=train_se, fmt="o-", color="#c65f16", + capsize=2.5) + axes[0].set_xscale("log", base=2) + axes[0].set_yscale("log") + axes[0].set_xlabel("hidden width") + axes[0].set_ylabel("FA train MSE - BP train MSE") + axes[0].set_title("Optimization gap: soft, monotone, never zero") + axes[0].grid(alpha=0.18, which="both") + + axes[1].errorbar(widths, test_gap, yerr=test_se, fmt="o-", color="#174f91", + capsize=2.5) + axes[1].axhline(0.0, color="black", linewidth=1) + axes[1].set_xscale("log", base=2) + axes[1].set_xlabel("hidden width") + axes[1].set_ylabel("FA test MSE - BP test MSE") + axes[1].set_title("Generalization gap: task-dependent, changes sign") + axes[1].grid(alpha=0.18) + + fig.suptitle("Teacher task (MLP teacher, SGD, T=8000): optimization cost vs test effect", + y=1.02) + fig.tight_layout() + out = args.run_dir / "teacher_train_vs_test_gap.png" + fig.savefig(out, bbox_inches="tight") + print(f"plot: {out}") + + +if __name__ == "__main__": + main() -- cgit v1.2.3