1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
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()
|
