Stochastic rounding boosts low-precision transformer accuracy in key parts
Stochastic Rounding in Low-Precision Transformer Inference: A Variable-Precision Emulation Study of a Small GPT-2
Machine LearningComputation and Language
Summary
When computers run transformer language models like GPT-2 using fewer bits to save memory and speed, rounding numbers can cause errors. The authors studied two rounding methods: stochastic rounding, which adds randomness, and round-to-nearest, which chooses the closest number deterministically. They found that stochastic rounding works better in some parts of the model, especially in the multilayer perceptron layers, while round-to-nearest is better in the language-model head. Combining these methods in different parts reduces errors significantly compared to using just one rounding method.
What this means in practice
- •For machine learning engineers: Implement mixed rounding methods in low-precision transformer inference to reduce language model errors.
- •For hardware designers: Design hardware supporting variable stochastic rounding at fine-grained levels to improve transformer inference accuracy at low precision.
Authors
Yohan Chatelain, Pablo de Oliveira Castro
Abstract
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.