Stable FP4 Training via Transposition-Invariant Block Quantization

TL;DR: FP4 training keeps blowing up, and this paper offers a satisfyingly concrete diagnosis: with standard 1D microscaling (MXFP4/NVFP4-style row blocks), the same tensor gets different quantization scales in the forward and backward passes, because backprop uses the transpose and transposing shuffles values into different scaling blocks. Forward and backward are effectively computed on two different quantized networks, which biases gradients. The fix is almost embarrassingly simple — use square 2D blocks, which map to themselves under transposition — plus overflow-proof scale selection and stochastic rounding. With that, they train dense models to 7B and a 30B MoE for 100B tokens in (emulated) FP4 within ~1% of BF16 loss, with the gap shrinking as models get bigger.

Why FP4 training is the current frontier of precision reduction

Every halving of numerical precision roughly doubles matmul throughput and halves memory bandwidth for the affected tensors. BF16 is the default; FP8 training is now practical on H100-class hardware; FP4 is the next rung, and Blackwell-generation accelerators ship native FP4 tensor cores. But FP4 (the E2M1 format: 1 sign bit, 2 exponent bits, 1 mantissa bit) gives you only 16 representable magnitudes per sign. You cannot cover a tensor's dynamic range with a single scale factor.

The community's answer is microscaling: chop each tensor into small blocks — typically 1\times 32 (MXFP4) or 1\times 16 (NVFP4) contiguous values along a row — and give each block its own shared scale S, chosen from the block's max absolute value. Each value is stored as P_i = \mathrm{round}_{\mathrm{FP4}}(X_i / S) and reconstructed as P_i \cdot S. This keeps local dynamic range under control while spending only a few extra bits per block on the scale.

That handles quantization error. But quantization error alone doesn't explain why FP4 training runs diverge or drift where FP8 runs don't. This paper's claim is that the dominant failure mode is somewhere else entirely.

The aha: transposition scrambles the blocks

A linear layer during training touches its tensors in multiple orientations:

where X is the activation, W the weight, and \nabla denotes gradients of the loss. The forward pass consumes W one way; the backward pass consumes it in the transposed orientation (and similarly \nabla Y appears both plain and transposed). For hardware efficiency you want quantization blocks aligned with the GEMM's reduction dimension — which means when the tensor is transposed, you re-block it along the other axis.

Here's the problem. With 1D blocks, a given value x_{ij} shares a scale with its row neighbors in the forward pass and with its column neighbors in the backward pass. Those two groups have different maxima, hence different scales, hence the value rounds to a different quantized number in each pass:

Why does this matter for optimization rather than just being extra noise? Training with quantization relies on the straight-through estimator (STE): the quantizer's derivative is approximated as identity, i.e. \partial\mathcal{L}/\partial X \approx \partial\mathcal{L}/\partial \hat{X}. That approximation is tolerable when forward and backward see the same quantized tensor \hat{X}. When the scales differ, the backward pass computes gradients of a different quantized objective than the one whose loss you evaluated. The mismatch isn't zero-mean noise that averages out — it's a systematic, per-value bias that accumulates over updates. The paper's Figure 1 visualizes this as scale-difference maps: large mismatch everywhere for 1D blocks, near-zero for 2D.

The reframing is the intellectually interesting part: FP4 instability is cast not as a representational fidelity problem (values too coarse) but as a consistency problem (forward and backward disagree about which network is being trained). Fidelity problems you fight with smaller blocks and fancier formats; consistency problems you fix structurally.

The fix, and its supporting cast

2D blocks. Partition tensors into b\times b square blocks (they use 32\times 32 to align with GEMM tiling). Block B_{i,j} of X becomes block B_{j,i} of X^{\top} — same values, same max, same scale. Transposition invariance is exact by construction. This is applied to weights and to both gradient tensors (\nabla W and \nabla X paths).

Activations stay 1D. They keep a 1\times-block layout for activations, arguing that square blocks would group "semantically unrelated" hidden channels. This is a real asterisk: X also appears transposed in the weight-gradient GEMM, so the consistency argument doesn't fully apply to activations — the paper accepts a residual inconsistency there as a statistical-structure trade-off. The main benefit evidently comes from making weights and gradients consistent.

Truncation-free scaling. Instead of scaling by exactly the block max (which puts the max right at the format boundary and risks clipping after rounding), the scale is rounded up to the next power of two:

where M is the block's max absolute value and Q_p, Q_n are FP4's representable bounds. Every value is guaranteed in-range; nothing ever clips. In FP4's tiny dynamic range, even rare clipping of large-magnitude gradients destabilizes training — their ablation shows removal causes overflow-driven instability.

