Source-linked AI summary
Compressing Streaming Neural Audio Encoders via Latent-Space Distillation
Prasanth Yadla, Mohammad Samragh, Dongseong Hwang, Mingbin Xu, Yuanyuan Zhang, Chung-Cheng Chiu, Yongqiang Wang, Yuan Liu, Zhen Huang, Xiaodan Zhuang
TL;DR
An always-on audio tokenizer must share on-device resources with the foundation model, making encoder size consequential. The paper distills the teacher's pre-quantizer latent, and at 2.8× compression five of six students stay within 1.9% relative WER of their teachers without fine-tuning.
Problem
An always-on tokenizer shares a fixed on-device memory and compute budget with the foundation model, while its full-size encoder leaves the model underprovisioned.
Source: §1 · pp. 1–2Method
The student encoder regresses the teacher's per-frame pre-quantizer latent with squared error, using an affine layer to reconcile widths.
Source: §1 · p. 2Results
Takeaways & Limitations
Abstract
from arXiv · showhide
System-wide Dictation on Apple devices runs entirely on-device, and the speech it transcribes reaches the foundation model through a tokenizer: an encoder that maps short windows of waveform onto the representation the language model reads. Because that model is sparsely activated under Instruction-Following Pruning, only a small subset of its experts occupies DRAM at any time, so the always-on tokenizer competes for the same memory, and its parameter count bears directly on power and latency. In this work we study how to compress such a tokenizer by distillation, taking as the supervision target neither the discrete token nor the output distribution but the pre-quantizer latent the model actually consumes - the last representation the two token interfaces share. We train only the student encoder to regress the teacher's per-frame latent under a squared-error objective, with a single affine layer absorbing the teacher-student width mismatch. Because the target precedes both the quantizer and the language-model bridge, one recipe covers both token interfaces we support, and applies both to a tokenizer pretrained alone and to one jointly trained with a language model. At 2.8x compression the distilled student stays within 1.9% relative WER of its teacher on five of six teacher-student pairs without any fine-tuning, and improves on an independently trained tokenizer of identical capacity by 3.9% relative.
1 Introduction
An always-on audio tokenizer shares a fixed on-device memory and compute budget with the foundation model, motivating a smaller encoder. The paper distills the pre-quantizer latent and reports teacher-comparable recognition results at 2.8× compression.
1 Introduction
- An always-on tokenizer must share memory and compute with the on-device foundation model, and the full-size encoder leaves that model underprovisioned.The encoder's resident footprint directly constrains the admissible language-model size.Source: §1 · pp. 1–2
- The pre-quantizer latent is the common distillation target because discrete tokens require handling a non-differentiable arg min and decoder outputs would target components absent from deployment.The authors motivate this target by construction, not by ablation against alternatives.Source: §1 · p. 2
- The student encoder regresses the teacher's per-frame latent with squared error and an affine layer for width mismatch, using teachers pretrained alone or jointly with a language model.The same target precedes either token interface.Source: §1 · p. 2
- At 2.8× compression or more, five of six students remain within 1.9% relative WER of their teacher without fine-tuning, and beat an equally sized independently trained tokenizer by 3.9% relative.The teacher comparison spans both distillation stages, both token interfaces, and three joint models.Source: §1 · p. 2
2 Related Work
Prior work establishes discrete and continuous audio-token interfaces and studies representation distillation and model compression. This work differs by shrinking an encoder whose output boundary is fixed by a downstream quantizer or bridge, without labels.
2 Related Work
- SpeechTokenizer and RepCodec also distill into tokenizer latents, but change token meaning; this work distills to shrink the encoder while keeping the token inventory fixed.SpeechTokenizer aligns its first residual-quantization level with HuBERT representations, while RepCodec learns a codec over speech representations rather than waveforms.Source: §2 · p. 2
- Discrete audio tokens can enter a language-model vocabulary, while a continuous interface projects encoder output into the model's embedding space.The continuous interface is used by AudioPaLM for discrete tokens and by Fathullah et al. and SLAM-ASR for projected vectors.Source: §2 · p. 2
- Intermediate-feature distillation matches representations across width differences, while speech encoder studies report that hidden-layer matching transfers more signal than output matching alone.Related examples include FitNets (Romero et al., 2015; Gou et al., 2021), TinyBERT (Jiao et al., 2020), and DistilHuBERT (Chang et al., 2022).Source: §2 · p. 3
- Unlike pruning, elastic-supernet, and pseudo-label compression approaches, this setting fixes the encoder output at a downstream quantizer or bridge and requires no labels.Distil-Whisper requires transcripts and distills a decoder that this work discards.Source: §2 · p. 3
3 Distilled Tokenizer
The distilled tokenizer reproduces a teacher encoder's per-frame latent, then uses the teacher-side interface block. Its shared pre-quantizer target supports either token interface and teachers from either training stage.
3 Distilled Tokenizer
- A distilled tokenizer trains a student encoder to reproduce the teacher's per-frame latent and pairs it with the teacher-side quantizer or bridge.The target is the pre-quantizer latent rather than discrete tokens or decoder outputs.Source: §3 · p. 3
3.1 Tokenizer–Model Interface
- Both interfaces branch from the same latent, so targeting it lets one student pair with either the discrete quantizer or continuous bridge.The continuous interface avoids quantization loss but carries a tokenizer-specific parameter block with the language model when extracted from a joint checkpoint.Source: §3.1 · p. 4
3.2 Distillation Objective
- The objective trains the student encoder against the frozen teacher's per-frame latent using masked squared error, with an affine map reconciling their widths.Only the student encoder and affine layer receive gradients; the student inherits the teacher-side quantizer or bridge, and stage-0 decoder heads are not distilled.Source: §3.2 · pp. 4–5
- Latent matching avoids supervising codebook assignments through a non-differentiable arg min and uses no temperature, softmax, or label smoothing.The authors state that downstream accuracy depends on the latent the model consumes.Source: §3.2 · p. 5
3.3 Student Encoder
- The student trades width for depth; narrowing yields a 2.8× parameter reduction and a matching 2.8× reduction in per-frame multiply–accumulate cost.The affine layer accounts for well under 1% of the student encoder and is folded into the downstream interface or dropped at inference.Source: §3.3 · p. 5
3.4 Distillation Stages
- A frozen tokenizer teacher from either stage supplies the target, and the resulting smaller tokenizer replaces it without student fine-tuning.Stage-0 targets a language-model-agnostic representation; stage-1 targets one co-adapted with a specific language model, tying the student to that model.Source: §3.4 · pp. 5–6
4 Experiments and Results
The experiments use unlabeled audio and evaluate students distilled from teachers at two training stages. The setup compares six teacher–student pairs through their full recognition stacks, without student fine-tuning.
4.1 Data
- Stage-0 runs use five speech datasets, while stage-1 runs use a different evaluation suite; their results are tabulated separately and should not be compared across tables.The stage-0 suite covers read, prepared, parliamentary, meeting, and accented long-form speech; stage 1 uses aggregate multilingual evaluators.Source: §4.1 · p. 6
4.2 Distillation Setup
- Students are trained for 500k steps on unlabeled audio, with the teacher frozen and only its encoder instantiated; the student’s EMA weights are evaluated.Training uses float32, AdamW, a peak learning rate of 10^-3, and a global batch size of 1280.Source: §4.2 · p. 6, §4.1 · p. 6
- Three stage-0 teachers vary in architecture and capacity, and three stage-1 joint models are built on teachers B or C, with J3 restricting gradient updates.The stage-0 set includes a Conformer, a Transformer, and a depth-pruned Conformer; each teacher defines its own compression ratio against the common student capacity.Source: §4.2 · p. 6
- Six teacher–student pairs are evaluated without student fine-tuning through the teacher’s fixed transcription decoder or the joint language model; same-capacity independent-training controls are also included.The stage-0 decoder has 262M parameters and the stage-1 joint model has 2.8B parameters; both controls inherit four initialized sub-blocks from a distilled model.Source: §4.2 · p. 6
4.3 Results on Stage-0 Distillation
Stage-0 distillation keeps student recognition performance close to the teacher across three pairs, with occasional better results on individual datasets.
4.3 Results on Stage-0 Distillation
- The relative WER differences are +0.4% for teacher A, +1.9% for teacher B, and +0.8% for teacher C.The corresponding absolute WER differences are +0.06, +0.22, and +0.10, respectively.Source: §4.3 · p. 6
- Across all three stage-0 pairs, distilled students track their teachers within 1.9% relative WER; Table 1 reports lower-is-better WER.Recognition uses the same 262M-parameter transcription decoder across rows, and students are evaluated without fine-tuning.Source: §4.3 · p. 6, Table 1 · p. 7
- Students outperform their teachers on AMI in all three pairs and on VoxPopuli in two, consistent with a small regularization effect from regressing a smoothed target.The proposed explanation contrasts regression against direct optimization of the recognition objective.Source: §4.3 · p. 7
4.4 Results on Stage-1 Distillation
Stage-1 results are less uniform than stage-0 results: student performance ranges from better than the teacher to the largest degradation observed.
4.4 Results on Stage-1 Distillation
- J2’s student improves average WER from 12.33 to 11.68, a 5.3% relative improvement, with a third of the encoder parameters.The student is better than its teacher on every test set.Source: §4.4 · p. 7
- J1’s student stays within 1.5% relative WER of its teacher.Source: §4.4 · p. 7
- J3’s student is 7.7% relatively worse than its teacher, the largest degradation observed.Source: §4.4 · p. 7
4.5 Ablation Studies
The ablations find an advantage over independent training, a strong relationship between teacher and student quality, and similar transfer across token interfaces. Stage-0 distillation is more reusable, while stage 1 is tied to its joint model.
4.5.1 Distillation vs. Independent Training
- At identical capacity, the distilled student reaches 11.90 average WER versus 12.36 for an independently trained student, a 3.9% relative gain.The frozen-decoder independent-training variant reaches 12.31 average WER; the gap is widest on Earnings-22 and TED-LIUM and nearly closed on clean read speech.Source: §4.5.1 · p. 7
4.5.2 Effect of Teacher Quality
- Across six pairs, teacher WER is the strongest predictor of student WER, and the authors conclude that teacher improvement transfers almost one-for-one.Teacher A is 25% relatively worse than teacher B, and its student inherits nearly the whole gap: 14.70 versus 11.90, a difference of 2.80.Source: §4.5.2 · pp. 7–8
4.5.3 Discrete vs. Continuous Interfaces
- The same pre-quantizer recipe transfers across interfaces: the discrete-interface pair is +0.4% relative, versus +1.9% and +0.8% for the continuous-interface pairs.The reported absolute difference between discrete- and continuous-interface rows reflects teacher differences, since quantization loss is incurred on both sides and cancels.Source: §4.5.3 · p. 8
4.5.4 Effect of the Distillation Stage
- Stage-0 students degrade by +0.4% to +1.9% relative, whereas stage-1 students span −5.3% to +7.7%; stage 0 is more reusable because its representation is language-model-agnostic.The authors default to stage 0 and use stage 1 when the tokenizer must match an already-trained joint model.Source: §4.5.4 · p. 8
5 Conclusion
The paper presents latent-space distillation as a way to compress tokenizer encoders while retaining teacher-comparable recognition performance. The recipe uses the shared pre-quantizer latent across interfaces and teacher training stages.
5 Conclusion
- Across six teacher–student pairs, five students stay within 1.9% relative WER of their teacher without fine-tuning, and the distilled student beats an equal-capacity independently trained tokenizer by 3.9% relative.These pairs cover both distillation stages, both token interfaces, and three joint models.Source: §5 · p. 8
- The pre-quantizer latent supports one distillation recipe across discrete and continuous interfaces and teachers pretrained alone or jointly with a language model.The paper reports that student quality depends chiefly on teacher quality and that the interface does not affect transfer.Source: §5 · p. 8
- A 2.8× smaller encoder can fit the shared on-device memory and latency budget, with lower compute and no downstream retraining.The production configuration differs in several respects from the research configuration; reported WERs use public benchmarks and the research evaluation protocol.Source: §5 · p. 8
A Formal Description of the Objective
The appendix formalizes latent matching between teacher and student encoders, showing how the shared pre-quantizer representation supports either downstream interface. It also specifies masking, width reconciliation, and training-graph choices.
A.2 The encoder as a composition of operators
- Both encoders process 10 ms filterbank frames through a causal subsampler, Conformer stack, and further downsampling; only the student adds an affine layer.The temporal reductions total 8×, yielding one output vector per 80 ms; the student is narrower than every teacher.Source: §A.2 · pp. 11–12
- Both downstream routes consume the encoder latent, so matching that representation supplies a common target for the quantized and continuous interfaces.The student affine layer makes its latent directly comparable with the teacher’s width.Source: §A.2 · p. 12, §A.3 · p. 12
A.3 The masked objective
- The objective is masked squared-error latent matching, optimized over student parameters while the teacher remains fixed.Per-frame differences are zeroed on padded frames, squared, summed over teacher-width channels, then averaged over batch-time positions.Source: §A.3 · p. 12
- The loss sums across channels, so its scale depends on teacher width; that matters when comparing widths or weighting it alongside another loss.The passage recommends gradient-scale normalization for a future multi-term objective.Source: §A.3 · p. 13
- Because padded positions remain in the batch-time denominator, the gradient scale varies with padding ratio and mildly upweights densely packed batches under variable-length bucketing.Under a fixed batching strategy, the padding-ratio factor is constant.Source: §A.3 · p. 13
- An ℓ1 loss could be less sensitive to occasional large frame deviations, but the reported results use squared error and do not compare the alternative.The passage notes that large deviations may occur on frames where the teacher is uncertain.Source: §A.3 · p. 13
A.4 Why matching the latent suffices
- For the continuous interface, the linear bridge bounds the perturbation entering the language model by the latent error; discrete tokens are unchanged below a codebook margin.These are sufficient conditions, not an explanation of measured differences; the discrete condition can fail near Voronoi boundaries.Source: §A.4 · p. 13
A.5 The affine layer
- When widths differ, a single affine map inside the student converts its output to teacher width and can absorb a change of basis.It is omitted when widths match and costs well under 1% of student parameters.Source: §A.5 · pp. 13–14
B Training Configuration
The training configuration is shared across reported runs, and the trainer inherits its data and evaluation setup from the teacher experiment. Downstream evaluation uses the student’s exponential moving average weights.
B Training Configuration
- The trainer inherits the data pipeline, sharding, and evaluation harness by cloning the teacher’s experiment configuration.The model and learner are replaced.Source: B Training Configuration · p. 14
- The distillation training configuration is identical for every reported run.The supplied passage identifies Table 4 but does not state its individual settings.Source: Table 4 · p. 14
- Downstream consumers load the student’s EMA weights; evaluating raw weights will not reproduce the reported numbers.Source: B Training Configuration · p. 14
C The Recipe in Full
The recipe extracts or loads a frozen teacher encoder and trains a student encoder to match its latent on unlabeled batches. Changing the teacher requires only updating its configuration name and checkpoint path.
C The Recipe in Full
- The procedure supports reuse across six teacher–student pairs and three joint models by requiring changes only to the teacher configuration name and checkpoint path.Source: C The Recipe in Full · p. 14
- For a joint tokenizer–language-model checkpoint, the recipe extracts the tokenizer, freezes the teacher, and enables its pre-quantizer output.It instantiates the student without decoders and adds an affine layer only when student and teacher widths differ.Source: C The Recipe in Full · p. 15
- Training draws unlabeled batches and returns the student’s EMA weights rather than its raw weights.Source: C The Recipe in Full · p. 15