Back to Moonshine

Quantization

docs/models/quantization.md

0.1.22.3 KB
Original Source

Quantization

We typically quantize our models to eight-bit weights across the board, and eight-bit calculations for heavy operations like MatMul. This is all post-training quantization, using a combination of OnnxRuntime's tools and my Onnx Shrink Ray utility. The only anomaly in the process is the treatment of the frontend, which uses convolution layers to generate features, which produces results similar to the more traditional MEL spectrogram preprocessing, but in a learned way with standard ML operations. The inputs to this initial stage correspond to 16-bit signed integers from the raw audio data (though they're encoded as floats) so we've found it necessary to leave the convolution operations in at least B16 float precision.

We give each output channel of a weight its own scale factor rather than sharing one across the whole tensor. This matters far more than it sounds like it should. Our frontend convolutions are trained with weight normalization, which by construction learns a separate magnitude per output channel — on Tiny Streaming the largest channel is 17x the smallest. A single scale for the whole tensor has to cover the largest channel, so the smallest ones get only a handful of the 256 available levels. Measured on LibriSpeech test-clean, moving to per-channel scales took Tiny Streaming from 7.57% to 4.83% WER, Small from 3.03% to 2.61%, and Medium from 2.37% to 2.17%, while adding 0.5% to the download. Roughly 90% of that gain came from the frontend alone.

If you quantize your own Moonshine-style model, this is the first thing to check. It is also why we do not fold the weight-normalization out of the exported graph even though doing so would save a little work at runtime: folding makes the per-channel magnitudes more extreme, not less, so it must be paired with per-channel scales rather than done on its own.

You can see the options we use for the conversions in scripts/quantize-streaming-model.sh.

Re-quantized models are published to a new dated directory on our CDN (currently quantized_26_07_30) instead of overwriting the previous files. That way a given version of this library always resolves the exact weights it was tested against, and your local model cache never mixes old and new files.