Trace v1: a draft model trained on the target's probabilities

Trace is our line of draft models: small models trained to predict one large one, so speculative decoding accepts more of their guesses. Trace v0 (note 15) imitated the 7B's text and gained about 3% acceptance. Trace v1 learns from the 7B's probabilities instead.

Prediction (docs/predictions.md, committed before the runs): held-out top-1 agreement plateaus between 0.78 and 0.80, clearing the +3 point bar on the full 35,000 prompts but not on the first 8,000. The partial-data half held (no run on 9,000 prompts cleared the bar); the plateau did not reach 0.78, and the full-data half is untested.

Method

The 7B Q4_0 answered each training prompt greedily, and llama-server recorded its top 20 candidates and their log-probabilities at every token (experiments/drafter/teacher.py). The draft trains on those distributions, renormalized, with positions where the stock draft's top guess differs from the 7B's weighted 3 times (train_trace.py, ALPHA=2). AdamW, one epoch, held-out top-1 agreement on 180 validation prompts every 200 steps. The trainer reproduces the stock draft's agreement on both machines (0.7606 M5 Pro, 0.7616 Strix Halo).

A second method, on-policy distillation (train_onpolicy.py): the draft generates, and the bf16 7B scores the draft's own tokens, with total-variation loss (as in DistillSpec, Zhou et al. 2023).

Results: agreement

At 1e-5 the first checkpoint fell below stock (0.7564 at step 200) and stayed there, as the MLX trainer did on the Mac. At 3e-6 every checkpoint beat stock and the curve was still rising slowly at the end of the epoch. On-policy declined steadily; the likely reason is that it distilled from the bf16 7B while agreement is measured against the Q4_0 7B's answers.

Results: speculative decoding

The best v1 (off-policy, LR 3e-6, 0.7685) and the stock draft, both converted to Q4_0 through the same pipeline, against plain decoding on 40 held-out prompts (20 chat, 20 code), 256 tokens, two passes with the second reversed, chip (APU package) energy, idle 16.3 W. Passes agreed within 1.3%, except plain decoding at 6%.

1. Higher agreement, no faster

Trace v1's 0.7 points of agreement show up as higher acceptance (0.734 against 0.716 at four tokens, 0.680 against 0.638 at eight), and buy nothing else. At four tokens it is 3% slower and costs 3% more energy per token than the stock draft; at eight it is level. A draft's acceptance has to rise by more than this before fewer target passes pay for its own cost.

2. On this machine, speculation saves energy, not just time

With either draft, speculative decoding is 1.5 times faster than plain decoding and costs 34% less chip energy per token (1,502 against 2,287 mJ). On the M5 Pro it saved time but cost 8 to 11% more chip energy on chat (note 15). The Strix Halo package draws about the same power either way, 116 W plain and 118 W speculating (energy per token times tokens per second), because decode here waits on memory; finishing in fewer passes is nearly all saving.

What to take from this

1. Off-policy distillation from the target's probabilities gave +0.7 points of agreement on 9,000 prompts at a low learning rate, and lost agreement at a higher one. It did not make a faster draft. Trace v1 does not ship; scaling to all 35,000 prompts is the next test.
2. On-policy distillation as run here moved the draft away from the target, most likely because the teacher was the bf16 7B and the target the Q4_0 7B. A fair test needs the quantized target's own probabilities.
3. On the Strix Halo, speculative decoding with any reasonable draft cuts energy per token by a third. Where decode waits on memory, speculation is an energy technique.

Limits

One target, one draft size. Agreement is measured on 180 held-out prompts. The 35,000-prompt run is not done yet. On-policy used a bf16 teacher against a Q4_0 target. Strix Halo energy is the whole APU package.

Reproduce

python experiments/drafter/teacher.py results/drafter/train-prompts.jsonl data/drafter/teacher.jsonl
LR=3e-6 ALPHA=2 python experiments/drafter/train_trace.py data/drafter/teacher.jsonl models/trace-0.5b-v1
sh experiments/drafter/convert.sh models/trace-0.5b-v1 models/quant/drafter-v1
python experiments/drafter/evaluate.py results/drafter/eval-v1-halo.json plain,stock-4-p0.5,v1-4-p0.5,stock-8-p0.5,v1-8-p0.5

Raw results: results/drafter/eval-v1-halo.json, results/drafter/trace-v1.md.

Source, scripts and raw data: phase/notes/2026-10-trace-v1.md