🤖 machine learning

Sphere Retraction Normalizations

This paper generalizes Geodesic Normalization by demonstrating that all retraction maps for spherical residual streams reduce to a single scalar parameter controlling the rotation angle, leading to the proposed $p$ -SpheretNorm family which outperforms existing methods on nanoGPT by identifying that the exponential map is merely a limiting case rather than the optimal choice.

Original authors: Jie Zhang, Cheng-Fang Su, Yi-Jui Huang, Min-Te Sun

Published 2026-08-05

📖 6 min read🧠 Deep dive

CC BY 4.0

Original authors: Jie Zhang, Cheng-Fang Su, Yi-Jui Huang, Min-Te Sun

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 build a skyscraper out of LEGO bricks, but every time you add a new floor, the whole tower starts to wobble, twist, or collapse. This is the daily struggle for scientists building "deep neural networks"—the super-smart computer brains that power everything from chatbots to image generators. These networks are made of layers stacked on top of each other, like floors in a building. To make them work, engineers use a special trick called a "residual connection," which is like a slide that lets information skip from one floor to the next, preventing the tower from getting too shaky.

For a long time, the best way to keep these towers stable was to use a method called "LayerNorm," which acts like a strict architect, constantly checking the size of the bricks to make sure they don't get too big or too small. But even this strict architect has flaws: sometimes the tower gets too tall and the top floors stop learning, or the bottom floors get too excited and the whole thing shakes apart. Recently, a new idea called "Geodesic Normalization" (or GeoNorm) tried to fix this by imagining the tower isn't built on flat ground, but on the surface of a giant, invisible sphere. On this sphere, the information travels along the shortest possible path, like a plane flying over the Earth, which keeps the size of the data perfectly constant. It was a brilliant idea, but the scientists who invented it used a very specific, complex mathematical tool to draw that path, assuming it was the only way to do it.

This paper, titled "Sphere Retraction Normalizations," asks a simple but revolutionary question: "Is that specific tool the only way to walk on the sphere, or are there other, better paths?" The authors, Jie Zhang, Cheng-Fang Su, and their team, decided to stop treating the sphere as a rigid, single-path track and instead looked at the entire family of ways you can move across a curved surface. They discovered that the complex tool everyone was using (the exponential map) is actually just one extreme end of a spectrum. By exploring the middle ground, they found two new, simpler ways to move across the sphere that are just as good at keeping the tower stable, but much easier to calculate.

The team introduced three new methods, which they collectively call "SpheretNorms." Think of the sphere as a giant globe, and the information as a traveler. The old method (GeoNorm) was like a traveler who had to follow a very specific, winding road that required complex navigation tools. The authors showed that you can actually take a shortcut. They found two new "retractions" (mathematical shortcuts for moving on a sphere): one called Proj-SpheretNorm, which is like projecting a shadow straight down to the ground and then stepping back up to the surface, and another called Cay-SpheretNorm, which uses a clever geometric trick to rotate the traveler to the next spot.

To prove these ideas work, the researchers didn't just do math on paper; they built digital towers using a popular, lightweight AI model called "nanoGPT." They tested their new methods on two massive datasets of text: OpenWebText and FineWeb-Edu. They built models with different heights, ranging from 12 layers up to 36 layers, and trained them on about 6.55 billion tokens (a token is a chunk of text, like a word or part of a word).

The results were surprising. While the old "GeoNorm" method was good, it wasn't the best. In fact, the authors found that the "perfect" path isn't at the extreme end of the spectrum where the old method lives, but somewhere in the middle. They created a flexible family of methods called p-SpheretNorm, where a single number, p, controls how the traveler moves.

When p = 1, you get the "shadow projection" method (Proj-SpheretNorm).
When p = 2, you get the "geometric trick" method (Cay-SpheretNorm).
When p gets very large, you get the old GeoNorm method.
When p gets very small, you get the standard, flat-ground method.

