| |
| """Phase-1 LR-schedule ablation. Every number read from p1ev_*/summary.csv.""" |
| import csv,os,math |
| import matplotlib; matplotlib.use("Agg") |
| import matplotlib.pyplot as plt |
| S={"surface":"#fcfcfb","ink":"#0b0b0b","ink2":"#52514e","grid":"#e4e4e1"} |
| C={"abl_const":"#b3243b","abl_cosine":"#eb6834","p1final":"#12508f","base680":"#9dc3ea"} |
| LAB={"abl_const":"const 8e-5 (Muon+AdamW)","abl_cosine":"cosine 8e-5->8e-6 (Muon+AdamW)", |
| "p1final":"p1final 1.8e-4 WSD (AdamW)","base680":"base680 NO training"} |
| def sty(ax): |
| ax.set_facecolor(S["surface"]) |
| for s in ax.spines.values(): s.set_color(S["grid"]) |
| ax.grid(True,color=S["grid"],lw=.7,alpha=.9); ax.tick_params(colors=S["ink2"],labelsize=9) |
| def wilson(k,n=100,z=1.96): |
| p=k/n; d=1+z*z/n; c=(p+z*z/(2*n))/d |
| h=z*math.sqrt(p*(1-p)/n+z*z/(4*n*n))/d |
| return 100*max(0,c-h),100*min(1,c+h) |
| PL=[4096,8192,16360]; D={} |
| for pl in PL: |
| r=list(csv.reader(open(f"p1ev_{pl}/summary.csv"))); h=r[0] |
| row=[x for x in r[1:] if x and x[0]=="niah_single"][0] |
| for k,v in zip(h[1:],row[1:]): D.setdefault(k,{})[pl]=float(v) |
|
|
| fig,(a,b,c)=plt.subplots(1,3,figsize=(17.6,5.6),facecolor=S["surface"]) |
| X=range(len(PL)) |
| for m in ("base680","p1final","abl_const","abl_cosine"): |
| ys=[D[m][p] for p in PL] |
| lo=[y-wilson(round(y))[0] for y in ys]; hi=[wilson(round(y))[1]-y for y in ys] |
| a.errorbar(list(X),ys,yerr=[lo,hi],color=C[m],lw=2.4,ms=8,marker="o",capsize=3,label=LAB[m]) |
| sty(a); a.set_xticks(list(X)); a.set_xticklabels(["4k","8k","16k"]); a.set_ylim(60,105) |
| a.set_xlabel("prompt length",color=S["ink2"],fontsize=10) |
| a.set_ylabel("niah_single (%)",color=S["ink2"],fontsize=10) |
| a.set_title("A. All four models, Wilson 95% CI",color=S["ink"],fontsize=11.5,loc="left",pad=9) |
| a.legend(frameon=False,fontsize=8,labelcolor=S["ink2"],loc="lower left") |
|
|
| d=[D["abl_const"][p]-D["abl_cosine"][p] for p in PL] |
| b.bar(list(X),d,color=["#b3243b" if x>0 else "#eb6834" for x in d],width=.5) |
| b.axhline(0,color=S["ink2"],lw=1) |
| b.axhspan(-11,11,color="#f0f0ee",zorder=0) |
| b.text(1,8.4,"Wilson 95% CI band at N=100",ha="center",fontsize=8.4,color=S["ink2"],style="italic") |
| sty(b); b.set_xticks(list(X)); b.set_xticklabels(["4k","8k","16k"]); b.set_ylim(-14,14) |
| b.set_ylabel("const - cosine (pp)",color=S["ink2"],fontsize=10) |
| b.set_title("B. THE ABLATION: schedule shape -- null",color=S["ink"],fontsize=11.5,loc="left",pad=9) |
| for i,v in enumerate(d): b.text(i,v+(0.7 if v>=0 else -1.6),f"{v:+.0f}",ha="center",fontsize=10,color=S["ink2"]) |
|
|
| names=["base680","p1final","abl_const","abl_cosine"] |
| tot=[sum(D[m][p] for p in PL) for m in names] |
| c.bar(range(4),tot,color=[C[m] for m in names]) |
| for i,v in enumerate(tot): c.text(i,v+3,f"{v:.0f}",ha="center",fontsize=10.5,color=S["ink2"],weight="bold") |
| sty(c); c.set_xticks(range(4)) |
| c.set_xticklabels(["base680\n(none)","p1final\n1.8e-4 AdamW","const\n8e-5 Muon","cosine\n8e-5 Muon"],fontsize=8.4) |
| c.set_ylim(0,310); c.set_ylabel("sum over 4k/8k/16k (pp)",color=S["ink2"],fontsize=10) |
| c.set_title("C. Both new arms sit BELOW the untrained base",color=S["ink"],fontsize=11.5,loc="left",pad=9) |
| fig.text(.005,.005,"680M, 4k->16k, 1,200 steps, identical data (ablmix16, token-identical to p1mix) and identical haystack " |
| "(199,754,694 B). N=100/point. CAVEAT: the two new arms differ from p1final in BOTH optimizer (Muon vs AdamW) and LR " |
| "(8e-5 vs 1.8e-4) -- this run cannot separate them.",fontsize=7.8,color=S["ink2"]) |
| fig.tight_layout(); fig.savefig("plots_abl/phase1_ablation.png",dpi=150,bbox_inches="tight") |
| print(" wrote plots_abl/phase1_ablation.png") |
|
|
