Whisper large-v3 Medusa Heads (N=5) for AWS Trainium

Medusa-1 speculative-decoding heads for openai/whisper-large-v3, trained on AWS Trainium (PyTorch Native, torch-neuronx). These heads let the Whisper decoder propose several future tokens per step, which a verify pass then accepts/rejects — reducing the number of sequential decoder calls at batch size 1 while producing byte-identical greedy output.

What this is: a small set of prediction heads that sit on top of a frozen whisper-large-v3. They add no new vocabulary and change no transcripts — they only speculate. Output is bit-for-bit identical to greedy whisper-large-v3 regardless of how accurate the heads are (the verify step rejects wrong guesses).

Status: a lightly-trained reference / demonstration head, published to make the speculative-decoding framework testable end-to-end. It is not an accuracy-tuned production head — see Limitations.

At a glance

How it works

Standard autoregressive decoding emits one token per decoder forward pass. Medusa adds N lightweight heads that read the decoder's last hidden state at position t and each predict a future token (head_i → token t + i + 2). A single verify pass then runs the real decoder over the proposed tokens and keeps the longest prefix that matches what greedy decoding would have produced, rejecting at the first mismatch.

Because every accepted position equals the target model's own argmax and decoding rejects at the first disagreement, the emitted text is identical to plain greedy whisper-large-v3 — the heads only change how many decoder calls it takes, never what tokens come out. The heads are Medusa-1: their output projection is tied to the frozen proj_out (LM head), so they add only the small per-head ResBlock.

Data

• Dataset: LibriSpeech (public, CC BY 4.0) — English read speech. Audio only.
• Pseudo-labels. The heads are trained to predict whisper-large-v3's own greedy continuation (obtained by decoding each clip with the frozen model), not the reference transcript. Speculative acceptance measures "does the head predict what whisper-large-v3 would say next," so the training target is exactly that.
• Head-0 top-1 accuracy of this checkpoint: ~40.6% on held-out audio.

Train your own heads on AWS Trainium

Both scripts are in this repo: harvest_whisper_pseudo.py (build the pseudo-label dataset) and train_medusa_heads_whisper.py (train the heads on Trainium). Run on a trn2 instance with the AWS Neuron PyTorch Native environment.

1 — Download the scripts and this checkpoint:

pip install -U "huggingface_hub[cli]"
huggingface-cli download jburtoft/whisper-large-v3-medusa-heads \
  harvest_whisper_pseudo.py train_medusa_heads_whisper.py \
  medusa_heads_whisper_large_v3.pt \
  --local-dir ./whisper-medusa
cd ./whisper-medusa

2 — Export the frozen LM-head (proj_out) weight the heads tie to (Medusa-1):

python -c "from transformers import WhisperForConditionalGeneration as W; import torch; \
m=W.from_pretrained('openai/whisper-large-v3'); \
torch.save(m.proj_out.weight.detach().cpu(), 'proj_out.pt')"

3 — Harvest pseudo-labels (frozen whisper-large-v3 greedy decode → (hidden, token) pairs). Uses LibriSpeech by default; runs on CPU/GPU:

python harvest_whisper_pseudo.py \
  --out-dir ./data \
  --dataset openslr/librispeech_asr \
  --n-samples 1000

4 — Train the heads on Trainium:

python train_medusa_heads_whisper.py \
  --data-dir ./data \
  --proj-out proj_out.pt \
  --out-dir ./ckpt \
  --n-heads 5 \
  --lr 3e-4 \
  --max-steps 4000 \
  --batch-size 8192

Validation top-1 accuracy per head is logged every --val-every steps; the best checkpoint is written to ./ckpt/best_heads.pt in the format described below. More data and steps raise acceptance (and the achievable speedup).

Files

Checkpoint format

State dict keyed medusa_head_{i}.{0}.linear.{weight,bias} for i in 0..4 (each weight is [1280, 1280], bias [1280], bf16), wrapped as:

{"heads": <state_dict>, "step": 1800, "head0_acc": 0.4065}

No output-projection tensors are stored — the heads are Medusa-1 and tie to the base model's proj_out weight at load time.

Loading the heads

import torch

ckpt = torch.load("medusa_heads_whisper_large_v3.pt", map_location="cpu")
heads_sd = ckpt["heads"]          # medusa_head_{i}.0.linear.{weight,bias}
print(ckpt["step"], ckpt["head0_acc"])   # 1800 0.4065

# Each head_i: h -> h + SiLU(Linear(1280,1280)(h)), then tie to the frozen
# whisper-large-v3 proj_out (LM head) for the vocab projection (Medusa-1).

The heads are designed to plug into a Whisper Medusa speculative-decoding runner that reads the decoder's last hidden state and runs the verify/accept step. Wiring depends on your serving stack; the training script documents the exact head module (ResBlock / MedusaHeads) and the shift = i + 2 convention so the heads can be reproduced or re-integrated.

Intended use

• Speculative decoding for whisper-large-v3 at batch size 1, to reduce single-stream transcription latency without changing the transcript.
• As a reference implementation / starting checkpoint for training your own, higher-acceptance Whisper Medusa heads (fine-tune on more of your own audio's pseudo-labels).

Limitations

• Lightly trained (demonstration). Head-0 accuracy ~40.6% and mean ~1.63 accepted tokens per verify. Whether this yields a net wall-clock speedup depends on your decoder/verify implementation's per-step overhead: you need mean-accepted-per-verify to exceed your verify-step / decode-step cost ratio. A more heavily trained head (more data, more steps) raises acceptance and the achievable speedup.
• Greedy. Trained and validated for greedy decoding. Sampling modes are not evaluated here.
• English, read speech. Trained on English pseudo-labels from read-speech audio. Acceptance will vary on other languages, noisy, or conversational audio.
• Correctness is guaranteed by the verify step, not by head accuracy. Low-accuracy heads simply get rejected more often (less speedup) — they never corrupt the output.

License

Apache-2.0. Base model openai/whisper-large-v3 is MIT-licensed; LibriSpeech is CC BY 4.0.