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
75
76
77
78
79
80
81
82
83
84
85
86
|
#!/usr/bin/env python3
"""Compare capacity-transition curves at different SGD horizons."""
from __future__ import annotations
from pathlib import Path
import matplotlib.pyplot as plt
import pandas as pd
RUNS = [
("T=3000, 256 traj", Path("outputs/phase_transition_sgd_256_lr001_T3000")),
("T=10000 probe", Path("outputs/phase_transition_sgd_long_lr001_T10000_probe")),
("T=30000 probe", Path("outputs/phase_transition_sgd_verylong_lr001_T30000_probe")),
]
def main() -> None:
outdir = Path("outputs/phase_transition_training_time_comparison")
outdir.mkdir(parents=True, exist_ok=True)
fig, ax = plt.subplots(figsize=(9.2, 5.7), dpi=180)
for label, run_dir in RUNS:
df = pd.read_csv(run_dir / "width_summary.csv").sort_values(
"fa_capacity_margin", ascending=False
)
ax.plot(
df["fa_capacity_margin"],
df["train_gap_mean"],
marker="o",
linewidth=2.0,
label=label,
)
ax.axvline(0, color="black", linestyle="--", linewidth=1.1, label="hard FA margin=0")
ax.axhline(0, color="black", linewidth=1.0)
ax.set_xlim(1100, -430)
ax.set_xlabel("hard FA capacity margin P - K_FA - N*out (capacity decreases ->)")
ax.set_ylabel("FA train MSE - BP train MSE")
ax.set_title("Training time shifts the apparent FA capacity transition")
ax.legend(frameon=True)
ax.grid(alpha=0.18)
fig.tight_layout()
path = outdir / "train_gap_vs_margin_by_training_time.png"
fig.savefig(path)
plt.close(fig)
fig2, ax2 = plt.subplots(figsize=(9.2, 5.7), dpi=180)
for label, run_dir in RUNS:
df = pd.read_csv(run_dir / "width_summary.csv").sort_values(
"fa_capacity_margin", ascending=False
)
ax2.plot(
df["fa_capacity_margin"],
df["bp_train_mean"],
marker="o",
linewidth=1.8,
linestyle="--",
label=f"BP {label}",
)
ax2.plot(
df["fa_capacity_margin"],
df["fa_train_mean"],
marker="o",
linewidth=2.0,
label=f"FA {label}",
)
ax2.axvline(0, color="black", linestyle="--", linewidth=1.1)
ax2.set_xlim(1100, -430)
ax2.set_yscale("log")
ax2.set_xlabel("hard FA capacity margin P - K_FA - N*out (capacity decreases ->)")
ax2.set_ylabel("train MSE")
ax2.set_title("BP/FA train losses across training time")
ax2.legend(frameon=True, fontsize=8, ncol=2)
ax2.grid(alpha=0.18)
fig2.tight_layout()
path2 = outdir / "bp_fa_train_loss_by_training_time.png"
fig2.savefig(path2)
plt.close(fig2)
print(f"plot: {path}")
print(f"plot: {path2}")
if __name__ == "__main__":
main()
|