AbstractPhil commited on
Commit
20e3f8f
·
verified ·
1 Parent(s): c0deed5

Create ar_differentiation_bed.py

Browse files
Files changed (1) hide show
  1. ar_differentiation_bed.py +493 -0
ar_differentiation_bed.py ADDED
@@ -0,0 +1,493 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ """ar_differentiation_bed.py — THE FOCUS (2026-07-09 redirect, Phil verbatim):
2
+ "refining the autoregressive techniques for differentiation rather than attempting
3
+ to just mash numbers together."
4
+
5
+ Differentiation is cultivated by PREDICTIVE pressure along the sequence — the
6
+ address parameterizing the next-byte distribution (Law 2: chain-rule advantage pays
7
+ ONLY where the composed address directly parameterizes the predictive distribution).
8
+ This bed puts the aleph in the autoregressive gradient path and measures what
9
+ differentiates. It is the Law-2 construction (codebook-pressure C3) + Tree 3d in
10
+ one harness; the Jun-19 "discuss before building" gate was resolved by the redirect.
11
+
12
+ Byte-level causal LM on wikitext-2-raw (HF parquet, CDN-fast), block 256. ARMS:
13
+ sdpa — standard causal transformer control (matched trunk).
14
+ hub — attention replaced by CAUSAL HUB: linear attention whose feature map
15
+ is the 2K-oriented aleph address, prefix-sum memories (no selection
16
+ event; O(n*K*d)). Differentiation cultivated INSIDE attention.
17
+ addr_head — sdpa trunk, but the OUTPUT HEAD reads ONLY the signed aleph
18
+ coefficient vector w_k = sinh(u_k)/sum_j cosh(u_j) of the final
19
+ hidden state (K -> 256 logits). The address MUST carry every bit of
20
+ next-byte information — the hardest Law-2 bottleneck.
21
+
22
+ JUDGED BY: val bits-per-byte per arm (task) + CULTIVATION VITALS on every aleph
23
+ codebook (readouts, never losses): axis aliveness/hppl, drift-from-init +
24
+ binding fraction @0.29154, winner-|cos| saturation (sign-code emergence), shadow
25
+ path diversity (fixed high-bits hash). Never by recon.
26
+
27
+ Riders: pure Adam wd=0; no BN/Dropout/GAP on geometric paths; orthogonal init;
28
+ Colab-cell-safe (paste-ahead imports, no bare argparse, no __file__ reliance);
29
+ GPU-only for verdict runs; data_root OUTSIDE the mind repo.
30
+
31
+ Terminal: python ar_differentiation_bed.py # shapes/parse smoke
32
+ python ar_differentiation_bed.py --train # verdict run
33
+ Colab: paste geolip_vitals.py cell, then this file (smoke auto-runs),
34
+ then train(steps=2000, data_root="/content/data") in the next cell.
35
+
36
+ Author: AbstractPhil + Claude
37
+ Home: https://huggingface.co/AbstractPhil
38
+ License: MIT
39
+
40
+ """
41
+ from __future__ import annotations
42
+ import math
43
+ import torch
44
+ import torch.nn as nn
45
+ import torch.nn.functional as F
46
+
47
+ if "anchor_drift" not in globals():
48
+ try:
49
+ from geolip_vitals import anchor_drift, axis_aliveness, path_diversity
50
+ except ImportError:
51
+ _here = globals().get("__file__")
52
+ if _here is not None:
53
+ import sys, pathlib
54
+ sys.path.insert(0, str(pathlib.Path(_here).parent))
55
+ from geolip_vitals import anchor_drift, axis_aliveness, path_diversity
56
+ else:
57
+ raise ImportError(
58
+ "geolip_vitals not found — paste/run its cell first, or "
59
+ "hf_hub_download tools/geolip_vitals.py from AbstractPhil/claude-mind.")
60
+
61
+ VOCAB = 256 # bytes
62
+
63
+
64
+ # ------------------------------------------------------------------ aleph address
65
+ def _super_fibonacci_s3(n: int) -> torch.Tensor:
66
+ """Near-uniform unit quaternions (Alexa CVPR'22; constants per canon) —
67
+ starts the codebook INSIDE the RP^3 attractor basin. D=4 only."""
68
+ PHI, PSI = math.sqrt(2.0), 1.533751168755204288118041
69
+ i = torch.arange(n, dtype=torch.float64)
70
+ s = (i + 0.5) / n
71
+ r, R = torch.sqrt(s), torch.sqrt(1.0 - s)
72
+ a, b = 2 * math.pi * i / PHI, 2 * math.pi * i / PSI
73
+ q = torch.stack([r * torch.sin(a), r * torch.cos(a),
74
+ R * torch.sin(b), R * torch.cos(b)], dim=-1)
75
+ return F.normalize(q, dim=-1).float()
76
+
77
+
78
+ class AlephAddress(nn.Module):
79
+ """Closed-form aleph over 2K oriented half-axes (canon/aleph_core.md).
80
+ signed(x): (..., K) w_k = sinh(u_k)/sum_j cosh(u_j) — the Law-2 head feature.
81
+ oriented(x): ((..., K), (..., K)) positive halves of the 2K softmax — HUB map."""
82
+
83
+ def __init__(self, K: int, D: int, tau: float = 0.1, init: str = "random"):
84
+ super().__init__()
85
+ self.K, self.D, self.tau = K, D, tau
86
+ if init == "fibonacci":
87
+ assert D == 4, "fibonacci init lives on S^3 (D=4)"
88
+ A = _super_fibonacci_s3(K)
89
+ else:
90
+ A = F.normalize(torch.randn(K, D), dim=-1)
91
+ self.codebook = nn.Parameter(A)
92
+ self.register_buffer("home", self.codebook.detach().clone())
93
+
94
+ def _u(self, x):
95
+ A = F.normalize(self.codebook, dim=-1)
96
+ return (F.normalize(x, dim=-1) @ A.transpose(-1, -2)) / self.tau
97
+
98
+ def oriented(self, x):
99
+ u = self._u(x)
100
+ m = u.abs().amax(dim=-1, keepdim=True)
101
+ ep, en = torch.exp(u - m), torch.exp(-u - m)
102
+ Z = (ep + en).sum(dim=-1, keepdim=True)
103
+ return ep / Z, en / Z
104
+
105
+ def signed(self, x):
106
+ u = self._u(x)
107
+ m = u.abs().amax(dim=-1, keepdim=True)
108
+ ep, en = torch.exp(u - m), torch.exp(-u - m)
109
+ return (ep - en) / (ep + en).sum(dim=-1, keepdim=True)
110
+
111
+ def signed_at(self, x, taus):
112
+ """Multi-tau stroboscope (rule of 3): signed coefficients at several
113
+ temperatures, concatenated — softer taus keep the vector dense while a
114
+ hard tau supplies the sign-code sharpness. v2 refinement (b)."""
115
+ A = F.normalize(self.codebook, dim=-1)
116
+ cos = F.normalize(x, dim=-1) @ A.transpose(-1, -2)
117
+ outs = []
118
+ for t in taus:
119
+ u = cos / t
120
+ m = u.abs().amax(dim=-1, keepdim=True)
121
+ ep, en = torch.exp(u - m), torch.exp(-u - m)
122
+ outs.append((ep - en) / (ep + en).sum(dim=-1, keepdim=True))
123
+ return torch.cat(outs, dim=-1)
124
+
125
+ def m_hat(self, x):
126
+ """Closed-form soft read (decoders read M_hat, never M). v2 control (c)."""
127
+ u = self._u(x)
128
+ m = u.abs().amax(dim=-1, keepdim=True)
129
+ ep, en = torch.exp(u - m), torch.exp(-u - m)
130
+ A = F.normalize(self.codebook, dim=-1)
131
+ return ((ep - en) @ A) / (ep + en).sum(dim=-1, keepdim=True)
132
+
133
+ def m_hard_ste(self, x):
134
+ """Canon hard mode: M_hard = sign(cos_win) * A[win], straight-through to
135
+ the soft read — forward fully discrete SIGN CODE, backward soft gradient.
136
+ Legal per theme A (reconstructive sign code, not a one-hot roster pick)."""
137
+ u = self._u(x)
138
+ soft = self.m_hat(x)
139
+ win = u.abs().argmax(dim=-1)
140
+ A = F.normalize(self.codebook, dim=-1)
141
+ sign = torch.sign(torch.gather(u, -1, win.unsqueeze(-1))).squeeze(-1)
142
+ hard = sign.unsqueeze(-1) * A[win]
143
+ return hard + soft - soft.detach()
144
+
145
+ @torch.no_grad()
146
+ def vitals(self, x_sample) -> dict:
147
+ u = self._u(x_sample.reshape(-1, x_sample.shape[-1]))
148
+ p, n = self.oriented(x_sample.reshape(-1, x_sample.shape[-1]))
149
+ two_k = torch.cat([p, n], dim=-1)
150
+ win = two_k.argmax(dim=-1)
151
+ cos_win = (u.abs().amax(dim=-1) * self.tau) # winner |cos| — sign-code sat.
152
+ d = anchor_drift(self.codebook, self.home)
153
+ return {"drift": round(d["mean"], 4),
154
+ "binding_frac": round(d["binding_fraction"], 4),
155
+ "aliveness": axis_aliveness(two_k),
156
+ "win_cos_mean": round(cos_win.mean().item(), 4),
157
+ "paths": path_diversity(win)}
158
+
159
+
160
+ # ------------------------------------------------------------------------- blocks
161
+ class CausalSDPA(nn.Module):
162
+ def __init__(self, d: int, heads: int = 4):
163
+ super().__init__()
164
+ self.h = heads
165
+ self.qkv = nn.Linear(d, 3 * d, bias=False)
166
+ self.o = nn.Linear(d, d, bias=False)
167
+ nn.init.orthogonal_(self.qkv.weight); nn.init.orthogonal_(self.o.weight)
168
+
169
+ def forward(self, x):
170
+ B, n, d = x.shape
171
+ q, k, v = self.qkv(x).chunk(3, dim=-1)
172
+ q, k, v = (t.view(B, n, self.h, d // self.h).transpose(1, 2) for t in (q, k, v))
173
+ y = F.scaled_dot_product_attention(q, k, v, is_causal=True)
174
+ return self.o(y.transpose(1, 2).reshape(B, n, d))
175
+
176
+
177
+ class CausalHUB(nn.Module):
178
+ """Causal aleph linear attention: prefix-sum memories over the two K-wide
179
+ halves of the oriented address; 2K never materialized; no selection event."""
180
+
181
+ def __init__(self, d: int, K: int = 32, D: int = 4, tau: float = 0.1):
182
+ super().__init__()
183
+ self.addr = AlephAddress(K, D, tau)
184
+ self.q = nn.Linear(d, D, bias=False)
185
+ self.k = nn.Linear(d, D, bias=False)
186
+ self.v = nn.Linear(d, d, bias=False)
187
+ self.o = nn.Linear(d, d, bias=False)
188
+ for m in (self.q, self.k, self.v, self.o):
189
+ nn.init.orthogonal_(m.weight)
190
+
191
+ def forward(self, x):
192
+ qp, qn = self.addr.oriented(self.q(x)) # (B, n, K)
193
+ kp, kn = self.addr.oriented(self.k(x))
194
+ v = self.v(x) # (B, n, d)
195
+ Sp = torch.cumsum(torch.einsum("bnk,bnd->bnkd", kp, v), dim=1)
196
+ Sn = torch.cumsum(torch.einsum("bnk,bnd->bnkd", kn, v), dim=1)
197
+ zp = torch.cumsum(kp, dim=1)
198
+ zn = torch.cumsum(kn, dim=1)
199
+ num = torch.einsum("bnk,bnkd->bnd", qp, Sp) + torch.einsum("bnk,bnkd->bnd", qn, Sn)
200
+ den = (qp * zp).sum(-1, keepdim=True) + (qn * zn).sum(-1, keepdim=True)
201
+ return self.o(num / den.clamp_min(1e-12))
202
+
203
+
204
+ class MslRelay(nn.Module):
205
+ """Depth-composition unit (chain-rule probe): multi-slot M_hat read entering
206
+ the trunk as a NEAR-ZERO gated residual (gate init -3.0, sigma~0.047 — theme D:
207
+ geometry enters as a nudge and grows only if it earns gradient)."""
208
+
209
+ def __init__(self, d: int, n_slots: int = 16, K: int = 64):
210
+ super().__init__()
211
+ self.n_slots = n_slots
212
+ self.proj = nn.Linear(d, n_slots * 4, bias=False)
213
+ self.out = nn.Linear(n_slots * 4, d, bias=False)
214
+ nn.init.orthogonal_(self.proj.weight)
215
+ nn.init.orthogonal_(self.out.weight)
216
+ self.addr = AlephAddress(K, 4)
217
+ self.gate = nn.Parameter(torch.tensor(-3.0))
218
+
219
+ def forward(self, x):
220
+ B, n, _ = x.shape
221
+ slots = self.proj(x).view(B, n, self.n_slots, 4)
222
+ m = self.addr.m_hat(slots).reshape(B, n, -1)
223
+ return x + self.gate.sigmoid() * self.out(m)
224
+
225
+
226
+ class Block(nn.Module):
227
+ def __init__(self, d: int, attn: nn.Module):
228
+ super().__init__()
229
+ self.n1, self.n2 = nn.LayerNorm(d), nn.LayerNorm(d)
230
+ self.attn = attn
231
+ self.mlp = nn.Sequential(nn.Linear(d, 4 * d), nn.GELU(), nn.Linear(4 * d, d))
232
+
233
+ def forward(self, x):
234
+ x = x + self.attn(self.n1(x))
235
+ return x + self.mlp(self.n2(x))
236
+
237
+
238
+ class ByteLM(nn.Module):
239
+ def __init__(self, arm: str, d: int = 192, layers: int = 4, block: int = 256,
240
+ K: int = 32, D: int = 4):
241
+ super().__init__()
242
+ # "<arm>_tri" suffix = trigram byte embedding (AlephLM byte_emb x3 lineage):
243
+ # token embedding is the sum of embeddings of bytes t, t-1, t-2.
244
+ self.trigram = arm.endswith("_tri")
245
+ if self.trigram:
246
+ arm = arm[:-4]
247
+ # "_fib" = super-Fibonacci S^3 codebook init (basin test: starts INSIDE
248
+ # the RP^3 attractor; primary observable is init->final geodesic drift).
249
+ self.fib = arm.endswith("_fib")
250
+ if self.fib:
251
+ arm = arm[:-4]
252
+ # "relay*" = stacked addresses in depth: MslRelay after every block.
253
+ # relay -> sdpa trunk + standard head; relay_msl64 -> + addressed head.
254
+ self.use_relay = arm.startswith("relay")
255
+ if arm == "relay":
256
+ arm = "sdpa"
257
+ elif arm == "relay_msl64":
258
+ arm = "addr_msl64"
259
+ self.arm, self.block = arm, block
260
+ self.emb = nn.Embedding(VOCAB, d)
261
+ if self.trigram:
262
+ self.emb1 = nn.Embedding(VOCAB, d)
263
+ self.emb2 = nn.Embedding(VOCAB, d)
264
+ self.pos = nn.Parameter(torch.zeros(1, block, d) + 0.01 * torch.randn(1, block, d))
265
+ mk_attn = (lambda: CausalHUB(d, K, D)) if arm == "hub" else (lambda: CausalSDPA(d))
266
+ self.blocks = nn.ModuleList([Block(d, mk_attn()) for _ in range(layers)])
267
+ if self.use_relay:
268
+ self.relays = nn.ModuleList([MslRelay(d) for _ in range(layers)])
269
+ self.nf = nn.LayerNorm(d)
270
+ if arm == "addr_head":
271
+ self.head_addr = AlephAddress(K, d) # v1: codebook in model dim — COLLAPSED
272
+ self.head = nn.Linear(K, VOCAB, bias=True)
273
+ elif arm in ("addr_d4", "addr_3tau", "addr_mhat"):
274
+ # v2 refinements: LOW-D HOME — learned projection to the canon D=4 home
275
+ # before addressing (mirrors the healthy HUB arms), K=64.
276
+ self.head_proj = nn.Linear(d, 4, bias=False)
277
+ nn.init.orthogonal_(self.head_proj.weight)
278
+ self.head_addr = AlephAddress(64, 4)
279
+ if arm == "addr_d4":
280
+ self.head = nn.Linear(64, VOCAB, bias=True) # w alone, D=4 home
281
+ elif arm == "addr_3tau":
282
+ self.taus = (0.05, 0.1, 0.3) # rule-of-3 strobe
283
+ self.head = nn.Linear(64 * 3, VOCAB, bias=True)
284
+ else: # addr_mhat
285
+ self.head = nn.Linear(4, VOCAB, bias=True) # tightest: M_hat
286
+ elif arm.startswith("addr_msl"):
287
+ # v3: MULTI-SLOT heads — the 16s funnel widening: P parallel D=4 slots
288
+ # over a SHARED codebook. addr_msl consumes the reconstructive M_hat per
289
+ # slot (Px4 dims); addr_msl_w consumes signed w per slot (Px64) — tests
290
+ # whether slot-parallel consumption alone rescues the coefficient path.
291
+ # addr_msl<P> = slot-count dose-response. addr_mslh<P> = HARD sign-code
292
+ # consumption (straight-through M_hard per slot).
293
+ self.hard = arm.startswith("addr_mslh")
294
+ if arm in ("addr_msl", "addr_msl_w"):
295
+ self.n_slots = 16
296
+ else:
297
+ self.n_slots = int(arm[len("addr_mslh" if self.hard else "addr_msl"):])
298
+ self.head_proj = nn.Linear(d, self.n_slots * 4, bias=False)
299
+ nn.init.orthogonal_(self.head_proj.weight)
300
+ self.head_addr = AlephAddress(
301
+ 64, 4, init="fibonacci" if self.fib else "random")
302
+ width = self.n_slots * (64 if arm == "addr_msl_w" else 4)
303
+ self.head = nn.Linear(width, VOCAB, bias=True)
304
+ elif arm == "addr_3tau_mhat":
305
+ # v3: combine the two v2 winners — 3-tau stroboscope + reconstructive read.
306
+ self.head_proj = nn.Linear(d, 4, bias=False)
307
+ nn.init.orthogonal_(self.head_proj.weight)
308
+ self.head_addr = AlephAddress(64, 4)
309
+ self.taus = (0.05, 0.1, 0.3)
310
+ self.head = nn.Linear(64 * 3 + 4, VOCAB, bias=True)
311
+ else:
312
+ self.head = nn.Linear(d, VOCAB, bias=True)
313
+ self._last_h = None
314
+
315
+ def forward(self, idx):
316
+ x = self.emb(idx)
317
+ if self.trigram: # past-only shifts — causality preserved
318
+ x = x + self.emb1(F.pad(idx, (1, 0), value=0)[:, :-1]) \
319
+ + self.emb2(F.pad(idx, (2, 0), value=0)[:, :-2])
320
+ x = x + self.pos[:, : idx.shape[1]]
321
+ if self.use_relay:
322
+ for b, r in zip(self.blocks, self.relays):
323
+ x = r(b(x))
324
+ else:
325
+ for b in self.blocks:
326
+ x = b(x)
327
+ h = self.nf(x)
328
+ self._last_h = h.detach()
329
+ if self.arm == "addr_head":
330
+ return self.head(self.head_addr.signed(h))
331
+ if self.arm == "addr_d4":
332
+ return self.head(self.head_addr.signed(self.head_proj(h)))
333
+ if self.arm == "addr_3tau":
334
+ return self.head(self.head_addr.signed_at(self.head_proj(h), self.taus))
335
+ if self.arm == "addr_mhat":
336
+ return self.head(self.head_addr.m_hat(self.head_proj(h)))
337
+ if self.arm.startswith("addr_msl"):
338
+ B, n, _ = h.shape
339
+ slots = self.head_proj(h).view(B, n, self.n_slots, 4)
340
+ if self.arm == "addr_msl_w":
341
+ feats = self.head_addr.signed(slots).reshape(B, n, -1)
342
+ elif getattr(self, "hard", False):
343
+ feats = self.head_addr.m_hard_ste(slots).reshape(B, n, -1)
344
+ else:
345
+ feats = self.head_addr.m_hat(slots).reshape(B, n, -1)
346
+ return self.head(feats)
347
+ if self.arm == "addr_3tau_mhat":
348
+ p = self.head_proj(h)
349
+ feats = torch.cat([self.head_addr.signed_at(p, self.taus),
350
+ self.head_addr.m_hat(p)], dim=-1)
351
+ return self.head(feats)
352
+ return self.head(h)
353
+
354
+ @torch.no_grad()
355
+ def vitals(self) -> dict:
356
+ out = {}
357
+ if self.arm == "hub":
358
+ for i, b in enumerate(self.blocks):
359
+ if self._last_h is not None:
360
+ out[f"L{i}"] = b.attn.addr.vitals(b.attn.q(self._last_h[:2]))
361
+ elif self.arm == "addr_head" and self._last_h is not None:
362
+ out["head"] = self.head_addr.vitals(self._last_h[:2])
363
+ elif self.arm in ("addr_d4", "addr_3tau", "addr_mhat",
364
+ "addr_3tau_mhat") and self._last_h is not None:
365
+ out["head"] = self.head_addr.vitals(self.head_proj(self._last_h[:2]))
366
+ elif self.arm.startswith("addr_msl") and self._last_h is not None:
367
+ slots = self.head_proj(self._last_h[:2])
368
+ out["head"] = self.head_addr.vitals(
369
+ slots.reshape(*slots.shape[:-1], self.n_slots, 4))
370
+ if self.use_relay and self._last_h is not None:
371
+ for i, r in enumerate(self.relays):
372
+ s = r.proj(self._last_h[:2])
373
+ v = r.addr.vitals(s.reshape(*s.shape[:-1], r.n_slots, 4))
374
+ out[f"relay{i}"] = {"gate": round(r.gate.sigmoid().item(), 4),
375
+ "drift": v["drift"],
376
+ "binding_frac": v["binding_frac"],
377
+ "ppl": round(v["aliveness"]["usage_ppl"], 1)}
378
+ return out
379
+
380
+
381
+ # --------------------------------------------------------------------------- data
382
+ def _wikitext_bytes(data_root: str):
383
+ """wikitext-2-raw as flat uint8 tensors via the HF parquet CDN."""
384
+ from huggingface_hub import hf_hub_download
385
+ import pyarrow.parquet as pq
386
+
387
+ def load(split):
388
+ p = hf_hub_download("Salesforce/wikitext",
389
+ f"wikitext-2-raw-v1/{split}-00000-of-00001.parquet",
390
+ repo_type="dataset", local_dir=data_root)
391
+ text = "".join(pq.read_table(p).column("text").to_pylist())
392
+ return torch.frombuffer(bytearray(text.encode("utf-8")), dtype=torch.uint8).clone()
393
+
394
+ return load("train"), load("validation")
395
+
396
+
397
+ def _batch(data: torch.Tensor, batch: int, block: int, device, g: torch.Generator):
398
+ ix = torch.randint(0, data.numel() - block - 1, (batch,), generator=g)
399
+ x = torch.stack([data[i:i + block] for i in ix]).long().to(device)
400
+ y = torch.stack([data[i + 1:i + block + 1] for i in ix]).long().to(device)
401
+ return x, y
402
+
403
+
404
+ # -------------------------------------------------------------------- train/smoke
405
+ def train(arms=("sdpa", "hub", "addr_head"), steps: int = 2000, batch: int = 32,
406
+ block: int = 256, device: str = "cuda", data_root: str = "./data",
407
+ seed: int = 0, eval_every: int = 500, save: bool = True):
408
+ """Verdict run — GPU only. Pure Adam wd=0. Reports val bits-per-byte + vitals.
409
+ save=True writes {data_root}/ar_ckpts/{arm}_s{seed}_t{steps}.pt per arm —
410
+ the cultivated codebooks are SPECIMENS for the projective reading instruments."""
411
+ import os
412
+ if device == "cuda" and not torch.cuda.is_available():
413
+ raise RuntimeError("Verdict runs are GPU-only (never CPU-train for accuracy).")
414
+ ckpt_dir = os.path.join(data_root, "ar_ckpts")
415
+ os.makedirs(ckpt_dir, exist_ok=True)
416
+ tr, va = _wikitext_bytes(data_root)
417
+ print(f"data ready: train {tr.numel():,} bytes, val {va.numel():,} bytes", flush=True)
418
+ results = {}
419
+ for arm in arms:
420
+ torch.manual_seed(seed)
421
+ g = torch.Generator().manual_seed(seed)
422
+ model = ByteLM(arm, block=block).to(device)
423
+ n_params = sum(p.numel() for p in model.parameters())
424
+ opt = torch.optim.Adam(model.parameters(), lr=3e-4, weight_decay=0.0)
425
+ for step in range(1, steps + 1):
426
+ x, y = _batch(tr, batch, block, device, g)
427
+ logits = model(x)
428
+ loss = F.cross_entropy(logits.reshape(-1, VOCAB), y.reshape(-1))
429
+ opt.zero_grad(set_to_none=True)
430
+ loss.backward()
431
+ opt.step()
432
+ if step % eval_every == 0 or step == steps:
433
+ model.eval()
434
+ with torch.no_grad():
435
+ losses = []
436
+ for _ in range(20):
437
+ xv, yv = _batch(va, batch, block, device, g)
438
+ lv = F.cross_entropy(model(xv).reshape(-1, VOCAB),
439
+ yv.reshape(-1))
440
+ losses.append(lv.item())
441
+ bpb = sum(losses) / len(losses) / math.log(2)
442
+ print(f"[{arm}] step {step} val_bpb={bpb:.4f} vitals={model.vitals()}",
443
+ flush=True)
444
+ model.train()
445
+ results[arm] = {"val_bpb": bpb, "params": n_params, "vitals": model.vitals()}
446
+ if save:
447
+ path = os.path.join(ckpt_dir, f"{arm}_s{seed}_t{steps}.pt")
448
+ torch.save({"arm": arm, "seed": seed, "steps": steps, "val_bpb": bpb,
449
+ "state_dict": {k: v.cpu() for k, v in
450
+ model.state_dict().items()}}, path)
451
+ print(f"saved specimen: {path}", flush=True)
452
+ print(results, flush=True)
453
+ return results
454
+
455
+
456
+ def smoke():
457
+ """Shapes/parse only — no accuracy claims."""
458
+ x = torch.randint(0, VOCAB, (2, 64))
459
+ for arm in ("sdpa", "hub", "addr_head"):
460
+ m = ByteLM(arm, d=96, layers=2, block=64, K=16)
461
+ logits = m(x)
462
+ assert logits.shape == (2, 64, VOCAB)
463
+ logits.sum().backward()
464
+ # causality check: future byte must not affect past logits
465
+ with torch.no_grad():
466
+ a = m(x)[0, 10]
467
+ x2 = x.clone(); x2[0, 40] = (x2[0, 40] + 7) % 256
468
+ b = m(x2)[0, 10]
469
+ assert torch.allclose(a, b, atol=1e-4), f"{arm} leaks future context"
470
+ print(f"{arm}: OK params={sum(p.numel() for p in m.parameters()):,} "
471
+ f"vitals={m.vitals()}", flush=True)
472
+ print("OK — AR bed smoke passed (verdict run: train() on GPU)", flush=True)
473
+
474
+
475
+ def _in_notebook() -> bool:
476
+ try:
477
+ get_ipython() # type: ignore[name-defined] # noqa: F821
478
+ return True
479
+ except NameError:
480
+ return False
481
+
482
+
483
+ if __name__ == "__main__":
484
+ if _in_notebook():
485
+ smoke()
486
+ print("Notebook mode: call train(steps=2000) in the next cell (GPU).")
487
+ else:
488
+ import argparse
489
+ ap = argparse.ArgumentParser()
490
+ ap.add_argument("--train", action="store_true")
491
+ ap.add_argument("--steps", type=int, default=2000)
492
+ a, _ = ap.parse_known_args()
493
+ train(steps=a.steps) if a.train else smoke()