On Stochastic Rounding with Few Random Bits

TL;DR

This paper examines a number of possible implementations of few- bit stochastic rounding (FBSR), and shows how several natural implementations can introduce sometimes significant bias into the rounding process, which are not present in the case of infinite- bit, infinite-precision examinations of these implementations.

Abstract

Large-scale numerical computations make increasing use of low-precision (LP) floating point formats and mixed precision arithmetic, which can be enhanced by the technique of stochastic rounding (SR), that is, rounding an intermediate high-precision value up or down randomly as a function of the value's distance to the two rounding candidates. Stochastic rounding requires, in addition to the high-precision input value, a source of random bits. As the provision of high-quality random bits is an additional computational cost, it is of interest to require as few bits as possible while maintaining the desirable properties of SR in a given computation, or computational domain. This paper examines a number of possible implementations of few-bit stochastic rounding (FBSR), and shows how several natural implementations can introduce sometimes significant bias into the rounding process, which are not present in the case of infinite-bit, infinite-precision examinations of these implementations. The paper explores the impact of these biases in machine learning examples, and hence opens another class of configuration parameters of which practitioners should be aware when developing or adopting low-precision floating point. Code is available at http://github.com/graphcore-research/arith25-stochastic-rounding.

Figures (4)

• Figure 1: Bias in stochastic rounding with finite randomness. For each $X$ value, 5000 samples are rounded with 2 bits of randomness, and the mean of the rounded values is plotted. (Left, SRFF): Mean rounded values are below the line $Y=X$, indicating bias. The computed bias of $-0.124 \approx-2^{-3}$, as predicted by the calculation of §\ref{['sec:bias-srff']}. (Right, SRF): Mean rounded values are symmetric around $Y=X$, with bias $\approx$ 0, following §\ref{['sec:unbiased-srf']}.
• Figure 2: Bias in stochastic rounding with finite randomness, finite-precision inputs. For each $X$ value in BFloat16, 100,000 samples are rounded with 2 bits of randomness, and the mean of the rounded values is plotted. (Left, SRFF): Mean rounded values are below the line $Y=X$, indicating bias. The computed bias is $\approx$ -0.046875, as predicted by the calculation of §\ref{['sec:fbbias-srff']} and §\ref{['asec:fbbias-srff']}. (Middle, SRF): Mean rounded values are not symmetric around $Y=X$, with bias $\approx$ 0.015625, following §\ref{['asec:fbbias-srf']}. (Right, SRC): Mean rounded values are symmetric around $Y=X$, with empirical bias $\approx$ 0.000, following the correction of §\ref{['sec:fbbias-srff']}.
• Figure 3: Experiments on language model training (nanoGPT, small model). Binary16 baseline: achieves minimum validation loss at about 1900 iterations, after which training loss (not shown) reduces, but validation loss increases. Binary8P4 RTNE: initial loss reduction followed by "stagnation"; SRFF: initial loss reduction followed by "divergence"; SRF, SRC: convergence to stable values. Although SRF and SRC converge to better validation errors than Binary16, this regularization behaviour is expected to apply only to such small-scale experiments, and should not be taken as an indication of any superiority of Binary8+SR over Binary16.
• Figure 4: Experiments on language model training (nanoGPT, GPT-2 (350M params)). Binary16 baseline: validation loss follows training loss (not shown) very closely. Binary8P4 SRFF: initial loss reduction followed by "divergence"; SRF, SRC: convergence to stable values. In this case SRF and SRC converge to better validation errors than SRFF, but not to the same loss as Binary16.