diff options
| -rw-r--r-- | design/experiment-spec.md | 17 | ||||
| -rw-r--r-- | design/preregistration.md | 26 | ||||
| -rw-r--r-- | src/config.py | 3 | ||||
| -rw-r--r-- | src/eval.py | 4 | ||||
| -rw-r--r-- | src/models/rnn.py | 54 | ||||
| -rw-r--r-- | src/train.py | 9 | ||||
| -rw-r--r-- | tests/test_models.py | 9 |
7 files changed, 94 insertions, 28 deletions
diff --git a/design/experiment-spec.md b/design/experiment-spec.md index c814057..77d1606 100644 --- a/design/experiment-spec.md +++ b/design/experiment-spec.md @@ -125,3 +125,20 @@ cross the gap from pattern matching to computation. --- > Copied from the research repo `prime-grokking/main.md` @ 546dc2c (provenance: `~/Projects/research`, remote ssh://meru/~/projects/research.git). + + +--- + +## Prior art (appended 2026-08-14, Gemini 3.6 Flash design review) + +No prior grokking work on next-prime / primality found. Canonical refs: + +- Power et al. 2022 (arXiv:2201.02177) — grokking, modular arithmetic +- Nanda et al. 2023 (arXiv:2301.05217) — grokking progress measures +- Liu et al. 2022 (arXiv:2205.10343) — empirical grokking study +- Varma et al. 2023 (arXiv:2309.02390) — grokking via circuit efficiency +- Graves 2016 (arXiv:1603.08983) — ACT +- Banino et al. 2021 (arXiv:2107.05407) — PonderNet +- Giannou et al. 2023 (arXiv:2301.13196) — looped transformers +- Xu et al. ICLR 2020 (arXiv:1905.13211) — algorithmic alignment +- Xu et al. 2021 (arXiv:2009.11848) — extrapolation / GNNs diff --git a/design/preregistration.md b/design/preregistration.md index d879912..f800461 100644 --- a/design/preregistration.md +++ b/design/preregistration.md @@ -66,3 +66,29 @@ Caveat: single seed — all architecture comparisons are seed-0 anecdotes until 2. No hyperparameter tuning on the val set. The wd sweep is a separate experiment run only after v1 results, at jayrup's call. 3. Probe interpretation locked above; NOTES.md must compare outcomes against this matrix verbatim (cite codes). 4. Training-range caveat recorded: for n ≤ 100 the sieve only needs divisors {2, 3, 5, 7}; "grokking the algorithm" in-range does not imply the general sieve. + +--- + +## Addendum 1 (2026-08-14, pre-launch — setup amendments only) + +Trigger: Gemini 3.6 Flash design review (experiment repo `design/reviews/gemini-design-review.md`). +The interpretation matrix (O/H/P codes) above is NOT amended; only setup details changed. Original lock commit: 00c696d. + +1. **Fully-tied cell (was: un-tied GRU decoder).** All recurrence — input read-in, K compute steps, AND output-digit decoding — now runs through the SAME 2-layer cell. The earlier draft's GRU decoder would have masked whether the tied cell solved the task. RNN param count ≈ 36.6k (was 168.7k). Regression test added (no GRU/LSTM/RNN modules). +2. **Recurrent input read-in (was: masked mean-pool).** Digits are read through the tied cell with sinusoidal positional encoding. The mean-pool blurred place value ("10" and "100" share the token multiset {1,0}). +3. **λ schedule: linear ramp 1000→5000 steps (was: hard switch at step 1000).** Avoids a discontinuous loss jump late in training. +4. **Future wd sweep revised to {0.01, 0.1, 0.3, 1.0, 3.0}** (10.0 dropped: at lr=1e-3 with AdamW, λ=10 decays weights ~1%/step). Seed-0 default wd=1.0 unchanged. +5. **ACT verification note:** aggregation is Graves (2016) standard — w_t = p_t·Π_{s<t}(1−p_s); Σw_t + remainder = 1 by construction. A reviewer initially flagged this as non-normalized; verified correct, comment + covered by tests. +6. **Prior art appended** (below). No prior grokking work on next-prime / primality prediction found — the closest literature is grokking on modular arithmetic (different: group structure) and ACT/PonderNet halting work. + +## References + +- Power et al. (2022), "Grokking: Generalization Beyond Overfitting on Small Algorithmic Datasets", arXiv:2201.02177 +- Nanda et al. (2023), "Progress measures for grokking via mechanistic interpretability", arXiv:2301.05217 +- Liu et al. (2022), "Towards Understanding Grokking: An Empirical Study", arXiv:2205.10343 +- Varma et al. (2023), "Explaining Grokking Through Circuit Efficiency", arXiv:2309.02390 +- Graves (2016), "Adaptive Computation Time for Recurrent Neural Networks", arXiv:1603.08983 +- Banino et al. (2021), "PonderNet: Learning to Ponder", arXiv:2107.05407 +- Giannou et al. (2023), "Looped Transformers as Programmable Computers", arXiv:2301.13196 +- Xu et al. (ICLR 2020), "What Can Neural Networks Reason About?", arXiv:1905.13211 +- Xu et al. (2021), "How Neural Networks Extrapolate: From Feedforward to Graph Neural Networks", arXiv:2009.11848 diff --git a/src/config.py b/src/config.py index 74cb4d3..c0532bc 100644 --- a/src/config.py +++ b/src/config.py @@ -19,7 +19,8 @@ class Config: halting: bool = True # False -> fixed-K ablation halt_eps: float = 0.05 # ACT cumulative threshold (documented; K small so no early break) halt_penalty: float = 0.01 # lambda on mean steps - halt_warmup_steps: int = 1000 # penalty off before this (anti-collapse) + halt_warmup_steps: int = 1000 # penalty = 0 before this (anti-collapse) + halt_ramp_end_steps: int = 5000 # penalty ramps linearly 0 -> halt_penalty between warmup and here min_steps: int = 2 # ACT: halt prob forced to 0 before this step n_layers: int = 2 n_heads: int = 4 diff --git a/src/eval.py b/src/eval.py index 9d55437..4485523 100644 --- a/src/eval.py +++ b/src/eval.py @@ -91,8 +91,8 @@ def halting_report(model, cfg: Config) -> dict: gaps, steps = [], [] for n in range(cfg.range_start, cfg.range_end + 1): x = torch.tensor(encode_int(n, cfg), dtype=torch.long).unsqueeze(0) - h = model._initial_state(x) - _, s = model._run_cell(h) + h = model._encode(x) + _, s = model._run_compute(h) gaps.append(next_prime(n, primes) - n) steps.append(float(s.mean())) mean = float(np.mean(steps)) diff --git a/src/models/rnn.py b/src/models/rnn.py index 9947c66..19123b7 100644 --- a/src/models/rnn.py +++ b/src/models/rnn.py @@ -1,14 +1,15 @@ -"""Weight-tied RNN: one 2-layer cell applied K times, ACT learned halting, GRU digit decoder. +"""Weight-tied RNN: ONE 2-layer cell reused for input read-in, K compute steps, and output decoding. Spec pseudocode (design/experiment-spec.md): state = embed(input_number) for step in range(max_steps): - state = step_module(state) # same weights every iteration + state = step_module(state) # same weights, every iteration if halt_condition(state): break output = project(state) -Initial state: masked mean-pool of (digit embedding + sinusoidal positional encoding) -passed through a small MLP, so digit ORDER reaches the tied cell. +Everything recurrent is the SAME cell (per design review: no un-tied GRU decoder, +no mean-pool blur — order reaches the cell via sinusoidal position added per digit). +ACT halting applies only to the K compute steps. """ import torch import torch.nn as nn @@ -24,15 +25,16 @@ class TiedRNN(PrimeModel): self.cfg = cfg d = cfg.d_model self.embed = nn.Embedding(cfg.vocab + 1, d) # +1 row = pad - self.in_proj = nn.Sequential(nn.Linear(d, d), nn.GELU(), nn.Linear(d, d)) - self.ln1 = nn.LayerNorm(d) self.cell_ln = nn.LayerNorm(d) self.cell_w1 = nn.Linear(d, d) self.cell_w2 = nn.Linear(d, d) + self.ln1 = nn.LayerNorm(d) self.halt_head = nn.Linear(d, 1) - self.decoder = nn.GRUCell(d, d) self.out_head = nn.Linear(d, cfg.vocab) + def _cell_step(self, h: torch.Tensor) -> torch.Tensor: + return h + self.cell_w2(F.gelu(self.cell_w1(self.cell_ln(h)))) + @staticmethod def _sinusoidal(T: int, d: int) -> torch.Tensor: pe = torch.zeros(T, d) @@ -42,35 +44,39 @@ class TiedRNN(PrimeModel): pe[:, 1::2] = torch.cos(pos / 10000 ** (2 * i[:, 1::2] / d)) return pe - def _initial_state(self, x: torch.Tensor) -> torch.Tensor: + def _encode(self, x: torch.Tensor) -> torch.Tensor: + """Read input digits through the tied cell (pad positions inject nothing).""" B, T = x.shape - mask = (x != self.cfg.pad_id).float().unsqueeze(-1) # (B,T,1) - pos = self._sinusoidal(T, self.cfg.d_model).to(x.device) # (T,d) - e = self.embed(x) + pos.unsqueeze(0) # (B,T,d) - h = (e * mask).sum(1) / mask.sum(1).clamp(min=1) # (B,d) - return self.in_proj(h) + d = self.cfg.d_model + pos = self._sinusoidal(T, d).to(x.device) # (T,d) + e = self.embed(x) # (B,T,d) + mask = (x != self.cfg.pad_id).float().unsqueeze(-1) # (B,T,1) + h = torch.zeros(B, d, device=x.device) + for t in range(T): + h = self._cell_step(h + (e[:, t] + pos[t]) * mask[:, t]) + return h - def _run_cell(self, h0: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: - """Tied cell x K steps. Returns (final_state (B,d), mean_steps (B,)).""" + def _run_compute(self, h0: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: + """K tied compute steps with ACT learned halting. Returns (final_state, mean_steps).""" cfg = self.cfg B = h0.shape[0] device = h0.device if not cfg.halting: h = h0 for _ in range(cfg.max_steps): - h = h + self.cell_w2(F.gelu(self.cell_w1(self.cell_ln(h)))) + h = self._cell_step(h) steps = h0.new_full((B,), float(cfg.max_steps)) return h, steps - # ACT: run all K steps, accumulate weighted average (K=20 -> no early break needed) h_list, p_list = [], [] h = h0 for t in range(cfg.max_steps): - h = h + self.cell_w2(F.gelu(self.cell_w1(self.cell_ln(h)))) + h = self._cell_step(h) p = torch.sigmoid(self.halt_head(self.ln1(h))).squeeze(-1) # (B,) if t < cfg.min_steps: p = p * 0.0 h_list.append(h) p_list.append(p) + # ACT aggregation (Graves 2016): w_t = p_t * prod_{s<t}(1-p_s); weights + remainder = 1 final = torch.zeros_like(h0) steps = torch.zeros(B, device=device) remaining = torch.ones(B, device=device) @@ -86,13 +92,13 @@ class TiedRNN(PrimeModel): def forward(self, x: torch.Tensor, y_in: torch.Tensor) -> dict: cfg = self.cfg - h = self._initial_state(x) - h, steps = self._run_cell(h) - # autoregressive digit decoder, teacher-forced during training - e = self.embed(y_in) # (B,T_out,d) + h = self._encode(x) + h, steps = self._run_compute(h) + # decode output digits through the SAME tied cell (teacher-forced during training) + e_out = self.embed(y_in) # (B,T_out,d) outs = [] for t in range(y_in.shape[1]): - h = self.decoder(e[:, t], h) + h = self._cell_step(h + e_out[:, t]) outs.append(self.out_head(h)) - logits = torch.stack(outs, dim=1) # (B,T_out,vocab) + logits = torch.stack(outs, dim=1) # (B,T_out,vocab) return {"logits": logits, "halt_steps": steps} diff --git a/src/train.py b/src/train.py index f0c63e9..975598d 100644 --- a/src/train.py +++ b/src/train.py @@ -133,7 +133,14 @@ def main() -> None: loss_tokens = ce(logits.reshape(-1, cfg.vocab), y_safe.reshape(-1)).reshape( logits.shape[0], -1) * batch["y_mask"].float() loss_tokens = loss_tokens.sum() / batch["y_mask"].sum().clamp(min=1) - lam = cfg.halt_penalty if (cfg.halting and step >= cfg.halt_warmup_steps) else 0.0 + if cfg.halting and step >= cfg.halt_warmup_steps: + if step >= cfg.halt_ramp_end_steps: + lam = cfg.halt_penalty + else: + frac = (step - cfg.halt_warmup_steps) / max(1, cfg.halt_ramp_end_steps - cfg.halt_warmup_steps) + lam = cfg.halt_penalty * frac + else: + lam = 0.0 halt = out["halt_steps"] penalty = halt.float().mean() * lam if halt is not None and lam > 0 else 0.0 loss = loss_tokens + penalty diff --git a/tests/test_models.py b/tests/test_models.py index 7994e86..b6fec4d 100644 --- a/tests/test_models.py +++ b/tests/test_models.py @@ -48,6 +48,15 @@ def test_transformer_integers_mode(): assert out["logits"].shape == (4, b["y"].shape[1], cfg.vocab) +def test_rnn_fully_tied_no_gru_decoder(): + """Regression: design review fix — everything recurrent must be the ONE tied cell.""" + import torch.nn as nn + cfg = Config(model="rnn") + m = build_model(cfg) + assert not any(isinstance(mod, nn.GRUCell) for mod in m.modules()) + assert not any(isinstance(mod, (nn.GRU, nn.LSTM, nn.RNN)) for mod in m.modules()) + + def test_param_counts_logged_not_gated(): r = build_model(Config(model="rnn")).param_count() t = build_model(Config(model="transformer")).param_count() |