💬 NLP

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

This paper introduces GeoNorm, a novel normalization method that unifies Pre-Norm and Post-Norm approaches by interpreting layer outputs as geodesic updates on a manifold with layer-wise decay, achieving consistent performance improvements in Transformer models with negligible computational overhead.

Original authors: Chuanyang Zheng, Jiankai Sun, Yihang Gao, Chi Wang, Yuehao Wang, Jing Xiong, Liliang Ren, Bo Peng, Qingmei Wang, Xiaoran Shang, Mac Schwager, Anderson Schneider, Yuriy Nevmyvaka, Xiaodong Liu

Published 2026-01-30

📖 4 min read☕ Coffee break read

CC BY 4.0

Original authors: Chuanyang Zheng, Jiankai Sun, Yihang Gao, Chi Wang, Yuehao Wang, Jing Xiong, Liliang Ren, Bo Peng, Qingmei Wang, Xiaoran Shang, Mac Schwager, Anderson Schneider, Yuriy Nevmyvaka, Xiaodong Liu

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you are trying to teach a robot to write a story. To do this, the robot uses a complex machine called a Transformer. Inside this machine, there are many layers of "workers" (like attention workers and feed-forward workers) that pass information along a chain.

Every time a worker passes a message to the next, the message needs to be "normalized." Think of normalization as a volume knob or a ruler. It ensures the message isn't too loud (too big) or too quiet (too small) before it moves to the next step.

For a long time, engineers had two main ways to turn this volume knob:

Post-Norm: The worker does their job, adds their message to the chain, and then someone checks the volume and adjusts it.
Pre-Norm: Someone checks the volume before the worker starts, so the worker always starts with a perfectly sized message.

Both methods work, but they have flaws. Post-Norm can sometimes cause the volume to spike wildly (like a feedback squeal), while Pre-Norm can make the early workers shout too loudly compared to the later ones.

The New Idea: GeoNorm

The authors of this paper, GeoNorm, say: "Why are we just checking the volume with a ruler? Let's think about the path the message takes."

They realized that the way these messages move is like walking on the surface of a globe (a sphere), not walking on a flat floor.

The Old Way (Projection): Imagine you are walking on a globe. You take a step in a straight line (like a laser beam). If you keep walking straight, you eventually fall off the globe into space. To fix this, the old methods say, "Oh no, you fell! Let's just pull you straight back down to the surface." This is like Post-Norm. It's a clumsy, abrupt correction that can distort your direction.
The New Way (Geodesic): The authors suggest that instead of walking in a straight line and falling off, you should walk along the curved surface of the globe itself. In math, the shortest path on a curved surface is called a geodesic (like how airplanes fly in arcs, not straight lines, to get from New York to London).

GeoNorm replaces the clumsy "pull back down" method with a smooth "walk along the curve" method. It uses a mathematical tool called an exponential map to ensure the message stays on the "globe" naturally, following the curve without ever falling off or needing a jarring correction.

The "Decay" Trick

The paper also introduces a smart way to handle the "step size" (how big a step the robot takes).

Imagine you are hiking up a mountain. At the bottom, you can take big, confident strides. But as you get higher and the terrain gets trickier, you should take smaller, more careful steps.
GeoNorm applies this logic. It starts with larger updates and gradually makes the steps smaller as the message moves through the deeper layers of the network. This prevents the robot from stumbling or getting confused near the end of the chain.

What Did They Find?

The authors tested this new method on different "hiking trails" (datasets like Arxiv and Books3) and with different sizes of robots (model sizes from small to huge).

Better Performance: The robot using GeoNorm consistently learned faster and ended up with a better "score" (lower error rate) than robots using the old methods.
Stability: The old methods sometimes had "loss spikes" (sudden crashes in performance), especially when the robot was very deep or the training was long. GeoNorm was much smoother and didn't crash.
No Extra Cost: The best part? This new way of walking on the globe doesn't require the robot to carry a heavier backpack. It doesn't add extra memory or slow down the computer; it just changes how the robot moves.

In a Nutshell

The paper argues that instead of forcing AI messages to fit into a box (the old normalization methods), we should let them flow naturally along the curved surface they are meant to travel on. By doing this, the AI becomes more stable, learns better, and doesn't need any extra computing power to achieve it.

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

Problem Statement

