Radial Suppression Accelerates Algorithmic Generalization:
A Geometric Analysis of Delayed Generalization

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=Wx$, where $W\in\mathbb{R}^{d\times d_{in}}$ is the weight matrix and $x$ is the input. The standard cross-entropy loss $L_{CE}$ operates on $h$. Let $\ell(h)$ denote the loss mapping from pre-activations to the scalar loss.

By the chain rule, the gradient with respect to the weights is $\nabla_{W}L_{CE}=(\nabla_{h}\ell)x^{T}$. The Hessian $H_{W}$ with respect to the weights—omitting the second-derivative tensor terms of the network architecture for the local linear approximation—is dominated by:

Unconstrained cross-entropy optimization naturally drives the magnitude of features outward to maximize softmax margins, driving $\|h\|_{2}\to\infty$. Under the mapping $h=Wx$, this radial inflation necessitates unconstrained growth in the spectral norm of $W$, specifically along the principal components of the input covariance $\Sigma_{x}$. Consequently, the maximal eigenvalue of the Hessian, $\lambda_{max}$, grows proportionally cohen2021gradient.

The Edge of Stability dictates that gradient descent becomes unstable when $\lambda_{max}>2/\eta$, where $\eta$ is the learning rate. Our norm penalty, $L_{norm}=\frac{1}{d}(\|h\|_{2}-\sqrt{d})^{2}$, introduces an opposing restoring force. By locking $\|h\|_{2}\approx\sqrt{d}$, we artificially bound the spectral norm of $W$. Because $\lambda_{max}$ is a function of this bounded weight norm, the local sharpness is forcibly held below the $2/\eta$ threshold, explaining the dramatic reduction in Hessian trace observed in our experiments.

Proposition 2. The activation norm penalty acts as a continuous Lagrangian relaxation of a Riemannian gradient flow on the hypersphere $\mathbb{S}^{d-1}(\sqrt{d})$, where the penalty multiplier $\lambda$ dictates the stiffness of the manifold retraction.

Derivation. Consider the continuous-time gradient flow of the activations $\dot{h}=-\nabla_{h}L_{total}$, where the total objective is $L_{total}=L_{CE}(h)+\frac{\lambda}{d}(\|h\|_{2}-\sqrt{d})^{2}$. The gradient evaluates to:

Let $P_{r}=\frac{hh^{T}}{\|h\|_{2}^{2}}$ be the radial projection matrix. We decompose the continuous flow into radial ($\dot{h}_{rad}$) and tangential ($\dot{h}_{tan}$) components:

In the asymptotic limit as $\lambda\to\infty$, the penalty strictly dominates the radial dynamics, correcting any deviation from the target radius infinitely fast. Thus, $\dot{h}_{rad}\to 0$ and $\|h\|_{2}\to\sqrt{d}$. In this limit, the optimization trajectory reduces exactly to:

This equation is the exact formulation of Riemannian gradient descent restricted to the manifold $\mathbb{S}^{d-1}(\sqrt{d})$ absil2008optimization. Because we operate at a finite $\lambda$, 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.

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:

where $\sigma_{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 li2020towards. Specifically, $L_{CE}$ inflates the dominant singular value $\sigma_{1}$ exponentially faster than the tail, driving $s(A)\to 1$.

Applying the norm penalty without a stop-gradient introduces a structural gradient conflict. The primary task attempts to maximize $\sigma_{1}$ to increase logit margins. Simultaneously, the radial gradient of the norm penalty, $-\frac{2\lambda}{d}(\|h\|_{2}-\sqrt{d})\frac{h}{\|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 $\sigma_{1}$. By actively bounding $\sigma_{1}$ via this continuous “antagonistic regularization,” the optimization energy is distributed across the trailing singular values ($\sigma_{2},\dots,\sigma_{k}$). This formally explains the empirical preservation of stable rank, ensuring the latent space maintains sufficient effective dimensionality for complex $\mathcal{O}(d)$ Fourier features to emerge prior to the phase transition.

Our evidence is consistent with the following account of algorithmic phase transitions:

1. 1.

   Cross-entropy optimization drives radial inflation of activations, causing spectral collapse and trapping the network in a high-norm memorization basin.

2. 2.

   Constraining activations to a $\sqrt{d}$-radius hypersphere suppresses radial gradient components, redirecting optimization to angular (tangential) updates.

3. 3.

   Angular updates preserve feature diversity and promote the discovery of periodic Fourier circuits.

4. 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., $R^{2}>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.

LayerNorm enforces a hard per-token constraint (zero mean, unit variance) and restores expressivity via learned affine parameters $\gamma,\beta$. 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 $\lambda$ is worthwhile. We argue that it is, for three reasons: (i) the penalty achieves $41\times$ 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.

We track the assembly of Fourier circuits over training by projecting the hidden activations onto the complete 4-dimensional Fourier basis $\{\sin(2\pi ka/P),\,\cos(2\pi ka/P),\,\sin(2\pi kb/P),\,\cos(2\pi kb/P)\}$ for each frequency $k\in\{1,\dots,48\}$ and computing the $R^{2}$ 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 $R^{2}{>}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 $\sigma_{\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\times$ 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.

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\in\{1,\dots,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\pi ka/P),\,\cos(2\pi ka/P)\}$, taking the larger of the two as the selectivity score. The resulting $512\times 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\approx 10.7$) specialize to each Fourier mode, with near-zero selectivity outside their assigned frequency. This neuron-cluster-per-frequency organization is the circuit motif identified by nanda2023progress as the hallmark of a well-formed modular-arithmetic Fourier circuit. The unpenalized baseline, whose radial inflation suppresses $\tilde{\Phi}_{\mathrm{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.

The penalty consistently and substantially accelerates grokking across all moduli and data fractions tested. Speedups relative to Strong WD range from ${\sim}5\times$ at $f_{\text{train}}{=}0.6$ to ${\sim}5\times$ at $f_{\text{train}}{=}0.4$. Notably, at $f_{\text{train}}{=}0.3$ the penalty continues to induce grokking (26,980–12,260 epochs depending on $P$) while Strong WD fails entirely for $P{\in}\{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.

Penalty strength $\lambda$. We sweep $\lambda\in\{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 $\lambda{=}0.05$ optimal at 2,460 epochs. Very low $\lambda$ ($0.001$) delays onset to 8,200 epochs; very high $\lambda$ ($1.0$) over-constrains angular updates, slowing onset to 9,900 epochs. The method is robust across an order of magnitude ($\lambda\in[0.01,0.1]$).

Target radius. Testing $c\sqrt{d}$ for $c\in\{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 $\mathcal{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.

To verify the penalty does not pathologically degrade standard feature learning, we applied it ($\lambda{=}0.05$) to a 500K-parameter character-level Transformer (4 layers, 4 heads, $d{=}128$) on Tiny Shakespeare.

The penalty does not accelerate language modeling—as expected, since character-level LM lacks a sharp memorization$\to$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.