🤖 machine learning

Stable FP4 Training via Transposition-Invariant Block Quantization

This paper introduces a stable FP4 training framework for large language models that resolves optimization instability caused by tensor transposition in conventional microscaling approaches by enforcing transposition-invariant 2D block quantization, achieving performance comparable to BF16 with minimal degradation across models up to 30B parameters.

Original authors: Mehdi Rahimifar, Amin Darabi, Mehran Taghian Jazi, Xing Huang, Yao Wang, Zhijun Tu, Yufei Cui, Yunke Peng, Hongliang Li

Published 2026-07-29

📖 7 min read🧠 Deep dive

CC BY 4.0

Original authors: Mehdi Rahimifar, Amin Darabi, Mehran Taghian Jazi, Xing Huang, Yao Wang, Zhijun Tu, Yufei Cui, Yunke Peng, Hongliang Li

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 the world of artificial intelligence as a massive, bustling library where giant robots (called Large Language Models) are trying to learn to read, write, and reason. To do this, they need to process mountains of information, which requires a tremendous amount of energy and memory. It's like trying to run a supercomputer on a single AA battery; it just doesn't work well. Scientists have been trying to make these robots more efficient by teaching them to think in "simpler" numbers. Instead of using complex, high-precision decimals (like measuring a distance to the millionth of a millimeter), they try to use rougher, 4-bit numbers (like measuring in whole inches). This saves a huge amount of space and energy. However, there's a catch: when the robots try to learn using these rough numbers, they often get confused, make mistakes, and stop learning entirely. It's like trying to navigate a maze with a map that keeps changing its own rules.

The paper you're about to read tackles this specific confusion. It asks: Why does teaching these AI robots with ultra-simple numbers cause them to crash? The authors discovered that the problem isn't just that the numbers are simple; it's that the way the robots organize these numbers changes when they switch from "learning forward" (reading a sentence) to "learning backward" (checking their mistakes). It's as if the robot reads a book left-to-right, but when it checks its notes, it suddenly reads them right-to-left, scrambling the page numbers and losing the plot. The paper proposes a clever fix: a new way of organizing the numbers so that the rules stay the same, no matter which direction the robot is looking. This allows the robots to learn stably, even with the simplest numbers, making them faster and cheaper to train without losing their smarts.

The Great Number Shuffle: Fixing AI's Memory Glitch

So, you want to build a super-smart AI, but your computer is running out of juice and memory. The usual trick is to tell the AI to use "low-precision" math. Think of it like this: normally, the AI measures things with a super-precise ruler that has tiny, tiny markings (like BF16 or FP8). But to save space, we want it to use a ruler with only four big, chunky marks (FP4). It's much faster and takes up less room. But here's the trouble: when the AI tries to learn with these chunky marks, it often falls apart. It starts making wild guesses, and the training crashes.

For a long time, scientists thought the problem was just that the chunky marks weren't detailed enough. They tried to fix it by grouping numbers together and giving each group a shared "scale" (a way to stretch or shrink the numbers to fit). This worked okay, but the new paper by Mehdi Rahimifar and his team says, "Wait a minute! We found the real culprit."

The culprit is a sneaky little glitch called transposition.

The "Left-to-Right vs. Right-to-Left" Problem

Imagine you have a grid of sticky notes on a wall, arranged in rows. You write a number on each note and put a label on the whole row to say how big the numbers are. This is how the AI usually organizes its data.

Now, imagine the AI is learning. First, it reads the notes from left to right (the "forward pass"). It sees the numbers and the labels, and it learns. But then, to check its work and fix mistakes, it has to flip the grid over (this is called "transposition" in math). Suddenly, the rows become columns.

Here's the disaster: In the old way of doing things (called 1D block quantization), when the grid flips, the sticky notes get shuffled into new groups. A note that was in "Group A" with a specific label is now in "Group B" with a totally different label. The AI thinks, "Wait, I thought this number was small, but now the label says it's huge!" This mismatch between what the AI saw when it learned and what it sees when it checks its work creates a confusing, biased signal. It's like trying to follow a recipe where the ingredients list changes every time you turn the page. The result? The AI gets dizzy, the training becomes unstable, and it fails to learn.

The 2D Square Solution

The authors of this paper had a brilliant idea: Don't let the groups change when you flip the grid.

