Shared rounding thresholds speed training when errors are corrected
Efficient Quantization Training with Stochastic Rounding via Structured Subspace Randomness · arXiv:2610.36587
![Overview of INT8 (S8) quantization training with S3R and its structured error correction S3R-EC. On the VPU, S3R broadcasts a compact seed 𝐮t∈[0,1]B×S×1, one threshold per token, to 𝐑t=Φ(𝐮t) inside…](https://arxiv.org/html/2610.36587v1/INT8_quant_flow.svg)
In plain words
With 8-bit training, randomly rounding gradient values can preserve small updates, but generating a random number for every value can slow an accelerator.
Structured Subspace Stochastic Rounding (S3R) reuses one random threshold across the values for each token. Its error-correction method removes the shared part of the rounding error.
In one 12-layer language-model training run, corrected S3R ended with validation loss close to independent rounding, though three runs could not settle the small difference.
The research
Low-precision training can reduce computation, but generating an independent random threshold for every gradient value can become costly on accelerators. It remains unclear whether sharing thresholds can reduce that cost while preserving the training quality independent rounding provides.
Structured Subspace Stochastic Rounding (S3R) shares random thresholds to cut rounding work; S3R-EC corrects the shared error.
S3R reduces random-number work by sharing thresholds across gradient values, but that sharing correlates their rounding errors. The authors pair it with S3R-EC, which subtracts the known structured error from the gradients before optimization.
In one 12-layer C4 training run, two-way S3R-EC’s mean final validation loss was 4.919, close to 4.929 for independent rounding; three seeds could not resolve that difference. Full-step timing also found S3R-EC faster than both BF16 and independent rounding for three tested dense models on TPU v5p.
The quality comparison covers one 12-layer C4 training workload and three seeds per configuration. The authors say this comparison can resolve differences of about 0.05 validation-loss units or more, so it cannot determine whether corrected S3R and independent rounding differ by a smaller amount.
What the evidence shows
- The cost of random rounding
Structured Subspace Stochastic Rounding (S3R) is a method to reduce the random-number cost of rounding gradients during low-precision transformer training. Independent stochastic rounding uses a separate random threshold for every gradient value.
On accelerators where random-number generation runs on vector units, that work can slow training even when matrix operations can use faster dedicated hardware.
- How thresholds are shared
S3R quantizes an upstream gradient once for both backward matrix products and shares a random threshold across the hidden-dimension values for each token. Each value remains unbiased, but values sharing a threshold can have correlated rounding errors. The method reduces random draws by trading independent errors for this shared structure.
- Where sharing adds error
The authors’ analysis locates the added error in the activation gradient, which is passed to earlier layers. For per-token sharing, their analysis finds that the weight-gradient error keeps the variance of independent rounding, while correlations between channel errors add variance to the activation gradient. This identifies why correcting the weight gradient alone would leave the incoming error in place.
- Correction restores close quality
In one 12-layer Llama-style model trained on English C4 for 2,000 steps, uncorrected per-token S3R had a mean final validation loss of 5.055 over three seeds, versus 4.929 for independent rounding. Two-way S3R-EC averaged 4.919.
Researchers used each configuration’s best learning rate, selected on seed 0, and the same model, data, steps, and AdamW optimizer.
- Speed in tested full steps
In full-step timing on TPU v5p, two-way S3R-EC was 1.24 to 1.32 times as fast as BF16 and about 1.11 times as fast as independent rounding for three tested dense models. These figures use median times across 20 steps after 12 warm-ups.
The quality and speed results together support the method in these evaluated workloads, with the quality comparison limited by its three-seed resolution.
Deep
- Original paper authors
- Zeyu Wang, Yiheng Duan, Tao Chen, Donghyun Koo, Sihan Liu
- Submitted to arXiv
- License
- CC BY 4.0