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.
@article{arxiv.2601.22095,
title = {GeoNorm: Unify Pre-Norm and Post-Norm with Geodesic Optimization},
author = {Chuanyang Zheng and Jiankai Sun and Yihang Gao and Chi Wang and Yuehao Wang and Jing Xiong and Liliang Ren and Bo Peng and Qingmei Wang and Xiaoran Shang and Mac Schwager and Anderson Schneider and Yuriy Nevmyvaka and Xiaodong Liu},
journal= {arXiv preprint arXiv:2601.22095},
year = {2026}
}