The placement of normalization layers within Transformer architectures—specifically the choice between Pre-Norm and Post-Norm—remains a critical design decision influencing training stability, model performance, and gradient flow. While Pre-Norm has become the predominant choice in modern Large Language Models (LLMs) due to its stability, it introduces imbalanced gradients where lower layers receive larger updates than higher ones. Conversely, Post-Norm, originally designed to facilitate deep network training, is prone to training instability and loss spikes. Existing solutions like DeepNorm and SandwichNorm improve stability through parameter initialization or additional normalization layers but lack a unified theoretical framework to analyze these mechanisms.

Methodology

The authors propose GeoNorm, a novel normalization method grounded in manifold optimization. The core theoretical insight is to reinterpret the standard Transformer update operations ( $x + \text{FFN}(x)$ and $x + \text{Attention}(x)$ ) as optimization steps on a sphere. In this view:

Token embeddings act as the initial state on a manifold.
Attention and FFN modules generate update directions (estimated gradients).
Normalization is viewed as a projection back onto a constraint set (the sphere).

Standard normalization (Post-Norm) performs an unconstrained update followed by a Euclidean projection back to the sphere. The authors argue that this projection discards geometric information and can distort update directions. Instead, GeoNorm replaces the projection step with a geodesic update using the exponential map.

For a point $x$ on a sphere and a tangent vector $v$ (the projected update direction), the update is performed intrinsically along the geodesic:
$\text{exp}_x(v) = \cos\left(\frac{\|v\|}{\|x\|}\right)x + \|x\|\sin\left(\frac{\|v\|}{\|x\|}\right)\frac{v}{\|v\|}$

This formulation ensures that the representation moves smoothly along the manifold, preserving the intrinsic geometry of the feature space. Furthermore, the authors introduce a layer-wise update decay mechanism for the step size ( $\alpha_k$ ), analogous to learning rate schedules. They evaluate three decay strategies:

Harmonic: $\alpha_k = \alpha / k$
Square Root: $\alpha_k = \alpha / \sqrt{k}$
Linear: $\alpha_k = \alpha (T - k) / T$

The paper also theoretically demonstrates that Pre-Norm can be interpreted as a special case of GeoNorm with non-adaptive, depth-dependent angles, thereby unifying the two approaches under a single geometric framework.

Key Contributions

Theoretical Framework: The authors provide a unified theoretical perspective interpreting Transformer layers as iterative optimization steps on a spherical manifold, bridging the gap between Pre-Norm and Post-Norm.
GeoNorm Algorithm: They introduce GeoNorm, which replaces standard projection-based normalization with geodesic updates via the exponential map. This method requires no additional trainable parameters and adds negligible computational overhead.
Layer-wise Decay: A novel proposal for layer-wise update decay for attention and FFN components to stabilize training dynamics.
Empirical Validation: Extensive experiments across varying model sizes (125M to 1.3B), datasets (Arxiv, Books3, FinWeb-Edu), and training lengths demonstrate the method's efficacy.

Experimental Results

The paper reports that GeoNorm consistently outperforms existing normalization methods (Pre-Norm, Post-Norm, DeepNorm, SandwichNorm) across all tested configurations:

Performance: On the Arxiv and Books3 datasets with a 125M model, GeoNorm achieved the lowest validation loss (1.8792 on Arxiv; 3.4040 on Books3) compared to all baselines.
Scalability: The performance advantage persists and widens as model size increases. For a 1.3B model, GeoNorm achieved a loss of 3.1759, outperforming the best baseline (Pre-Norm at 3.2067).
Stability: GeoNorm demonstrated superior training stability, particularly in long training runs (up to 4096 sequence length). While Post-Norm exhibited loss spikes and training collapse at longer lengths, GeoNorm maintained stable convergence.
Decay Mechanism: Among the decay strategies, the Harmonic decay ( $\alpha/k$ ) yielded the best performance, establishing new benchmarks on the tested datasets.

Significance and Claims

The paper claims that GeoNorm offers a seamless integration into standard Transformer architectures, requiring no modification to core blocks (Attention or FFN) and introducing no additional trainable parameters. The authors assert that by leveraging geodesic optimization, GeoNorm enables representations to move more coherently along the spherical manifold, leading to improved information propagation and training stability.

The work positions itself as a foundational step in understanding normalization mechanisms through the lens of Riemannian geometry. It suggests that the empirical success of Pre-Norm can be explained by its implicit step-size decay, a mechanism that GeoNorm makes explicit and adaptive. The authors conclude that this unified view provides a robust foundation for future innovations in Transformer architecture design, particularly for scaling models to greater depths and complexities.