TinyMOSS-Diarize

Inspiration

Speaker-aware transcription is still a stack: an ASR model, a diarizer, an aligner, sometimes a language model on top. Each hand-off drops context, and the pile is too large and slow to keep on a phone. MOSS-Transcribe-Diarize 0.9B does the whole job in one pass, words, speakers, and timestamps, because a small language-model backbone can carry all three. The remaining question was whether that backbone could be quantized hard enough to live on mobile and edge devices.

What it does

TinyMOSS-Diarize is a quantized build of MOSS-Transcribe-Diarize 0.9B that transcribes, emits timestamps, and labels speakers in a single pass, across nine languages. The weights are quantized to an average of 2.911 bits per value, 1,817,113,576 bytes down to 330,968,414 bytes (5.49x) and, critically, they stay packed at inference. We wrote NEON kernels for every precision in the package, so nothing is expanded back to BF16 in memory:

Accuracy. 16 datasets, 553 utterances, 0 inference errors and 0 transcript-format errors under the grammar-constrained decoder. The strongest result is multi-speaker audio, where a single-pass diarizing model has a real advantage: on synthetic Chinese mixtures we score 7.50% cpWER where whisper-large-v3-turbo produces 33.87% CER of transcription-only text and whisper-small collapses to 120.53%; on AliMeeting, 21.86% cpWER versus 42.92%.

On-device. On an Apple M4 MacBook Air, the packed build runs a 10.4 s clip in 18.228 s versus 24.087 s for the BF16 model, with TTFT 9.686 s versus 19.704 s, model load 0.541 s versus 14.680 s (27x), and peak RSS 3,998 MiB versus 7,011 MiB. On a 4 GB Raspberry Pi 5 the packed build is the only configuration that finishes an utterance at all: BF16 needed 177 s just to load and was still in its first token when killed at 628 s, and FP32 completed a single token before timing out with 2.8 GiB of swap in use.

How we built it

A llama.cpp-style mmap container (MQF1) holds the packed tensors with a 64-byte-aligned tensor table, and a fused PyTorch operator multiplies straight out of it. The kernels are strictly weight-only: activations stay FP32, dequantization is fused into the matmul, and no dequantized weight tensor is ever materialized. There is a blocked GEMM path for prefill that unpacks each 128-column weight group once into an L1 tile and reuses it across all activation rows (4x4 NEON FMA microkernel, NR=8 x MR=4), and a separate GEMV path for decode. Runtime feature detection picks the kernel per device via getauxval(AT_HWCAP) on Linux/Android and sysctlbyname on macOS.

Everything is measured with one metric contract on every device: one discarded warmup plus three timed runs, median reported; ttft_seconds from a separate max_new_tokens=1 generation; decode_tokens_per_second = (tokens - 1) / (total - ttft). Every number in the repository is backed by a JSON file under results/.

Challenges we ran into

The hard problem was low precision, not packing. At 1.375 bits on the language backbone, post-training quantization alone drops speaker turns and transcript grammar. We recovered that with quantization-aware training plus knowledge distillation: the original BF16 MOSS-Transcribe-Diarize checkpoint teaches a fake-quant student, so the low-bit decoder matches the teacher's token distribution instead of a brittle hard-label loss.

That still leaves an exposure-bias gap, where the student trains on teacher states, then at inference visits its own. We closed it partially with VESPO (Variational Sequence-level Soft Policy Optimization) and used the original teacher as the on-policy distillation (OPD) teacher: the quantized student rolls out transcripts, and the BF16 teacher scores those student-induced states. VESPO's sequence-level importance weights stay stable when the rollout worker lags the learner, which is the usual failure mode of async low-precision training.

Supervision had to come from somewhere. Clean single-speaker corpora do not teach overlap. We extracted ground truth from synthetic multi-speaker mixtures: speaker identity and timestamps known by construction, and from real-world audio, including the audio tracks of FineVideo, mixed into meetings after teacher transcription and acoustic gating. The 0.9B model then has to do, in one forward pass, what today's pipelines still split across many models. Extreme quantization of the language-model backbone is what made that single packed checkpoint small enough to run on Arm.

Accomplishments that we're proud of

The packed build sometime beats its own BF16 baseline end-to-end on a laptop-class Arm chip while using 57% of the memory, and it turns a 4 GB Raspberry Pi from "cannot run this model" into "runs this model".

What we learned

Low-bit weight-only inference is not automatically faster; it is faster only if unpack cost is amortized over the M dimension and the threading granularity matches the op size. Two of our three biggest wins came from removing parallelism and removing work, not from adding SIMD. We also learned to distrust a speedup that arrives without a correctness check: our largest apparent gain was a cache bug that skipped computation.

What's next for TinyMOSS-Diarize

Closing the remaining decode gap against vendor BF16 GEMV; an i8mm GEMM path that does not require quantizing activations; and pushing the audio encoder below 4 bits, since W4 audio and embedding tensors are 70.7% of the packed bytes while the 1.375-bit language projections hold 48.5% of all values in only 22.9% of the bytes.

Built With

• android
• apple-silicon
• arm-neon
• c++
• cmake
• python
• pytorch
• raspberry-pi
• transformers