implement Addendum 8: 6-arm token-space recurrence suite (arms A, B, C, D1, D2, D3), tests, and jobs/e8.csv

7 files changed, 319 insertions, 36 deletions

diff --git a/jobs/e8.csv b/jobs/e8.csv
new file mode 100644
index 0000000..0c393a0
--- /dev/null
+++ b/jobs/e8.csv
@@ -0,0 +1,7 @@
+job,model,seed,flags
+e8-armA-none,transformer,0,--range_end 1000 --max_steps 32 --device cuda --weight_decay 0.1 --max_train_steps 200000 --scratch_mode none
+e8-armB-structured,transformer,0,--range_end 1000 --max_steps 32 --device cuda --weight_decay 0.1 --max_train_steps 200000 --scratch_mode structured
+e8-armC-filler,transformer,0,--range_end 1000 --max_steps 32 --device cuda --weight_decay 0.1 --max_train_steps 200000 --scratch_mode filler --scratch_len 16
+e8-armD1-randlearn,transformer,0,--range_end 1000 --max_steps 32 --device cuda --weight_decay 0.1 --max_train_steps 200000 --scratch_mode random_learned --scratch_len 16
+e8-armD2-randfroz,transformer,0,--range_end 1000 --max_steps 32 --device cuda --weight_decay 0.1 --max_train_steps 200000 --scratch_mode random_frozen --scratch_len 16
+e8-armD3-randnoise,transformer,0,--range_end 1000 --max_steps 32 --device cuda --weight_decay 0.1 --max_train_steps 200000 --scratch_mode random_noise --scratch_len 16
diff --git a/src/config.py b/src/config.py
index 945c259..915cdae 100644
--- a/src/config.py
+++ b/src/config.py
@@ -25,6 +25,9 @@ class Config:
     min_steps: int = 2               # ACT: halt prob forced to 0 for the first min_steps-1 steps
     n_layers: int = 2
     n_heads: int = 4
+    # scratchpad / token-space recurrence (Addendum 8, E8)
+    scratch_mode: str = "none"       # "none" | "structured" | "filler" | "random_learned" | "random_frozen" | "random_noise"
+    scratch_len: int = 16            # length of filler / random token sequence
     # training
     lr: float = 1e-3
     lr_decay: bool = False           # cosine 1e-3 -> 1e-4 over the run (E3)
@@ -45,7 +48,11 @@ class Config:
 
     @property
     def vocab(self) -> int:
-        """Token count. Integers mode: value tokens 0..next_prime(range_end), plus EOS (+1) and pad (+1)."""
+        """Token count.
+        - Integers mode: value tokens 0..next_prime(range_end) + EOS + pad.
+        - Digits mode (none): 10 digits + EOS = 11.
+        - Digits mode (scratchpad): 10 digits + EOS + SEP + PAUSE + 16 random (a-p) + 4 structured symbols (c,=,d,:) + 1 noise slot = 34.
+        """
         if self.vocab_mode == "integers":
             # inline mini-sieve: next prime above range_end bounds the target domain
             limit = self.range_end + 100
@@ -57,15 +64,21 @@ class Config:
                         is_prime[m] = False
             np_ = next(i for i in range(self.range_end + 1, limit + 1) if is_prime[i])
             return np_ + 2                    # values 0..np_, EOS=np_+1, pad=np_+2
-        return 11                              # digits 0-9 + EOS
+        if self.scratch_mode != "none":
+            return 34                         # expanded vocabulary for scratchpad/filler/random tokens
+        return 11                             # digits 0-9 + EOS
 
     @property
     def eos_id(self) -> int:
-        return self.vocab - 1
+        if self.vocab_mode == "integers":
+            return self.vocab - 1
+        if self.scratch_mode != "none":
+            return 10                         # fixed ID 10 for EOS in digits mode
+        return self.vocab - 1                 # 10 for digits mode none
 
     @property
     def pad_id(self) -> int:
-        return self.vocab              # one extra embedding row reserved for pad
+        return self.vocab                     # one extra embedding row reserved for pad
 
     def to_json(self) -> dict:
         d = {f.name: getattr(self, f.name) for f in fields(self)}