In their experiments, the Proj-SpheretNorm (where p = 1) turned out to be the champion. On the medium and large models, it achieved the lowest training and validation losses, beating not only the old GeoNorm but also other popular methods like Pre-LN, Peri-LN, and Keel. For example, on the FineWeb-Edu dataset with a 24-layer model, the new method achieved a validation perplexity (a score where lower is better) of 16.81, compared to 19.09 for the previous best method, Peri-LN. On the OpenWebText dataset, it similarly outperformed the competition.

Perhaps most importantly, the new methods were incredibly stable. In some tests, the older methods like Pre-LN and Keel completely failed, with their validation scores exploding to numbers like 111.70 or even 2081.50, meaning the AI stopped learning and started guessing randomly. In contrast, all three of the new SpheretNorm methods converged smoothly in every single setting, even in the deepest 36-layer models.

The authors also tested these models on "downstream tasks," which are like giving the AI a series of quizzes to see how smart it really is. They used 11 different benchmarks covering things like reading comprehension, logic puzzles, and science questions. The 1-SpheretNorm model (the p=1 version) scored the highest on 11 out of 12 metrics on the FineWeb-Edu training data and 10 out of 12 on the OpenWebText data. It particularly shined in "perplexity" tests, where it reduced the confusion score on the LAMBADA dataset from 59.55 (the best baseline) down to 46.74, a massive improvement.

The paper concludes that the exponential map used in GeoNorm is not the "holy grail" of spherical movement; it is merely one end of a spectrum. By choosing a different point on that spectrum (specifically p=1), we can build deeper, more stable, and more accurate AI models without needing complex calculations. The authors suggest that this opens up a new way of thinking about how we connect layers in neural networks, moving away from rigid, single-path solutions to a flexible family of geometric choices. While they note that these methods might not work for every single type of AI architecture (like those that need to preserve specific symmetries), they have successfully shown that on the hypersphere, there is a whole world of better paths waiting to be explored.

Technical Summary: Sphere Retraction Normalizations

Problem Statement
Training deep neural networks stably remains a significant challenge. Existing methodologies, primarily relying on residual connections combined with normalization (e.g., Pre-LN and Post-LN in Transformers), exhibit inherent limitations. Pre-LN architectures suffer from the "curse of depth," where deeper layers contribute diminishingly to learning, and allow unbounded residual streams that accumulate massive activations, leading to attention sinks. Conversely, Post-LN architectures face large gradients near the output layer at initialization, necessitating careful learning-rate warm-up schedules that slow optimization. While recent work like Geodesic Normalization (GeoNorm) attempts to stabilize training by confining hidden states to a hypersphere via Riemannian geometry, it relies exclusively on the exponential map for updates. The authors posit that the exponential map is merely one member of a broader family of retraction maps, raising the question of whether alternative retractions on the hypersphere could offer superior performance or stability.

Methodology
The paper proposes a unified framework for deep connections based on retraction maps on a hypersphere, termed Sphere Retraction Normalizations (SpheretNorms).

Unified Framework: The authors recast standard Residual Connections (RC) and GeoNorm within a single computational skeleton defined on a manifold $\mathcal{M}$ . The update rule is abstracted as:
$h_l = R_{h_{l-1}}\left[ \alpha_l \cdot \text{proj}_{T_{h_{l-1}}\mathcal{M}}(F_l(h_{l-1})) \right]$
where $R$ is a retraction map. While GeoNorm selects the hypersphere ( $S^{d-1}$ ) and the exponential map, this framework allows for other retractions.
Specific Instantiations:
- Proj-SpheretNorm: Instantiates the framework using the metric projection retraction (also known as the gnomonic retraction). This involves normalizing the vector sum of the current state and the update to lie on the hypersphere.
- Cay-SpheretNorm: Uses the Cayley retraction (stereographic retraction), derived via the Cayley transform of a skew-symmetric matrix generated by the current state and the update vector. This yields a closed-form algebraic update.
Generalization (p-SpheretNorm): The authors observe that both Proj-SpheretNorm and Cay-SpheretNorm share a planar-rotation structure where the update magnitude is converted into a rotation angle. They introduce a one-parameter family, p-SpheretNorm, defined by a rotation angle $\theta_l = p \arctan(\beta_l / p)$ , where $\beta_l$ is the normalized step size.
- $p=1$ recovers Proj-SpheretNorm.
- $p=2$ recovers Cay-SpheretNorm.
- $p \to \infty$ converges to GeoNorm (exponential map).
- $p \to 0^+$ converges to the identity mapping.
- Crucially, for finite $p$ , the rotation angle saturates, preventing unbounded growth, unlike the exponential map.
Implementation Details: To ensure stability, the authors employ a decay factor $\lambda_l = 1/\sqrt{l}$ for the learnable step size $\alpha_l$ and apply a modified ScaleNorm in the first layer to initialize the hidden state radius.

