Where, When, and How Many Bits? Scaling Laws for SGD with Stochastic Rounding
Authors: Zhichao Wang, Michael W. Mahoney
Organizations: ICSI and Department of Statistics University of California, Berkeley, USA · ICSI, LBNL and Department of Statistics University of California, Berkeley, USA
Low-precision training reduces computation but introduces additional rounding error. We theoretically study how stochastic rounding (SR) affects prediction risk as model size, training time, and precision change. We consider a two-layer linear network with a fixed random first layer, trained by mini-batch SGD on fresh Gaussian data with power-law spectra. At each SGD step, the features used in forward and backward computations are rounded independently. We prove matched upper and lower bounds on prediction error for the final model and for uniform parameter averaging, with general learning-rate and precision schedules chosen before training. For fixed precisions and constant learning rates with parameter averaging, we derive how the smallest prediction error decreases with compute, and we identify model sizes and training times that achieve this optimal-compute scaling. Rounding only the forward features preserves the compute scaling of exact arithmetic, while rounding backward features changes it. Averaging parameters can change which computation needs higher precision and when higher precision is useful. Our results also give bit requirements for a given accuracy tolerance and show how backward rounding can change learning-rate schedule design. At a fixed compute budget, lower precision can improve prediction when the benefit of a larger model or longer training outweighs the additional SR-induced error. Model simulations illustrate these tradeoffs, and experiments with an MLP and a small transformer examine rounding choices beyond the theoretical model.
Figures & tables
Figure 1: (a) We round the features either at forward pass or backward pass in SGD dynamics ( 5 ) with one-bit precision and compare the SR-induced error over the corresponding risk without rounding. Solid lines evaluate the final model (last iterate); dashed lines evaluate the model obtained by averaging parameters over training (uniform average). Forward rounding causes more error for last iterate, while backward causes more error for the average. (b) Using three bits in the first half of training and one bit in the second, or the reverse, with nearly equal average precision. Each bar divides the SR-induced risk with higher precision early by that with higher precision late. Values above one indicate higher precision at late phase of training outperforms the case with higher precision at early phase of training. (c) Each point selects the smallest tested precision that keeps the final prediction risk no more than 1% higher than its value without rounding, at the same width and duration. Both locations use the same precision from 1 to 4 bits. Appendix F gives experimental details.
Feature precision (forward, backward)
Risk
Width M
Optimization time τ
Exact, exact
F−a+1b−1
Fa+11
Fa+1a
Finite, exact
F−a+1b−1
Fa+11
Fa+1a
Exact, finite
F−a+bb−1
Fa+bb
Fa+ba
Finite, finite
F−2ab−1
F1/2
F1/2
Table 1: Compute-optimal prediction risk Ravg(T) with fixed precisions, η,B , and compute budget F . Precisions pfw,pbw may differ when they are finite. Columns of M and τ show the orders attaining the compute-optimal risk. Here we only focus on the scaling order as F→∞ . 1<b<a .
Last iterate
Uniform average
Exact-arithmetic risk scale Dout
M1−b+ρsM
M1−b+ρsMM1−a
Forward coefficient Afw
ρMa−b
ρM1−b
Backward coefficient Abw
ρ(sM+M1−b)
ρ(sM+M1−b)
Joint coefficient A×
ρ2(sM+M1−b)Ma−1
ρMa−b
Table 2: Following Section 3.2 with τ≍Ma . Scaling definitions in ( 28 ) with sM defined in ( 19 ).
Figure 2: (a) With five mantissa bits in total, more forward bits reduce last-iterate error; averaged-output error is smallest at (pfw,pbw)=(2,3) . Errors are relative to the corresponding exact-arithmetic risks. (b) The orange schedule accounts for backward-rounding error using ( 31 ); the blue schedule uses the exact-arithmetic design ( 30 ). Both use one-bit backward SR; the dotted line is a constant learning rate. (c) Full last-iterate prediction risk under the precision-dependent design minus that under the exact-SGD design, at the same backward precision. Positive values favor the exact-SGD design. In (c), points average eight independent regression problems. Forward features are exact in (b,c). Appendix F.2 gives details.
Figure 3: Prediction risk of uniformly averaged outputs under a compute budget. (a,b) Exact-arithmetic and joint one-bit SR training ( pfw=pbw=1 ) at several widths. Black curves show the smallest prediction risks among tested widths and training durations. Red dashed curves use the optimal-compute scaling in Table 1 for exact-arithmetic training. (c) At a fixed budget F , joint SR permits a wider model and longer training with lower prediction risk. We average over three sketches. Appendix F.3 gives details.
Figure 4: Prediction risk under a fixedcompute budget with fresh and reused data. Uniformly averaged outputs at M=256 , using exact arithmetic or joint two-bit feature SR. (a) Full prediction risks at F=227 . (b) The difference between SR and exact prediction risks across budgets; negative values indicate SR is beneficial. (c) Reuse-minus-fresh risk differences at F=227 for different data sizes. For more details, see Appendix F.3 .
Figure 5: Rounding locations, shared activations and test-time prediction in neural networks. (a) Final transformer validation loss; the combined case rounds both activation computations, incoming gradients and forward weights to two bits, with seven-bit parameter storage. (b) Shared SR reuses the same forward-rounded activation in the weight-gradient computation. Points show final loss differences from same-seed training with SR and without SR. Negative values indicate lower loss than this baseline. (c) Each trained MLP is tested twice with its parameters fixed: with forward-activation SR disabled, and with SR enabled at its training precision. Both tests use the same examples. Appendix F.5 gives details and final performance values.
Appendix figures & tables3 assets
Supplementary material from the paper’s appendix.
Appendix
Figure 6: Supplementary precision comparisons. Figure 6 (a,b) shows SR-induced risk increases for last-iterate and uniformly averaged outputs at fixed duration, with fresh or reused data. The four categories use feature rounding alone, add temporary prediction-weight rounding, add stored-parameter rounding, or round at all six locations, as defined below. The vertical axes use different logarithmic ranges; bars show twice the standard error across eight task means. Figure 6 (c) compares four integer precision schedules with constant two-bit precision at the same feature location. Mean bit counts are shown; the realized budgets differ.
Training configuration
Validation loss
No added SR
4.1846±0.0049
Independent activations, 1 bit
4.2452±0.0034
Independent activations, 2 bits
4.2016±0.0082
Shared activations, 1 bit
4.2238±0.0019
Shared activations, 2 bits
4.1976±0.0036
Backpropagated gradients, 2 bits
4.1852±0.0051
Appendix
Table 3: Final transformer validation loss with added SR disabled at evaluation. The table covers Figure 5 (a,b) and a three-bit storage comparison.
Training configuration
Test loss SR off
Test loss SR on
Accuracy (%) SR off
No added SR
1.6198±0.0032
1.6198±0.0032
56.409±0.147
Forward activations, 1 bit
1.6237±0.0041
1.6721±0.0056
56.541±0.109
Forward activations, 2 bits
1.6210±0.0027
1.6338±0.0030
56.458±0.112
Forward activations, 3 bits
1.6212±0.0039
1.6242±0.0039
56.381±0.124
Independent activations, 1 bit
1.6208±0.0034
1.6706±0.0041
56.501±0.120
Independent activations, 2 bits
1.6223±0.0032
1.6356±0.0035
56.403±0.131
Appendix
Table 4: Final MLP test performance for Figure 5 (b,c). Entries are means ± two standard errors across 16 seeds. The two loss columns evaluate the same final parameters without or with SR; accuracy uses evaluation without SR.