They proposed using 2D square blocks. Imagine instead of long rows of sticky notes, you arrange them in perfect squares (like a 32x32 checkerboard). When you flip a square grid over, the squares stay squares. The sticky notes that were in the top-left square are still in a square that corresponds to the same group, just flipped. The label (the scale) stays exactly the same for those numbers, whether the AI is reading forward or backward.

By enforcing this transposition-invariant rule, the AI no longer gets confused. The numbers it sees when learning are the same numbers it sees when correcting. This simple change fixes the "dizziness" and allows the AI to train stably using the ultra-simple FP4 numbers.

The Safety Net: No Clipping and Random Guessing

But wait, there's more! Just fixing the groups wasn't enough. Because FP4 numbers are so simple, they can easily get too big (overflow) or get rounded in a way that introduces a bias (always rounding up, for example).

To fix this, the team added two safety features:

Truncation-Free Scaling: Instead of chopping off numbers that are too big (which distorts the data), they adjust the "ruler" itself so that every number fits comfortably inside the range. It's like stretching the ruler so the tallest person in the room can still fit without their head getting cut off.
Stochastic Rounding: When a number falls between two marks on the ruler, instead of always rounding to the nearest one (which creates a bias), the AI flips a coin. Sometimes it rounds up, sometimes down. Over time, the average stays perfectly accurate. It's like a fair game where the house doesn't cheat.

The "Special Sauce" for Attention

There is one tricky part of the AI brain called "Attention," where the model decides which words in a sentence are important. This part is super sensitive. If you use the simple FP4 numbers here, the AI gets confused.

So, the authors came up with a mixed-precision strategy. They use the super-efficient FP4 numbers for the heavy lifting (the main math), but they use a slightly more precise format (MXFP8) just for the "Attention" part. It's like using a rough, fast shovel for digging a big hole, but switching to a precise trowel for the delicate flower bed in the middle. This keeps the whole system fast and efficient while protecting the most sensitive parts.

Does It Actually Work?

