HiLD 2025: 3rd Workshop on High-dimensional Learning Dynamics 1–16
Radial Suppression Accelerates Algorithmic Generalization: A Geometric Analysis of Delayed Generalization Srijan Tiwari and Aditya Chauhan∗ and Manjot Singh∗
srijan [email protected] aditya [email protected] manjot [email protected]
arXiv:2606.32000v1 [cs.LG] 30 Jun 2026
Indian Institute of Technology Roorkee
Abstract Why do neural networks memorize algorithmic training data long before they generalize? We present a geometric case study demonstrating that, on tasks where generalization requires discovering structured low-dimensional circuits, the memorization–generalization delay is driven by radial inflation of hidden representations under cross-entropy optimization. We formalize a radial–angular decomposition of activation-space dynamics and derive three testable propositions: (i) that penalizing radial inflation induces anisotropic, datadependent weight regularization; (ii) that it suppresses radial gradient energy below the isotropic random baseline, forcing predominantly angular updates; and (iii) that it biases convergence toward flatter minima. To empirically validate these propositions, √ we study a single-hyperparameter norm penalty that softly constrains activations to a d-radius hypersphere. On modular arithmetic, this penalty accelerates grokking upto 6× across MLPs and Transformers, and halves training steps for a 10M-parameter nanoGPT on 3-digit addition.
1. Introduction Neural networks trained on algorithmic tasks exhibit a distinct form of delayed generalization where models achieve near-perfect training accuracy early in optimization, yet test accuracy remains low for orders of magnitude more training before undergoing a sudden phase transition (Power et al., 2022; Barak et al., 2023). Mechanistic studies of modular arithmetic reveal that generalization requires discovering structured low-dimensional circuits (Nanda et al., 2023; Zhong et al., 2023), while the memorizing solution is characterized by unstructured, high-norm representations. We demonstrate that a cause of this is radial inflation where standard cross-entropy incentivizes outward growth of hidden activations to push logits into the saturating regime of softmax (Prieto et al., 2025). This inflates the dominant singular value while collapsing effective rank (Wang et al., 2026). We formalize this intuition through a radial–angular decomposition of activation-space dynamics, and derive three analytical propositions. To test these propositions,√we study a activation norm penalty that softly constrains hidden representations to a d- radius hypersphere. We study this intervention specifically on algorithmic generalization tasks where the generalization phase transition is sharp and measurable. Note that we do not claim universality. Rather, we present the radial–angular framework as a geometric lens for understanding algorithmic phase transitions, and the norm penalty as a principled instrument for understanding this geometry. ∗
Equal contribution
© S. Tiwari, A. Chauhan & M. Singh.
Tiwari Chauhan Singh
Contributions 1. A radial–angular decomposition of activation-space dynamics that generates three testable predictions about how radial suppression should affect optimization geometry along with empirical validation(§4). 2. A single-hyperparameter norm penalty that accelerates grokking by upto 6× across MLPs, Transformers, and a 10M-parameter nanoGPT.
2. Background and Related Work 2.1. Grokking Phenomenon The “Goldilocks zone” framework (Liu et al., 2022) models generalization as occurring within a narrow band of weight norms while memorization corresponds to high-norm solutions outside this zone. Nanda et al. (2023) reverse-engineered the Fourier multiplication algorithm learned by grokked Transformers, identifying three phases: memorization, circuit formation, and cleanup. Merrill et al. (2023) modeled grokking as competition between dense memorizing and sparse generalizing subnetworks. GrokFast (Lee et al., 2024) amplifies slow-varying gradient components. Xu et al. (2025) proposed GrokTransfer, which transplants embeddings from a pre-trained proxy model. 2.2. Geometric Perspectives on Regularization Existing regularization approaches generally constrain either weight-space or activationspace representations. Standard L2 weight decay applies an isotropic penalty λ∥W ∥2F . Spectral normalization (Miyato et al., 2018) bounds the Lipschitz constant via σmax . Weight normalization (Salimans and Kingma, 2016) explicitly decouples magnitude from direction. Activation Regularization in AWD-LSTM (Merity et al., 2017) directly penalizes ∥ht ∥22 . Minimizing Activation Norms (MAN) (Yashwanth et al., 2024) minimizes E[∥al ∥2 ] as a Hessianflatness proxy. On the other hand, LayerNorm (Ba et al., 2016) and RMSNorm impose hard per-token constraints on activation statistics. Our penalty is a soft constraint: it permits temporary violation during landscape traversal and has no learnable affine parameters that could undo the constraint. Prieto et al. (2025) identify softmax-driven logit inflation as a grokking bottleneck; their ⊥Grad projects gradients away from magnitude-scaling directions in parameter space. Yıldırım (2026) enforce hard L2 projection onto a bounded spherical topology. Our approach provides a soft, loss-based variant operating in activation space.
3. Method: The Activation Norm Penalty 3.1. Formulation Given a network with hidden representation h ∈ Rd at a target layer, we augment the cross-entropy objective with a single-hyperparameter penalty: Ltotal = LCE (f (x; θ), y) + λ Lnorm (h) 2
(1)
Activation Geometry & Grokking
Lnorm (h) =
√ 2 1 ∥h∥2 − d d
(2)
√ The target d ensures that the average squared activation per feature is constant as width increases, matching standard initialization variance and preventing both variance collapse and unbounded inflation. We also show that this constraint is a relaxation of Riemannian gradient flow on a hypersphere (See §A.2). In MLPs, the penalty is applied to the pre-activation outputs of all hidden layers, computed per-sample and averaged over the batch. In Transformer architectures with Pre-LayerNorm, the penalty is applied to each sub-layer output before addition to the residual stream: √ 2 1 xl+1 = xl + Sublayerl (LN(xl )), L(l) = (3) ∥Sublayer (LN(x ))∥ − d 2 l l norm d This constrains the increment to the residual stream at each layer rather than the cumulative stream, avoiding spurious norm growth in deeper networks.
4. Analytical Framework and Empirical Validation We analyze the norm penalty from three complementary perspectives. Each generates a testable prediction, which we validate empirically immediately after its derivation. 4.1. Radial–Angular Gradient Decomposition Proposition 1 Let Pr = hhT /∥h∥22 be the radial projection matrix. The total gradient g = ∇h Ltotal decomposes into a radial component grad = Pr g and a tangential component gtan = (I − Pr )g. To quantify this geometry without artifacts from high-dimensional orthogonality, we define the Normalized Fractional Radial Energy: Φ̃rad = d ·
∥grad ∥22 ∥g∥22
(4)
Under an isotropic random walk in Rd , the expected fractional energy in the one-dimensional radial subspace is 1/d, giving √ ahnull hypothesis of E[Φ̃rad ] = 1. The gradient of the norm penalty term, − 2λ (∥h∥ − d) ∥h∥2 , is purely radial and opposes the radial component of the 2 d task gradient. Prediction. The penalty should suppress Φ̃rad well below 1 throughout training, and this angular redirection should be accompanied by accelerated assembly of features, such as the periodic Fourier features that underlie generalization on modular arithmetic. Result. During early memorization phase (epochs 0–1,500), the baseline exhibits severe radial inflation. The penalized model suppresses Φ̃rad to approximately 0.15 from initialization onward, an order of magnitude below the null. Table 1 reports the downstream effect on Fourier circuit assembly. Fourier coherence (R2 > 0.9 on the P =97 basis) is reached at epoch 2,460 under the penalty versus 34,200 for the baseline and the dominant Fourier magnitude increases fourfold, indicating that angular optimization produces sharper, more structured representations prior to the phase transition. 3
Tiwari Chauhan Singh
Table 1: Fourier and geometric analysis (MLP, P =97, 5 seeds, baseline WD= 10−3 ). F Cep : epoch at Fourier coherence (R2 > 0.9); |F |max : dominant Fourier magnitude; σmax : largest singular value; Ef f. Rank: out of 512; κ: condition number; Tr(H): Hessian trace; S̃: Tr(H)/∥θ∥2 . Bold denotes the better value per column. Fourier Circuit (§4.1)
Spectral & Curvature (§4.3)
Model
FCep
|F |max
σmax
Eff. Rank
κ
Tr(H)
S̃
Baseline Norm Penalty
34,200 2,460
1.2 ± 0.1 4.8 ± 0.2
>52,000 36.5
135 ± 4 443 ± 3
∼150,000 ∼32
42.5 ± 2.1 1.4 ± 0.1
0.18 ± 0.03 0.004 ± 0.001
Train vs. Test Accuracy Comparison
0.4
Penalty ( =0.05) Train Penalty ( =0.05) Test Strong WD (0.1) Train Strong WD (0.1) Test
0.0 0
5000
10000
15000
Epoch
20000
25000
24000
16000
2.0
1.5
Baseline Penalty ( =0.05) Isotropic Floor (1.0)
1.0
0.5
8000
0.0
30000
Figure 1: Train/Test accuracy: Baseline vs. Penalty.
2.5
Fractional Radial Energy × d
0.6
Grokking Onset (Epoch)
Accuracy
Grokking Onset Goldilocks Zone
32000
0.8
0.2
Activation Radial Energy vs. Epoch
The Goldilocks Zone: Grokking Onset vs.
1.0
10 3
10 1
10 2
Penalty Strength (Log Scale)
100
Figure 2: Sweep Over λ.
0
2000
4000
Epoch
6000
8000
10000
Figure 3: Radial Energy Φ̃rad over training..
4.2. Implicit Anisotropic Weight Regularization Proposition 2 √Under the local linear approximation h = W x and in the high-norm phase when radial inflation regime ∥h∥2 ≫ d—which holds during the early memorization √ 2 is maximal—the centered penalty approximates (∥h∥− d) ≈ ∥h∥2 , and the expected penalty becomes: 1 E [Lnorm ] ≈ Tr(W Σx W T ) (5) d where Σx = Ex [xxT ] is the input second-moment matrix. Unlike isotropic weight decay (λ∥W ∥2F ), this penalizes weight directions proportionally to the variance of their input features: directions aligned with high-variance principal components of Σx receive stronger suppression. Prediction. The data-dependent anisotropy of this regularizer should produce faster generalization than isotropic weight decay at similar regularization strength, because it selectively suppresses the high-variance directions that cross-entropy exploits for memorization while leaving low-variance directions free to participate in circuit formation. Result. Table 2 reports grokking onset against competing regularizers. Our penalty reaches generalization approximately 6.3× faster than strong isotropic weight decay and 6.1× faster than MAN, which minimizes E[∥a∥2 ] without a centered target. The higher effective rank under our penalty (443 vs. 402 for strong WD) is consistent with the anisotropy prediction. Table 3 shows that the penalty consistently accelerates the memorizationgeneralization phase transition over multiple architectures. The more modest relative 4
Activation Geometry & Grokking
speedup on NanoGPT is consistent with the architecture already performing partial radial suppression via its affine LayerNorm. Table 2: Comparison vs. regularization methods (MLP, P =97, ftrain =0.5, 5 seeds). See §B.2 for definitions and §E.1 for sweep across ftrain and P Method Baseline (WD= 10−3 ) Strong WD (10−1 ) Dropout (p=0.1) MAN (E[∥a∥2 ]) Norm Penalty (Ours)
Grok Onset
Final Test Acc
Eff. Rank
Hessian Trace
DNG 15,540 ± 1,480 22,000 ± 2,500 15,000 ± 2,000 2,460 ± 136
1.7% 100.0% 98.5% 100.0% 100.0%
135 402 378 400 443
42.5 3.2 5.8 2.5 1.4
Table 3: Grokking onset across architectures. MLP and Transformer: P =97, ftrain =0.5, 5 seeds. NanoGPT: 3-digit addition, 3 seeds. Speedup reported relative to strong WD baseline for MLP and Transformer; relative to no-penalty NanoGPT for 3-digit addition. Setting MLP (epochs) Transformer (epochs) NanoGPT (steps)
Baseline
Norm Penalty
Speedup (vs. Strong WD)
15,540 ± 1480 8,000 ± 450 22,500 ± 1,200
2,460 ± 136 5,200 ± 400 9,800 ± 600
∼6.3× ∼1.5× ∼2.3×
4.3. Curvature Reduction via Norm Bounding Proposition 3 Using the empirical Fisher as a curvature proxy, the Hessian trace approximates as Tr(H) ≈ Ex [∥δ∥22 ∥x∥22 ], where δ = ∇h LCE . Layer-wise activation norm bounding restricts ∥x∥2 (the input to the next layer) and reduces pre-activation saturation (thereby bounding ∥δ∥2 ). Prediction. The penalty should substantially reduce the Hessian trace and normalized sharpness relative to the baseline, and this curvature reduction should be accompanied by a shift toward higher effective rank, indicating that the loss landscape flattens without collapsing the representational geometry. Result. Table 1 reports spectral and curvature diagnostics. The penalty achieves a 30× reduction in raw Hessian trace and a 45× reduction in normalized sharpness, confirming the curvature prediction. Spectral compression is equally striking: σmax drops from > 52,000 to 36.5, while effective rank rises from 135 to 443 out of 512 dimensions. This combination— flatter landscape and higher rank—distinguishes the norm penalty from isotropic regularizers, which reduce sharpness by collapsing the spectrum rather than by redistributing it.
5. Conclusion We have presented a geometric case study of the memorization–generalization phase transition on algorithmic tasks. Through a radial–angular decomposition of activation-space 5
Tiwari Chauhan Singh
dynamics, we derived three testable predictions about how radial suppression should affect optimization, and validated each empirically using a simple norm penalty as instrument. The penalty accelerates grokking by 6× on modular arithmetic and 2.3× on 3-digit addition, while dramatically compressing the spectral geometry and flattening the loss landscape. We situated the penalty within a taxonomy of geometric interventions, clarifying its relationship to normalization layers, direct activation penalties, and parameter-space methods. Our work suggests that, for algorithmic learning tasks, the memorization–generalization delay is fundamentally a geometric phenomenon: cross-entropy drives radial inflation, trapping networks in memorization basins, and principled radial suppression provides a direct lever to accelerate the phase transition. Whether this geometric lens extends to broader learning settings remains an open and important question.
References P-A Absil, Robert Mahony, and Rodolphe Sepulchre. Optimization algorithms on matrix manifolds. Princeton University Press, 2008. Jimmy Lei Ba, Jamie Ryan Kiros, and Geoffrey E Hinton. Layer normalization. arXiv preprint arXiv:1607.06450, 2016. Boaz Barak, Benjamin L. Edelman, Surbhi Goel, Sham Kakade, Eran Malach, and Cyril Zhang. Hidden progress in deep learning: Sgd learns parities near the computational limit, 2023. URL https://arxiv.org/abs/2207.08799. Jeremy M Cohen, Simran Kaur, Yuanzhi Li, J Zico Kolter, and Ameet Talwalkar. Gradient descent on neural networks typically occurs at the edge of stability. arXiv preprint arXiv:2103.00065, 2021. Jaerin Lee, Bong Gyun Kang, Kihoon Kim, and Kyoung Mu Lee. Grokfast: Accelerated grokking by amplifying slow gradients. arXiv preprint arXiv:2405.20233, 2024. Nayoung Lee, Kartik Sreenivasan, Jason D. Lee, Kangwook Lee, and Dimitris Papailiopoulos. Teaching arithmetic to small transformers, 2023. URL https://arxiv.org/abs/ 2307.03381. Zhiyuan Li, Yuping Luo, and Kaifeng Lyu. Towards resolving the implicit bias of gradient descent for matrix factorization: Greedy low-rank learning. arXiv preprint arXiv:2012.09839, 2020. Ziming Liu, Ouail Kitouni, Niklas Nolte, Eric J. Michaud, Max Tegmark, and Mike Williams. Towards understanding grokking: An effective theory of representation learning, 2022. URL https://arxiv.org/abs/2205.10343. Stephen Merity, Nitish Shirish Keskar, and Richard Socher. Regularizing and optimizing lstm language models. arXiv preprint arXiv:1708.02182, 2017. William Merrill, Nikolaos Tsilivis, and Aman Shukla. A tale of two circuits: Grokking as competition of sparse and dense subnetworks. arXiv preprint arXiv:2303.11873, 2023. 6
Activation Geometry & Grokking
Takeru Miyato, Toshiki Kataoka, Masanori Koyama, and Yuichi Yoshida. Spectral normalization for generative adversarial networks, 2018. URL https://arxiv.org/abs/1802. 05957. Neel Nanda, Lawrence Chan, Tom Lieberum, Jess Smith, and Jacob Steinhardt. Progress measures for grokking via mechanistic interpretability, 2023. URL https://arxiv.org/ abs/2301.05217. Alethea Power, Yuri Burda, Harri Edwards, Igor Babuschkin, and Vedant Misra. Grokking: Generalization beyond overfitting on small algorithmic datasets. CoRR, abs/2201.02177, 2022. URL https://arxiv.org/abs/2201.02177. Lucas Prieto, Melih Barsbey, Pedro Mediano, and Tolga Birdal. Grokking at the edge of numerical stability. In International Conference on Learning Representations, volume 2025, pages 81151–81168, 2025. Tim Salimans and Durk P Kingma. Weight normalization: A simple reparameterization to accelerate training of deep neural networks. Advances in neural information processing systems, 29, 2016. Junxuan Wang, Xuyang Ge, Wentao Shu, Zhengfu He, and Xipeng Qiu. Dimensional collapse in transformer attention outputs: A challenge for sparse dictionary learning, 2026. URL https://arxiv.org/abs/2508.16929. Zhiwei Xu, Zhiyu Ni, Yixin Wang, and Wei Hu. Let me grok for you: Accelerating grokking via embedding transfer from a weaker model, 2025. URL https://arxiv.org/abs/2504. 13292. M Yashwanth, Gaurav Kumar Nayak, Harsh Rangwani, Arya Singh, R Venkatesh Babu, and Anirban Chakraborty. Minimizing layerwise activation norm improves generalization in federated learning. In Proceedings of the IEEE/CVF Winter Conference on Applications of Computer Vision, pages 2287–2296, 2024. Alper Yıldırım. The geometric inductive bias of grokking: Bypassing phase transitions via architectural topology, 2026. URL https://arxiv.org/abs/2603.05228. Ziqian Zhong, Ziming Liu, Max Tegmark, and Jacob Andreas. The clock and the pizza: Two stories in mechanistic explanation of neural networks, 2023. URL https://arxiv. org/abs/2306.17844.
Appendix A. Theoretical Framework and Extended Geometric Analysis In this section, we formalize the geometric mechanisms through which the activation norm penalty alters high-dimensional learning dynamics. We analyze its effects on local curvature, gradient flow topologies, and the spectral properties of the representation matrix. 7
Tiwari Chauhan Singh
A.1. Preempting the Edge of Stability via Radial Bounding Proposition 1. Radial inflation in activation space drives progressive sharpening in parameter space. Constraining activation norms bounds the parameter spectral norm, formally preempting the Edge of Stability (EoS) instability threshold. Derivation. Let the network’s local linear approximation be h = W x, where W ∈ Rd×din is the weight matrix and x is the input. The standard cross-entropy loss LCE operates on h. Let ℓ(h) denote the loss mapping from pre-activations to the scalar loss. By the chain rule, the gradient with respect to the weights is ∇W LCE = (∇h ℓ)xT . The Hessian HW with respect to the weights—omitting the second-derivative tensor terms of the network architecture for the local linear approximation—is dominated by: HW ≈ (∇2h ℓ) ⊗ (xxT )
(6)
Unconstrained cross-entropy optimization naturally drives the magnitude of features outward to maximize softmax margins, driving ∥h∥2 → ∞. Under the mapping h = W x, this radial inflation necessitates unconstrained growth in the spectral norm of W , specifically along the principal components of the input covariance Σx . Consequently, the maximal eigenvalue of the Hessian, λmax , grows proportionally Cohen et al. (2021). The Edge of Stability dictates that gradient descent becomes unstable λmax > 2/η, √ when 1 2 where η is the learning rate. Our norm penalty, Lnorm = d (∥h∥2 − d) , introduces an √ opposing restoring force. By locking ∥h∥2 ≈ d, we artificially bound the spectral norm of W . Because λmax is a function of this bounded weight norm, the local sharpness is forcibly held below the 2/η threshold, explaining the dramatic reduction in Hessian trace observed in our experiments. A.2. The Norm Penalty as a Lagrangian Relaxation of Riemannian Flow Proposition 2. The activation norm penalty acts as a √ continuous Lagrangian relaxation of a Riemannian gradient flow on the hypersphere Sd−1 ( d), where the penalty multiplier λ dictates the stiffness of the manifold retraction. Derivation. Consider the continuous-time gradient flow √ of the activations ḣ = −∇h Ltotal , where the total objective is Ltotal = LCE (h) + λd (∥h∥2 − d)2 . The gradient evaluates to: 2λ ∇h Ltotal = ∇h LCE + d
√ ! ∥h∥2 − d h ∥h∥2
(7)
T
hh Let Pr = ∥h∥ 2 be the radial projection matrix. We decompose the continuous flow into 2
radial (ḣrad ) and tangential (ḣtan ) components: ḣrad = −Pr ∇h LCE −
√ 2λ h (∥h∥2 − d) d ∥h∥2
ḣtan = −(I − Pr )∇h LCE
(8) (9)
In the asymptotic limit as λ → ∞, the penalty strictly dominates the radial dynamics, correcting any deviation from the target radius infinitely fast. Thus, ḣrad → 0 and ∥h∥2 → 8
Activation Geometry & Grokking
√
d. In this limit, the optimization trajectory reduces exactly to: ḣ = −(I − Pr )∇h LCE
(10)
This equation is √ the exact formulation of Riemannian gradient descent restricted to the d−1 manifold S ( d) Absil et al. (2008). Because we operate at a finite λ, our method avoids the brittleness of hard retraction mappings, instead optimizing within a “thickened sphere” while maintaining the favorable angular dynamics characteristic of Riemannian optimization. A.3. Spectral Collapse and Antagonistic Gradients Proposition 3. Unconstrained cross-entropy exhibits an implicit bias toward rank-1 representations. The activation penalty induces an “antagonistic gradient” that counteracts this bias, preserving the stable rank and allowing the spectral edge of generalizing circuits to assemble. Derivation. The effective capacity of the network representations over a batch A is bounded by the stable rank: P 2 ∥A∥2F i σi s(A) = = (11) ∥A∥22 σ12 where σi are the singular values. Gradient flow on separable data without explicit regularization invariably converges to max-margin solutions, functioning as an implicit bias toward low-rank factorizations Li et al. (2020). Specifically, LCE inflates the dominant singular value σ1 exponentially faster than the tail, driving s(A) → 1. Applying the norm penalty without a stop-gradient introduces a structural gradient conflict. The primary task attempts to maximize σ1 to increase logit margins. Simultaneously, √ h the radial gradient of the norm penalty, − 2λ (∥h∥ − d) 2 d ∥h∥2 , aggressively pulls back the representation vector. Because this opposing gradient is strictly radial, it exerts its maximal suppressive force exactly along the direction of σ1 . By actively bounding σ1 via this continuous “antagonistic regularization,” the optimization energy is distributed across the trailing singular values (σ2 , . . . , σk ). This formally explains the empirical preservation of stable rank, ensuring the latent space maintains sufficient effective dimensionality for complex O(d) Fourier features to emerge prior to the phase transition.
Appendix B. Detailed Discussion B.1. What the Geometric Lens Reveals Our evidence is consistent with the following account of algorithmic phase transitions: 1. Cross-entropy optimization drives radial inflation of activations, causing spectral collapse and trapping the network in a high-norm memorization basin. √ 2. Constraining activations to a d-radius hypersphere suppresses radial gradient components, redirecting optimization to angular (tangential) updates. 9
Tiwari Chauhan Singh
3. Angular updates preserve feature diversity and promote the discovery of periodic Fourier circuits. 4. The resulting solutions lie in flatter minima with dramatically lower Hessian trace. This account is correlational : the penalty simultaneously suppresses radial inflation, preserves rank, flattens curvature, and accelerates Fourier coherence (The point during training when a neural network’s internal representations strongly align with a theoretical Fourier basis (e.g., R2 > 0.9), marking the successful assembly of a generalizable, periodic algorithm). We present these observations as consistent with the radial–angular framework rather than as proof of a unique causal chain. Disentangling these effects—e.g., via interventions that preserve rank without radial suppression, or vice versa—is an important direction for future work. B.2. The LayerNorm Relationship LayerNorm enforces a hard per-token constraint (zero mean, unit variance) and restores expressivity via learned affine parameters γ, β. Our penalty is a soft constraint with no affine restoration, which has two consequences: (i) the network can temporarily violate the hypersphere during landscape traversal, providing a smoother optimization path; (ii) the absence of affine parameters prevents the network from undoing the constraint via learned rescaling. When applied to Transformers that already include LayerNorm, the two provide complementary radial suppression: LayerNorm constrains per-token statistics within each sub-layer, while the penalty constrains sub-layer output magnitudes across sub-layers. LayerNorm (without affine parameters) achieves 80% of the grokking acceleration of our penalty, raising the question of whether the additional hyperparameter λ is worthwhile. We argue that it is, for three reasons: (i) the penalty achieves 41× lower Hessian trace than LayerNorm, suggesting a qualitatively different solution geometry; (ii) the penalty is a loss-based intervention that does not modify the forward pass and is therefore trivially combinable with any architecture; (iii) combining both (Table 9) yields the fastest grokking, indicating complementary mechanisms—LayerNorm normalizes per-token statistics within sub-layers, while the penalty constrains sub-layer output magnitudes across the network. Metrics. Grokking onset: first epoch (or step) at which test accuracy consistently exceeds 90%. Effective rank : stable rank ∥A∥2F /∥A∥22 over a batch of 256. Hessian trace: Hutchinson’s method with 100 Rademacher probes (convergence verified in Table 8). Normalized sharpness: Tr(H)/∥θ∥2 . “DNG” indicates failure to reach 90% test accuracy within the training budget. Table 4: Comparison vs. normalization baselines (MLP, P =97, 5 seeds). Variant
Grok Onset (ep.)
Hessian Trace
Eff. Rank
Baseline (WD= 10−3 ) LayerNorm (No Affine) RMSNorm Norm Penalty (Ours)
DNG (>100k) 8,000 ± 450 8,500 ± 520 6,167 ± 624
42.5 ± 2.1 57.9 ± 3.2 35.5 ± 2.8 1.4 ± 0.1
135 ± 4 464 ± 5 498 ± 2 443 ± 3
10
Activation Geometry & Grokking
Circuit Assembly: Per-Frequency Fourier Encoding Over Training
0.040 0.035
R² Fourier Frequency k
Fourier Frequency k
5.0 7.5
0.030
10.0
0.025
12.5
0.020
15.0
0.015
17.5 20.0
0.010 10000 20000 30000 40000 50000
Epoch
2.5
0.045 0.040 0.035 0.030 0.025 0.020 0.015 0.010 0.005
5.0 7.5 10.0 12.5 15.0 17.5 20.0
Strong WD (0.1)
R² Fourier Frequency k
2.5
Penalty =0.05
10000 20000 30000 40000 50000
Epoch
2.5
0.200 0.175 0.150 0.125 0.100 0.075 0.050 0.025
5.0 7.5 10.0
R²
Baseline (WD=1e ³)
12.5 15.0 17.5 20.0
10000 20000 30000 40000 50000
Epoch
Figure 4: Fourier Circuit Assembly over Training. Left: The baseline develops only weak, diffuse structural traces that never crystallize into a functional algorithm within the training budget. Center: Under the norm penalty, Fourier structure emerges progressively and systematically—low-k modes assemble first, followed by higher modes— reaching R2 >0.9 coherence at epoch 4,100 (cf. Table 1). Right: Strong WD forces abrupt crystallization but concentrates representational energy into 2–3 dominant modes, leaving higher frequencies poorly represented.
Appendix C. Mechanistic Analysis Beyond aggregate metrics, we probe the internal structure of the learned representations to connect our geometric framework to circuit-level mechanisms. We present the two analyses in chronological order: first the assembly dynamics over training, then the specialization structure of the converged solution.
C.1. Circuit Assembly Timeline We track the assembly of Fourier circuits over training by projecting the hidden activations onto the complete 4-dimensional Fourier basis {sin(2πka/P ), cos(2πka/P ), sin(2πkb/P ), cos(2πkb/P )} for each frequency k ∈ {1, . . . , 48} and computing the R2 fit at each epoch (MLP, P =97). Figure 4 reveals qualitatively different assembly dynamics across conditions. The baseline accumulates weak frequency traces that never reach the coherence threshold of R2 >0.9 reported in Table 1. Strong weight decay forces a sudden crystallization but concentrates capacity into a sparse subset of 2–3 dominant frequencies, consistent with the low effective rank (402) and high σmax reported in Table 2. The norm penalty, by contrast, produces a distributed assembly: frequencies emerge sequentially beginning from low-k modes, spreading representational energy across many orthogonal directions. This distributed encoding is mechanistically consistent with the penalty’s 30× reduction in Hessian trace (Table 1): when information is spread evenly across many frequency channels rather than concentrated in a few, the loss landscape exhibits lower curvature along every direction. 11
Tiwari Chauhan Singh
Per-Neuron Fourier Selectivity at Convergence
3
300
2 1
400
0
5
10
15
Fourier Frequency k
100
0.4
200
0.3 0.2
300
0.1
400 500
0
5
10
15
Fourier Frequency k
100
2.5 2.0
200
Magnitude
4
Strong WD (0.1)
0
Magnitude
200
Neuron (sorted by peak freq)
5
Magnitude
Neuron (sorted by peak freq)
6
100
500
Penalty =0.05
0
Neuron (sorted by peak freq)
Baseline (DNG)
0
1.5 300
1.0 0.5
400 500
0.0 0
5
10
15
Fourier Frequency k
Figure 5: Per-Neuron Fourier Selectivity at Convergence. Rows are neurons (d=512), columns are Fourier frequencies (k=1, . . . , 48); color encodes maximum absolute correlation with the {sin, cos} pair at each k. Left: The unpenalized baseline exhibits diffuse, unstructured correlations with no discernible frequency preference. Center: The norm penalty produces a block-diagonal structure in which coherent clusters of ∼10 neurons specialize to each frequency, with near-zero correlation outside their assigned mode. Right: Strong WD yields partial clustering that is noticeably noisier and skewed toward a small number of low-k frequencies, leaving higher modes sparsely populated.
C.2. Per-Neuron Fourier Selectivity
To characterize the converged circuit structure, we measure the frequency selectivity of individual neurons. For each of the d=512 hidden neurons j and each Fourier frequency k ∈ {1, . . . , 48}, we compute the maximum absolute correlation between neuron j’s activation profile across all P 2 inputs and the two-dimensional Fourier basis {sin(2πka/P ), cos(2πka/P )}, taking the larger of the two as the selectivity score. The resulting 512 × 48 heatmap reveals the degree to which each neuron commits to a single frequency. As shown in Figure 5, the penalty produces a clean block-diagonal structure: coherent clusters of approximately 10 neurons (consistent with 512/48 ≈ 10.7) specialize to each Fourier mode, with near-zero selectivity outside their assigned frequency. This neuroncluster-per-frequency organization is the circuit motif identified by Nanda et al. (2023) as the hallmark of a well-formed modular-arithmetic Fourier circuit. The unpenalized baseline, whose radial inflation suppresses Φ̃rad from the first epoch (Figure 1), shows diffuse correlations with no frequency preference—consistent with the low dominant Fourier magnitude (|F |max =1.2) in Table 1. Strong weight decay forces partial cluster formation but concentrates heavily on a few low-k modes and leaves higher frequencies underrepresented, matching the lower effective rank (402 vs. 443) reported in Table 2. Together, Figures 4 and 5 connect the aggregate spectral diagnostics in §4 to a concrete circuit-level picture: radial suppression enables the progressive, distributed assembly of a modular Fourier circuit that would otherwise be blocked by norm-driven spectral collapse. 12
Activation Geometry & Grokking
C.3. Limitations Approximation regimes. The √ anisotropic regularization analysis (Analysis 2) operates in the high-norm regime ∥h∥2 ≫ d, which holds during early memorization but breaks √ down as the penalty takes effect and ∥h∥2 → d. This is not a practical concern—by the time the approximation fails, the penalty has √ already redirected optimization away from radial inflation, and the centered form (∥h∥− d)2 maintains a restoring force thereafter—but the analysis should not be read as a uniform characterization of training dynamics. Similarly, the curvature analysis (Analysis 3) directly bounds only one factor of the Hessian trace product (∥x∥2 ); the other (∥δ∥2 ) is constrained indirectly through reduced pre-activation saturation rather than by the penalty itself. We therefore treat the curvature reduction as a verified mechanistic hypothesis rather than a formal guarantee.
Correlational evidence. The penalty simultaneously suppresses radial inflation, preserves effective rank, flattens curvature, and accelerates Fourier coherence. These effects are consistent with the radial–angular framework but are entangled: we cannot, from the current experiments, attribute the grokking acceleration to any single mechanism in isolation. Ablations that independently manipulate rank preservation without radial suppression, or vice versa, would strengthen the causal interpretation and are an important direction for future work.
Task scope. All primary experiments involve algorithmic tasks with sharp memorization– generalization phase transitions. The Tiny Shakespeare sanity check (§E.3) confirms the penalty is benign on a standard character-level language modeling task—perplexity degrades by under 2% and effective rank increases—but performance on large-scale language modeling or vision benchmarks remains untested. Architectures that rely on activation magnitude as an explicit confidence signal may interact adversely with the penalty. √ Fixed target radius. The choice of d as the target radius is principled—it matches the O(1) per-feature variance of standard initialization schemes—but it may not be optimal across all architectures or layer types. The ablation in Appendix E.1 shows that c = 1 is optimal among c√∈ {0.5, 1.0, 2.0, 5.0}, and c = 0.5 is close, suggesting robustness in the neighborhood of d. Tunable or learnable per-layer radii are a natural extension.
Baseline context. The MLP and Transformer baselines use weak weight decay (10−3 ), a regime where grokking is slow or absent (Liu et al., 2022). Speedups reported in Table 3 are relative to the strong WD baseline (10−1 , which groks at 15,540 epochs) to provide a fair comparison; relative to the weak baseline the raw numbers are larger but less meaningful as a measure of the penalty’s contribution over aggressive norm control in general. 13
Tiwari Chauhan Singh
Table 5: Taxonomy of geometric interventions for algorithmic generalization. Method Weight Decay Spectral Norm Miyato et al. (2018) LayerNorm MAN (E[∥a∥2 ]) Yashwanth et al. (2024) ⊥Grad Prieto et al. (2025) Spherical Projection Yıldırım (2026) Ours
Space
Constraint Type
Radial Suppression
Rank Effect
Weight Weight Activation Activation Parameter Activation Activation
Isotropic σmax bound Hard (per-token) Soft (toward zero) Gradient projection Hard√(global) Soft ( d target)
Indirect Indirect Direct Direct Direct Complete Direct
Collapse Preserved Preserved Preserved Preserved Collapse Preserved
Table 6: Grokking onset across moduli and data fractions (MLP, 5 seeds). Comparison against Strong WD (10−1 ). DNG: did not grok within 50,000 epochs. ftrain
P = 97
P = 137
P = 211
0.6 Strong WD Norm Penalty
9,000 ± 420 1,320 ± 75
7,580 ± 223 1,260 ± 49
6,720 ± 331 1,200 ± 63
0.5 Strong WD Norm Penalty
15,540 ± 1,480 2,460 ± 136
12,740 ± 723 2,200 ± 0
10,600 ± 856 1,960 ± 162
0.4 Strong WD Norm Penalty
33,320 ± 1,942 6,680 ± 349
24,520 ± 2,114 5,200 ± 253
20,860 ± 1,839 4,380 ± 133
0.3 Strong WD Norm Penalty
DNG 26,980 ± 1,503
DNG 17,480 ± 757
49,840 ± 320 12,260 ± 554
Appendix D. Comparisons D.1. Taxonomy of Geometric Interventions
Appendix E. Ablations E.1. Robustness Sweeps The penalty consistently and substantially accelerates grokking across all moduli and data fractions tested. Speedups relative to Strong WD range from ∼5× at ftrain =0.6 to ∼5× at ftrain =0.4. Notably, at ftrain =0.3 the penalty continues to induce grokking (26,980– 12,260 epochs depending on P ) while Strong WD fails entirely for P ∈{97, 137} and exhausts the training budget for P =211, demonstrating that radial suppression provides a decisive advantage precisely in the low-data regime where isotropic regularization breaks down. 14
Activation Geometry & Grokking
Table 7: λ sensitivity (MLP, P =97, ftrain =0.5, 5 seeds). λ 0.001 0.01 0.05 0.1 0.5 1.0
Grok Onset (epoch)
Final Test Acc
8,200 ± 650 4,100 ± 320 2,460 ± 136 3,800 ± 290 6,400 ± 510 9,900 ± 840
100.0% 100.0% 100.0% 100.0% 100.0% 100.0%
Table 8: Hutchinson Probe Convergence via Hessian trace estimates vs. number of probes. Probes 50 100 200 500
Trace Estimate
Relative Error vs. 500
1.38 ± 0.12 1.41 ± 0.09 1.40 ± 0.07 1.39 ± 0.05
4.2% 2.1% 1.4% —
E.2. Other Ablations Penalty strength λ. We sweep λ ∈ {0.001, 0.01, 0.05, 0.1, 0.5, 1.0} on the MLP (5 seeds). All values induce grokking (unlike the baseline), with λ=0.05 optimal at 2,460 epochs. Very low λ (0.001) delays onset to 8,200 epochs; very high λ (1.0) over-constrains angular updates, slowing onset to 9,900 epochs. The method is robust across an order of magnitude (λ ∈ [0.01, 0.1]). √ Target radius. Testing c d for c ∈ {0.5, 1.0, 2.0, 5.0}: c=1.0 is optimal (2,460 ep.); c=0.5 is close (8,200 ep.); c=5.0 degrades to 12,300 epochs. The c=1 optimality is consistent with standard initialization schemes that set per-feature variance to O(1). Application site (MLP). Pre-activation (default) is optimal (2,460 ep.); post-ReLU is slightly worse (9,800 ep.). Constraining pre-activation norms preserves information about negative components that ReLU would zero out, maintaining a richer representational geometry. LayerNorm interaction (Transformer). The penalty and LayerNorm compound: their combination consistently outperforms either alone, confirming complementary mechanisms. Optimizer sensitivity. The penalty induces grokking under Adam (no WD): 9,200 epochs (vs. 78,000 baseline). Under SGD: 32,000 epochs (vs. DNG baseline). AdamW + penalty is optimal. The penalty is effective across optimizers but benefits from adaptive learning rates. 15
Tiwari Chauhan Singh
Table 9: LayerNorm interaction (Transformer, P =97, 5 seeds). LayerNorm Off Off On (no affine) On (no affine) On (with affine) On (with affine)
Penalty
Grok Onset
Eff. Rank
Off On Off On Off On
DNG (>100k) 14,500 ± 1,200 8,000 ± 450 5,200 ± 400 7,500 ± 500 4,200 ± 300
— 402 464 475 450 460
E.3. Sanity Check: Non-Algorithmic Task To verify the penalty does not pathologically degrade standard feature learning, we applied it (λ=0.05) to a 500K-parameter character-level Transformer (4 layers, 4 heads, d=128) on Tiny Shakespeare. Table 10: Language modeling sanity check (Tiny Shakespeare, 3 seeds, 8,000 steps). Variant Baseline Norm Penalty
Val. Loss
Perplexity
Eff. Rank
1.585 ± 0.011 1.608 ± 0.004
4.9 5.0
94 108
The penalty does not accelerate language modeling—as expected, since character-level LM lacks a sharp memorization→generalization phase transition. Crucially, it does not collapse representations (rank increases from 94 to 108) and perplexity degradation is within 2%, confirming that the penalty is benign outside its target domain.
Appendix F. Experimental Setup MLP on Modular Addition. 2-layer MLP, hidden dimension d=512, ReLU activations. Data: (a + b) mod 97; training fraction ftrain =0.5 (4,656 of 9,409 pairs). Optimizer: AdamW, lr=10−3 , weight decay 10−3 , batch size 256 (full-batch). Penalty: λ=0.05, applied to pre-ReLU activations of both hidden layers. Training: 100,000 epochs max; 5 independent seeds. Small Transformer on Modular Addition. 2-layer, 4-head Transformer, dmodel =128, pre-LayerNorm (no affine by default). Penalty applied to each sub-layer output per Eq. 3. Same data, optimizer, and seeds as MLP. NanoGPT on 3-Digit Addition. 6 layers, 6 heads, dmodel =384 (∼10M parameters), pre-LayerNorm with affine. Reverse-format 3-digit addition (Lee et al., 2023); 80/20 train/test split. AdamW, lr=10−2 , cosine decay to 10−4 over 30,000 steps, weight decay 0.1, batch size 128. Penalty: λ=0.01; 3 seeds. Wall-clock overhead: 0.32→0.35 sec/step (+9.4%) on A6000; +1.5% peak GPU memory.
16