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
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
|
"""Aggregate results/*.json into summary tables (+ figures if matplotlib is
present). Run on any box with access to the shared NFS home.
Usage: python experiments/analyze.py [results_dir]
"""
import glob
import json
import os
import sys
from collections import defaultdict
RES = sys.argv[1] if len(sys.argv) > 1 else "results"
try:
import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt
HAVE_MPL = True
except Exception:
HAVE_MPL = False
def load_all(res):
runs = []
for p in sorted(glob.glob(os.path.join(res, "*.json"))):
try:
with open(p) as f:
d = json.load(f)
d["_path"] = p
d["_tag"] = os.path.basename(p)[:-5]
runs.append(d)
except Exception as e:
print("skip", p, e)
return runs
def mean_final_cos(run):
f = run.get("final", {})
c = f.get("cos_r_negg")
return sum(c) / len(c) if c else None
def group(runs, keyfn):
g = defaultdict(list)
for r in runs:
k = keyfn(r)
if k is not None:
g[k].append(r)
return g
def agg(vals):
vals = [v for v in vals if v is not None]
if not vals:
return (None, None, 0)
m = sum(vals) / len(vals)
sd = (sum((v - m) ** 2 for v in vals) / len(vals)) ** 0.5
return (m, sd, len(vals))
def fmt(m, sd, n):
if m is None:
return " n/a "
return f"{m:.3f}±{sd:.3f}(n{n})"
def table_main(runs):
print("\n================ MAIN COMPARISON (mnist depth3 width256) ================")
sel = [r for r in runs if r["args"].get("depth") == 3 and r["args"].get("width") == 256
and r["args"].get("dataset") == "mnist" and r["args"].get("nuis_rho", 0) == 0
and r["_tag"].startswith(("main_", "A_main_"))]
by_mode = group(sel, lambda r: r["args"]["mode"])
print(f"{'mode':6s} {'test_acc':22s} {'mean_cos(r,-g)':22s}")
for mode in ["bp", "dfa", "sdil"]:
rs = by_mode.get(mode, [])
acc = agg([r["final"].get("test_acc") for r in rs])
cos = agg([mean_final_cos(r) for r in rs])
print(f"{mode:6s} {fmt(*acc):22s} {fmt(*cos):22s}")
def table_nuisance(runs):
print("\n================ RESIDUALIZATION UNDER NUISANCE ================")
print("(sdil; does subtracting the soma-predictable baseline preserve alignment as rho grows?)")
sel = [r for r in runs if r["_tag"].startswith(("nuis_", "A_nuis_"))]
by = group(sel, lambda r: (r["args"].get("nuis_rho"), r["args"].get("use_residual")))
rhos = sorted({k[0] for k in by})
print(f"{'rho':>5s} | {'residual=1 acc':22s} {'cos':16s} | {'residual=0 acc':22s} {'cos':16s}")
for rho in rhos:
r1 = by.get((rho, 1), [])
r0 = by.get((rho, 0), [])
a1 = agg([r["final"].get("test_acc") for r in r1]); c1 = agg([mean_final_cos(r) for r in r1])
a0 = agg([r["final"].get("test_acc") for r in r0]); c0 = agg([mean_final_cos(r) for r in r0])
print(f"{rho:5.1f} | {fmt(*a1):22s} {fmt(*c1):16s} | {fmt(*a0):22s} {fmt(*c0):16s}")
def table_trainedA(runs):
print("\n================ TRAINED vs FIXED APICAL PATHWAY ================")
fixed = [r for r in runs if "fixedA" in r["_tag"]]
af = agg([r["final"].get("test_acc") for r in fixed]); cf = agg([mean_final_cos(r) for r in fixed])
print(f"{'fixed-A (DFA-like)':22s} acc {fmt(*af):22s} cos {fmt(*cf):16s}")
tr = [r for r in runs if "trainedA_nd" in r["_tag"]]
by = group(tr, lambda r: r["args"].get("pert_ndirs"))
for nd in sorted(by):
rs = by[nd]
a = agg([r["final"].get("test_acc") for r in rs]); c = agg([mean_final_cos(r) for r in rs])
print(f"trained-A ndirs={nd:<3d} acc {fmt(*a):22s} cos {fmt(*c):16s}")
def table_depth(runs):
print("\n================ DEPTH SCALING ================")
sel = [r for r in runs if r["_tag"].startswith("depth_")]
if not sel:
return
by = group(sel, lambda r: (r["args"]["mode"], r["args"]["depth"]))
depths = sorted({k[1] for k in by})
print(f"{'depth':>6s} | " + " | ".join(f"{m:>26s}" for m in ["bp", "dfa", "sdil"]))
for d in depths:
cells = []
for m in ["bp", "dfa", "sdil"]:
rs = by.get((m, d), [])
a = agg([r["final"].get("test_acc") for r in rs])
c = agg([mean_final_cos(r) for r in rs])
cells.append(f"a{fmt(*a)[:5]} c{('%.2f'%c[0]) if c[0] is not None else ' n/a'}")
print(f"{d:6d} | " + " | ".join(f"{c:>26s}" for c in cells))
def plots(runs):
if not HAVE_MPL:
print("\n[matplotlib not available -> skipping figures]")
return
figdir = os.path.join(RES, "figs")
os.makedirs(figdir, exist_ok=True)
# Fig 1: alignment trajectory sdil vs dfa (seed0 main)
plt.figure(figsize=(6, 4))
for tag, lab in [("A_main_sdil_s0", "SDIL (trained A)"), ("A_main_dfa_s0", "DFA (fixed A)"),
("main_sdil_mnist_s0", "SDIL"), ("main_dfa_mnist_s0", "DFA")]:
r = next((x for x in runs if x["_tag"] == tag), None)
if not r:
continue
xs = [s["step"] for s in r["steps"] if "cos_r_negg" in s]
ys = [sum(s["cos_r_negg"]) / len(s["cos_r_negg"]) for s in r["steps"] if "cos_r_negg" in s]
if xs:
plt.plot(xs, ys, label=lab)
plt.xlabel("step"); plt.ylabel("mean cos(r, -grad)"); plt.title("Gradient alignment over training")
plt.legend(); plt.grid(alpha=.3); plt.tight_layout()
plt.savefig(os.path.join(figdir, "alignment_traj.png"), dpi=120); plt.close()
# Fig 2: acc & alignment vs rho (residual vs not)
sel = [r for r in runs if r["_tag"].startswith(("nuis_", "A_nuis_"))]
if sel:
by = group(sel, lambda r: (r["args"].get("nuis_rho"), r["args"].get("use_residual")))
rhos = sorted({k[0] for k in by})
fig, ax = plt.subplots(1, 2, figsize=(10, 4))
for ur, lab in [(1, "residual (SDIL)"), (0, "raw apical")]:
accs = [agg([r["final"].get("test_acc") for r in by.get((rho, ur), [])])[0] for rho in rhos]
coss = [agg([mean_final_cos(r) for r in by.get((rho, ur), [])])[0] for rho in rhos]
ax[0].plot(rhos, accs, marker="o", label=lab)
ax[1].plot(rhos, coss, marker="o", label=lab)
ax[0].set_xlabel("nuisance rho"); ax[0].set_ylabel("test acc"); ax[0].set_title("Accuracy vs nuisance")
ax[1].set_xlabel("nuisance rho"); ax[1].set_ylabel("mean cos(r,-grad)"); ax[1].set_title("Alignment vs nuisance")
for a in ax:
a.legend(); a.grid(alpha=.3)
plt.tight_layout(); plt.savefig(os.path.join(figdir, "nuisance.png"), dpi=120); plt.close()
# Fig 3: alignment vs depth
sel = [r for r in runs if r["_tag"].startswith("depth_")]
if sel:
by = group(sel, lambda r: (r["args"]["mode"], r["args"]["depth"]))
depths = sorted({k[1] for k in by})
plt.figure(figsize=(6, 4))
for m in ["dfa", "sdil"]:
ys = [agg([mean_final_cos(r) for r in by.get((m, d), [])])[0] for d in depths]
plt.plot(depths, ys, marker="o", label=m)
plt.xlabel("hidden depth"); plt.ylabel("mean cos(r,-grad)"); plt.title("Alignment vs depth")
plt.legend(); plt.grid(alpha=.3); plt.tight_layout()
plt.savefig(os.path.join(figdir, "depth.png"), dpi=120); plt.close()
print(f"\n[figures written to {figdir}]")
def main():
runs = load_all(RES)
print(f"loaded {len(runs)} runs from {RES}")
table_main(runs)
table_nuisance(runs)
table_trainedA(runs)
table_depth(runs)
plots(runs)
if __name__ == "__main__":
main()
|
