← Posts · Home

SSRQ: Structured Stochastic Rounding for
Quantization Aware Post Training on FP3

Tsun-Yi Yang · 17 Jul 2026
Supporting AI model: Grok 4.5 High · Supporting agent: Cursor

Most low-bit quantization stops at FP4 / INT4. SSRQ is the first to close the FP3–FP16 gap: packed FP3 is ~81% smaller than FP16 (\(5.33\times\)), one-shot PTQ collapses (71.3% on mlp9), and a short SSRQ post-training step recovers 94.2% — within ~1.5 pp of FP16 float (95.7%) — with a plain RTNE FP3 deploy checkpoint.

Deploy (RTNE) test accuracy on MNIST mlp9, seeds 161–164. One-shot PTQ-FP3 collapses (~71%). Scratch SRC recovers some ground. Warm-starting from FP16 @ 5000 and continuing with SSRQ closes most of the gap to FP16 float (~95.7%), beating post-FP16 SRC by ~2–4 pp across \(R\in\{2,3,4,6,8\}\).

Why FP3 is a different regime

Compact formats are cheap to store and multiply, but many weights are not exactly representable on the lattice. Train in high precision and snap once, and you often lose accuracy. Train for that lattice — update wide, round back each step — and you can find better deploy weights. That is the setting of Graphcore’s stochastic-rounding work and ARITH 2025. Push the lattice one step coarser — to FP3 — and the story gets harder: PTQ variance explodes, and scratch SRC alone leaves a clear gap to FP16.

The storage win is real. Format alias fp3 = p3109_k3p2sf is a 3-bit signed finite lattice (7 values). Against an FP16 master (2 bytes/param), a pure packed FP3 dump is \(16/3 \approx 5.33\times\) smaller — about 81% less weight storage — before counting a tiny per-tensor scale header:

Sizes count every dim ≥ 1 parameter (same coverage as deploy RTNE). FP3 column is bit-packed at 3 bits/param; a float32 scale per tensor adds <0.3 KiB. Activations / optimizer state are not included.

On the same mlp9 probe (depth 9, hidden 64), after FP16 training for 5000 steps:

Same test split (seed 161) for every run.

FP8/FP6 PTQ is essentially free. FP4 already hurts. FP3 is a cliff: one-shot RTNE is not a deployable policy — you pay the 81% size cut, then lose most of the accuracy unless you train for that lattice. That is the setting where quantization-aware training with stochastic rounding matters.

SRC: centered stochastic rounding

Round-to-nearest-even (RTNE) is the obvious deploy snap, but during training it stalls on a coarse lattice: once updates fall below half a bin, the rounded weight stops moving. Stochastic rounding gets optimization unstuck: ideal SR rounds \(x\) up or down with probability given by the fractional part (\(\operatorname{SR}(x,u)=\lfloor x+u\rfloor\), \(u\sim U[0,1]\)), so \(\mathbb{E}_u[\operatorname{SR}(x,u)]=x\).

Hardware does not hand you infinite-precision \(u\). With an \(R\)-bit integer \(K\in\{0,\ldots,2^R-1\}\), the floor form \(\lfloor x+K/2^R\rfloor\) is biased for small \(R\). Graphcore’s fix is to center the comparator — call it SRC. On a weight matrix, ordinary SRC draws a fresh independent \(K_{ij}\) at every entry each step:

Optimizer math stays wide; the forward / backward see the rounded weights. At the end of training, drop the randomness and deploy with one RTNE snap — same footprint as PTQ.

Same story, more detail: Graphcore blog · ARITH 2025 · arith25-stochastic-rounding.

Motivation: post-train the lattice, not the network from scratch

On FP3, one-shot PTQ collapses and scratch QAT plateaus below FP16. Continuing from an FP16 master with a short quantization-aware post-training step recovers most of the gap — so the hard work is settling weights onto the lattice, not rediscovering the task under a brutal codebook. Stochastic rounding (SRC) already works in that post-training role, but leaves a few points on the table.

The mental model is where the noise comes from:

So SSRQ keeps Graphcore’s SRC rounder unchanged and only changes how \(K\) is proposed — used as a short continuation after the high-precision run:

Inference footprint matches PTQ / SRC: ordinary FP3 weights, no learned dither, no lookup table. The only change is which FP3 points the masters settle on before the final snap.

SSRQ: same rounder, a structured draw of \(K\)

Structured Stochastic Rounding (SSRQ) does not invent a new FP3 cast. It keeps \(\operatorname{SRC}(\,\cdot\,;K)\) exactly. The only change is how \(K\) is proposed: mix fresh Uniform noise \(\xi\) with a learned row–column preference, then feed that \(K\) into the unchanged Graphcore cast:

For each weight matrix \(W\in\mathbb{R}^{m\times n}\), the vectors \(a\in\mathbb{R}^{m}\) and \(b\in\mathbb{R}^{n}\) are learnable logits — one scalar per row and per column — trained only during QAT (not part of deploy). Entry \((i,j)\) shares structure through the rank-1 field \(a_i+b_j\): rows with large \(a_i\) (and columns with large \(b_j\)) get a systematically different rounding bias. That is cheap: \(m+n\) extras instead of \(mn\). \(\sigma\) turns those logits into a rounding preference in \((0,1)\): near \(0\) or \(1\) pushes the dither toward rounding down or up; at init \(a=b=\mathbf{0}\) gives \(\sigma(0)=\tfrac12\) (no preference). Defaults \(\alpha=\beta=0.5\); \(\xi\) is redrawn each step. Forward / backward see the rounded weights; the STE is identity. At deploy, drop \(a\), \(b\), and \(\xi\), and snap with \(W^{\mathrm{deploy}}=\operatorname{RTNE}(\operatorname{scaled}(W))\).

