🤖 AI

Rethinking Shrinkage Bias in LLM FP4 Pretraining: Geometric Origin, Systemic Impact, and UFP4 Recipe

This paper identifies a fundamental "Shrinkage Bias" in non-uniform E2M1 FP4 formats that causes training instability, and proposes UFP4, a uniform 4-bit (E1M2/INT4) training recipe that leverages Random Hadamard Transform to achieve superior pretraining performance across various model scales compared to existing E2M1-based approaches.

Original authors: Qian Zhao, Kunlong Chen, Changxin Tian, Zhonghui Jiang, Haitao Zhang, Chaofan Yu, Peijie Jiang, Mingliang Gong, Jia Liu, Ziqi Liu, Zhiqiang Zhang, Jun Zhou

Published 2026-06-19

📖 4 min read☕ Coffee break read

CC BY 4.0

Original authors: Qian Zhao, Kunlong Chen, Changxin Tian, Zhonghui Jiang, Haitao Zhang, Chaofan Yu, Peijie Jiang, Mingliang Gong, Jia Liu, Ziqi Liu, Zhiqiang Zhang, Jun Zhou

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 giant, incredibly complex LEGO castle (a Large Language Model). To make the construction faster and cheaper, you decide to stop using the standard, heavy, precise bricks and switch to a new, ultra-lightweight set of tiny 4-bit bricks.

The problem is that the current "standard" lightweight bricks (called E2M1) have a weird, lopsided shape. The authors of this paper discovered that these bricks have a hidden flaw: they naturally shrink things slightly every time you use them.

Here is the breakdown of their discovery and solution, explained simply:

1. The Problem: The "Shrinking" Bricks

Think of the standard E2M1 bricks as having uneven spacing. Some gaps between the bricks are tiny, while others are huge.

The Flaw: Because of this uneven spacing, when you try to snap a piece of data into one of these bricks, the math forces it to round down slightly more often than up.
The Result: This creates a "Shrinkage Bias." It's like a leaky bucket; every time you pass a message through a layer of the model, the message gets just a tiny bit smaller and weaker.
The Compounding Effect: In a deep model with hundreds of layers, this tiny shrinkage happens over and over. By the time the message reaches the end, it has shrunk so much that the model starts to lose its mind (training becomes unstable).

2. The "Magic Spin" That Made It Worse

To fix the problem of data being too big or too weird (outliers), engineers use a trick called Random Hadamard Transform (RHT).

The Analogy: Imagine you have a pile of sand where most of it is in one giant mound (outliers). The RHT is like a mixer that spins the sand around so it spreads out evenly across the whole tray.
The Conflict: When you mix the sand (RHT) and then try to put it into the lopsided E2M1 bricks, the sand lands in the worst possible gaps. The uneven spacing of the bricks grabs the mixed-up sand and shrinks it even more than before. The paper found that the standard recipe actually makes the "leaky bucket" problem worse by using this mixer.

3. The Solution: The "Even" Bricks (UFP4)

The authors propose a new recipe called UFP4. Instead of using the lopsided E2M1 bricks, they use a new set of bricks called E1M2/INT4.

The Difference: These new bricks are perfectly uniform. The gaps between them are all exactly the same size.
Why it Works: Because the gaps are even, there is no "shrinkage bias." When the "mixed sand" (the RHT data) lands in these uniform bricks, it fits perfectly without losing any size.
The Strategy: With these new even bricks, the authors realized they could safely use the "mixer" (RHT) on every single part of the construction process (forward pass, backward pass, and weight updates). Previously, they had to be careful and only use the mixer in one spot to avoid breaking the model. Now, they can use it everywhere.

4. The Proof

The team tested this new recipe on three different sizes of AI models (small, medium, and huge).

The Result: The models built with the new UFP4 (even bricks) stayed much closer to the performance of the heavy, standard bricks (BF16) than the models built with the old E2M1 (lopsided bricks).
The Takeaway: The "even" grid allows the model to train longer and more stably without losing its signal.

Summary

The paper argues that the industry has been trying to fix a leaky bucket by patching the holes, but the bucket itself is the wrong shape. By switching to a bucket with perfectly even spacing (Uniform Grids), they can finally use the powerful mixing tools (RHT) that were previously too dangerous, leading to faster, cheaper, and more stable AI training.

The Bottom Line: Don't just patch the old, lopsided 4-bit format; switch to a new, uniform 4-bit format to stop the signal from shrinking away.

Technical Summary: Rethinking Shrinkage Bias in LLM FP4 Pretraining

Problem Statement

As Large Language Models (LLMs) scale, the computational and memory costs of pretraining necessitate a shift toward lower-precision formats. While FP8 has been widely adopted, emerging hardware (e.g., NVIDIA Blackwell/Rubin, AMD MI350) supports native FP4 computation. Current FP4 training recipes, such as NVFP4, predominantly rely on the E2M1 format (2 exponent bits, 1 mantissa bit) combined with the Random Hadamard Transform (RHT) to mitigate outlier sensitivity.

However, these E2M1-based recipes suffer from training instability and significant loss degradation relative to BF16 baselines. The paper identifies a fundamental, previously overlooked limitation: Shrinkage Bias. This is a systematic negative rounding error inherent to non-uniform quantization formats like E2M1. The geometric asymmetry of E2M1's representable bins causes a consistent "shrinkage" of magnitudes under Round-to-Nearest-Even (RTNE) rounding. Crucially, the paper argues that while RHT is designed to disperse outliers and improve codebook utilization, it inadvertently exacerbates this bias in E2M1 by shifting tensor mass into the format's most asymmetric rounding bins, leading to multiplicative signal decay across deep networks.

Methodology

1. Theoretical Analysis of Shrinkage Bias

