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
|
### 1. Overall Setup & Feasibility for Next-Prime Grokking
- **Verdict**: **ISSUE** (High risk of pure memorization without grokking).
- **Domain size vs. Capacity**: 69 training examples in $[2, 100]$ with ~50k–140k parameters is overparameterized by orders of magnitude. The network can memorize all 69 pairs in a few hundred steps.
- **Representational Mismatch**: Grokking (Power et al., 2022) typically occurs on algebraic group operations ($x + y \pmod p$) where Fourier representations can compactly solve the task. Next-prime mapped from base-10 digits lacks smooth continuous group symmetries, giving weight decay little structural leverage to discover a compact closed-form algorithm.
- **Mean-Pooling Blur**: Masked mean-pooling of digit embeddings blurs place value. For example, $10$ and $100$ share token sets $\{1, 0\}$, making exact spatial/positional decoding difficult unless positional embeddings dominate.
---
### 2. Soundness of ACT Setup
- **Verdict**: **ISSUE** (Flawed state aggregation and penalty schedule).
- **Non-Standard State Aggregation**: Standard ACT (Graves, 2016; `arXiv:1603.08983`) requires halting weights to sum to 1 ($\sum_{t=1}^N w_t = 1$, using a remainder $R(x)$ at step $N$). Simply taking $\sum p_t h_t$ without normalization allows the state vector norm to scale arbitrarily with step count, causing uncalibrated gradients.
- **Gradient Shock from Step Warmup**: Hard-switching $\lambda$ from $0$ to $0.01$ at step 1000 introduces a discontinuous loss jump. At batch size 32 on 69 examples (~2 steps/epoch), step 1000 is epoch ~460—by which time training loss is near zero. Turning on $\lambda=0.01$ suddenly causes severe optimization instability.
- **Better Alternative**: Consider PonderNet (Banino et al., 2021; `arXiv:2107.05407`), which formulates halting probabilistic distributions with a Geometric prior KL penalty, eliminating discrete step-floor collapse. Otherwise, use a smooth linear warmup for $\lambda$ across 5,000 steps.
---
### 3. Optimizer & Weight Decay Tuning
- **Verdict**: **MIXED** (WD=1.0 is a good baseline; sweep upper bound is too high).
- **Baseline ($WD = 1.0$)**: **OK**. Strong weight decay is essential to drive network weights out of high-norm memorization regimes into generalizable representations (Power et al., 2022; Nanda et al., 2023).
- **Sweep Range $\{0.3, 1.0, 3.0, 10.0\}$**: **ISSUE**. For AdamW with $\text{lr}=1\times 10^{-3}$, weight updates decay by $(1 - \eta \cdot \lambda)$ per step. At $\lambda=10.0$, the parameter weight decays by $1\%$ *per step*, which will cause underfitting or complete gradient divergence.
- **Recommended Sweep**: Shift range to $\{0.01, 0.1, 0.3, 1.0, 3.0\}$ to include lower control baselines.
---
### 4. Prior Art & Theoretical Context
- **Verdict**: **OK** (Relevant papers identified).
- **Canonical Grokking Literature**:
- Power et al. (2022), *"Grokking: Generalization Beyond Overfitting on Small Algorithmic Datasets"*, [`arXiv:2201.02177`](https://arxiv.org/abs/2201.02177).
- Nanda et al. (2023), *"Progress measures for grokking via mechanistic interpretability"*, [`arXiv:2301.05217`](https://arxiv.org/abs/2301.05217).
- Liu et al. (2022), *"Towards Understanding Grokking: An Empirical Study"*, [`arXiv:2205.10343`](https://arxiv.org/abs/2205.10343).
- Varma et al. (2023), *"Explaining Grokking Through Circuit Efficiency"*, [`arXiv:2309.02390`](https://arxiv.org/abs/2309.02390).
- **Adaptive Computation & Recurrent Looping**:
- Graves (2016), *"Adaptive Computation Time for Recurrent Neural Networks"*, [`arXiv:1603.08983`](https://arxiv.org/abs/1603.08983).
- Banino et al. (2021), *"PonderNet: Learning to Ponder"*, [`arXiv:2107.05407`](https://arxiv.org/abs/2107.05407).
- Giannou et al. (2023), *"Looped Transformers as Programmable Computers"*, [`arXiv:2301.13196`](https://arxiv.org/abs/2301.13196).
- **Algorithmic Alignment**:
- Xu et al. (2020), *"What Can Neural Networks Reason About? Reasoning Algorithmic Alignment"*, [`arXiv:1905.13211`](https://arxiv.org/abs/1905.13211) (ICLR 2020) — shows sample complexity and generalization depend on structural alignment between architecture steps and target algorithm steps (e.g., dynamic programming / trial division steps).
- Xu et al. (2021), *"How Neural Networks Extrapolate: From Feedforward to Graph Neural Networks"*, [`arXiv:2009.11848`](https://arxiv.org/abs/2009.11848).
- **Primality Domain Note**: Past grokking work focuses on *prime modulo arithmetic* ($x+y \pmod p$), not predicting *next-prime* from base-10 representations.
---
### 5. GRU Decoder Architectural Confound
- **Verdict**: **ISSUE** (Breaks strict weight-tying comparison).
- **Confound**: Introducing an un-tied 1-layer GRU decoder on Arm A creates an architectural asymmetry against Arm B (GPT-style causal transformer). The GRU decoder contains enough un-tied parameters to perform sequential pattern lookup on output tokens, masking whether the tied recurrent cell actually solved the problem.
- **Cleaner Alternatives**:
1. **Direct Tied Linear Readout**: Project the final state $h_{\text{final}}$ directly to vocabulary logits via a single linear layer (or shared embedding matrix).
2. **Recurrent Cell Decoding**: Feed previous output tokens back into the *same tied cell* during generation rather than using a separate GRU block.
---
### 6. Interpretation Logic & Metrics
- **Verdict**: **MIXED** (Valid in-domain criteria, but needs extrapolation check).
- **In-Domain Sharp Transition**: Train $\text{EM} \ge 0.95 \to$ Val $\text{EM}$ jumping $0.2 \to 0.9$ within 5 evals correctly identifies grokking dynamics on $[2, 100]$.
- **Out-of-Domain Generalization ($[101, 200]$)**: High val EM on $[2, 100]$ alone does *not* prove a primality algorithm (sieve/trial division) was grokked. If validation accuracy on $[101, 200]$ remains near 0%, the network merely grokked a bounded interpolation table for $n \le 100$.
---
### Top 3 Things to Fix Before Launching Runs
1. **Remove the GRU Decoder Confound**: Replace the un-tied GRU decoder in Arm A with a single shared linear projection layer to strictly isolate the effect of weight-tied recurrence.
2. **Fix the ACT Formulation & Warmup Schedule**: Normalize state weighting ($\sum w_t = 1$) following Graves (2016) or switch to PonderNet (`arXiv:2107.05407`). Replace the abrupt step-1000 penalty turn-on with a smooth linear schedule for $\lambda$ across 5,000 steps.
3. **Fix Input Representation & Alignment**: Replace mean-pooling with explicit sequence concatenation/causal prefix encoding to preserve place-value structure. Lower the AdamW weight decay sweep upper bound from $10.0$ to $3.0$.
|
