jburtoft/Voxtral-Mini-3B-2507-draft-4layer

Voxtral-Mini-3B-2507 -- 4-layer distilled draft for speculative decoding

A shallow distilled draft of mistralai/Voxtral-Mini-3B-2507, intended as an assistant/draft model for speculative decoding. Not intended for standalone use.

Not a general-purpose ASR model. This model is a much smaller, faster imitator of the full Voxtral-Mini-3B-2507 target. It reproduces the target's greedy output well enough to serve as a draft in speculative decoding, but its standalone transcripts are lower quality than the target's. Always pair it with the target inside a spec-decode loop.

What it is

• Architecture: identical to Voxtral-Mini-3B-2507 in every component except the LLM decoder is pruned from 30 layers to 4 layers (kept: indices [0, 10, 20, 29]).
• Shared verbatim with target (frozen, not trained): audio encoder (32 layers, Whisper-derived), multi_modal_projector, embed_tokens, final norm, lm_head. Same tokenizer, same audio path, same output vocab.
• Trained: only the 4 kept Llama decoder layers (~428M trainable params of the ~3B total).
• Compute: 4-layer decoder is ~13% of the target's decoder FLOPs (4/30).

Results

Measured on LibriSpeech dev-other (public, CC BY 4.0): 100 randomly sampled clips (seed 0), 1.8-21.2 s each, 10.4 minutes total, 2,155 target tokens. The draft was trained on dev-clean + test-clean + test-other, so these clips are held out. They are, however, the same domain as the training data (read English audiobooks). Expect lower agreement on conversational, accented, noisy or domain-specific audio, and measure on your own data before relying on these numbers.

Method. Target greedy transcript (bf16, CPU), then a teacher-forced forward pass of both models over the prompt + transcript. Agreement is the fraction of transcript positions where the draft's argmax equals the target's argmax. For greedy speculative decoding this match sequence fully determines how many draft tokens are accepted, so the tokens-per-iteration figures below are computed directly from it (up to bf16 numerical differences between a single teacher-forced pass and incremental decoding). They are not wall-clock speedups, which depend on the draft/target cost ratio of your serving stack.

Metric	Value
Top-1 agreement (token-weighted)	78.4% (95% bootstrap CI over clips: 75.8-80.9%)
Top-1 agreement (mean per clip / median per clip)	80.1% / 81.8%
Per-clip agreement, 10th-90th percentile	65.5%-94.7%

Expected target tokens produced per target forward pass (greedy speculative decoding with K draft tokens per iteration; 1.0 = no speculation):

K	1	2	3	4	6	8
tokens / iteration	1.72	2.25	2.72	3.05	3.59	4.01

• Byte-identical greedy correctness: yes, by construction. The target verifies every draft token, so the final output always equals the target's greedy output.
• Choosing K: larger K means fewer target passes but more draft work per iteration. The best K depends on the cost of one draft step relative to one target pass on your hardware; sweep it on your deployment.
• Reproduce / measure on your data: training_scripts/measure_agreement.py --root <dir of .flac files> --n 100 --out results.jsonl (CPU; about 1 minute per clip). training_scripts/validate_draft.py does the same for a single clip.

Intended use

from transformers import VoxtralForConditionalGeneration, AutoProcessor
import torch

target = VoxtralForConditionalGeneration.from_pretrained(
    "mistralai/Voxtral-Mini-3B-2507",
    dtype=torch.bfloat16,
).cuda().eval()
draft = VoxtralForConditionalGeneration.from_pretrained(
    "jburtoft/Voxtral-Mini-3B-2507-draft-4layer",
    dtype=torch.bfloat16,
).cuda().eval()
processor = AutoProcessor.from_pretrained("mistralai/Voxtral-Mini-3B-2507")

inputs = processor.apply_transcription_request(
    language="en",
    audio="my_clip.wav",
    model_id="mistralai/Voxtral-Mini-3B-2507",
).to("cuda", dtype=torch.bfloat16)

# Standard HuggingFace speculative decoding
out = target.generate(
    **inputs,
    max_new_tokens=256,
    do_sample=False,
    assistant_model=draft,
)
transcript = processor.tokenizer.decode(
    out[0, inputs["input_ids"].shape[1]:], skip_special_tokens=True,
)

Known caveat: as of transformers==5.14.1, HF's built-in assistant_model= path exhibits inefficiencies with VoxtralForConditionalGeneration (multimodal target). The final output is byte-identical to greedy baseline, but per-iteration acceptance-rate reporting is incorrect and wall-clock speedup is often below the theoretical ceiling. Use training_scripts/validate_draft.py to measure the draft's teacher-forced agreement with the target directly; a correctly-implemented spec-decode loop is needed to realize the speedup that agreement implies.

