understudy-lm

A 109M-parameter draft model for speculative decoding, trained from scratch only on text its target model wrote, with no access to the target's logits or hidden states. This checkpoint's target is Qwen/Qwen3.5-4B; it uses the same tokenizer, so its guesses can be checked by the target directly.

The checkpoint is a standard LlamaForCausalLM and loads with plain transformers. Three extra lookahead heads, stored separately in lookahead_heads.safetensors, let one forward pass guess four tokens.

Code and training pipeline: github.com/SamratDuttaOfficial/understudy-lm

Results

Greedy speculative decoding replayed over 200 held-out target answers (90K rounds, four guesses per round, answer tokens only):

Tokens per target pass	1.58
Accepted guesses per round	0.58
Draft round (4 guesses, CUDA graph)	1.77 ms
Target decode step (Qwen3.5-4B, Ollama, 4-bit)	16.6 ms
Projected speedup	1.43x

The speedup is projected as tokens_per_round / (1 + draft_time / target_time), not measured end to end. Details are in eval_results.json.

Pretraining loss	Acceptance per guess
Tokens produced per round	Projected speedup

Usage

As an assistant model in transformers

The draft's vocabulary size matches the target's, so it works with assisted generation as is. In this mode transformers asks the draft for one token at a time and does not use the lookahead heads.

import torch
from transformers import AutoModelForCausalLM, AutoTokenizer

tok = AutoTokenizer.from_pretrained("Qwen/Qwen3.5-4B")
target = AutoModelForCausalLM.from_pretrained("Qwen/Qwen3.5-4B", dtype=torch.bfloat16, device_map="auto")
draft = AutoModelForCausalLM.from_pretrained("samratduttaofficial/understudy-lm", dtype=torch.bfloat16).to(target.device)

messages = [{"role": "user", "content": "Explain what a hash table is."}]
inputs = tok.apply_chat_template(messages, add_generation_prompt=True, enable_thinking=False,
                                 return_tensors="pt", return_dict=True).to(target.device)
out = target.generate(**inputs, assistant_model=draft, max_new_tokens=256, do_sample=False)
print(tok.decode(out[0, inputs["input_ids"].shape[1]:], skip_special_tokens=True))

Four guesses from one pass, with the lookahead heads

Head j maps the final hidden state h to h + silu(W_j h), and the shared output layer turns it into a guess for token t+1+j.

import torch
import torch.nn.functional as F
from huggingface_hub import hf_hub_download
from safetensors.torch import load_file
from transformers import AutoModelForCausalLM

repo = "samratduttaofficial/understudy-lm"
draft = AutoModelForCausalLM.from_pretrained(repo).eval()
heads = load_file(hf_hub_download(repo, "lookahead_heads.safetensors"))

@torch.no_grad()
def guess4(input_ids):
    h = draft.model(input_ids).last_hidden_state[:, -1]
    hs = [h] + [h + F.silu(h @ heads[f"heads.{j}.weight"].T) for j in range(3)]
    return draft.lm_head(torch.stack(hs, 1)).argmax(-1)   # tokens t+1 .. t+4

On its own the draft writes fluent but often wrong text. It is meant to be checked by the target, not used by itself.

Model

Architecture	Llama-style decoder: RMSNorm, RoPE, grouped-query attention, SwiGLU
Layers / width	8 / 384
Attention heads	6 query, 2 key-value (head size 64)
Feed-forward	1,088
Context	1,024 tokens
Vocabulary	Qwen3.5 tokenizer, 248,320 entries (tied input and output embedding)
Parameters	95.4M embedding + 13.2M transformer + 0.44M lookahead heads

Training

1. Data from the target. The target answered 40,000 prompts drawn from SmolTalk, the Tulu 3 SFT mixture and FineWeb-Edu (temperature 0.7, top-p 0.95, up to 512 tokens, thinking off): 13.4M answer tokens.
2. Corpus. 305M tokens of complete SmolTalk and Tulu 3 conversations, excluding those whose prompt was used in step 1.
3. Pretraining. 12,000 steps of 65,536 tokens, 75% corpus and 25% target answers. AdamW, peak learning rate 3e-4, cosine decay. Validation loss on the target's held-out answers fell to 3.16. An earlier run trained on the target's answers alone overfit after about 13 passes (best 3.68).
4. Post-training. Three lookahead heads added, 3,000 steps at 1e-4 on the target's answers only (prompt tokens masked), loss L1 + 0.4 * (L2 + L3 + L4).

Measured in llama.cpp

In llama.cpp (the build bundled with Ollama 0.40.8, --spec-type draft-simple, up to four drafted tokens, greedy, three prompts, RTX A3000), the draft averaged 1.57 to 1.82 tokens per target pass, in line with the replay above. Yet throughput fell from 69 tok/s with the target alone to 27 to 36 tok/s with the draft. llama.cpp drafts one token per draft pass and does not use the lookahead heads, and Qwen3.5's hybrid layers cannot drop rejected tokens from their state, so the server restores a checkpoint after each rejection. The projected speedup needs a runtime that uses the heads and rolls back cheaply. Ollama 0.40.8 does not enable a plain draft model like this one.

Limitations

• The 1.43x figure is projected from replayed acceptance and measured draft and target timings. The one end-to-end measurement so far, in llama.cpp, was slower than the target alone (see above).
• It assumes checking five positions costs the target about as much as one decode step. If it costs 10% more, the projection drops to about 1.31x.
• Results cover greedy decoding. With sampling, acceptance is usually lower.
• Acceptance was measured against sampled (temperature 0.7) answers; against the target's greedy answers it is slightly higher.

License

Apache-2.0.