diff --git a/src/data.py b/src/data.py
index 26fb6b9..411ba7e 100644
--- a/src/data.py
+++ b/src/data.py
@@ -7,6 +7,18 @@ import torch
 from src.config import Config
 
 
+# Special token constants for digits mode (Addendum 8, E8)
+EOS_ID = 10
+SEP_ID = 11            # '#' separator between scratchpad and final target
+PAUSE_ID = 12          # '<p>' pause/filler token (Arm C)
+RAND_START_ID = 13     # 'a' (token IDs 13..28 represent 'a'..'p' for Arms D1/D2)
+SYM_C = 29             # 'c'
+SYM_EQ = 30            # '='
+SYM_D = 31             # 'd'
+SYM_COLON = 32         # ':'
+NOISE_SLOT_ID = 33     # '<noise>' placeholder for Arm D3 continuous dynamic noise
+
+
 def sieve_primes(limit: int) -> list[int]:
     """All primes <= limit (inclusive)."""
     if limit < 2:
@@ -33,15 +45,51 @@ def encode_int(n: int, cfg: Config) -> list[int]:
     return [int(d) for d in str(n)]
 
 
+def make_structured_trace(n: int, target_prime: int, primes_list: list[int]) -> list[int]:
+    """Generates trace tokens for candidate search and trial division:
+    For each candidate c in [n+1 .. target_prime]:
+      c, =, digits(c), d, digits(p), :, 0/1, ...
+    """
+    tokens = []
+    for c in range(n + 1, target_prime + 1):
+        tokens.append(SYM_C)
+        tokens.append(SYM_EQ)
+        tokens.extend([int(d) for d in str(c)])
+        for p in primes_list:
+            if p * p > c:
+                break
+            tokens.append(SYM_D)
+            tokens.extend([int(d) for d in str(p)])
+            tokens.append(SYM_COLON)
+            if c % p == 0:
+                tokens.append(0)  # divisible -> composite found, halt checks for c
+                break
+            else:
+                tokens.append(1)  # not divisible -> check next prime
+    return tokens
+
+
 def decode_tokens(ts, cfg: Config) -> int:
-    """Decode a token sequence, stopping at EOS. -1 if nothing decodable."""
+    """Decode a token sequence, stopping at EOS. If SEP_ID is present, decodes digits AFTER SEP_ID.
+    -1 if nothing decodable."""
     if cfg.vocab_mode == "integers":
         return int(ts[0]) if len(ts) else -1
+
+    if cfg.scratch_mode != "none":
+        if SEP_ID in ts:
+            idx = ts.index(SEP_ID)
+            ts = ts[idx + 1:]
+        else:
+            return -1
+
     digits = []
     for t in ts:
         if t == cfg.eos_id:
             break
-        digits.append(int(t))
+        if 0 <= t <= 9:
+            digits.append(int(t))
+        else:
+            break
     return int("".join(map(str, digits))) if digits else -1
 
 
@@ -67,7 +115,6 @@ def get_splits(cfg: Config) -> tuple[list[int], list[int]]:
     train, val = sorted(inputs[n_val:]), sorted(inputs[:n_val])
     if cfg.train_frac < 1.0:
         n_tr = max(1, round(len(train) * cfg.train_frac))
-        # deterministic subsample: seeded shuffle, take first n_tr
         sub = random.Random(cfg.seed + 1000)   # distinct stream from split shuffle
         sub.shuffle(train)
         train = sorted(train[:n_tr])
@@ -75,29 +122,79 @@ def get_splits(cfg: Config) -> tuple[list[int], list[int]]:
 
 
 def build_examples(inputs: list[int], cfg: Config) -> list[tuple[list[int], list[int]]]:
-    """[(input_tokens, target_tokens+EOS), ...]. task_mode selects the target function."""
+    """[(input_tokens, target_tokens+EOS), ...]. task_mode and scratch_mode select format."""
     margin = max(100, int(cfg.range_end * 0.05) + 50)