The intended deployment target for this draft is AWS Trainium.

Standalone quality

The draft is not intended for standalone transcription. Run alone, it produces topically related, grammatically English output with word substitutions relative to the target. This is a fundamental property of a highly-shrunk distilled model. Inside a spec-decode loop the target's verification pass corrects these mispredictions, so the final output is byte-identical to the target's greedy output.

Training

Dataset

• LibriSpeech train subsets (dev-clean + test-clean + test-other, public, CC BY 4.0). Sampled 2000 clips, mean duration 6.95 s, total ~230 minutes of English audio.
• Pseudo-labeling: for each clip, ran Voxtral-Mini-3B-2507 end-to-end greedy on CPU to obtain the "teacher transcript" token IDs. These are the distillation targets (not the ground-truth LibriSpeech transcripts). Rationale: we want the draft to match the target's idiosyncrasies (its capitalization, punctuation, word choice), not to be a good ASR model on its own.

Distillation recipe

• Framework: torch-neuronx custom training loop on AWS Trainium (trn2.3xlarge, LNC=2, TP=1).
• Loss: cross-entropy on teacher's argmax tokens (teacher forcing). No KL on soft logits, no auxiliary losses.
• Optimizer: AdamW, betas (0.9, 0.95), weight decay 0.01.
• Learning rate: peak 1e-4, linear warmup 100 steps, cosine decay to 10% of peak over remaining 2900 steps.
• Batch size: 2. Max sequence length: 512 (padded).
• Precision: bfloat16 model, fp32 loss computation.
• Gradient clipping: max norm 1.0.
• Total steps: 3000. Approximately 3 epochs over 2000 clips at batch=2.
• Frozen: all params outside the 4 kept decoder layers.

Convergence

Validation top-1 agreement rose steadily over training and plateaued in the last ~600 of 3000 steps; the best checkpoint (step 2800) is the one published here. The training data appears to be the bottleneck at this scale; training on more LibriSpeech clips (e.g. train-clean-100) is the most likely way to raise acceptance. Held-out results are in Results.

Trainium-specific workaround (XLA path only)

What breaks: On the torch-neuronx XLA path (the default), Voxtral's masked_scatter call in VoxtralModel.forward -- which splices the audio-encoder output into inputs_embeds at the audio-token positions -- lowers to an HLO sort op on a s32[1176576] index tensor (for the standard [1, 383, 3072] inputs_embeds shape). Neuron's XLA backend rejects this with [NCC_EVRF029] Operation sort is not supported on trn2. The compiler's suggestion is to use TopK or an NKI kernel instead, but no such lowering exists in torch_xla for masked_scatter at this size as of neuronx-cc 2.26.6360.0+6f180f47.

What is fine: The audio encoder (audio_tower + multi_modal_projector) compiles cleanly on Neuron via both the XLA path and PyTorch Native (TorchNeuron) eager mode. This is not a hardware or model-architecture limitation.

Workaround (used by the training scripts here): precompute audio embeddings on CPU once per training clip, splice them into inputs_embeds also on CPU in the DataLoader collate function, and only send the audio-conditioned inputs_embeds to Neuron. The audio encoder + projector never runs on Neuron in the training loop. Since the audio path is frozen (never trained), the audio-encoder forward pass only needs to run once per clip anyway, so moving it to CPU costs nothing in the amortized case.

What PyTorch Native (TorchNeuron) does differently: in eager mode, masked_scatter dispatches directly to ATen kernels on the Neuron device, with no XLA compilation involved, so Voxtral runs end-to-end on Neuron without any CPU workaround: full model(**inputs) including audio encoder + projector + splice + 30-layer decoder runs in one call on torch.device("neuron:0"). The XLA failure mode does not occur there.

For training on either stack, we still use the CPU-precompute workaround: it's more efficient than re-running the frozen audio encoder every training step, and it makes the training loop's static-shape graph simpler (no variable-length audio inside the compiled region). See training_scripts/README.md for details.

For inference on PyTorch Native, no workaround is needed -- move the whole model to neuron:0 and call .generate(**inputs) normally.

Reproducibility

Training scripts and full recipe: see the training_scripts/ directory alongside this model card. Also linked from the file browser above.

Quick reproduction on a fresh trn2.3xlarge (SDK 2.31 DLAMI 20260708 or later):

# --- 0. Setup
source /opt/aws_neuronx_venv_pytorch_2_9/bin/activate
pip install -q -U "transformers>=4.54" "mistral-common[audio]>=1.8.1" librosa soundfile

# --- 1. Sanity check that the LLM decoder loads on trn2
python training_scripts/sanity_voxtral_xla_v2.py

