Stochastic Rounding in Low-Precision Transformer Inference: A Variable-Precision Emulation Study of a Small GPT-2
The paper asks whether stochastic rounding (SR) can improve low-precision transformer inference over round-to-nearest (RN), and shows that the benefit is site-dependent: SR outperforms RN in long MLP reductions because its $O(\sqrt{n}u)$ probabilistic error envelope beats the $O(nu)$ deterministic drift of RN, while RN is superior at the language-model head because SR's non-uniform variance is heavily penalized by the Fisher-weighted cross-entropy. The authors support this with a fast variable-precision emulation tool (PRISM), a vector Bienaymé–Chebyshev bound, a second-order loss decomposition, and experiments on DistilGPT-2 that motivate mixed-precision rounding recipes.
The paper rigorously isolates the rounding rule from quantization mapping and delivers both analytical and empirical evidence that site-specific rounding matters, but its conclusions are severely limited by a single tiny model, a minuscule evaluation set, and software emulation that ignores exponent-range constraints of real narrow formats. As the authors themselves caution, "The complete sweeps use one six-block model and one leading text prefix. The measured site ordering and composition effects need replication across texts, larger models, and broader tasks." (Section 6.4).
The VPSR algorithm and its PRISM implementation are correct and substantially faster than prior emulators, and the two theoretical analyses provide clear, interpretable insight. The vector Bienaymé–Chebyshev bound gives an $O(\sqrt{n}u)$ error envelope that avoids the exponential blow-up of Azuma–Hoeffding at low precision, and the loss decomposition neatly explains why MLP noise is largely a uniform logit shift that the softmax ignores while head noise is not. These mechanisms align with the data: at $t=6$ in the MLP, "SR holds perplexity degradation to 15% over the full-precision reference against 121% for RN" (Section 6.2), and at the head "SR's excess loss is driven by its non-uniform variance penalty, whereas deterministic RN carries zero variance penalty" (Abstract).
The experimental scope is too small to support broad claims about transformer inference. The evaluation uses only 1020 scored tokens from a single WikiText-2 prefix, so the perplexity numbers are effectively point estimates with no cross-text replication, and the mixed-rounding recipe is tested only on this narrow slice.
The theoretical story also weakens at very low precision, where the second-order decomposition's remainder is large. In Appendix B, Table 4, the head remainder at $t=4$ is $R=-5.860$ nats, comparable in magnitude to the variance term $V=23.204$, showing that the Taylor approximation is not dominating there. Moreover, the study emulates reduced significand precision while keeping the binary32 exponent range, so it "does not model the exponent-range restrictions of narrow hardware formats such as FP8 or NVFP4" (Section 6.4). Transfer to native hardware is therefore speculative.
The evidence strongly supports the site-level ordering within the experiment: the MLP down-projection shows a clear SR advantage ($66.24$ vs $79.64$ perplexity at $t=6$) and the head a clear RN advantage ($5054$ vs $519$), with the variance penalty $V_{\text{SR}}$ accounting for essentially all of SR's excess loss at the head. As the authors note, "Over $t=6,\dots,9$, the variance penalty $V_{\text{SR}}$ accounts for $100$–$105\%$ of SR's excess loss" (Section 5.3). However, the theoretical bounds bound forward error of a single reduction and are only "consistent with the observations," not predictive of end-to-end perplexity. Comparisons to related work are generally fair — the authors clearly distinguish their dynamic rounding from AdaRound-style offline weight rounding and from outlier-mitigation quantization — though they overstate the categorical nature of prior opposition to forward-pass SR, since cited works such as Chmiel et al. (2023) and NVFP4 (Abecassis et al., 2025) focus on training rather than inference.
Code, scripts, and aggregated data are archived on Zenodo and GitHub, which is commendable. However, independent reproduction is blocked by the heavy emulation overhead: PRISM SR is "$\times 8.5$–$32.1$" slower than native execution (Section 4.3, Table 1), and the full pipeline requires recompiling PyTorch through Verificarlo. The paper reports random seeds and virtual precision $t$, but does not fully specify the deterministic RN tie-breaking behavior beyond "untied" rounding away from zero, and the evaluation harness uses a fixed custom text prefix rather than a standard test-set split, which complicates direct comparison. The model is off-the-shelf, so no training hyperparameters are needed, yet reproducing the mixed-rounding gains would require rebuilding the entire instrumented stack.
Should low-precision transformer inference use stochastic rounding (SR) or round-to-nearest (RN)? The answer depends on where in the network you look. We isolate this effect by holding the numerical format fixed and varying only the rounding rule at individual operation sites. To enable experiments at freely chosen precisions, we extend the PRISM vectorized rounding library to arbitrary virtual precision via a variable-precision stochastic rounding (VPSR) algorithm, proving that the rounding decision is evaluated exactly in hardware floating point. We develop two analyses providing complementary insight into this site-level trade-off. First, a probabilistic forward-error bound for linear projections shows that SR's error envelope grows as $O(\sqrt{n} u)$ in reduction length $n$, versus $O(n u)$ for RN, a gap that widens rapidly at low precision and is most pronounced in the long multilayer perceptron (MLP) down-projection. Second, a second-order decomposition of expected cross-entropy loss change at the output softmax into signed drift, drift curvature, and a Fisher-weighted variance penalty reveals why the two sites behave oppositely: MLP noise is predominantly a uniform logit shift to which softmax is invariant, so SR's variance is largely discounted; head noise is non-uniform across the vocabulary and is not. On DistilGPT-2 at $t=6$ significand bits, observations match theory: SR in the MLP raises perplexity to 1.15x the full-precision reference, versus 2.21x for RN. At the language-model head, the ordering reverses because SR introduces non-uniform variance, whereas deterministic RN carries none. In a mixed-precision configuration (MLP output at $t=6$), assigning SR to the MLP and RN to the head brings perplexity within 1.10x of the full-precision reference, a 28% reduction over matched-bit RN.
Pick a starting point or write your own. Challenges run in the background, so you can keep reading while the AI investigates.
No challenges yet. Disagree with the review? Ask the AI to revisit a specific claim.