---
license: apache-2.0
base_model: mistralai/Voxtral-Mini-3B-2507
base_model_relation: quantized
tags:
- audio
- automatic-speech-recognition
- speculative-decoding
- voxtral
- neuron
- trainium
- distillation
language:
- en
library_name: transformers
pipeline_tag: automatic-speech-recognition
datasets:
- librispeech_asr
---

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

A shallow distilled draft of [`mistralai/Voxtral-Mini-3B-2507`](https://huggingface.co/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. Practical latency ratio measured on GPU L40S: **draft is ~4.4x faster per generated token** than the target.

## Results

Reference measurements on a held-out corpus of 18 short-to-medium English audio clips (0.4-27 seconds, TED-derived speech), never seen during training:

- **Mean top-1 argmax agreement with target**: 66.67%. Equivalent to the draft's per-position "would this token match the target's greedy pick?" rate under teacher forcing.
- **Analytical GPU speedup at K=2 candidates/iteration**: 1.416x mean (1.381x on clips >= 15 s).
- **Byte-identical greedy correctness**: yes, by construction of the greedy speculative-decoding algorithm.

Per-K analytical speedup on L40S (target: 24.2 ms/tok; draft: 5.5 ms/tok):

| K | mean accepted per iter | mean speedup | speedup (long clips >= 15s) |
|---|---:|---:|---:|
| 2 | 1.13 | **1.416x** | **1.381x** |
| 3 | 1.37 | 1.364x | 1.351x |
| 4 | 1.57 | 1.300x | 1.280x |
| 6 | 1.80 | 1.148x | 1.122x |
| 8 | 1.95 | 1.017x | 0.992x |

K=2 is optimal on GPU L40S given the ~4.4x cost ratio; larger K adds draft cost that isn't repaid by the incremental acceptance gain. On different hardware (e.g., AWS Neuron/Trainium with a fused-K draft NEFF), the optimal K will differ.

## Intended use

```python
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. The training-time analytical measurement (66.67% top-1 agreement) reflects the model's true quality; a correctly-implemented spec-decode wrapper should approach the K=2 -> 1.4x speedup shown above. See `measurement/README.md` in the [companion training scripts](#reproducibility) for a working analytical measurement approach.

The intended production deployment for this draft is on **AWS Neuron/Trainium** via a custom NxD Inference spec-decode wrapper (planned separately). GPU is useful for validation but is not the target platform.

## Standalone quality (illustrative)

The draft is **not** intended for standalone transcription. If run alone on a 9.67 s clip about nuclear reactor refueling, it emits:

> "That's right, you're always refiling the revolution. So I have lots of people and lots of conditions that can go wrong. That thing where you'm opening things in and moving things in and move."

The target on the same clip emits:

> "That's right. Today, you're always refueling the reactor. So you have lots of people and lots of controls that can go wrong. That thing where you're opening it up and moving things in and out, that's..."

The draft's output is topically related and grammatically English but has word substitutions ("refueling the reactor" -> "refiling the revolution", "controls" -> "conditions"). This is a fundamental property of a highly-shrunk distilled model. When used inside a spec-decode loop, the target's verification pass corrects these mispredictions on the fly, 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, ~5-day capacity block; actual training used ~20 minutes).
- **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

Held-out validation top-1 agreement (measured every 200 steps on a validation set of unrelated short English audio clips):

| Step | Val top-1 agreement |
|---:|---:|
| 200 | 12.82% |
| 400 | 19.23% |
| 600 | 25.05% |
| 800 | 31.36% |
| 1000 | 36.59% |
| 1200 | 41.62% |
| 1400 | 49.01% |
| 1600 | 51.78% |
| 1800 | 53.75% |
| 2000 | 57.40% |
| 2200 | 62.43% |
| 2400 | 64.40% |
| 2600 | 61.93% |
| **2800** | **66.67%** (best) |
| 3000 | 65.78% |

Monotonically increasing until ~step 2400; plateau at 62-67% for the last 600 steps. The training data appears to be the bottleneck at this scale; scaling to ~10k+ LibriSpeech clips should push val top-1 into the 75-85% range and improve downstream spec-decode speedup accordingly.

### 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 and Beta 4 eager paths. 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 Beta 4 does differently:** In eager mode Beta 4 dispatches `masked_scatter` directly to ATen kernels on the Neuron device, no XLA compilation involved. It runs natively and Voxtral works end-to-end on Neuron without any CPU workaround. Verified 2026-07-31: full `model(**inputs)` including audio encoder + projector + splice + 30-layer decoder runs in one call on `torch.device("neuron:0")`, ~1150 ms steady state. The same failure mode does NOT reproduce on Beta 4.

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 Beta 4, 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](./training_scripts) 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):

```bash
# --- 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: val top-1 climbs from ~4% (step 200) to ~65-68% (step 2800), 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](#reproducibility)) |
| `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 on GPU.** Documented above; use a custom spec-decode loop or wait for Neuron NxD Inference integration.
5. **4-layer choice was not swept.** Alternative layer selections (5-6 layers, different indices) may give better accepance/cost tradeoffs. Reproducibility recipe supports easy re-training with `--keep-layers`.

## Citation

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

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

And LibriSpeech for the training corpus:

```bibtex
@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.