The Geometric Cost of Normalization: Affine Bounds on the Bayesian Complexity of Neural Networks

1. Problem Statement

Layer Normalization (LayerNorm) and RMS Normalization (RMSNorm) are widely treated as interchangeable architectural choices in deep learning. However, they impose fundamentally different geometric constraints on the data manifold passed to subsequent weight matrices.

• The Core Question: Does the specific geometry of the normalization output (e.g., a flat hyperplane vs. a sphere) structurally alter the Bayesian complexity (effective parameter count) of the neural network before training begins?
• The Gap: While normalization is known to aid training stability, its impact on the Local Learning Coefficient (LLC)—a measure from Singular Learning Theory (SLT) quantifying the effective number of parameters seen by the Bayesian posterior—has not been rigorously quantified or theoretically derived.

2. Methodology

The authors employ a combination of Singular Learning Theory (SLT) and controlled empirical experiments to isolate the geometric effects of normalization.

• Theoretical Framework:

  • LLC & RLCT: The paper utilizes the Local Learning Coefficient ($\lambda$), which measures the effective complexity of a model. A lower $\lambda$ indicates higher degeneracy (symmetries) in the loss landscape.
  • Holonomic Constraints: Normalization is modeled as a holonomic constraint that confines data to a specific submanifold $M \subset \mathbb{R}^d$.
  • Symmetry Space: The authors define a symmetry space $S_M$ consisting of weight matrix directions orthogonal to the span of the data manifold. Directions in $S_M$ are "invisible" to the loss function, creating continuous symmetries that reduce the LLC.
  • Key Metric: The reduction in LLC ($\Delta \lambda$) is derived as $\Delta \lambda = \frac{m(d - d_s)}{2}$, where $m$ is the output dimension, $d$ is the input dimension, and $d_s$ is the dimension of the linear span of the data manifold.

• Experimental Setup:

  • Architecture: Minimal single-layer teacher-student setups ($W \in \mathbb{R}^{m \times d}$) with Mean Squared Error (MSE) loss. This eliminates confounding factors like multi-layer scaling symmetries ($GL(r)$).
  • Estimation: The wrLLC framework is used to estimate LLC via Stochastic Gradient Langevin Dynamics (SGLD).
  • Variables: Experiments vary output dimension ($m$), input dimension ($d$), and manifold geometry (flat hyperplanes, paraboloids, hyperboloids, saddles, and simplices).

3. Key Contributions

A. The Geometric Threshold for LLC Reduction

The paper proves that the reduction in LLC is determined by whether the data manifold is affinely flat.

• Theorem 1: For a linear layer $W$ operating on data confined to a manifold $M$ with linear span dimension $d_s$, the LLC reduction is exactly $\Delta \lambda = \frac{m(d - d_s)}{2}$.
• Binary Threshold: The drop is binary.
  • Affinely Flat: If the manifold is a linear hyperplane (e.g., $1^\top x = 0$), $d_s = d-1$, resulting in a drop of $\Delta \lambda = m/2$.
  • Curved: If the manifold has any non-zero curvature (positive or negative), it spans the full ambient space ($d_s = d$), resulting in no LLC drop ($\Delta \lambda = 0$).

B. Specific Bounds for Normalization Layers

• LayerNorm: Enforces zero mean ($1^\top x = 0$). This projects data onto a linear hyperplane.
  • Result: Reduces the LLC of the subsequent weight matrix by exactly $m/2$.
• RMSNorm: Enforces unit RMS without mean-centering. Data lies on a hypersphere.
  • Result: The sphere has full linear span ($d_s = d$). No LLC reduction occurs.

C. The "Smuggled Bias" in Softmax

The authors analyze the Simplex-Bias Duality.

• Mechanism: Softmax outputs lie on a simplex ($1^\top x = 1$), which is an affine hyperplane not passing through the origin.
• Strictly Linear Layer ($y=Wx$): The linear span is full ($d_s=d$); no symmetry exists; $\Delta \lambda = 0$.
• Affine Layer ($y=Wx+b$): The explicit bias $b$ becomes indistinguishable from the "smuggled bias" ($b_{smuggled} = \frac{1}{d}W\mathbf{1}$) inherent in the simplex geometry. This creates an $m$-dimensional symmetry space.
• Result: An explicit bias paired with Softmax data induces an LLC drop of $m/2$.

D. Finite-Sample Phase Transition

While the algebraic theory predicts a sharp binary threshold, experiments with finite sample sizes ($N=2000$) reveal a smooth crossover.

• Observation: The SGLD estimator only "sees" the curvature if the data distribution significantly overlaps with the curved region of the manifold.
• Implication: Narrow, high-amplitude curvature may be missed by the sampler (appearing flat), while wide, low-amplitude curvature is detected. The effective RLCT depends on the data-curvature overlap, not just the existence of curvature.

4. Experimental Results

1. LayerNorm vs. RMSNorm:

   • Experiments varying $m$ (output dim) confirmed that LayerNorm reduces LLC by $\approx m/2$ (e.g., at $m=4$, $\Delta \lambda \approx 2.0$).
   • RMSNorm showed $\Delta \lambda \approx 0$ across all dimensions, confirming it preserves the full parameter count.
   • Varying input dimension $d$ showed the drop is independent of $d$, depending solely on the lost dimension ($d-d_s$).

2. Curvature Threshold:

   • Flat Hyperplane: $\Delta \lambda \approx 2.0$ (matches theory).
   • Curved Manifolds (Paraboloid, Hyperboloid, Saddle): $\Delta \lambda \approx 0$. The sign of curvature (positive/negative) was irrelevant; only the presence of curvature mattered.

3. Softmax Degeneracy:

   • Linear Layer: $\Delta \lambda \approx 0$.
   • Affine Layer: $\Delta \lambda \approx 2.26$ (theoretical $m/2 = 2.0$), confirming the "smuggled bias" hypothesis.

5. Significance and Implications

• Structural Complexity Reduction: The paper establishes that LayerNorm is not just a stabilizer but a structural reducer of model complexity. It effectively removes $m/2$ degrees of freedom per layer before training starts. RMSNorm avoids this "cost," preserving the full expressivity of the weight matrix.
• Attention Sinks: The authors hypothesize that the $m/2$ reduction in LayerNorm might structurally bias weight matrices (like $W_V$ in Transformers) toward low-rank solutions, potentially explaining the "attention sink" phenomenon. However, empirical validation in this specific context was inconclusive.
• Grokking Dynamics: The reduction in degrees of freedom by LayerNorm might facilitate the "grokking" transition (sudden generalization) by removing parameters associated with memorization, making generalizing circuits more accessible.
• Theoretical Rigor: The work provides the first exact algebraic derivation of normalization's impact on Bayesian complexity, moving beyond heuristic observations to precise geometric bounds.

Conclusion

The paper demonstrates that the choice between LayerNorm and RMSNorm has a precise, quantifiable consequence for Bayesian complexity. LayerNorm reduces the effective parameter count by $m/2$ per layer due to the affine flatness of its output manifold, whereas RMSNorm preserves the full parameter count due to the spherical curvature of its output. This geometric cost is structurally guaranteed and independent of training dynamics, offering a new lens for understanding model capacity and generalization.