-    primes = sieve_primes(cfg.range_end + margin) if cfg.task_mode != "is_prime" else []
+    primes = sieve_primes(cfg.range_end + margin) if (cfg.task_mode != "is_prime" or cfg.scratch_mode == "structured") else []
     out = []
     for n in inputs:
+        x_toks = encode_int(n, cfg)
         if cfg.task_mode == "is_prime":
             p = 1 if is_prime_n(n) else 0
+            ans_toks = encode_int(p, cfg) + [cfg.eos_id]
+            if cfg.scratch_mode == "none":
+                out.append((x_toks, ans_toks))
+            elif cfg.scratch_mode == "filler":
+                filler = [PAUSE_ID] * cfg.scratch_len
+                y_toks = filler + [SEP_ID] + ans_toks
+                out.append((x_toks, y_toks))
+            else:
+                out.append((x_toks, ans_toks))
         else:
             p = next_prime(n, primes)
-        out.append((encode_int(n, cfg), encode_int(p, cfg) + [cfg.eos_id]))
+            ans_toks = encode_int(p, cfg) + [cfg.eos_id]
+            if cfg.scratch_mode == "none":
+                out.append((x_toks, ans_toks))
+            elif cfg.scratch_mode == "structured":
+                trace = make_structured_trace(n, p, primes)
+                y_toks = trace + [SEP_ID] + ans_toks
+                out.append((x_toks, y_toks))
+            elif cfg.scratch_mode == "filler":
+                filler = [PAUSE_ID] * cfg.scratch_len
+                y_toks = filler + [SEP_ID] + ans_toks
+                out.append((x_toks, y_toks))
+            elif cfg.scratch_mode in ("random_learned", "random_frozen"):
+                rng = random.Random(cfg.seed + n * 37)
+                rand_toks = [RAND_START_ID + rng.randint(0, 15) for _ in range(cfg.scratch_len)]
+                y_toks = rand_toks + [SEP_ID] + ans_toks
+                out.append((x_toks, y_toks))
+            elif cfg.scratch_mode == "random_noise":
+                noise_toks = [NOISE_SLOT_ID] * cfg.scratch_len
+                y_toks = noise_toks + [SEP_ID] + ans_toks
+                out.append((x_toks, y_toks))
+            else:
+                out.append((x_toks, ans_toks))
     return out
 
 
 def _global_lengths(cfg: Config) -> tuple[int, int]:
-    """(in_max, out_max): fixed global lengths so batch layout == singleton layout (codex BLOCKER fix)."""
+    """(in_max, out_max): fixed global lengths so batch layout == singleton layout."""
     if cfg.vocab_mode == "integers":
-        return 1, 2                      # [value], [value, EOS]
+        return 1, 2
+    in_max = len(str(cfg.range_end))
     if cfg.task_mode == "is_prime":
-        return len(str(cfg.range_end)), 2  # digit token "1"/"0" + EOS
-    margin = max(100, int(cfg.range_end * 0.05) + 50)
-    primes = sieve_primes(cfg.range_end + margin)
-    max_target = next_prime(cfg.range_end, primes)
-    return len(str(cfg.range_end)), len(str(max_target)) + 1   # digits + EOS
+        base_out = 2
+    else:
+        margin = max(100, int(cfg.range_end * 0.05) + 50)
+        primes = sieve_primes(cfg.range_end + margin)
+        max_target = next_prime(cfg.range_end, primes)
+        base_out = len(str(max_target)) + 1
+
+    if cfg.scratch_mode == "none":
+        return in_max, base_out
+    elif cfg.scratch_mode in ("filler", "random_learned", "random_frozen", "random_noise"):
+        return in_max, cfg.scratch_len + 1 + base_out
+    elif cfg.scratch_mode == "structured":
+        margin = max(100, int(cfg.range_end * 0.05) + 50)
+        primes = sieve_primes(cfg.range_end + margin)
+        # sample max trace length across the entire range
+        max_trace_len = 0
+        for n in range(cfg.range_start, cfg.range_end + 1):
+            p = next_prime(n, primes)
+            tr = make_structured_trace(n, p, primes)
+            if len(tr) > max_trace_len:
+                max_trace_len = len(tr)
+        return in_max, max_trace_len + 1 + base_out
+    return in_max, base_out
 
 
 def pad_inputs(x: torch.Tensor, cfg: Config) -> torch.Tensor:
@@ -110,16 +207,19 @@ def pad_inputs(x: torch.Tensor, cfg: Config) -> torch.Tensor:
 
 
 def make_batch(examples, cfg: Config) -> dict[str, torch.Tensor]:
-    """Fixed global layout (not batch-max): x left-padded to in_max, y right-padded to out_max."""
-    xs, ys = zip(*examples)
+    """Fixed global layout: x left-padded to in_max, y right-padded to out_max."""
     in_max, out_max = _global_lengths(cfg)
     B = len(examples)
     x = torch.full((B, in_max), cfg.pad_id, dtype=torch.long)
     y = torch.full((B, out_max), cfg.pad_id, dtype=torch.long)
+    loss_mask = torch.zeros((B, out_max), dtype=torch.bool)
     for i, (xi, yi) in enumerate(examples):
         x[i, in_max - len(xi):] = torch.tensor(xi, dtype=torch.long)
         y[i, : len(yi)] = torch.tensor(yi, dtype=torch.long)
-    # teacher-forced decoder input: BOS(=EOS reuse) then shifted y
+        if cfg.scratch_mode in ("filler", "random_learned", "random_frozen", "random_noise"):
+            loss_mask[i, cfg.scratch_len : len(yi)] = True
+        else:
+            loss_mask[i, : len(yi)] = True
     y_in = torch.cat([torch.full((B, 1), cfg.eos_id, dtype=torch.long), y[:, :-1]], dim=1)
     y_mask = y != cfg.pad_id
