💬 NLP

Gated Normalization Removal and Scale Anchoring in Pre-Norm Transformers

This paper introduces TaperNorm, a gated approach for pre-norm transformers that gradually removes internal normalization layers to improve inference throughput and reveals that final normalization primarily serves to anchor the scale of pre-logit representations, a role that can be replaced by fixed-target scaling to enable fully norm-free models.

Original authors: Andrei Kanavalau, Carmen Amo Alonso, Sanjay Lall

Published 2026-05-21

📖 4 min read☕ Coffee break read

CC BY 4.0

Original authors: Andrei Kanavalau, Carmen Amo Alonso, Sanjay Lall

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 a Transformer model (the brain behind modern AI chatbots) as a massive, multi-story factory. Every time a piece of data (a word) moves through the factory, it passes through several rooms. In every single room, there is a "Quality Control Inspector" (called a Normalization Layer) that checks the data, adjusts its size, and makes sure it doesn't get too wild or too small before moving to the next room.

For years, engineers thought these inspectors were absolutely necessary at every step, all the time. But this paper asks a simple question: Do we really need these inspectors to be active after the factory has finished its training?

Here is the breakdown of their findings using simple analogies:

1. The "Tapering" Trick (Turning Inspectors into Doorways)

The authors created a new method called TaperNorm. Think of it like a dimmer switch for the Quality Control Inspectors.

During Training: At the start, the inspectors work hard, checking every piece of data just like normal.
The Transition: As training finishes, the authors slowly "dim" the inspectors. They don't just fire them; they teach them to stop checking the data and just let it pass through with a simple, fixed rule (like "multiply by 1.5").
The Result: By the end, the inspectors are gone. They have been replaced by a simple, static door. Because the door is just a fixed rule, the factory can "fold" it into the next room's machinery. This means the computer doesn't have to do any extra math to check the data anymore; it just moves it along.

Why does this matter?
The paper tested this on a small model (TinyStories) and a famous model (GPT-2). They found that removing these internal inspectors didn't hurt the model's intelligence much (only a tiny drop in performance). However, because the inspectors were removed, the factory became faster.

The Speed Boost: When they tested how fast the model could write text, it was 1.18 times faster. That's like a car going from 60 mph to 70 mph just by removing a few unnecessary speed bumps.

2. The "Anchor" Problem (Why the Last Inspector Must Stay)

Here is the most interesting discovery. While they could remove the inspectors in the middle of the factory, they found that the very last inspector (right before the AI gives its final answer) plays a special, unique role.

The Analogy of the Anchor:
Imagine the AI's final thought is a balloon floating in the wind.

With the Last Inspector: The inspector acts like an anchor. It holds the balloon at a specific height. No matter how hard the wind blows (the math trying to make the answer more confident), the balloon stays at a steady level. This keeps the AI stable.
Without the Last Inspector: If you remove that last anchor, the balloon is free to drift. The AI's math naturally tries to make the balloon fly higher and higher (making the answers more "confident" or extreme) just to lower its error score. This is called "Logit Chasing." The AI gets unstable because it keeps inflating its own confidence without a limit.

The Solution:
If you must remove that last inspector (to make the model even faster), you have to replace the anchor with something else. The authors used a "Fixed-Target Scale Loss."

The Metaphor: Imagine a tether that gently pulls the balloon back to a specific height if it tries to float too high or sink too low. This tether acts as a "virtual anchor," stopping the AI from going crazy, allowing them to remove the last inspector safely.

3. The Bottom Line

The paper concludes with two main takeaways:

Internal Inspectors are Optional: You can train the AI with inspectors, then turn them into simple doors at the end. This makes the AI faster (up to 18% faster in their tests) with almost no loss in smarts.
The Last Inspector is Special: The final check is crucial because it stops the AI from getting "overconfident" and unstable. If you remove it, you must add a "tether" (a specific math rule) to hold the AI's confidence in check.

In short: You can strip out the middlemen to make the AI run faster, but you have to be careful about the very last step, or the AI might get carried away by its own confidence.

Technical Summary: Gated Normalization Removal and Scale Anchoring in Pre-Norm Transformers

Problem Statement

Normalization layers (e.g., LayerNorm, RMSNorm) are standard components in transformer architectures, stabilizing activations and enabling the training of deep networks. However, these layers introduce sample-dependent, nonlinear computations that incur significant overhead during inference. It remains unclear whether these per-token statistics are strictly necessary throughout the entire inference phase, particularly in pre-norm transformers. While previous works have attempted to remove normalization via initialization tricks or sequential fine-tuning, there is a lack of understanding regarding the distinct functional roles of internal versus final normalization layers, specifically concerning how they control the scale of the pre-logit representation.

Methodology

The authors propose TaperNorm, a gated normalization-removal approach designed for pre-norm transformers. This method allows the model to transition from standard normalization to learned, sample-independent linear or affine maps during training, enabling the removal of normalization operations at inference.

Core Mechanism: TaperNorm