Key Contributions

Theoretical Unification: The paper unifies Euclidean residual connections and GeoNorm as retraction-based updates, demonstrating that GeoNorm is a specific case of a Riemannian residual connection driven by the exponential map.
Novel Normalization Layers: It introduces two new, norm-preserving normalization layers (Proj-SpheretNorm and Cay-SpheretNorm) that rely solely on algebraic operations, avoiding the computational cost of trigonometric functions required by the exponential map.
Parametric Family: It generalizes these methods into p-SpheretNorm, a single-parameter family that bridges the gap between the identity map, algebraic retractions, and the exponential map. The authors prove that all members of this family are second-order retractions with identical first-order tangent dynamics.
Empirical Superiority: Extensive experiments on nanoGPT (up to 36 layers, 6.55B tokens) show that 1-SpheretNorm (corresponding to $p=1$ ) achieves the lowest pre-training and validation losses and the highest zero-shot accuracy across OpenWebText and FineWeb-Edu datasets, outperforming Pre-LN, Pre-DyT, Peri-LN, Keel, and GeoNorm. Notably, 1-SpheretNorm converges in settings where Pre-LN and Keel diverge.

Results

Pre-training: On medium (24 layers) and large (36 layers) models, 1-SpheretNorm achieved the best validation perplexity. For instance, on FineWeb-Edu at 24 layers, it reduced validation perplexity from 19.09 (Peri-LN) to 16.81.
Stability: Unlike Pre-LN (which diverged on OpenWebText at medium scale) and Keel (which diverged on FineWeb-Edu at small scale), all SpheretNorm variants converged in all tested settings.
Downstream Tasks: On the large-size models trained on FineWeb-Edu, SpheretNorm variants achieved the best score on 11 of 12 metrics; on OpenWebText, they achieved the best score on 10 of 12 metrics. The gains were most pronounced in perplexity metrics (e.g., LAMBADA and WikiText-2), where 1-SpheretNorm significantly outperformed GeoNorm and other baselines.
Parameter Sensitivity: The results indicate that the optimal performance is attained at finite $p$ (specifically $p=1$ ), suggesting that the exponential map (the limit as $p \to \infty$ ) is not the preferred retraction for spherical residual streams but merely one end of a spectrum.

Significance and Claims
The paper claims that the choice of retraction map is a critical design choice for deep connections on hyperspheres. By moving beyond the exponential map to algebraic retractions (specifically metric projection), the authors demonstrate that it is possible to achieve norm preservation and training stability with superior performance. The work suggests that the "exponential map" is not a fundamental necessity for stable deep learning on manifolds but rather a specific point in a broader spectrum of valid geometric updates. The authors conclude that designing deep connection mechanisms is a mathematical problem at the intersection of graph theory, optimization, and dynamical systems, and that identifying the optimal trade-off between information propagation and Jacobian conditioning remains an open direction for future work. They acknowledge limitations regarding specialized architectures (e.g., Equivariant Neural Networks) and the need for evaluation on larger-scale models beyond the current scope.

Technical Summary: Sphere Retraction Normalizations

More like this