# --- 2. Download LibriSpeech dev-clean (~340 MB) + test-clean + test-other (~700 MB)
cd /mnt/data
for split in dev-clean test-clean test-other; do
  curl -sSL -o ${split}.tar.gz https://www.openslr.org/resources/12/${split}.tar.gz
  tar xzf ${split}.tar.gz
done

# --- 3. Build a training manifest (500 dev-clean clips)
cd - && cd training_scripts
python build_librispeech_manifest.py \
  --librispeech-root /mnt/data/LibriSpeech/dev-clean \
  --output /mnt/data/manifest_train.csv --limit 500

# Optional: combine multiple splits for more data
python build_librispeech_combined_manifest.py \
  --librispeech-roots /mnt/data/LibriSpeech/test-clean /mnt/data/LibriSpeech/test-other \
  --output /mnt/data/manifest_extra.csv --limit 1500 \
  --exclude-existing-dir /mnt/data/pseudo_labels_train

# --- 4. Precompute pseudo-labels (~20 min per 500 clips on CPU)
python precompute_pseudo_labels.py \
  --manifest /mnt/data/manifest_train.csv --dataset-root / \
  --output-dir /mnt/data/pseudo_labels_train

# --- 5. Precompute a small validation set of your own audio for validation
#     (any short English audio clips work; ~10-20 clips is sufficient)
python precompute_pseudo_labels.py \
  --manifest /path/to/your/val_manifest.csv --dataset-root /path/to/val \
  --output-dir /mnt/data/pseudo_labels_val

# --- 6. Train (~16 minutes on trn2.3xlarge for 3000 steps at batch=2)
python train_distill_v2.py \
  --train-data-dir /mnt/data/pseudo_labels_train \
  --val-data-dir /mnt/data/pseudo_labels_val \
  --output-dir /mnt/drafts/my-voxtral-draft \
  --keep-layers 0,10,20,29 \
  --max-seq-len 512 --batch-size 2 --lr 1e-4 \
  --warmup-steps 100 --total-steps 3000 \
  --save-every 500 --validate-every 200

# --- 7. Validate on one clip
python validate_draft.py \
  --draft /mnt/drafts/my-voxtral-draft/best \
  --audio /mnt/data/LibriSpeech/dev-clean/2035/147960/2035-147960-0015.flac

Expected outcome: validation top-1 agreement climbs steadily and plateaus in the last few hundred steps, and standalone draft output on any English clip is coherent (not gibberish).

Files in this repo

File	Purpose
model.safetensors	4-layer draft weights (~3.6 GB in bfloat16)
config.json	Voxtral config with text_config.num_hidden_layers=4 and pruned layer_idx values on the 4 kept layers
generation_config.json	Copied from target
processor_config.json, tekken.json	Copied from target (same tokenizer + processor)
training_state.pt	Training state at best-val step: losses, val-top1 history, hyperparameters
training_scripts/*.py	Full reproduction pipeline (see Reproducibility)
training_scripts/measure_agreement.py	Held-out agreement + expected tokens/iteration (the Results measurement)
training_scripts/README.md	Detailed script-by-script guide

Limitations

1. Not a standalone ASR model. See "Standalone quality" above.
2. Small training corpus. Only 2000 LibriSpeech clips (~230 minutes English audio). Domain shift for non-English or non-conversational audio (medical, financial, etc.) will lower acceptance rate. Retraining on domain-matched pseudo-labels is straightforward and recommended.
3. English only. Trained on LibriSpeech English. Untested on other languages. Voxtral supports multilingual audio; the draft's other-language behavior is undefined.
4. HuggingFace assistant_model= inefficiency. Documented above; use a custom spec-decode loop for best results.
5. 4-layer choice was not swept. Alternative layer selections (5-6 layers, different indices) may give better acceptance/cost tradeoffs. Reproducibility recipe supports easy re-training with --keep-layers.

Citation

If you use this draft, please cite the base Voxtral model:

@misc{voxtral2025,
  title={Voxtral},
  author={Mistral AI},
  year={2025},
  url={https://huggingface.co/mistralai/Voxtral-Mini-3B-2507}
}

And LibriSpeech for the training corpus:

@inproceedings{panayotov2015librispeech,
  title={Librispeech: an ASR corpus based on public domain audio books},
  author={Panayotov, Vassil and Chen, Guoguo and Povey, Daniel and Khudanpur, Sanjeev},
  booktitle={2015 IEEE international conference on acoustics, speech and signal processing (ICASSP)},
  pages={5206--5210},
  year={2015},
  organization={IEEE}
}

License

Apache 2.0. This model is a distilled variant of mistralai/Voxtral-Mini-3B-2507 (Apache 2.0), retaining the base model's tokenizer, config schema, audio encoder weights, projector weights, embed_tokens, lm_head, and final norm. Only the 4 kept decoder layers were trained by us on LibriSpeech.