TaperNorm dynamically interpolates between a normalized branch and a fixed scaling branch using a global gate $g \in [0, 1]$ . For a token vector $h$ :
$\text{TaperNorm}(h; g) = g \cdot \frac{h}{r(h)} D\gamma + (1 - g) \cdot c \cdot h D\tilde{\gamma}$
Where:

$r(h)$ is the RMSNorm scaling factor.
$D\gamma$ and $D\tilde{\gamma}$ are learnable diagonal scaling matrices.
$c$ is a per-layer scalar.
$g$ is a global gate that starts at 1 (standard normalization) and is cosine-decayed to 0 over training.

Training Procedure:

Warmup: The gate $g$ is held at 1. The model trains as a standard pre-norm transformer.
Calibration: At the end of warmup, the scalar $c$ is computed to minimize the least-squares difference between the normalized branch and the scaling branch, ensuring a smooth transition without distribution shift.
Tapering: The gate $g$ is decayed to 0. The scaling branch parameters ( $\tilde{\gamma}$ ) remain trainable.
Inference Folding: Once $g=0$ , the TaperNorm layer becomes a simple linear/affine map ( $c \cdot h D\tilde{\gamma}$ ). This operation can be mathematically folded into adjacent linear projections (e.g., QKV projections or MLP weights), removing the explicit normalization operation from the inference graph.

Scale Anchoring Analysis

The paper identifies a critical distinction between internal and final normalization layers:

Final Normalization as an Anchor: The final normalization layer acts as an implicit scale anchor. Mathematically, it is 0-homogeneous, meaning it removes the radial gradient component at the output. Without this anchor, minimizing cross-entropy can drive the model to increase logit magnitudes indefinitely ("logit chasing") to reduce loss, leading to instability.
Explicit Alternative: When the final normalization is removed (in "All-Taper" settings), the authors introduce a fixed-target scale loss ( $L_{aux}$ ). This auxiliary loss penalizes deviations of token scales from a fixed target $s_{tgt}$ , providing an explicit radial restoring force that counteracts the logit-chasing tendency.

Two experimental configurations are tested:

Internal-Taper: Internal block normalizers are tapered to linear maps, but the final normalization remains active (implicit anchor).
All-Taper: All normalizers, including the final one, are tapered. This requires the fixed-target scale loss to maintain stability (explicit anchor).

Key Contributions

Gated Normalization Removal: A method to taper RMSNorm/LayerNorm modules to fixed linear/affine rescalings, allowing per-token statistics to be eliminated at inference via weight folding.
Output Scale Anchoring Theory: An analysis demonstrating that the final normalization layer's primary distinct role is removing radial gradients. The paper shows that without this anchor, cross-entropy induces unbounded logit growth, and that a fixed-target scale loss can effectively substitute for this stabilizing role.
Empirical Validation: Experiments on TinyStories pre-training and GPT-2 fine-tuning show that internal normalization can be removed with minimal validation loss increases. The approach achieves up to 1.18× higher throughput in KV-cached autoregressive decoding after folding internal scalings.

Results

Pre-training (TinyStories):
- Internal-Taper (+aux): Achieved validation losses within 0.7–1.5% of the baseline across model sizes (1M to 30M parameters).
- All-Taper (+aux): Achieved losses within 0.8–1.8% of the baseline, demonstrating that a fully norm-free model is feasible with explicit scale anchoring.
- Dynamics: Models without a scale anchor exhibited rapid growth in logit magnitudes ("logit chasing"). The fixed-target scale loss successfully suppressed this behavior, aligning gradient magnitudes with the anchored baseline.
Inference Throughput:
- On an NVIDIA H100, the Internal-Taper (fused) configuration achieved 1.15–1.18× higher tokens-per-second throughput compared to the standard RMSNorm baseline. The Internal-Taper (unfused) configuration (explicit scaling operations) achieved 1.10–1.14× speedup.
Fine-tuning (GPT-2):
- The proposed Internal-TaperLN and All-TaperLN methods achieved cross-entropy losses competitive with the "FakeLN" approach (Baroni et al., 2025) and standard Vanilla models across OpenWebText and The Pile datasets.
- Unlike sequential removal methods, TaperNorm uses a single global gate to remove all selected layers simultaneously.

Significance and Claims

The paper claims that internal and final normalization layers play distinct roles in pre-norm transformers. While internal normalizers can be replaced by learned sample-independent linear maps with negligible performance degradation, the final normalization layer serves a unique function as an output scale anchor.

The significance of this work lies in:

Efficiency: It provides a practical pathway to remove computationally expensive normalization operations from the inference graph, yielding measurable throughput gains (up to 1.18×) without retraining from scratch.
Theoretical Insight: It clarifies that the stability provided by the final normalization layer is largely due to its removal of radial gradients. This insight allows for the design of "fully norm-free" models by replacing the implicit anchor with an explicit fixed-target scale loss.
Robustness: The approach maintains standard pre-norm dynamics early in training, only transitioning to linear maps later, which appears more stable than methods that attempt to train without normalization from the start.

The authors remain modest regarding the scope, noting that experiments are limited to models up to 30M parameters and specific hardware stacks. They acknowledge that robustness to activation outliers in billion-parameter regimes and the applicability to other objectives beyond cross-entropy remain open questions for future work.