summaryrefslogtreecommitdiff
path: root/ep_run/fig_primer_arch.py
diff options
context:
space:
mode:
authorYuren Hao <yurenh2@illinois.edu>2026-07-15 10:29:06 -0500
committerYuren Hao <yurenh2@illinois.edu>2026-07-15 10:29:06 -0500
commit82da86dbd49a093a24e0827331667e8e9e6217df (patch)
tree96c40877ab6888d66c6e7992a8d6bfebb0453ba2 /ep_run/fig_primer_arch.py
parentbf83f39828fb79428fe65fe0fef49768351afc05 (diff)
Alexi primer deck: 10 slides (arch op-by-op, EP phases, schedule/parallelism, 72M blowup from real logs, loop-gain + beta-window hypotheses, diagnostics table, in-flight outcomes, $50k ladder) + figure scripts
Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_014FAPDWQ49M5Ye3NpTndTpn
Diffstat (limited to 'ep_run/fig_primer_arch.py')
-rw-r--r--ep_run/fig_primer_arch.py135
1 files changed, 135 insertions, 0 deletions
diff --git a/ep_run/fig_primer_arch.py b/ep_run/fig_primer_arch.py
new file mode 100644
index 0000000..1090f45
--- /dev/null
+++ b/ep_run/fig_primer_arch.py
@@ -0,0 +1,135 @@
+"""Primer figure 1: stock architecture + two-phase gradient measurement + PC contrast.
+Vector PDF + PNG preview to ../assets/. Restrained style, no glow."""
+import matplotlib
+matplotlib.use('Agg')
+import matplotlib.pyplot as plt
+plt.rcParams['svg.fonttype'] = 'none' # keep text editable in PPT
+from matplotlib.patches import FancyBboxPatch
+
+INK = '#3a3a3a'
+GRAY = '#8a8a8a'
+BOXF = '#f4f4f4'
+BOXE = '#9a9a9a'
+BLUE = '#2c6fbb'
+BLUEF = '#eaf1fa'
+ORAN = '#d95f02'
+RED = '#b03a2e'
+CREAM = '#faf6ee'
+CREME = '#c9b895'
+
+fig = plt.figure(figsize=(13.2, 10.0))
+ax = fig.add_axes([0, 0, 1, 1]); ax.set_xlim(0, 100); ax.set_ylim(0, 105); ax.axis('off')
+
+def box(x, y, w, h, fc=BOXF, ec=BOXE, lw=1.0, r=0.6):
+ b = FancyBboxPatch((x, y), w, h, boxstyle=f'round,pad=0.25,rounding_size={r}',
+ fc=fc, ec=ec, lw=lw)
+ ax.add_patch(b); return b
+
+def txt(x, y, s, size=8.5, color=INK, ha='center', va='center', w='normal', style='normal'):
+ ax.text(x, y, s, fontsize=size, color=color, ha=ha, va=va,
+ fontweight=w, fontstyle=style)
+
+def arrow(x0, y0, x1, y1, color=INK, lw=1.2, style='-|>', ms=11):
+ ax.annotate('', xy=(x1, y1), xytext=(x0, y0),
+ arrowprops=dict(arrowstyle=style, color=color, lw=lw,
+ mutation_scale=ms, shrinkA=0.5, shrinkB=0.5))
+
+# ============================= (A) the model =============================
+txt(16, 102.3, '(A) The model — a stock OLMo2-style decoder', 10.5, w='bold', ha='center')
+txt(16, 99.6, 'Inference is one ordinary forward pass.', 8.5, GRAY)
+
+cx = 14
+txt(cx, 93.5, 'logits', 9.5, w='bold')
+arrow(cx, 88.6, cx, 91.8)
+box(cx - 10, 84.6, 20, 4.0); txt(cx, 86.6, r'linear readout $W_{\mathrm{out}}$ (untied)', 8.5)
+arrow(cx, 81.0, cx, 84.2)
+box(cx - 10, 77.0, 20, 4.0); txt(cx, 79.0, 'final RMSNorm', 8.5)
+arrow(cx, 71.6, cx, 76.6)
+
+box(cx - 12.5, 53.0, 25, 18.6, fc='#eeeeee', ec='#777777', lw=1.3)
+txt(cx, 68.4, r'transformer block ($=f_l$) $\times\,12$', 9, w='bold')
+txt(cx, 63.9, r'$h\ =\ z + \mathrm{RMSNorm}(\,\mathrm{Attn}(z)\,)$', 9)
+txt(cx, 59.4, r"$z'\ =\ h + \mathrm{RMSNorm}(\,\mathrm{SwiGLU}(h)\,)$", 9)
+txt(cx, 55.3, 'causal SDPA · QK-RMSNorm · RoPE', 7.5, GRAY)
+
+arrow(cx, 48.0, cx, 52.6)
+box(cx - 10, 44.0, 20, 4.0); txt(cx, 46.0, 'token embedding', 8.5)
+
+txt(16, 40.0, '42.75M (TinyStories 4k BPE) / 72.11M (FineWeb-Edu 32k BPE)', 7.8, GRAY)
+txt(16, 36.6, 'Nothing in this graph is modified for training: the trained checkpoint\n'
+ 'is indistinguishable in form from a conventionally trained model.', 8.2)
+
+# ======================== (B) free phase ========================
+bx = 47
+txt(bx, 102.3, '(B) Free phase (= inference)', 10.5, w='bold')
+txt(bx, 99.6, 'give each layer a state $z_l$ (3 of $L$ drawn); settle the disagreement energy', 8.5, GRAY)
+txt(bx, 95.9, r'$E(z)\ =\ \sum_l\ \frac{1}{2}\,\|\,z_l - f_l(z_{l-1})\,\|^2$', 10.5)
+
+def state_col(x0, ys, labels, notes, fc=BLUEF, ec=BLUE, note_c=GRAY):
+ for y, lab, note in zip(ys, labels, notes):
+ box(x0 - 9.5, y - 2.0, 19, 4.0, fc=fc, ec=ec, lw=1.1)
+ txt(x0, y, lab, 9, color=INK)
+ if note: txt(x0 + 11.0, y, note, 7.8, note_c, ha='left')
+ for ya, yb in zip(ys[1:], ys[:-1]):
+ arrow(x0, ya + 2.5, x0, yb - 2.6, lw=1.1)
+
+ys = [86, 74, 62]
+state_col(bx - 6, ys, [r'$z_3 = f_3(z_2)$', r'$z_2 = f_2(z_1)$', r'$z_1 = f_1(\mathrm{emb})$'],
+ ['term = 0', 'term = 0', 'term = 0'])
+arrow(bx - 6, 53.6, bx - 6, 59.0, lw=1.1)
+txt(bx - 6, 51.5, r'$\mathrm{emb}(x)$', 9)
+
+txt(bx, 44.6, 'The minimum is exact: $E=0$, states $\\equiv$ forward activations,', 8.4)
+txt(bx, 41.6, 'and one bottom-up pass reaches it.', 8.4)
+txt(bx, 37.6, 'The free phase adds nothing and changes nothing at inference.', 8.2, GRAY)
+
+# ======================== (C) nudged phase ========================
+nx = 81
+txt(nx, 102.3, '(C) Nudged phase (training only)', 10.5, w='bold')
+txt(nx, 99.6, r'add the loss at strength $\beta \ll 1$ and settle again:', 8.5, GRAY)
+txt(nx, 95.9, r'$E(z)\ +\ \beta\cdot\mathrm{CE}(\mathrm{logits}(z_3),\,y)$', 10.5)
+
+txt(nx + 2.0, 92.3, r'pull $-\beta\,\nabla\mathrm{CE}$', 8.6, RED, ha='left')
+txt(nx + 2.0, 90.0, '(the only place the label enters)', 7.3, RED, ha='left')
+arrow(nx + 1.0, 91.4, nx - 4.0, 88.6, color=RED, lw=1.5)
+
+nys = [86, 74, 62]
+state_col(nx - 6, nys,
+ [r'$z_3^{\beta} = z_3 + d_3$', r'$z_2^{\beta} = z_2 + d_2$', r'$z_1^{\beta} = z_1 + d_1$'],
+ [None, None, None], fc='#fdeee2', ec=ORAN)
+arrow(nx - 6, 53.6, nx - 6, 59.0, lw=1.1)
+txt(nx - 6, 51.5, r'$\mathrm{emb}(x)$', 9)
+
+for (ya, yb, lab) in [(84.0, 76.6, r'$d_2 = J_3^{\top} d_3$'),
+ (72.0, 64.6, r'$d_1 = J_2^{\top} d_2$')]:
+ arrow(nx + 5.4, ya, nx + 5.4, yb, color=ORAN, lw=1.5)
+ txt(nx + 7.0, (ya + yb) / 2, lab, 8.6, ORAN, ha='left')
+
+txt(nx, 44.6, 'The top state is pulled toward lower loss; each layer\'s mismatch $d_l$', 8.4)
+txt(nx, 41.6, 'transmits DOWN through the same weights, and the stack re-settles.', 8.4)
+txt(nx, 37.6, r'($J^{\top}$ = the transpose read a bidirectional physical device provides)', 8.2, GRAY)
+
+# ======================== (D) the update ========================
+box(2, 20.5, 96, 13.5, fc=CREAM, ec=CREME, lw=1.2)
+txt(4, 31.6, '(D) The update — a difference measurement between the two settled states', 10, w='bold', ha='left')
+txt(50, 27.9, r'$\hat{g}\ =\ \left[\ \partial_\theta E(z^{\beta})\ -\ \partial_\theta E(z^{0})\ \right]\,/\,\beta$'
+ r'$\qquad\qquad(\partial_\theta E(z^0)\equiv 0\ \mathrm{here,\ since}\ E=0)$', 10.5)
+txt(50, 24.5, r'per layer: $\Delta\theta_l\ \propto\ \langle\ d_l\ ,\ \partial f_l(z_{l-1})/\partial\theta_l\ \rangle\ /\ \beta$', 9.5)
+txt(50, 21.9, r'Each layer updates from its own boundary mismatch — no global backward graph, no global tape, no loss'
+ '\n'
+ r'derivatives except the top nudge. $\beta\to 0$ gives the exact gradient; bias is $O(\beta)$; a $\pm\beta$ two-sided read cancels it to $O(\beta^2)$.',
+ 8.2)
+
+# ======================== (E) not PC ========================
+box(2, 2.0, 96, 16.2, fc='#f7f7f7', ec=BOXE, lw=1.1)
+txt(4, 15.7, '(E) This is not predictive coding — same energy family, different measurement', 10, w='bold', ha='left')
+txt(4.5, 12.4, '· PC (as typically run): ONE settled phase with the target clamped hard; the update uses the raw errors of that single state → finite-clamp', 8.2, ha='left')
+txt(4.5, 10.1, ' bias. Its "exact-BP" results require freezing predictions during error transport (fixed-prediction) — backprop re-expressed in local variables.', 8.2, ha='left')
+txt(4.5, 7.5, '· EP (here): TWO phases and an infinitesimal nudge; the update is a difference quotient in β → bias is measured and controllable, and the', 8.2, ha='left')
+txt(4.5, 5.2, ' settle stays fully self-consistent — which is what physical hardware actually does.', 8.2, ha='left')
+txt(4.5, 3.0, '· The free phase is the zero-reference of the measurement: on analog hardware the subtraction cancels state-independent device offsets.', 8.2, ha='left')
+
+fig.savefig('/home/yurenh2/ept/assets/fig_primer_arch.pdf')
+fig.savefig('/home/yurenh2/ept/assets/fig_primer_arch.png', dpi=190)
+fig.savefig('/home/yurenh2/ept/assets/fig_primer_arch.svg')
+print('saved fig_primer_arch.{pdf,png,svg}')