Paper deep dive
Marginals Before Conditionals
Mihir Sahasrabudhe
Intelligence
Status: succeeded | Model: google/gemini-3.1-flash-lite-preview | Prompt: intel-v1 | Confidence: 94%
Last extracted: 3/13/2026, 1:00:26 AM
Summary
The paper investigates the 'Marginals Before Conditionals' phenomenon in neural networks, where models learn marginal distributions before resolving conditional dependencies. Using a controlled task with K-fold ambiguity resolved by a selector token, the authors demonstrate that models plateau at log K loss, with the duration of this plateau scaling with dataset size D. The transition to conditional learning is collective, sharp, and stabilized by gradient noise, which acts as an entropic force. The study also identifies internal circuit formation (selector-routing heads) and links the findings to directional asymmetry and the reversal curse.
Entities (5)
Relation Signals (3)
Dataset Size → determines → Plateau Duration
confidence 98% · plateau duration = f(D) (set by dataset size D, not K)
Gradient Noise → stabilizes → Marginal Solution
confidence 95% · Gradient noise stabilizes the marginal solution: higher learning rates monotonically slow the transition
Selector-routing head → enables → Conditional Learning
confidence 90% · Internally, a selector-routing head assembles during the plateau, leading the loss transition
Cypher Suggestions (0)
No Cypher suggestions yet.
Abstract
Abstract:We construct a minimal task that isolates conditional learning in neural networks: a surjective map with K-fold ambiguity, resolved by a selector token z, so H(A | B) = log K while H(A | B, z) = 0. The model learns the marginal P(A | B) first, producing a plateau at exactly log K, before acquiring the full conditional in a sharp, collective transition. The plateau has a clean decomposition: height = log K (set by ambiguity), duration = f(D) (set by dataset size D, not K). Gradient noise stabilizes the marginal solution: higher learning rates monotonically slow the transition (3.6* across a 7* {\eta} range at fixed throughput), and batch-size reduction delays escape, consistent with an entropic force opposing departure from the low-gradient marginal. Internally, a selector-routing head assembles during the plateau, leading the loss transition by ~50% of the waiting time. This is the Type 2 directional asymmetry of Papadopoulos et al. [2024], measured dynamically: we track the excess risk from log K to zero and characterize what stabilizes it, what triggers its collapse, and how long it takes.
Tags
Links
- Source: https://arxiv.org/abs/2603.10074v1
- Canonical: https://arxiv.org/abs/2603.10074v1
Trouble viewing inline? Open PDF directly →
Full Text
43,208 characters extracted from source content.
Expand or collapse full text
Marginals Before Conditionals: Staged Disambiguation in Gradient-Trained Transformers Mihir Sahasrabudhe School of Information Sciences University of Illinois Urbana-Champaign mihirrs2@illinois.edu Abstract We construct a minimal task that isolates conditional learning in neural networks: a surjective map with K-fold ambiguity, resolved by a selector token z, so H(A∣B)=logKH(A B)= K while H(A∣B,z)=0H(A B,z)=0. The model learns the marginal P(A∣B)P(A B) first, producing a plateau at exactly logK K, before acquiring the full conditional in a sharp, collective transition. The plateau has a clean decomposition: height =logK= K (set by ambiguity), duration =f(D)=f(D) (set by dataset size D, not K). Gradient noise stabilizes the marginal solution: higher learning rates monotonically slow the transition (3.6×3.6× across a 7×7× η range at fixed throughput), and batch-size reduction delays escape, consistent with an entropic force opposing departure from the low-gradient marginal. Internally, a selector-routing head assembles during the plateau, leading the loss transition by ∼50% 50\% of the waiting time. This is the Type 2 directional asymmetry of Papadopoulos et al. [2024], measured dynamically: we track the excess risk from logK K to zero and characterize what stabilizes it, what triggers its collapse, and how long it takes. Code, data-generation scripts, and trained checkpoints: https://github.com/mihirrs/synass-lens 1 Introduction Neural networks learn structured tasks in stages. Delayed generalization can follow memorization by orders of magnitude in training steps [Power et al., 2022]; hidden progress accumulates long before external metrics move [Barak et al., 2022]; and models routinely fail to exploit information present in their inputs [Berglund et al., 2024, Kitouni et al., 2024]. Mechanistic studies of grokking have identified specific circuit-formation events underlying delayed generalization [Nanda et al., 2023], including Fourier-basis representations for modular arithmetic. We study a related but distinct transition: from marginal to conditional prediction, rather than from memorization to generalization. This connects to two lines of work on directional asymmetry. Papadopoulos et al. [2024] showed that forward-trained language models outperform backward-trained ones (an “arrow of time”) and identified a Type 2 asymmetry where the model can represent the inverse but struggles to learn it. The reversal curse [Berglund et al., 2024] showed that models trained on “A is B” cannot infer “B is A.” Both can be viewed through the lens of excess risk in the inverse direction: ΔX→Y(t)=[−logqθ(Y∣X)]−H(Y∣X) _X→ Y(t)=E[- q_θ(Y X)]-H(Y X). The arrows-of-time result shows these excess risks differ in aggregate; the reversal curse shows ΔB→A _B→ A stays large for specific relations. We measure Δ(t) (t) as a function of training time and characterize its dynamics (plateau, snap, scaling, stabilization) in a controlled setting where every variable is isolated. We construct a task where Δ(t)=ℒ(t) (t)=L(t) starts at exactly logK K and drops to zero when the model learns to condition on a selector z. Across hundreds of runs, the plateau duration depends on dataset size D, not disambiguation complexity K. The marginal solution is stabilized by gradient noise (an entropic force in the sense of Ziyin et al. [2024]), and the transition is collective: zero of 200 base groups are solved at τ/2τ/2, then all snap simultaneously. Contributions. (i) A controlled task (a “wind tunnel” for studying conditional learning) with exact information-theoretic benchmarks (§2). (i) Plateau duration scales with dataset size D, not ambiguity K (§3.2). (i) Evidence for collective transition, entropic stabilization (3.6×3.6× LR effect, 1.8×1.8× batch-size residual in tokens processed), and internal circuit formation (§3, §4). (iv) A directional asymmetry connecting to the reversal curse: unambiguous A→BA→ B is 1.71.7–4.4×4.4× slower than structured (B,z)→A(B,z)→ A (§5). (v) Seven falsified candidate mechanisms constraining the space of viable explanations (Appendix M). 2 Task and Apparatus Task. A surjective map with nbn_b base strings B (6 characters from 36 alphanumeric symbols) and constant fiber size K: each B maps to K distinct targets A (4 characters). A selector z (2 characters) indexes into the fiber, making (B,z)↦A(B,z) A one-to-one. The model receives [BOS,B,SEP,z,SEP][ BOS,\,B,\, SEP,\,z,\, SEP] and predicts A autoregressively. Total dataset: D=nb×KD=n_b× K unique examples. The task has exact entropy decomposition: H(A∣B)=logKH(A B)= K, H(A∣B,z)=0H(A B,z)=0. A model ignoring z achieves loss logK K; a model using z achieves loss 0. Model. 4-layer Transformer (d=128d=128, 4 heads, dmlp=512d_mlp=512; ∼600 600K params), AdamW [Loshchilov and Hutter, 2019], batch size 128, η=10−3η=10^-3, cosine warmup 500 steps, via TransformerLens [Nanda and Bloom, 2022]. Diagnostics. The z-shuffle gap Δz=ℒz-shuffle−ℒclean _z=L_z-shuffle-L_clean detects when z affects the output (Δz=0 _z=0 during plateau; Δz≫0 _z 0 after). It is computed at each evaluation step by running inference twice on the same batch: once with original z tokens, once with z tokens randomly permuted within the batch (preserving z marginals but breaking the (B,z)→A(B,z)→ A mapping). Δz _z is the difference in per-example cross-entropy loss, averaged over the batch. Onset is defined as the first step where Δz>0.1nats _z>0.1\,nats for 3 consecutive evaluations. The waiting time τ is when loss first drops below 50%50\% of logK K; this threshold is robust (exponent varies <0.1<0.1 across thresholds 0.30.3–0.70.7; Appendix D). Seeds and replication. Unless otherwise noted, results are from single-seed runs (seed 42). Where multi-seed data is available (phase-boundary experiments, 3 seeds per condition; Appendix C), we report means and standard deviations. Individual τ measurements have substantial seed variance (CV 3030–70%70\% in phase-boundary runs), which is typical of stochastic processes near a bifurcation: the time to align with a shallow escape direction is inherently noisy. Single-seed sweeps derive statistical power from consistency across many conditions (10 D values, five batch sizes, four learning rates) rather than per-condition replication. 3 The Phase Transition 3.1 The marginal plateau Every run with K>1K>1 exhibits two regimes (Figure 1). Loss drops to ≈logKnats≈ K\,nats within a few hundred steps (Δz=0 _z=0: the model ignores z), then plateaus for thousands of steps before a sharp transition to near-zero loss. The plateau height tracks logK K with ratio 1.01±0.021.01± 0.02 across K∈3,…,36K∈\3,…,36\: the model converges to the exact uniform-over-candidates solution, then stalls at the information-theoretic limit of that partial solution. Figure 1: Staged disambiguation. Loss curves for K∈5,10,20,36K∈\5,10,20,36\ at nb=1,000n_b=1,000. Dashed lines mark logK K. The model converges to the marginal solution within a few hundred steps, then plateaus until a sharp transition to near-zero loss. 3.2 Duration depends on dataset size, not ambiguity Since D=nb×KD=n_b× K, varying K at fixed nbn_b changes both ambiguity and dataset size simultaneously. To isolate the two factors, we hold D constant by adjusting nb=D/Kn_b=D/K. At D=10,000D=10,000: τ≈1,600τ≈ 1,600 for all K∈5,10,20,36K∈\5,10,20,36\ (δ=−0.004δ=-0.004, CI [−0.04,0.04][-0.04,0.04]). At D=20,000D=20,000: τ≈3,950τ≈ 3,950 with δ=−0.04δ=-0.04 (CI [−0.16,0.04][-0.16,0.04]; Figure 5 in Appendix B). K has no effect on waiting time once D is controlled. Note that holding D constant while varying K necessarily varies nb=D/Kn_b=D/K (e.g., at D=10,000D=10,000: K=5K=5 gives nb=2,000n_b=2,000 while K=36K=36 gives nb≈278n_b≈278). The flat τ across these conditions means neither K nor nbn_b matters independently—only their product D. Since K contributes nothing to τ beyond its role in determining D, the 10-point K-sweep at nb=1,000n_b=1,000 (Appendix A) is equivalently a 10-point D-sweep over D∈3,000,…,36,000D∈\3,000,…,36,000\, spanning a full decade. The argument is two steps: (1) the fixed-D control establishes that K has no independent effect on τ; (2) therefore every τ value in the K-sweep is attributable to D=1,000KD=1,000K, not to K per se. Fitting τ versus D yields τ∝D1.19τ D^1.19 (R2=0.98R^2=0.98, 95% CI [1.14,1.31][1.14,1.31]; Figure 4, Table 3). The fixed-D control establishes K-independence; the K-sweep, reinterpreted as a D-sweep, provides the scaling law on 10 data points spanning a full decade. Central finding: plateau height depends on K (how many candidates per group), but plateau duration depends on D (how many examples the optimizer must process). The model must learn to route z across all D pairs; more pairs means more optimization work, regardless of how ambiguity is structured. Multi-seed experiments (3 seeds per condition) confirm that τ increases with D despite substantial seed-to-seed variability (CV 2222–71%71\%; Appendix C). The super-linear exponent remains unexplained. 3.3 The collective snap Per-group evaluation at K∈10,20K∈\10,20\ shows that the transition is collective, not incremental. At τ/2τ/2: zero of 200 sampled groups exceed 80%80\% accuracy (mean accuracy 15.4%15.4\% at K=10K=10, 8.3%8.3\% at K=20K=20). At τ, all groups snap together within a narrow window (99%99\% of groups ≥80%≥ 80\% at 1.5τ1.5τ; full data in Table 6). A shared circuit becomes operational for all groups simultaneously. The transition width, defined as the step interval over which the fraction of groups exceeding 80%80\% accuracy rises from <5%<5\% to >80%>80\%, spans approximately 0.5τ0.5τ: at K=10K=10, from step ∼1,400 1,400 to ∼2,400 2,400; at K=20K=20, from ∼3,500 3,500 to ∼5,200 5,200 (Table 6). No groups solve before 0.5τ0.5τ; the earliest solvers appear near 0.9τ0.9τ, with the bulk following within 0.5τ0.5τ thereafter. The transition is sharp relative to the plateau but occupies a measurable window, not a single step. Under independent per-group learning with mean transition time τ and typical seed variance (CV ∼30% 30\%), the expected transition window for nb=1,000n_b=1,000 groups would be negligible relative to τ. The observed 0.5τ0.5τ window suggests the transition propagates through the network over a measurable timescale, but is far narrower than what incremental group-by-group coverage would predict. 3.4 Entropic stabilization If the plateau were a local minimum, noise should facilitate escape. The opposite holds, consistent with an entropic force [Ziyin et al., 2024] stabilizing the marginal solution. We distinguish manipulations that alter gradient noise without changing the task from those that also degrade the conditional signal. Task-preserving noise manipulations. Two manipulations change gradient noise while leaving the (B,z)→A(B,z)→ A mapping intact. Batch size (fixed η=10−3η=10^-3). Smaller batches increase per-step gradient noise without altering the task or learning rate [Smith et al., 2018]. At K=20K=20, sweeping B∈32,64,128,256,512B∈\32,64,128,256,512\: the step-count ratio spans 23×23×, but smaller batches also reduce per-step throughput. In the appropriate currency—tokens processed (τtok=τsteps×B _tok= _steps× B)—the effect is a 1.8×1.8× delay (409409K–742742K tokens; Table 1, Figure 2a). This 1.8×1.8× residual is modest but real and in the predicted direction: increased per-step noise delays escape even after controlling for data throughput. Learning rate (fixed B=128B=128). Higher η increases both step size and effective gradient noise. τ increases monotonically with η: from 2,1002,100 steps (269269K tokens) at η=3×10−4η=3× 10^-4 to 7,5507,550 steps (966966K tokens) at η=2×10−3η=2× 10^-3; fails to converge at η=5×10−3η=5× 10^-3. Because batch size is constant, the 3.6×3.6× token-normalized ratio is entirely attributable to the change in η, not throughput. The monotonic direction (more noise → longer plateau, no U-shape) is inconsistent with barrier crossing regardless of the step-size confound. Table 1: Batch-size sweep at K=20K=20, η=10−3η=10^-3: steps vs. tokens. The 23×23× step ratio is predominantly a throughput effect; the token-normalized residual is 1.8×1.8×. B τsteps _steps τtok _tok (K) Step ratio Token ratio 32 23,200 742 23.2×23.2× 1.81×1.81× 64 7,200 461 7.2×7.2× 1.12×1.12× 128 3,200 410 3.2×3.2× 1.00×1.00× 256 1,600 410 1.6×1.6× 1.00×1.00× 512 1,000 512 1.0×1.0× 1.25×1.25× Task-degrading manipulation: label noise. Within-fiber label noise at rate p replaces the correct target with a random candidate from the same B-group. This both increases gradient stochasticity and reduces I(A;z∣B)I(A;z B): at rate p, a fraction p of training examples provide incorrect z→Az→ A mappings, making the conditional task itself harder. τtok _tok grows from 410410K (p=0p=0) to 640640K (p=0.1p=0.1) to 6,5506,550K (p=0.2p=0.2), a 16×16× ratio at fixed B=128B=128 (Figure 2b). The sharp elbow between p=0.1p=0.1 and p=0.2p=0.2 (10×10× jump) likely reflects both noise-induced stabilization and task degradation; we cannot separate the two. We include label noise for completeness but note that it is not a pure noise manipulation; the batch-size and learning-rate sweeps, which preserve the task structure, provide cleaner evidence for noise-induced stabilization. Interpretation. The marginal solution has low gradient norm because K competing directions cancel within each group; noise penalizes departure from this low-gradient state, acting as a restoring force. The batch-size sweep shows the direction is correct (noise → slower escape) with a 1.8×1.8× residual after throughput normalization. The learning-rate sweep—the cleaner test, since batch size and thus per-step throughput are held constant—confirms the direction with a 3.6×3.6× effect. Together, both task-preserving manipulations are inconsistent with barrier crossing and consistent with entropic stabilization [Ziyin et al., 2024]. Connection to SGD theory. Wang and Wu [2023] show that minibatch SGD noise preferentially projects onto high-curvature directions, causing escape along flat rather than sharp directions. In our setting, the marginal solution sits at a saddle with extreme anisotropy (λmax≈2.8 _max≈ 2.8, λmin≈−0.005 _min≈-0.005 at K=20K=20; §4): the dominant curvature does not lead to escape, while the shallow λmin _min direction does. Noise-curvature alignment implies that stochastic perturbations predominantly excite the non-escape direction, and increasing noise amplitude strengthens this misalignment. First-exit-time analyses [Nguyen et al., 2019] predict that exit times from metastable states depend on both local geometry and the noise structure. In our anisotropic saddle, the relevant escape is along a direction ∼500× 500× shallower than the dominant curvature, making exit times sensitive to the noise distribution. Adaptive noise injection [Gong et al., 2025] shows that tuning noise tails to the landscape accelerates escape from sharp regions; the stabilization we observe is plausibly the mirror image, where untuned stochastic noise stabilizes flat regions by failing to project onto the escape direction. Stronger causal tests of this mechanism would hold the task fixed while manipulating noise geometry directly: adding isotropic Gaussian noise to the gradient (isolating stochasticity from task structure), varying dropout rate (changing effective noise without altering inputs), or using gradient accumulation to decouple batch noise from per-step throughput. Measuring the empirical gradient covariance and its alignment with Hessian eigenvectors during the plateau would provide a definitive test of the noise-curvature account. We leave these to future work. (a) Batch-size sweep (1.8×1.8× in tokens; 23×23× in steps is mostly throughput). (b) Label noise (16×16× in tokens at fixed B=128B=128). Figure 2: Noise delays escape. (a) Batch-size effect: a modest 1.8×1.8× delay in tokens processed (the 23×23× step-count ratio is predominantly throughput), consistent with entropic stabilization [Ziyin et al., 2024]. (b) Label noise shows a 16×16× delay in tokens, but also degrades the conditional signal (§3.4); the LR sweep (3.6×3.6× at constant throughput) provides the cleanest evidence. 4 Mechanism Internal cascade. Across all K values, Δz _z onset consistently precedes loss onset, with mean lead fraction 0.52±0.140.52± 0.14 of τ (Figure 3a). At K=20K=20, the aggregate Δz _z exceeds detection at step ∼2,750 2,750, well before loss drops at step ∼8,700 8,700. Causal ablation identifies head L0H3 as the critical selector-routing component: zeroing it mid-transition increases loss by 1.72nats1.72\,nats at K=20K=20, while the next head contributes only 1.07nats1.07\,nats (Appendix I). Caveat: cascade timing was measured at fixed nbn_b, so the apparent K-dependence of lead times may reflect D-dependence. Saddle-point geometry. Hessian tracking reveals λmin<0 _min<0 during the plateau (saddle, not minimum). At K=20K=20, plateau eigenvalues are λmax≈2.8 _max≈ 2.8 and λmin≈−0.005 _min≈-0.005, giving |λmax/λmin|≈560| _max/ _min|≈ 560. Across K values, the mean plateau anisotropy ranges from 10,80010,800 to 17,70017,700 (Appendix J), reflecting variation in λmin _min across checkpoints; the escape direction is consistently 500500–1,000×1,000× shallower than the dominant curvature. Consecutive weight displacements have cosine similarity ≈0.04≈ 0.04 during the plateau (random walk), snapping to ∼0.8 0.8 at τ (Figure 3b). At convergence, λmin _min crosses to positive: the first genuine local minimum. The optimizer diffuses on a flat saddle until it aligns with the shallow escape direction; noise makes this alignment harder, explaining the entropic stabilization (§3.4). Both noise-curvature alignment [Wang and Wu, 2023] and first-exit-time scaling [Nguyen et al., 2019] predict that shallower saddles are more sensitive to noise-induced stabilization; the extreme anisotropy here places the system firmly in that regime. (a) Δz _z onset leads loss onset across all K. (b) Direction consistency: ≈0≈0 on plateau, ∼0.8 0.8 at transition. Figure 3: Internal precursors. (a) z-dependence leads loss across all K. (b) Transition is onset of coherent descent. 5 Directional Asymmetry The “backward” task (B,z)→A(B,z)→ A has K-fold ambiguity resolved by z. The “forward” task A→BA→ B is unambiguous but requires flat memorization with no shared group structure. The forward direction is 1.71.7–4.4×4.4× slower (Table 2): the backward task’s shared B-group structure scaffolds circuit formation, while A→BA→ B requires memorizing each pair independently. Transfer from forward pre-training helps at high K (1.31×1.31× at K=36K=36) but hurts at low K (0.75×0.75× at K=5K=5). This connects to the reversal curse [Berglund et al., 2024]: the information-collapsing direction is slower because it lacks the group structure that enables circuit reuse. Linear networks stay at logK K permanently in both directions, confirming that escape requires nonlinear conditional computation (Appendix G). Controls and caveats. All four task variants at each K use the same surjective map f and the same (A,B,z)(A,B,z) triples. Input sequence lengths differ: forward A→BA→ B has |A|+|z|+3=9|A|+|z|+3=9 input tokens; backward (B,z)→A(B,z)→ A has |B|+|z|+3=11|B|+|z|+3=11 input tokens. Output lengths also differ (|B|=6|B|=6 for forward, |A|=4|A|=4 for backward). These differences mean absolute τ comparisons across directions should be interpreted cautiously; the qualitative pattern (forward consistently slower across all four K values, transfer helping at high K) is robust. Directional asymmetry results are single-seed (seed 42); per-condition replication is not available, but consistency across four K values spanning a 7×7× range provides robustness through conditions. Table 2: Directional asymmetry. Unambiguous forward is consistently slower. K τfwd _fwd (A→BA→B) τbwd _bwd ((B,z)→A(B,z)→A) Ratio Transfer 5 2,650 600 4.4×4.4× 0.75×0.75× 10 5,400 1,400 3.9×3.9× 0.88×0.88× 20 9,500 3,200 3.0×3.0× 1.07×1.07× 36 13,050 7,600 1.7×1.7× 1.31×1.31× 6 Discussion What we showed. Transformers learn marginals before conditionals. The metastable regime decomposes into height (logK K, set by ambiguity) and duration (f(D)f(D), set by dataset size). The marginal is stabilized by gradient noise, the transition is collective, and an internal cascade (head commitment, synergy, loss drop) precedes it. Seven candidate mechanisms were tested; all falsified or inconclusive (Appendix M). Limitations and open questions. The scaling τ∝D1.2τ D^1.2 is empirical (single model scale, AdamW only). The D-sweep spans a single decade (33K–3636K) and the fixed-D control has only two D values; confirming the exponent over a wider range remains open. Unless otherwise noted, results are single-seed (seed 42); where multi-seed data is available (Appendix C), seed-to-seed CV ranges from 22%22\% to 71%71\%. The connection to language-model asymmetry is structural analogy, not demonstrated equivalence. Open: why super-linear (D1.2D^1.2, not D1.0D^1.0)? Can entropic stabilization be formalized via an effective potential [Ziyin, 2023] or noise-curvature alignment theory [Wang and Wu, 2023]? Does the staging generalize to natural-structure tasks [Papadopoulos et al., 2024]? The entropic-stabilization interpretation predicts that SGD noise covariance during the plateau projects predominantly onto high-curvature (non-escape) directions [Wang and Wu, 2023]. Directly measuring this alignment, by estimating the empirical gradient covariance across mini-batches and projecting onto Hessian eigenvectors, would provide a definitive test. A complementary experiment would add isotropic gradient noise while keeping labels intact, isolating stochastic stabilization from task degradation. From diagnosis to prescription. A natural next step is to test interventions that accelerate the conditional transition: auxiliary attention losses that reward attending to z, architectural inductive biases for routing (e.g., cross-attention over the selector position), or curriculum strategies that present low-D subsets before the full dataset. Such experiments would move from characterizing the metastable regime to engineering around it. References B. Barak, B. L. Edelman, S. Goel, S. Kakade, E. Malach, and C. Zhang (2022) Hidden progress in deep learning: SGD learns parities near the computational limit. In Advances in Neural Information Processing Systems, External Links: Link Cited by: §1. L. Berglund, M. Tong, M. Kaufmann, M. Balesni, A. C. Stickland, T. Korbak, and O. Evans (2024) The reversal curse: LLMs trained on “A is B” fail to learn “B is A”. In International Conference on Learning Representations, External Links: Link Cited by: §1, §1, §5. B. Gong, G. E. Batista, and P. Lafaye de Micheaux (2025) Adaptive heavy-tailed stochastic gradient descent. External Links: Link Cited by: §3.4. O. Kitouni, N. Nolte, A. Williams, M. Rabbat, and D. Bouchacourt (2024) The factorization curse: which tokens you predict underlie the reversal curse and more. In Advances in Neural Information Processing Systems, External Links: Link Cited by: §1. I. Loshchilov and F. Hutter (2019) Decoupled weight decay regularization. External Links: Link Cited by: §2. N. Nanda and J. Bloom (2022) TransformerLens. External Links: Link Cited by: §2. N. Nanda, L. Chan, T. Lieberum, J. Smith, and J. Steinhardt (2023) Progress measures for grokking via mechanistic interpretability. In International Conference on Learning Representations, External Links: Link Cited by: §1. T. H. Nguyen, U. Simsekli, M. Gurbuzbalaban, and G. Richard (2019) First exit time analysis of stochastic gradient descent under heavy-tailed gradient noise. External Links: Link Cited by: §3.4, §4. V. Papadopoulos, J. Wenger, and C. Hongler (2024) Arrows of time for large language models. In International Conference on Machine Learning, External Links: Link Cited by: §1, §6. A. Power, Y. Burda, H. Edwards, I. Babuschkin, and V. Misra (2022) Grokking: generalization beyond overfitting on small algorithmic datasets. In ICLR Workshop on MATH-AI, External Links: Link Cited by: §1. S. L. Smith, P. Kindermans, C. Ying, and Q. V. Le (2018) Don’t decay the learning rate, increase the batch size. In International Conference on Learning Representations, External Links: Link Cited by: §3.4. M. Wang and L. Wu (2023) A theoretical analysis of noise geometry in stochastic gradient descent. External Links: Link Cited by: §3.4, §4, §6, §6. L. Ziyin, M. Wang, H. Li, and L. Wu (2024) Parameter symmetry and noise equilibrium of stochastic gradient descent. In Advances in Neural Information Processing Systems, External Links: Link Cited by: §1, Figure 2, Figure 2, §3.4, §3.4. L. Ziyin (2023) Symmetry induces structure and constraint of learning. External Links: Link Cited by: §6. Appendix A Dataset-Size Sweep Table 3 reports τ versus K at fixed nb=1,000n_b=1,000 and η=10−3η=10^-3. Since D=nb×K=1,000KD=n_b× K=1,000K in this design, varying K simultaneously varies D; the sweep is equivalently a sweep of D from 3,0003,000 to 36,00036,000. A power-law fit τ∝Dδτ D^δ gives δ=1.19δ=1.19 (95% CI [1.14,1.31][1.14,1.31], n=10n=10). All runs reach the logK K plateau within a few hundred steps and subsequently transition to near-zero loss, confirming convergence across the full range. Table 3: K-sweep at nb=1,000n_b=1,000, equivalently a D-sweep. Since D=nb×K=1,000KD=n_b× K=1,000K and K has no independent effect on τ (§3.2), this table is interpretable as τ vs. D. Power-law fit: τ∝D1.19τ D^1.19 (R2=0.98R^2=0.98, 95% CI [1.14,1.31][1.14,1.31]). K 3 5 7 10 13 17 20 25 30 36 D 3K 5K 7K 10K 13K 17K 20K 25K 30K 36K τ 450 800 1,050 1,850 2,100 3,300 3,950 5,250 6,950 8,750 Figure 4: τ versus D on log-log axes. Blue circles: 10-point K-sweep at nb=1,000n_b=1,000, reinterpreted as a D-sweep (D=1,000KD=1,000K). Red squares: fixed-D control (D=10D=10K and 2020K, four K values each), confirming K-independence. Dashed line: power-law fit τ∝D1.19τ D^1.19. Appendix B Fixed-D Control Figure 5 shows that at fixed dataset size D, the waiting time τ is flat across K. At D=10,000D=10,000, τ≈1,600τ≈ 1,600 for all K∈5,10,20,36K∈\5,10,20,36\ (δ=−0.004δ=-0.004). Doubling the dataset to D=20,000D=20,000 roughly doubles τ to ≈3,950≈3,950, with no K-dependence (δ=−0.04δ=-0.04, CI [−0.16,0.04][-0.16,0.04]). Figure 5: Duration depends on D, not K. At fixed D, τ is flat across K. Doubling D roughly doubles τ. Appendix C Multi-Seed Variability To assess seed-to-seed variability, we draw on phase-boundary experiments that ran 3 independent seeds (7, 42, 123) at each (K,η)(K,η) condition. Table 4 reports all conditions where all 3 seeds converged within 50,000 steps. Since D=1,000KD=1,000K, each row is equivalently a point on the τ vs. D curve at a given η. The coefficient of variation (CV) ranges from 22%22\% to 71%71\%, with a median of 31%31\%. This substantial seed variability is expected for a saddle-escape process: the time to align with a shallow escape direction is inherently stochastic. Despite this variability, at every η where we have two or more D values, τ increases with D: the D-scaling is robust across seeds, even though individual τ values fluctuate. At the baseline η=10−3η=10^-3, the single available multi-seed condition (K=36K=36, D=36,000D=36,000) gives τ=10,867±3,349τ=10,867± 3,349 (mean ± std), consistent with the single-seed value of 8,7508,750 from the main D-sweep. The main 11-point D-sweep (Appendix A) uses seed 42; the phase-boundary data confirms that this seed produces τ values within the multi-seed distribution. Table 4: Multi-seed τ across D values. 3 seeds per condition from phase-boundary experiments. SE =std/3=std/ 3. D K η Mean τ Std τ CV 5,000 5 5×10−35× 10^-3 2,850 1,727 61% 5,000 5 7×10−37× 10^-3 3,250 2,295 71% 5,000 5 1×10−21× 10^-2 2,183 504 23% 10,000 10 3×10−33× 10^-3 5,950 2,299 39% 10,000 10 5×10−35× 10^-3 9,867 3,838 39% 10,000 10 7×10−37× 10^-3 8,533 1,893 22% 20,000 20 2×10−32× 10^-3 5,750 1,316 23% 20,000 20 3×10−33× 10^-3 30,300 8,497 28% 36,000 36 1×10−31× 10^-3 10,867 3,349 31% Appendix D Threshold Sensitivity The waiting time τ is defined as the first step where loss drops below a fraction α of logK K. We test robustness of the power-law exponent δ to the choice of α across the range 0.3,0.4,0.5,0.6,0.7\0.3,0.4,0.5,0.6,0.7\. Table 5 reports the fitted exponents. The exponent varies from 1.121.12 (α=0.7α=0.7) to 1.221.22 (α=0.3α=0.3), with mean 1.18±0.041.18± 0.04. This variation is small relative to the confidence interval on any single fit, confirming that the power-law scaling is not an artifact of a particular threshold choice. Table 5: Threshold sensitivity. Power-law exponent δ for τ∝Dδτ D^δ at different threshold fractions α of logK K. α δ 95% CI n 0.3 1.22 [1.16,1.34][1.16,1.34] 10 0.4 1.20 [1.13,1.32][1.13,1.32] 10 0.5 1.20 [1.14,1.30][1.14,1.30] 10 0.6 1.17 [1.11,1.31][1.11,1.31] 10 0.7 1.12 [1.07,1.23][1.07,1.23] 10 Appendix E Per-Group Evaluation To test whether the transition is incremental (groups escape one by one) or collective (all groups snap together), we evaluate 200 randomly sampled base groups at multiple checkpoints for K∈10,20K∈\10,20\. For each group, we measure the fraction of its K targets correctly predicted and report the percentage of groups exceeding 80% and 100% accuracy. The key finding is that at τ/2τ/2, zero groups exceed 80% accuracy in either condition. By 1.5τ1.5τ, nearly all groups have transitioned. This narrow window confirms a collective mechanism: the selector-routing circuit, once functional, applies to all groups simultaneously rather than being learned group by group. Table 6: Collective transition. 200 sampled base groups evaluated at four checkpoints per K value. K Step Frac. τ ≥80%≥80\% 100%100\% Mean acc. 10 500 27% 0.0% 0.0% 0.118 10 900 49% 0.0% 0.0% 0.154 10 1,800 97% 31.5% 1.5% 0.645 10 2,800 151% 99.0% 75.0% 0.969 20 1,900 48% 0.0% 0.0% 0.083 20 2,900 73% 0.0% 0.0% 0.192 20 3,800 96% 1.5% 0.0% 0.508 20 5,900 149% 87.0% 8.5% 0.870 Appendix F Selector Variants To verify that the phenomenon is not an artifact of the selector token format, we vary the selector length |z|∈1,2,3,4|z|∈\1,2,3,4\ at K=10K=10. Table 7 reports the results. The plateau height remains tightly locked to logK K across all selector lengths (ratio to logK K ranges from 1.0071.007 to 1.0371.037). The waiting time τ varies modestly from 1,4501,450 to 1,8501,850 steps, with no systematic trend. The staged-disambiguation phenomenon is robust to the representation of the selector. Table 7: Selector length variants at K=10K=10 (logK=2.30nats K=2.30\,nats). |z||z| Plateau (nats) Plateau / logK K τ 1 2.34 1.017 1,650 2 2.39 1.037 1,850 3 2.33 1.014 1,500 4 2.32 1.007 1,450 Appendix G Architecture Comparison We test three nonlinear architectures and one linear baseline. All nonlinear architectures exhibit the logK K plateau and subsequent sharp transition; the quantitative details differ. Nonlinear architectures. Transformers (δ=1.70δ=1.70, R2=0.97R^2=0.97), Gated MLPs (δ=1.62δ=1.62, R2=0.97R^2=0.97), and RNNs (δ=1.47δ=1.47, R2=0.98R^2=0.98) all show power-law scaling of τ with D. The exponents differ across architectures but all exceed 1, suggesting the super-linear scaling is a property of the optimization landscape rather than a specific architecture. Note: these runs used nb=1,000n_b=1,000 (so D=1,000KD=1,000K) and did not control the D/K confound; the exponent values should be interpreted with caution. Two-layer linear networks. A two-layer linear network (d=128d=128, same tokenizer and dataset) stays at logK K permanently with zero z-shuffle gap up to 30,00030,000 training steps. This confirms that the transition requires nonlinear conditional computation: the marginal solution P(A∣B)P(A B) is representable in linear models, but the conditional P(A∣B,z)P(A B,z) requires nonlinear interaction between B and z. Appendix H Phase Boundary For each K∈5,10,20,36K∈\5,10,20,36\, we sweep over learning rates with 3 seeds per (K,η)(K,η) pair. A learning rate is classified as “succeeding” if all 3 seeds converge within 50,00050,000 steps, and “failing” if any seed fails to converge. The critical learning rate η∗(K)η^*(K) is defined as the geometric mean of the largest all-succeed and smallest any-fail values. Table 8 reports the boundary. A power-law fit gives η∗(K)=0.048⋅K−0.83η^*(K)=0.048· K^-0.83 (R2=0.96R^2=0.96, δ=0.83±0.16δ=0.83± 0.16). Near the boundary, τ diverges: at K=36K=36, η=2×10−3η=2× 10^-3, τ=68,350τ=68,350 (i.e., 6.3×6.3× the baseline at η=10−3η=10^-3). Across all 24 near-boundary runs that did converge, zero reversions were observed: once the transition initiates, it runs to completion. Note: these runs used nb=1,000n_b=1,000, so D=1,000KD=1,000K; the boundary may be η∗(D)η^*(D) rather than η∗(K)η^*(K). Table 8: Phase boundary. Critical learning rate η∗η^* per K. K D Max all-succeed η Min any-fail η η∗η^* 5 5,000 1.0×10−21.0× 10^-2 1.5×10−21.5× 10^-2 1.2×10−21.2× 10^-2 10 10,000 7.0×10−37.0× 10^-3 1.0×10−21.0× 10^-2 8.4×10−38.4× 10^-3 20 20,000 3.0×10−33.0× 10^-3 5.0×10−35.0× 10^-3 3.9×10−33.9× 10^-3 36 36,000 1.0×10−31.0× 10^-3 2.0×10−32.0× 10^-3 1.4×10−31.4× 10^-3 Appendix I Head Ablation We perform zero-ablation (replacing a head’s output with zeros) for all 16 heads (4 layers × 4 heads) at three phases: pre-transition, mid-transition, and post-transition. Table 9 reports the top heads by loss increase at each phase for K=10K=10 and K=20K=20. Pre-transition, no single head has an outsized effect (max Δℒ<0.02nats <0.02\,nats). Mid-transition, head L0H3 emerges as the dominant contributor: ablating it at K=10K=10 increases loss by 1.41nats1.41\,nats, and at K=20K=20 by 1.72nats1.72\,nats. The remaining layer-0 heads also contribute substantially (1.01.0–1.3nats1.3\,nats each), while layer-1 heads contribute an order of magnitude less (0.10.1–0.14nats0.14\,nats). Post-transition, all layer-0 heads become critical (3.13.1–3.7nats3.7\,nats), reflecting their role in reading the input tokens. The concentration of mid-transition effects in L0H3 is consistent with it being the first head to specialize in routing selector information. Layer-1 heads likely perform downstream computation on the features that layer-0 heads extract. Caveat: zero-ablation creates out-of-distribution inputs for downstream layers; mean-replacement or rescaling ablations would provide more faithful causal estimates and are left to future work. Table 9: Head ablation. Loss increase (Δℒ , nats) when zeroing each head. Top 4 heads shown per phase. K=10K=10 K=20K=20 Phase Head Δℒ Head Δℒ Pre L0H3 0.019 L0H3 0.022 L0H1 0.012 L0H0 0.015 Mid L0H3 1.41 L0H3 1.72 L0H0 1.30 L0H0 1.07 L0H2 1.28 L0H2 1.05 L0H1 1.27 L0H1 1.03 Post L0H3 3.65 L0H3 3.69 L0H2 3.37 L0H0 3.61 L0H1 3.36 L0H2 3.41 L0H0 3.13 L0H1 3.28 Appendix J Hessian Details We compute the largest and smallest eigenvalues of the Hessian via power iteration (50 iterations, batch size 512) at checkpoints spanning the plateau and transition for K∈10,20,36K∈\10,20,36\. During the plateau, the loss landscape is a saddle: λmin<0 _min<0 for all K values. At K=20K=20, representative plateau eigenvalues are λmax≈2.8 _max≈ 2.8 and λmin≈−0.005 _min≈-0.005, giving a pointwise ratio |λmax/λmin|≈560| _max/ _min|≈ 560. However, λmin _min varies across checkpoints within the plateau; the mean anisotropy (averaged over plateau checkpoints) is 12,10012,100 at K=20K=20 and ranges from 10,80010,800 to 17,70017,700 across K values (Table 10). The escape direction is consistently 500500–1,000×1,000× shallower than the dominant curvature. At the transition, λmax _max spikes to 3030–8080, reflecting the rapid weight changes as the model reorganizes. Post-convergence, λmin _min crosses to a small positive value (≈+0.0001≈+0.0001): the first genuine local minimum encountered during training. The anisotropy does not scale systematically with K (γ=0.06γ=0.06, R2=0.05R^2=0.05), consistent with the falsification of the “curvature scales with K” hypothesis. Robustness checks. We recomputed λmax _max and λmin _min at K=20K=20 across batch sizes 256,512,1024,2048\256,512,1024,2048\ and iteration counts 25,50,100,200\25,50,100,200\. The coefficient of variation is <0.1<0.1 for λmax _max and <0.15<0.15 for λmin _min. The sign of λmin _min (negative on plateau, positive post-convergence) is consistent across all settings. Table 10: Hessian anisotropy (|λmax/λmin|| _max/ _min|) during the plateau for each K. K Anisotropy (mean) λmax _max 3 17,70017,700 9.83 5 11,70011,700 5.90 10 10,80010,800 4.59 20 12,10012,100 2.83 36 17,60017,600 0.95 Appendix K Hierarchical Task We tested a two-level hierarchical version of the task to see whether the model would produce a staircase of plateaus (one per disambiguation level). In this variant, each base string B maps to K1K_1 groups, each of which maps to K2K_2 targets, yielding K1×K2K_1× K_2 total targets per B. Two selectors z1z_1 and z2z_2 resolve the ambiguity at each level. Small hierarchy (K1=5K_1=5, K2=4K_2=4, D=20,000D=20,000): The model plateaus at log(20)≈3.0nats (20)≈ 3.0\,nats and transitions directly to near-zero loss. No intermediate plateau appears at log(5)≈1.6nats (5)≈ 1.6\,nats. The model does not learn to resolve one level of ambiguity before the other. Large hierarchy (K1=20K_1=20, K2=10K_2=10, D=200,000D=200,000): The model remains stuck at log(200)≈5.3nats (200)≈ 5.3\,nats for 100,000100,000 steps, likely due to the large dataset size. No intermediate plateau at log(20) (20) is observed during this window. Both selectors provide gradient signal from step one; the staging we observe in the main task is between marginal and conditional prediction, not between hierarchical levels of ambiguity. This negative result suggests the “marginals before conditionals” principle does not straightforwardly extend to hierarchical disambiguation. Appendix L Gradient Dissipation We measure the total gradient dissipation Q=∑t∈transitionη‖∇ℒ‖2Q= _t η\| \|^2 during the transition window (from τ/2τ/2 to 2τ2τ). At η=10−3η=10^-3, Q scales linearly with logK K (slope =4.29=4.29, R2=0.991R^2=0.991; Table 11). At η=3×10−4η=3× 10^-4, the slope increases to 9.919.91 (R2=0.913R^2=0.913), and at η=3×10−3η=3× 10^-3 the relationship degrades (slope =110.2=110.2, R2=0.732R^2=0.732). The logK K scaling of Q at moderate learning rates suggests that the optimizer must dissipate a quantity of gradient energy proportional to the information-theoretic gap logK K during the transition. The strong η-dependence of the slope is expected: larger steps dissipate more energy per update. Caution: these runs used fixed nb=1,000n_b=1,000, so the logK K dependence may partly reflect D-dependence. Table 11: Gradient dissipation Q during transition, by K and η. K Q (η=3×10−4η=3× 10^-4) Q (η=10−3η=10^-3) Q (η=3×10−3η=3× 10^-3) 5 2.72 3.96 7.05 10 5.92 7.42 3.48 20 16.4 9.90 159.8 Appendix M Falsification Summary We tested seven candidate mechanisms for the plateau. Table 12 summarizes each hypothesis, its prediction, the experimental result, and the verdict. Six are clearly falsified; one is inconclusive due to numerical instability. Table 12: Tested mechanisms and their outcomes. Hypothesis Prediction Result Verdict Gradient cancellation drives plateau τ∝τ number of groups G β≈0.02β≈ 0.02; no scaling with G Falsified Barrier crossing (noise helps) More noise → faster escape LR: 3.6×3.6× slower; BS: 1.8×1.8× (tokens) Falsified Incremental group coverage Groups solved independently over time 0% of groups ≥80%≥80\% at τ/2τ/2 Falsified Label noise breaks symmetry Noise perturbs away from flat solution 16×16× delay (p=0.2p=0.2); confounded∗ Falsified∗ K drives duration τ∝Kδτ K^δ at fixed D δ=−0.004δ=-0.004; no K effect Falsified Linear networks suffice Two-layer linear escapes Stays at logK K permanently Falsified Curvature scales with K |λmax/λmin|∝K| _max/ _min| K Anisotropy ∼104 10^4 for all K Inconclusive ∗Label noise also degrades I(A;z∣B)I(A;z B); the 16×16× delay conflates noise stabilization with task degradation (§3.4).