The authors formalize Shrinkage Bias as a negative expected error in the normalized magnitude space caused by the geometric asymmetry of RTNE bins.

Geometric Origin: In non-uniform grids like E2M1, the spacing between representable values changes (e.g., from 0.5 to 1.0 to 1.5). For a level $q_i$ , the left and right bin widths ( $\ell_i$ and $r_i$ ) are often unequal. If $r_i > \ell_i$ , the conditional expected error under uniform local density is negative ( $\frac{\ell_i - r_i}{2} < 0$ ).
Systemic Impact: Unlike zero-mean noise which cancels out, this bias acts as a systematic attenuation factor ( $\alpha < 1$ ) in matrix multiplications (GEMMs). Across $K$ layers, these factors accumulate multiplicatively ( $\prod \alpha_k \approx \exp(-\sum \delta_k)$ ), causing systematic signal decay.
RHT Interaction: RHT transforms tensors from being "dynamic-range-limited" to "local-resolution-limited." While this improves bucket utilization, it forces data into the mid-magnitude regions where E2M1's asymmetry is most pronounced, thereby amplifying the shrinkage bias.

2. The UFP4 Recipe

Based on the finding that uniform grids bypass this geometric error, the authors propose UFP4 (Uniform FP4), a training recipe built on an E1M2/INT4-style uniform grid.

Uniform Grid: Unlike E2M1, the E1M2 grid maintains equal spacing between levels ( $\ell_i = r_i$ ), eliminating the geometric source of Shrinkage Bias.
Full RHT Coverage: Because the uniform grid is unbiased, UFP4 safely applies RHT to all three training GEMM paths: Forward (FPROP), Data Gradient (DGRAD), and Weight Gradient (WGRAD). Previous E2M1 recipes restricted RHT to the weight-gradient path to avoid compounding errors in non-leaf paths.
Stochastic Rounding (SR): SR is applied only to the upstream gradient ($dY$) to preserve gradient expectations, while RTNE is used elsewhere.

Key Contributions

Identification of Shrinkage Bias: The paper formally defines and demonstrates that non-uniform formats (E2M1) inherently suffer from a systematic negative rounding error due to geometric asymmetry, distinct from distribution-induced quantization errors.
Unified Explanation of Instability: It provides a theoretical and empirical explanation for the training instability in existing E2M1-based recipes, showing how Shrinkage Bias accumulates multiplicatively and is exacerbated by RHT.
Proposal of UFP4: A new 4-bit training recipe utilizing uniform grids (E1M2/INT4) that enables pervasive RHT across all training GEMMs, converting improved bucket utilization into higher quantization quality.
Industrial-Scale Validation: The approach is validated on Dense 1.5B, MoE 7.9B, and MoE 124B models, demonstrating that uniform grids are practical first-class FP4 training formats.

Experimental Results

Quantization and GEMM Diagnostics

RHT Effect: On outlier-heavy tensors, RHT significantly improves the effective bucket ratio for both formats. However, the resulting Signal-to-Quantization-Noise Ratio (SQNR) diverges: E1M2 shows improved SQNR post-RHT, while E2M1 often degrades (negative $\Delta$ SQNR) because the shifted data lands in asymmetric bins.
GEMM Outputs: In single-GEMM tests, E1M2 preserves or improves output SQNR across forward and backward paths with RHT, whereas E2M1 suffers degradation, particularly in backward paths.

End-to-End Pretraining

Loss Degradation: On long-run pretraining tasks, UFP4 (E1M2) consistently achieves lower BF16-relative loss degradation compared to strong E2M1 baselines.
- Dense 1.5B: Relative error reduced from 1.2570% (E2M1) to 0.9673% (UFP4).
- MoE 7.9B: Reduced from 2.3596% to 1.8469%.
- MoE 124B: Reduced from 1.7308% to 1.3863%.
Scaling Laws: The advantage of UFP4 persists across model scales (10M to 324M parameters). The fitted loss curves show the E1M2 gap to BF16 remains smaller than the E2M1 gap, and the FP4 penalty does not grow with compute scale.
Ablation Studies:
- RHT Scope: Full RHT coverage (all three GEMMs) yields the best results for UFP4, whereas restricting RHT (as done in E2M1 recipes) limits performance.
- Range Restriction: Attempts to emulate a uniform grid by restricting the dynamic range of E2M1 (e.g., capping max value) failed to match the performance of native E1M2, sacrificing dynamic range without fully eliminating the bias.

Efficiency

Kernel Fusion: The authors demonstrate that RHT can be fused with FP4 quantization efficiently. On tested hardware (SM90/SM100), the fused RHT+quantization latency is only ~1.06x–1.07x that of standalone quantization, compared to ~1.4x–1.6x for unfused implementations.

Significance and Claims

The paper argues that the current industry standard of treating E2M1 as the sole viable FP4 training format is suboptimal due to its geometric limitations. The significance of this work lies in:

Challenging the Status Quo: It suggests that future ML accelerators should support E1M2/INT4-style uniform 4-bit grids as first-class training primitives alongside E2M1, rather than relying solely on E2M1.
Format-Regime Mismatch: It highlights that the effectiveness of outlier mitigation techniques like RHT is dependent on the underlying grid geometry. RHT is beneficial for uniform grids but detrimental for non-uniform grids in deep training regimes.
Practical Recipe: UFP4 offers a concrete, scalable recipe that reduces the gap between FP4 and BF16 training without requiring complex estimator modifications, simply by aligning the grid geometry with the post-RHT tensor distribution.

The authors conclude that while E2M1 remains suitable for range-limited workloads and inference, training hardware should evolve to support uniform grids to fully realize the potential of FP4 pretraining.