samratduttaofficial

understudy-lm

README

License: apache-2.0

Results

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

Table

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.

Table with columns: Pretraining loss, Acceptance per guess
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.

python

import torchfrom 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.

python

import torchimport torch.nn.functional as Ffrom huggingface_hub import hf_hub_downloadfrom safetensors.torch import load_filefrom 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

Table

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

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.
Corpus. 305M tokens of complete SmolTalk and Tulu 3 conversations, excluding those whose prompt was used in step 1.
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).
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.

Available on FriendliAI

Dedicated Endpoints

Run this model inference on single tenant GPU with unmatched speed and reliability at scale.

Learn more

Model Details

Model Provider

samratduttaofficial

Model Tree

Base

this model

Input Modalities

Text

Output Modalities

Text

Supported Functionality

Dedicated Endpoints

Explore FriendliAI today

Get started Talk to an engineer

README

License: apache-2.0

Results

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

Table

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.

Table with columns: Pretraining loss, Acceptance per guess
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.

python

import torchfrom 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))

import torchfrom 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.

python

import torchimport torch.nn.functional as Ffrom huggingface_hub import hf_hub_downloadfrom safetensors.torch import load_filefrom 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

import torchimport torch.nn.functional as Ffrom huggingface_hub import hf_hub_downloadfrom safetensors.torch import load_filefrom 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

Table

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

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.
Corpus. 305M tokens of complete SmolTalk and Tulu 3 conversations, excluding those whose prompt was used in step 1.
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).
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

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.