-    return {"x": x, "y": y, "y_in": y_in, "y_mask": y_mask}
+    return {"x": x, "y": y, "y_in": y_in, "y_mask": y_mask, "loss_mask": loss_mask}
diff --git a/src/model_api.py b/src/model_api.py
index 85b5d55..07d5d8a 100644
--- a/src/model_api.py
+++ b/src/model_api.py
@@ -39,7 +39,9 @@ def greedy_decode(model: PrimeModel, x: torch.Tensor, cfg: Config, max_len: int
     if x.device != device:
         x = x.to(device)
     x = pad_inputs(x, cfg)
-    max_len = max_len or cfg.max_out_len
+    from src.data import _global_lengths
+    _, out_max = _global_lengths(cfg)
+    max_len = max_len or max(getattr(cfg, "max_out_len", 6), out_max)
     B = x.shape[0]
     y_in = torch.full((B, 1), cfg.eos_id, dtype=torch.long, device=device)
     for _ in range(max_len):
diff --git a/src/models/transformer.py b/src/models/transformer.py
index 32e158b..2b986ce 100644
--- a/src/models/transformer.py
+++ b/src/models/transformer.py
@@ -32,7 +32,7 @@ class TransformerBaseline(PrimeModel):
         self.cfg = cfg
         d = cfg.d_model
         self.embed = nn.Embedding(cfg.vocab + 1, d)     # +1 row = pad
-        self.pos = nn.Embedding(64, d)                  # 3 in + 6 out max, generous
+        self.pos = nn.Embedding(256, d)                 # supports sequences up to 256
         self.blocks = nn.ModuleList([CausalBlock(d, cfg.n_heads) for _ in range(cfg.n_layers)])
         self.ln_f = nn.LayerNorm(d)
         self.head = nn.Linear(d, cfg.vocab)
@@ -44,14 +44,25 @@ class TransformerBaseline(PrimeModel):
         seq = torch.cat([x, y_in], dim=1)               # (B, T_in+T_out)
         T = seq.shape[1]
         device = seq.device
-        e = self.embed(seq) + self.pos(torch.arange(T, device=device)).unsqueeze(0)
+        
+        # Token embeddings
+        tok_embed = self.embed(seq)
+        if cfg.scratch_mode == "random_frozen":
+            # detach embeddings for tokens 'a'..'p' (IDs 13..28)
+            is_rand = (seq >= 13) & (seq <= 28)
+            tok_embed = torch.where(is_rand.unsqueeze(-1), tok_embed.detach(), tok_embed)
+        elif cfg.scratch_mode == "random_noise":
+            # inject fresh Gaussian noise at NOISE_SLOT_ID (ID 33)
+            is_noise = (seq == 33)
+            noise = torch.randn_like(tok_embed)
+            tok_embed = torch.where(is_noise.unsqueeze(-1), noise, tok_embed)
+
+        e = tok_embed + self.pos(torch.arange(T, device=device)).unsqueeze(0)
         # causal mask: input positions (j < T_in) fully visible; output positions causal
         # (True = blocked, per torch.nn.MultiheadAttention bool convention)
         blocked = torch.zeros(T, T, dtype=torch.bool, device=device)
-        for i in range(T):
-            for j in range(T):
-                if j >= T_in and j > i:
-                    blocked[i, j] = True
+        causal = torch.triu(torch.ones(T, T, dtype=torch.bool, device=device), diagonal=1)
+        blocked[:, T_in:] = causal[:, T_in:]
         key_pad = seq == cfg.pad_id                       # (B,T) True = ignore
         h = e
         for blk in self.blocks:
diff --git a/src/train.py b/src/train.py
index a0874c4..2523325 100644
--- a/src/train.py
+++ b/src/train.py
@@ -38,6 +38,7 @@ def make_gpu_dataset(examples: list, cfg: Config, device: torch.device, max_cach
             "y": torch.empty((0, out_max), dtype=torch.long, device=device),
             "y_in": torch.empty((0, out_max), dtype=torch.long, device=device),
             "y_mask": torch.empty((0, out_max), dtype=torch.bool, device=device),
+            "loss_mask": torch.empty((0, out_max), dtype=torch.bool, device=device),
             "size": 0,
             "is_cached": True,
         }
@@ -50,6 +51,7 @@ def make_gpu_dataset(examples: list, cfg: Config, device: torch.device, max_cach
         "y": batch["y"].to(device, non_blocking=True),
         "y_in": batch["y_in"].to(device, non_blocking=True),
         "y_mask": batch["y_mask"].to(device, non_blocking=True),
+        "loss_mask": batch["loss_mask"].to(device, non_blocking=True),
         "size": len(examples),
         "is_cached": True,
     }
@@ -77,13 +79,13 @@ def evaluate_gpu(model, gpu_ds: dict[str, torch.Tensor | int | bool], examples:
             x = gpu_ds["x"][bi: bi + eval_bs]
             y = gpu_ds["y"][bi: bi + eval_bs]
             y_in_eval = gpu_ds["y_in"][bi: bi + eval_bs]
-            mask = gpu_ds["y_mask"][bi: bi + eval_bs]
+            mask = gpu_ds["loss_mask"][bi: bi + eval_bs]
         else:
             batch = make_batch(chunk_examples, cfg)
             x = batch["x"].to(device)
             y = batch["y"].to(device)
             y_in_eval = batch["y_in"].to(device)
-            mask = batch["y_mask"].to(device)
+            mask = batch["loss_mask"].to(device)
 
         with torch.amp.autocast(device_type="cuda", dtype=torch.float16, enabled=use_amp):
             out = model(x, y_in_eval)
@@ -93,10 +95,9 @@ def evaluate_gpu(model, gpu_ds: dict[str, torch.Tensor | int | bool], examples:
         token_total += int(mask.sum().item())
 
         B = x.shape[0]
-        max_out_len = getattr(cfg, "max_out_len", 6)
         from src.data import _global_lengths
         _, out_max = _global_lengths(cfg)
-        max_len = max(max_out_len, out_max)
+        max_len = max(getattr(cfg, "max_out_len", 6), out_max)
 
         y_in = torch.full((B, 1), cfg.eos_id, dtype=torch.long, device=device)
         with torch.amp.autocast(device_type="cuda", dtype=torch.float16, enabled=use_amp):
@@ -106,10 +107,11 @@ def evaluate_gpu(model, gpu_ds: dict[str, torch.Tensor | int | bool], examples:
                 y_in = torch.cat([y_in, nxt[:, None]], dim=1)
         gen = y_in[:, 1:].cpu()
 
-        for i, (xrow, target) in enumerate(chunk_examples):
+        for i, ex in enumerate(chunk_examples):
+            xrow, target = ex[0], ex[1]
             pred = decode_tokens(gen[i].tolist(), cfg)
             target_n = decode_tokens(target, cfg)
-            ok = pred == target_n
+            ok = (pred == target_n) and (pred != -1)
             em_correct += int(ok)
             per_example.append((tuple(xrow), target_n, pred, ok))
 
@@ -268,7 +270,7 @@ def main() -> None:
                 x = train_gpu["x"][idx]
                 y = train_gpu["y"][idx]
                 y_in = train_gpu["y_in"][idx]
-                mask = train_gpu["y_mask"][idx]
+                mask = train_gpu["loss_mask"][idx]
 
                 with torch.amp.autocast(device_type="cuda", dtype=torch.float16, enabled=use_amp):
                     out = model(x, y_in)
@@ -317,7 +319,7 @@ def main() -> None:
                 x = batch["x"].to(device)
                 y = batch["y"].to(device)
                 y_in = batch["y_in"].to(device)
-                mask = batch["y_mask"].to(device)
+                mask = batch["loss_mask"].to(device)
 
                 with torch.amp.autocast(device_type="cuda", dtype=torch.float16, enabled=use_amp):
                     out = model(x, y_in)
diff --git a/tests/test_scratchpad.py b/tests/test_scratchpad.py
new file mode 100644
index 0000000..2103453
--- /dev/null
+++ b/tests/test_scratchpad.py
@@ -0,0 +1,148 @@
+"""Unit tests for Addendum 8: 6-arm token-space recurrence & decomposition."""
+import pytest
+import torch
+import torch.nn as nn
+
+from src.config import Config
+from src.data import (
+    EOS_ID,
+    NOISE_SLOT_ID,
+    PAUSE_ID,
+    RAND_START_ID,
+    SEP_ID,
+    SYM_C,
+    SYM_COLON,
+    SYM_D,
+    SYM_EQ,
+    build_examples,
+    decode_tokens,
+    make_batch,
+    make_structured_trace,
+    next_prime,
+    sieve_primes,
+)
+from src.models.transformer import TransformerBaseline
+
+
+def test_vocab_and_special_tokens():
+    cfg_none = Config(scratch_mode="none")
+    assert cfg_none.vocab == 11
+    assert cfg_none.eos_id == 10
+    assert cfg_none.pad_id == 11
+
+    cfg_scratch = Config(scratch_mode="filler")
+    assert cfg_scratch.vocab == 34
+    assert cfg_scratch.eos_id == 10
+    assert cfg_scratch.pad_id == 34
+
+
+def test_make_structured_trace():
+    primes = sieve_primes(50)
+    # n=14 -> next prime 17. Candidates 15 (comp: d3:0), 16 (comp: d2:0), 17 (prime: d2:1, d3:1)
+    trace = make_structured_trace(14, 17, primes)
+    assert SYM_C in trace
+    assert SYM_EQ in trace
+    assert SYM_D in trace
+    assert SYM_COLON in trace
+
+    # n=2 -> next prime 3. Candidate 3 has no primes with p*p <= 3, trace should be [SYM_C, SYM_EQ, 3]
+    trace_2_3 = make_structured_trace(2, 3, primes)
+    assert trace_2_3 == [SYM_C, SYM_EQ, 3]
+
+
+def test_decode_tokens_all_modes():
+    cfg_none = Config(scratch_mode="none")
+    assert decode_tokens([4, 3, EOS_ID], cfg_none) == 43
+
+    cfg_filler = Config(scratch_mode="filler")
+    # with separator #
+    assert decode_tokens([PAUSE_ID, PAUSE_ID, SEP_ID, 4, 3, EOS_ID], cfg_filler) == 43
+    # without separator (should fail / return -1)
+    assert decode_tokens([4, 3, EOS_ID], cfg_filler) == -1
+
+
+def test_build_examples_and_loss_mask():
+    modes = ["none", "structured", "filler", "random_learned", "random_frozen", "random_noise"]
+    for mode in modes:
+        cfg = Config(scratch_mode=mode, range_start=2, range_end=30, scratch_len=16)
+        exs = build_examples([14], cfg)
+        assert len(exs) == 1
+        x, y = exs[0]
+        batch = make_batch(exs, cfg)
+        loss_mask = batch["loss_mask"][0]
+
+        if mode == "none":
+            assert decode_tokens(y, cfg) == 17
+            assert all(loss_mask[: len(y)])
+        elif mode == "structured":
+            assert SEP_ID in y
+            assert decode_tokens(y, cfg) == 17
+            assert all(loss_mask[: len(y)])  # structured trace is fully supervised
+        elif mode in ("filler", "random_learned", "random_frozen", "random_noise"):
+            assert SEP_ID in y
+            assert decode_tokens(y, cfg) == 17
+            # first 16 tokens must have loss_mask == 0
+            assert all(not m for m in loss_mask[:16])
+            # tokens after separator must have loss_mask == 1
+            assert all(m for m in loss_mask[16 : len(y)])
+
+
+def test_random_tokens_no_digit_leakage():
+    for mode in ("random_learned", "random_frozen"):
+        cfg = Config(scratch_mode=mode, range_start=2, range_end=100, scratch_len=16)
+        exs = build_examples(list(range(2, 101)), cfg)
+        for _, y in exs:
+            random_segment = y[:16]
+            for tok in random_segment:
+                assert RAND_START_ID <= tok <= RAND_START_ID + 15
+                assert not (0 <= tok <= 9)  # no digits leakage
+
+
+def test_transformer_forward_all_arms():
+    modes = ["none", "structured", "filler", "random_learned", "random_frozen", "random_noise"]
+    for mode in modes:
+        cfg = Config(scratch_mode=mode, range_start=2, range_end=30, d_model=32, n_layers=1, n_heads=2)
+        exs = build_examples([14, 15, 16], cfg)
+        batch = make_batch(exs, cfg)
+        model = TransformerBaseline(cfg)
+        out = model(batch["x"], batch["y_in"])
+        logits = out["logits"]
+        assert logits.shape == (3, batch["y"].shape[1], cfg.vocab)
+
+        # compute masked loss
+        ce = nn.CrossEntropyLoss(reduction="none")
+        y_safe = batch["y"].clamp(max=cfg.vocab - 1)
+        loss_tokens = ce(logits.reshape(-1, cfg.vocab), y_safe.reshape(-1)).reshape(3, -1) * batch["loss_mask"].float()
+        loss = loss_tokens.sum() / batch["loss_mask"].sum().clamp(min=1)
+        assert not torch.isnan(loss)
+        loss.backward()
+
+
+def test_random_frozen_zero_gradients():
+    cfg = Config(scratch_mode="random_frozen", range_start=2, range_end=30, d_model=32, n_layers=1, n_heads=2)
+    exs = build_examples([14, 15], cfg)
+    batch = make_batch(exs, cfg)
+    model = TransformerBaseline(cfg)
+    out = model(batch["x"], batch["y_in"])
+    logits = out["logits"]
+    loss = logits.sum()
+    loss.backward()
+
+    # random token embeddings (13..28) must have zero grad because they were detached
+    grad = model.embed.weight.grad
+    assert grad[RAND_START_ID : RAND_START_ID + 16].abs().sum() == 0
+    # non-random tokens must have non-zero grad
+    assert grad[:10].abs().sum() > 0
+
+
+def test_random_noise_stochasticity():
+    cfg = Config(scratch_mode="random_noise", range_start=2, range_end=30, d_model=32, n_layers=1, n_heads=2)
+    exs = build_examples([14], cfg)
+    batch = make_batch(exs, cfg)
+    model = TransformerBaseline(cfg)
+    model.eval()
+
+    out1 = model(batch["x"], batch["y_in"])["logits"]
+    out2 = model(batch["x"], batch["y_in"])["logits"]
+    # Dynamic Gaussian noise should produce distinct logits between forward passes
+    assert not torch.allclose(out1, out2)