Fake-quant (this work) vs inplace cast

Graphcore’s SR story is often told in an inplace training setting: cast weights into the target format as you go, on hardware that can round in the datapath. This write-up uses the usual ML-engineering cousin — STE fake quantization: keep a wide master \(W\) (FP32/FP16), build a quantized view \(W^q\) each step via functional_call, and send gradients back to the master. Deploy is then one RTNE snap of that master onto FP3. Same rounding math; different bookkeeping — and it runs on ordinary PyTorch without custom inplace kernels.

Scope of the quantizer during QAT: weights only. Activations stay in full precision (FP16/FP32 matmuls). In the common weight/activation bit-width shorthand, training/deploy here is W3A16 — 3-bit FP3 weights, wide activations. We also post-hoc evaluate the same RTNE FP3 weights under W3A8 (fake-quant activations to INT8 or FP8 at every Linear/Conv input) — see W3A8. That is an inference stress test, not activation QAT.

Weights stay put · \(\sigma(a_i+b_j)\) learns structure

Mid-layer (5.weight, 64×64) during SSRQ post-FP16 QAT (\(R=8\), \(\alpha=\beta=0.5\), seed 161). Top: float master weights — almost unchanged from step 1→2000 (the FP16 basin is already good). Bottom: structured preference \(\sigma(a_i+b_j)\) — starts flat (~0.5), then grows a separable row/column grid that biases which FP3 neighbors get chosen. At deploy, \(a,b\) are discarded; only the settled weights remain.

Scratch SRC on FP3

Train from random init with IID SRC onto FP3, then RTNE-deploy. That already beats PTQ (~71%), but on this seed grid it plateaus around the high 80s to low 90s — still several points short of FP16 float (~95.7%).

That is why SSRQ prefers a warm start: the FP16 run already found a good basin; the short post-QAT only has to settle those weights onto FP3, not rediscover the task under a brutal lattice.

Post-FP16 results · mlp9

Load FP16 masters @ 5000, continue QAT on FP3 for 1000 or 2000 steps. Report deploy RTNE accuracy (mean ± std over 4 seeds). All cells below are W3A16: FP3 weights, float activations.

Ablation · \(\alpha\) / \(\beta\) at \(R=4\)

Recall the dither mix \(r_{ij}=(\alpha\,\xi_{ij}+\beta\,\sigma(a_i+b_j))\bmod 1\). Default is \(\alpha=\beta=0.5\). Same post-FP16 recipe (\(R=4\), 2000 steps, seeds 161–164), turn one term off on both probes:

Post-FP16 results · tf4

Same recipe on MNIST tf4 (4-layer transformer): FP16 masters @ 5000, then SRC/SSRQ post-QAT onto FP3 for 1000 or 2000 steps. Seeds 161–164. No scratch-SRC grid on this probe — chart is PTQ vs post-FP16 only.

Deploy (RTNE) test accuracy on MNIST tf4. PTQ-FP3 collapses (~54%). Post-FP16 SSRQ recovers to the mid-80s across \(R\in\{2,3,4,6,8\}\), beating post-FP16 SRC by ~10–15 pp and closing most of the gap to FP16 float (~88.2%).

Post-QAT cells below are W3A16: FP3 weights, float activations.

Post-hoc W3A8 · FP3 weights, 8-bit activations

Take the SSRQ post-FP16 @ 2000 RTNE masters (weight-only QAT, W3A16 train) and re-evaluate with fake-quantized activations on every Linear/Conv input: per-tensor symmetric INT8, or RTNE FP8. No activation training — inference stress test only. Full MNIST test set (10k); mean ± std over seeds 161–164.

W3A8 · mlp9 · SSRQ post-FP16 @ 2000

W3A8 · tf4 · SSRQ post-FP16 @ 2000

Conclusion

Ultra-narrow formats buy storage, but one-shot PTQ assumes a float solution already near a good lattice point — on FP3 that fails. Stochastic rounding keeps weights moving between bins; SSRQ adds the missing piece by drawing the dither with a cheap shared preference instead of IID noise, so training can bias neighbor choice without changing the deploy cast. An \(\alpha/\beta\) ablation at \(R=4\) on both mlp9 and tf4 confirms both pieces matter: structure alone collapses, noise alone ≈ SRC, and the default mix wins. The workable recipe is post-train the lattice: warm-start from a wide-precision basin, settle with a short structured QAT, then snap. Rounding structure is train-only; deploy stays a plain deterministic checkpoint. Post-hoc W3A8 is nearly free on these checkpoints — 8-bit acts do not undo the FP3-weight recovery. Natural next steps are richer shared rounding preferences, joint weight–activation QAT when narrower acts are a real target, and the same structured-dither idea on larger models or other ultra-narrow formats — including hardware paths that cast inplace.

Cite

If you use this writeup, please cite it as a blog post. Also listed on Google Scholar.

@misc{yang2026ssrq,
  title        = {SSRQ: Structured Stochastic Rounding for Quantization Aware Post Training on FP3},
  author       = {Yang, Tsun-Yi},
  year         = {2026},
  month        = {7},
  howpublished = {Blog post},
  url          = {https://shamangary.github.io/post/ssrq-fp3/}
}