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
|
"""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}')
|