The team tested this idea on some very large AI models, including ones with up to 30 billion parameters (that's a lot of brain cells!). They trained them on up to 100 billion tokens (words and symbols).

The results were impressive:

Stability: The models trained smoothly from start to finish. The old methods that didn't use this 2D square trick would crash early on.
Performance: The AI models trained with this new method performed almost exactly the same as models trained with the slow, heavy, high-precision numbers. The difference in quality was tiny—less than 1.3% in some tests.
Efficiency: They saved a massive amount of memory (about 65% less) and could theoretically run 3.5 times faster on future hardware designed for this.

The Bottom Line

This paper suggests that the biggest hurdle to making AI training super-efficient isn't just about using smaller numbers; it's about making sure those numbers behave consistently, no matter how the AI looks at them. By switching from messy, shifting rows to neat, stable squares, the authors found a way to make AI training both lightning-fast and rock-solid.

Of course, there are still some hurdles. The team had to simulate this on current computers because the special chips that can do this math natively aren't widely available yet. But the math checks out, and the path forward is clear: if we want AI to be faster and cheaper, we need to stop the number shuffle and keep the rules consistent.

Technical Summary: Stable FP4 Training via Transposition-Invariant Block Quantization

1. Problem Statement

Reducing training precision is a primary lever for improving the efficiency of Large Language Model (LLM) training. While FP8 training has become practical, extending this to 4-bit floating point (FP4) remains a significant challenge due to optimization instability. Existing microscaling approaches (e.g., MXFP4, NVFP4) utilize one-dimensional (1D) block quantization, where groups of values share a scaling factor.

The paper identifies a fundamental source of instability in these existing methods: scale inconsistency induced by tensor transposition. In conventional 1D block quantization, the forward pass and the backward pass assign different scaling factors to the same values after tensor transposition occurs during backpropagation. This mismatch causes the backward pass to compute gradients with respect to a different quantized objective than the one used in the forward pass. Consequently, this leads to biased gradients, distribution shifts, and unstable training dynamics, preventing stable end-to-end FP4 training.

2. Methodology

The authors propose a low-precision training framework centered on Transposition-Invariant Block Quantization, combined with specific rounding and scaling techniques to manage error.

2.1 Transposition-Invariant 2D Block Quantization

To resolve the scale inconsistency, the authors replace 1D block layouts with 2D square blocks (e.g., $32 \times 32$ ).

Mechanism: Tensors are partitioned into square blocks. When a tensor is transposed (as occurs in backpropagation for weight and gradient matrices), a block $B_{i,j}$ maps to $B_{j,i}$ . Because the set of values within the block remains the same (only their positions change), the block maximum and the resulting scaling factor are preserved ( $S(B_{i,j}) = S(B_{i,j}^\top)$ ).
Effect: This ensures that the scaling factor assigned to a value in the forward pass is identical to the one assigned in the backward pass, eliminating the systematic bias in gradient updates.
Application: 2D FP4 quantization is applied to linear layer weights and gradients. Activations retain a 1D layout to avoid mixing semantically unrelated values across channels, representing a trade-off between statistical structure and quantization consistency.

2.2 Truncation-Free Scaling and Stochastic Rounding

To further control quantization error and maintain unbiased gradients:

Truncation-Free Scaling: The scale factor $S$ is computed to ensure all values within a block fall strictly within the representable range of the FP4 format, preventing overflow and the systematic distortion of large-magnitude activations.
Stochastic Rounding: During backpropagation, values are rounded stochastically rather than deterministically. This preserves the expectation of the values ( $E[\text{round}(x)] = x$ ), ensuring that gradient estimates remain unbiased despite aggressive quantization.

2.3 Mixed-Precision Attention (MXFP8)

Attention mechanisms are identified as highly sensitive to quantization noise due to the dot-product interaction ( $QK^\top$ ) and softmax normalization.

Strategy: The authors adopt a mixed-precision design where Query ( $Q$ ) and Key ( $K$ ) projections are quantized using MXFP8, while all other transformer linear layers (MLP, Value/Output projections) use the proposed 2D block FP4.
Rationale: This preserves higher precision where numerical stability is most critical while retaining the memory and compute benefits of FP4 for the majority of the model's parameters.

3. Key Contributions

Identification of a Failure Mode: The paper identifies transposition-induced scale mismatch in 1D microscaling as a primary cause of instability in FP4 training, distinct from simple quantization error.
Transposition-Invariant Quantization: The proposal of 2D block FP4 quantization with truncation-free scaling and stochastic rounding to enforce consistent forward–backward behavior.
Practical Mixed-Precision Training: A demonstration that combining 2D block FP4 linear layers with MXFP8 attention enables stable end-to-end training for models up to 30B parameters, achieving performance close to BF16.

4. Experimental Results

The method was evaluated on dense models (OLMo-1B, OLMo-7B) and a sparse Mixture-of-Experts model (Qwen 30B MoE), trained on up to 100B tokens.

Training Stability: Naive FP4 quantization without microscaling diverged early. Microscaling-based methods remained stable, but the proposed 2D block FP4 method consistently converged closest to the BF16 baseline.
Performance Gap: Across all settings, the approach achieved stable training with less than 1.3% degradation in perplexity and downstream accuracy compared to BF16.
- For OLMo-1B, the average perplexity degradation was ~1.1%.
- For larger models (OLMo-7B, Qwen 30B), the gap decreased to below 1%, suggesting larger models are more tolerant to quantization noise once scaling consistency is enforced.
Ablation Studies: Removing truncation-free scaling led to overflow-induced instability; disabling stochastic rounding introduced bias and degraded convergence. Both components were necessary for optimal results.
Efficiency: Theoretical analysis on OLMo-1B suggests a ~65% reduction in model weight memory and up to ~3.6x idealized throughput improvement for linear operations, with only ~12.5% of parameters requiring FP8 precision.

5. Significance and Claims

The paper claims that the primary obstacle to practical FP4 training is not merely limited numerical precision or representational fidelity, but inconsistency in scaling across forward and backward computations.

By reframing the problem as one of computational consistency, the authors demonstrate that enforcing transposition-invariant scaling is sufficient to enable stable, end-to-end FP4 training at scale. The results suggest that once systematic bias is removed, FP4 training can closely match BF16 performance. This provides a simple and effective pathway toward more efficient LLM training, shifting the design focus from improving local representational fidelity to enforcing global computational consistency.

Limitations noted by the authors:

Experiments relied on software emulation as native FP4 tensor-core support was unavailable on their hardware.
The evaluation is limited to models up to 30B parameters and 100B tokens; scaling to frontier models may reveal new challenges.
The optimal block size and scaling factor overhead depend on specific architectures and workloads.