GeoNorm: Unify Pre-Norm and Post-Norm with Geodesic Optimization

The placement of normalization layers, specifically Pre-Norm and Post-Norm, remains an open question in Transformer architecture design. In this work, we rethink these approaches through the lens of manifold optimization, interpreting the outputs of the Feed-Forward Network (FFN) and attention layers as update directions in optimization. Building on this perspective, we introduce GeoNorm, a novel method that replaces standard normalization with geodesic updates on the manifold. Furthermore, analogous to learning rate schedules, we propose a layer-wise update decay for the FFN and attention components. Comprehensive experiments demonstrate that GeoNorm consistently outperforms existing normalization methods in Transformer models. Crucially, GeoNorm can be seamlessly integrated into standard Transformer architectures, achieving performance improvements with negligible additional computational cost.

💡 Research Summary

The paper tackles a long‑standing design question in Transformer architectures: where to place the normalization layer. While Pre‑Norm (normalization before the residual addition) is widely adopted for its training stability, it suffers from gradient imbalance across layers. Post‑Norm (normalization after the addition) was the original design but often leads to loss spikes and instability in deep models. Existing remedies such as DeepNorm or SandwichNorm improve empirical performance but lack a unified theoretical foundation.

The authors reinterpret the core Transformer update— x + FFN(x) and x + Attention(x) —as steps of an optimization algorithm on a spherical manifold (the set of vectors with constant ℓ₂‑norm). In this view, the FFN and Attention modules generate update directions (sₖ), while the normalization operator projects the updated vector back onto the sphere. This projection is mathematically equivalent to a projected gradient step onto a feasible set Ω (the sphere). However, projection discards curvature information and can distort the intended direction, especially on smooth manifolds.

To address this, the paper introduces GeoNorm, a normalization scheme that replaces the extrinsic projection with an intrinsic geodesic move using the exponential map on the sphere. Concretely, for a current point xₖ and an update direction sₖ, the method first projects sₖ onto the tangent space TₓₖΩ:
vₖ = sₖ – (xₖᵀsₖ / ‖xₖ‖²) xₖ.
Then it updates the representation via the closed‑form exponential map:
xₖ₊₁ = expₓₖ(αₖ vₖ) = cos(‖vₖ‖/‖xₖ‖) xₖ + sin(‖vₖ‖/‖xₖ‖) (‖xₖ‖/‖vₖ‖) vₖ,
where αₖ is a step‑size that can be constant or follow a decay schedule (e.g., polynomial decay αₖ = α₀ / √k). Because the sphere’s exponential map has a simple analytic form, the computational overhead is negligible—only a few extra vector operations and no additional trainable parameters.

The authors further show that Pre‑Norm can be seen as a special case of GeoNorm where the exponential map is applied with a fixed, non‑adaptive angle. Thus GeoNorm unifies both normalization styles under a single Riemannian framework while allowing adaptive step sizes that improve convergence.

Experimental evaluation covers three corpora (ArXiv, Books3, FinWeb‑Edu) and three model sizes (125 M, 350 M, 1.3 B parameters). Across all settings, GeoNorm consistently yields lower validation loss than Pre‑Norm, Post‑Norm, DeepNorm, and SandwichNorm. For example, on the ArXiv dataset with a sequence length of 512, GeoNorm achieves a loss of 1.8792 versus 1.9032 for the best baseline (Pre‑Norm). Similar margins appear on Books3 and on larger models, and downstream fine‑tuning tasks (summarization, QA) show 1–2 % absolute improvements. Importantly, training stability is enhanced: deeper networks (up to 48 layers) exhibit fewer gradient explosions or vanishing gradients, and the loss curve is smoother.

A cost analysis confirms that GeoNorm adds less than 3 % extra FLOPs compared to standard LayerNorm/RMSNorm and does not increase memory consumption. The method can be dropped into existing codebases by simply swapping the normalization call with the exponential‑map update.

The paper concludes that viewing Transformer layers as Riemannian flows on a sphere provides a principled way to design normalization. GeoNorm leverages this insight to deliver a theoretically grounded, computationally cheap, and empirically superior alternative to existing practices. Limitations include the current focus on spherical manifolds; extending the approach to other manifolds (e.g., Stiefel for orthogonal constraints) or to mixed‑precision training remains future work. Nonetheless, GeoNorm represents a significant step toward a unified theory of normalization in deep sequence models.