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

TL;DR

GeoNorm reframes Transformer normalization as a geodesic optimization problem on the spherical manifold, replacing projection-based normalization with intrinsic updates via the exponential map $\exp_{\mathbf{x}}(\cdot)$. By viewing FFN and Attention outputs as gradient directions and applying layer-wise update decays analogous to learning-rate schedules, GeoNorm unifies Pre-Norm and Post-Norm under a geometric framework. Empirical evaluation across multiple datasets, model sizes, and training lengths shows GeoNorm consistently outperforms existing normalization methods with negligible extra cost, and harmonic decay further enhances performance and stability. This work offers a principled, scalable approach to normalization in large Transformers, with implications for training stability and efficiency at extreme scales.

Abstract

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.

Figures (5)

• Figure 1: The performance of different methods on the Arxiv and Books3 dataset, with model parameter 125M.
• Figure 2: The performance of different baselines on the Books3 dataset, with model parameter 350M and 1.3B, training lengths of 512.
• Figure 3: The performance of different baselines on the Books3 dataset, with model size 125M and training length 2048 and 4096.
• Figure 4: The performance of different decay methods on the Arxiv and Books3 dataset, with model parameter 125M, training lengths of 512. Sqrt: $\frac{\alpha}{k^{0.5}}$, Linear: $\frac{\alpha(T-k)}{T}$, and Harnomic: $\frac{\alpha}{k}$, where $k$ is the current layer index and $T$ is the total layer number.
• Figure 5: The training dynamic with loss metrics, with model parameter of 125M, training lengths of 512 and the Books3 dataset.