Stochastic rounding in the backward pass. Deterministic round-to-nearest introduces correlated bias at 4 bits; stochastic rounding chooses between the two adjacent representable values with probabilities such that \mathbb{E}[\mathrm{round}_{\mathrm{stoch}}(x)] = x, keeping gradient estimates unbiased in expectation. Their ablation (Figure 5) shows both components are individually necessary: drop either and convergence degrades or destabilizes.

MXFP8 for queries and keys. Attention is the one place they retreat from FP4: QK^{\top} multiplies quantization noise and softmax amplifies it, so Q and K projections are quantized in MXFP8 instead. Everything else — value/output projections and all MLP layers, which hold most of the parameters — stays in 2D-FP4. Since only ~12.5% of parameters get FP8, the theoretical accounting on OLMo-1B still shows a 65.6% reduction in total weight memory and ~3.56× idealized linear-layer throughput vs. BF16.

The evidence

Three models trained from scratch on 100B tokens each (2K sequence length, AdamW, identical hyperparameters across precision variants): OLMo-1B, OLMo-7B, and a Qwen 30B MoE. Naive FP4 without microscaling diverges around 20B tokens. All microscaled variants are stable, but 2D-FP4 tracks BF16 tightest: ~1.1% final loss gap at 1B versus ~2% for competing 1D microscaling schemes, shrinking to ~1% at 7B and 30B.

The perplexity gaps are small and — notably — decrease with scale, consistent with the general lore that bigger models tolerate quantization noise better once systematic bias is removed. Downstream reasoning accuracy tells the same story:

The pure 2D-FP4 configuration matches or slightly beats BF16 on downstream averages at 7B and 30B (the authors gesture at a regularization effect from quantization noise; treat that as speculation given the task count and single seeds).

On block size (Table 4): 8\times 8 blocks give the smallest loss gap (~0.8–1.0%) but ~0.85× throughput from scale-metadata overhead; 64\times 64 trades the other way (~1.3–1.6% gap, ~1.05× throughput). 32\times 32 is the chosen sweet spot at 1.1% gap and full tiling-aligned throughput.

What to be skeptical about

Everything is software-emulated. No native FP4 tensor cores on their platform, so all throughput/memory numbers are theoretical ("idealized kernels," "parameter storage only"). The 3.56× figure is a format-arithmetic ceiling, not a measurement. The stability and quality claims are what the experiments actually support.

The MXFP8 attention story is internally inconsistent. The paper motivates FP8 queries/keys by claiming FP4 attention "often leads to instability" — yet their own pure-FP4 configuration trains stably in every experiment and outperforms the MXFP8 hybrid on every metric in Tables 2 and 3. If FP4 Q/K were genuinely destabilizing, we'd expect the opposite ordering somewhere. Either the instability appears in settings not shown (longer training? larger scale? different init?), or the MXFP8 component is defensive engineering without demonstrated necessity. The appendix concedes it's a "stability-oriented option." Note also the abstract's "<1.3% degradation" claim is generous: the +MXFP8 variant at 1B shows a 1.6% perplexity gap in their own Table 2.

Scale and duration are modest. 100B tokens at 2K sequence length is far short of frontier pretraining runs; slow-accumulating gradient bias is exactly the kind of pathology that surfaces late. The 30B MoE result is encouraging but is one run. Baseline comparisons against NVFP4-style training are summarized in a loss curve ("~2% for competing methods") rather than in full tables, and recent competing FP4 recipes (e.g., with Hadamard rotations or 2D-scaling variants — the paper's own citations suggest the 2D idea has neighbors) aren't head-to-head in the quality tables.

Residual inconsistencies remain. Activations are still 1D-blocked, optimizer states are BF16, and embeddings are BF16, so this is "FP4 for linear-layer GEMMs" rather than an end-to-end 4-bit pipeline in the strictest sense — which is also true of every other FP4 training paper, to be fair.

Why it matters anyway

The result that endures here is likely the diagnosis, not the specific recipe. "Forward–backward scale consistency, not representational fidelity, is the binding constraint for FP4 training" is a crisp, falsifiable claim with an obvious mechanism, near-zero-cost fix, and an ablation trail supporting it. If it holds up at frontier scale, it changes how low-precision formats should be designed: hardware and format specs (which currently standardize 1D microscaling blocks) should care about transposition invariance, and 2D block scaling becomes a natural request for next-generation tensor-core designs. FP4 GEMMs at ~1% quality cost would roughly double effective training compute relative to FP8 for the dominant matmuls — a big lever if realizable on real silicon.

Where to spend your time in the paper: Section 2's "Scale inconsistency under transposition" plus Figure 1 is the whole idea in two pages and is genuinely clarifying. Then skim Table 2/3 and the Figure 5 ablation. Appendix D has the clean formal statement of the invariance property if you want to convince yourself it's exact rather than approximate.