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

TL;DR

GeoNorm introduces a novel normalization approach for Transformers using geodesic optimization, unifying Pre-Norm and Post-Norm, leading to improved performance with minimal extra computation.

Contribution

The paper proposes GeoNorm, a new normalization method based on manifold optimization, unifying Pre-Norm and Post-Norm in Transformer architectures.

Findings

GeoNorm outperforms existing normalization methods in Transformers.

GeoNorm can be integrated seamlessly with standard Transformer models.

Performance improvements are achieved with negligible additional computational cost.

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.