Release understudy-lm v1.0 weights

.gitattributes

@@ -33,3 +33,4 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+tokenizer.json filter=lfs diff=lfs merge=lfs -text

README.md

@@ -0,0 +1,138 @@
+---
+language:
+- en
+library_name: transformers
+pipeline_tag: text-generation
+tags:
+- speculative-decoding
+- draft-model
+- assistant-model
+- qwen3.5
+- llama
+datasets:
+- HuggingFaceTB/smoltalk
+- allenai/tulu-3-sft-mixture
+- HuggingFaceFW/fineweb-edu
+---
+
+# understudy-lm
+
+A 109M-parameter draft model for speculative decoding with
+[Qwen/Qwen3.5-4B](https://huggingface.co/Qwen/Qwen3.5-4B). It was trained from
+scratch, only on text the 4B wrote, with no access to the 4B's logits or
+hidden states. It uses the Qwen3.5 tokenizer, so its guesses can be checked
+by the 4B 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](https://github.com/SamratDuttaOfficial/understudy-lm)
+
+## Results
+
+Greedy speculative decoding replayed over 200 held-out Qwen3.5-4B answers
+(90K rounds, four guesses per round, answer tokens only):
+
+| | |
+|---|---|
+| Tokens per 4B pass | **1.58** |
+| Accepted guesses per round | 0.58 |
+| Draft round (4 guesses, CUDA graph) | 1.77 ms |
+| 4B decode step (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 |
+|---|---|
+| ![Loss](assets/val_loss.png) | ![Acceptance](assets/acceptance.png) |
+| **Tokens produced per round** | **Projected speedup** |
+| ![Round tokens](assets/round_tokens.png) | ![Speedup](assets/speedup.png) |
+
+## Usage
+
+### As an assistant model in transformers
+
+The draft's vocabulary size matches the 4B'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 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`.
+
+```python
+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 4B, 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.** Qwen3.5-4B 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)`.
+
+## Limitations
+
+- The speedup is projected from replayed acceptance and measured draft and
+  target timings. It has not been measured end to end.
+- It assumes checking five positions costs the 4B 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 4B's greedy answers it is slightly higher.

chat_template.jinja

@@ -0,0 +1,154 @@
+{%- set image_count = namespace(value=0) %}
+{%- set video_count = namespace(value=0) %}
+{%- macro render_content(content, do_vision_count, is_system_content=false) %}
+    {%- if content is string %}
+        {{- content }}
+    {%- elif content is iterable and content is not mapping %}
+        {%- for item in content %}
+            {%- if 'image' in item or 'image_url' in item or item.type == 'image' %}
+                {%- if is_system_content %}
+                    {{- raise_exception('System message cannot contain images.') }}
+                {%- endif %}
+                {%- if do_vision_count %}
+                    {%- set image_count.value = image_count.value + 1 %}
+                {%- endif %}
+                {%- if add_vision_id %}
+                    {{- 'Picture ' ~ image_count.value ~ ': ' }}
+                {%- endif %}
+                {{- '<|vision_start|><|image_pad|><|vision_end|>' }}
+            {%- elif 'video' in item or item.type == 'video' %}
+                {%- if is_system_content %}
+                    {{- raise_exception('System message cannot contain videos.') }}
+                {%- endif %}
+                {%- if do_vision_count %}
+                    {%- set video_count.value = video_count.value + 1 %}
+                {%- endif %}
+                {%- if add_vision_id %}
+                    {{- 'Video ' ~ video_count.value ~ ': ' }}
+                {%- endif %}
+                {{- '<|vision_start|><|video_pad|><|vision_end|>' }}
+            {%- elif 'text' in item %}
+                {{- item.text }}
+            {%- else %}
+                {{- raise_exception('Unexpected item type in content.') }}
+            {%- endif %}
+        {%- endfor %}
+    {%- elif content is none or content is undefined %}
+        {{- '' }}
+    {%- else %}
+        {{- raise_exception('Unexpected content type.') }}
+    {%- endif %}
+{%- endmacro %}
+{%- if not messages %}
+    {{- raise_exception('No messages provided.') }}
+{%- endif %}
+{%- if tools and tools is iterable and tools is not mapping %}
+    {{- '<|im_start|>system\n' }}
+    {{- "# Tools\n\nYou have access to the following functions:\n\n<tools>" }}
+    {%- for tool in tools %}
+        {{- "\n" }}
+        {{- tool | tojson }}
+    {%- endfor %}
+    {{- "\n</tools>" }}
+    {{- '\n\nIf you choose to call a function ONLY reply in the following format with NO suffix:\n\n<tool_call>\n<function=example_function_name>\n<parameter=example_parameter_1>\nvalue_1\n</parameter>\n<parameter=example_parameter_2>\nThis is the value for the second parameter\nthat can span\nmultiple lines\n</parameter>\n</function>\n</tool_call>\n\n<IMPORTANT>\nReminder:\n- Function calls MUST follow the specified format: an inner <function=...></function> block must be nested within <tool_call></tool_call> XML tags\n- Required parameters MUST be specified\n- You may provide optional reasoning for your function call in natural language BEFORE the function call, but NOT after\n- If there is no function call available, answer the question like normal with your current knowledge and do not tell the user about function calls\n</IMPORTANT>' }}
+    {%- if messages[0].role == 'system' %}
+        {%- set content = render_content(messages[0].content, false, true)|trim %}
+        {%- if content %}
+            {{- '\n\n' + content }}
+        {%- endif %}
+    {%- endif %}
+    {{- '<|im_end|>\n' }}
+{%- else %}
+    {%- if messages[0].role == 'system' %}
+        {%- set content = render_content(messages[0].content, false, true)|trim %}
+        {{- '<|im_start|>system\n' + content + '<|im_end|>\n' }}
+    {%- endif %}
+{%- endif %}
+{%- set ns = namespace(multi_step_tool=true, last_query_index=messages|length - 1) %}
+{%- for message in messages[::-1] %}
+    {%- set index = (messages|length - 1) - loop.index0 %}
+    {%- if ns.multi_step_tool and message.role == "user" %}
+        {%- set content = render_content(message.content, false)|trim %}
+        {%- if not(content.startswith('<tool_response>') and content.endswith('</tool_response>')) %}
+            {%- set ns.multi_step_tool = false %}
+            {%- set ns.last_query_index = index %}
+        {%- endif %}
+    {%- endif %}
+{%- endfor %}
+{%- if ns.multi_step_tool %}
+    {{- raise_exception('No user query found in messages.') }}
+{%- endif %}
+{%- for message in messages %}
+    {%- set content = render_content(message.content, true)|trim %}
+    {%- if message.role == "system" %}
+        {%- if not loop.first %}
+            {{- raise_exception('System message must be at the beginning.') }}
+        {%- endif %}
+    {%- elif message.role == "user" %}
+        {{- '<|im_start|>' + message.role + '\n' + content + '<|im_end|>' + '\n' }}
+    {%- elif message.role == "assistant" %}
+        {%- set reasoning_content = '' %}
+        {%- if message.reasoning_content is string %}
+            {%- set reasoning_content = message.reasoning_content %}
+        {%- else %}
+            {%- if '</think>' in content %}
+                {%- set reasoning_content = content.split('</think>')[0].rstrip('\n').split('<think>')[-1].lstrip('\n') %}
+                {%- set content = content.split('</think>')[-1].lstrip('\n') %}
+            {%- endif %}
+        {%- endif %}
+        {%- set reasoning_content = reasoning_content|trim %}
+        {%- if loop.index0 > ns.last_query_index %}
+            {{- '<|im_start|>' + message.role + '\n<think>\n' + reasoning_content + '\n</think>\n\n' + content }}
+        {%- else %}
+            {{- '<|im_start|>' + message.role + '\n' + content }}
+        {%- endif %}
+        {%- if message.tool_calls and message.tool_calls is iterable and message.tool_calls is not mapping %}
+            {%- for tool_call in message.tool_calls %}
+                {%- if tool_call.function is defined %}
+                    {%- set tool_call = tool_call.function %}
+                {%- endif %}
+                {%- if loop.first %}
+                    {%- if content|trim %}
+                        {{- '\n\n<tool_call>\n<function=' + tool_call.name + '>\n' }}
+                    {%- else %}
+                        {{- '<tool_call>\n<function=' + tool_call.name + '>\n' }}
+                    {%- endif %}
+                {%- else %}
+                    {{- '\n<tool_call>\n<function=' + tool_call.name + '>\n' }}
+                {%- endif %}
+                {%- if tool_call.arguments is defined %}
+                    {%- for args_name, args_value in tool_call.arguments|items %}
+                        {{- '<parameter=' + args_name + '>\n' }}
+                        {%- set args_value = args_value | tojson | safe if args_value is mapping or (args_value is sequence and args_value is not string) else args_value | string %}
+                        {{- args_value }}
+                        {{- '\n</parameter>\n' }}
+                    {%- endfor %}
+                {%- endif %}
+                {{- '</function>\n</tool_call>' }}
+            {%- endfor %}
+        {%- endif %}
+        {{- '<|im_end|>\n' }}
+    {%- elif message.role == "tool" %}
+        {%- if loop.previtem and loop.previtem.role != "tool" %}
+            {{- '<|im_start|>user' }}
+        {%- endif %}
+        {{- '\n<tool_response>\n' }}
+        {{- content }}
+        {{- '\n</tool_response>' }}
+        {%- if not loop.last and loop.nextitem.role != "tool" %}
+            {{- '<|im_end|>\n' }}
+        {%- elif loop.last %}
+            {{- '<|im_end|>\n' }}
+        {%- endif %}
+    {%- else %}
+        {{- raise_exception('Unexpected message role.') }}
+    {%- endif %}
+{%- endfor %}
+{%- if add_generation_prompt %}
+    {{- '<|im_start|>assistant\n' }}
+    {%- if enable_thinking is defined and enable_thinking is false %}
+        {{- '<think>\n\n</think>\n\n' }}
+    {%- else %}
+        {{- '<think>\n' }}
+    {%- endif %}
+{%- endif %}

config.json

@@ -0,0 +1,35 @@
+{
+  "architectures": [
+    "LlamaForCausalLM"
+  ],
+  "attention_bias": false,
+  "attention_dropout": 0.0,
+  "bos_token_id": null,
+  "dtype": "float32",
+  "eos_token_id": [
+    248046,
+    248044
+  ],
+  "head_dim": 64,
+  "hidden_act": "silu",
+  "hidden_size": 384,
+  "initializer_range": 0.02,
+  "intermediate_size": 1088,
+  "max_position_embeddings": 1024,
+  "mlp_bias": false,
+  "model_type": "llama",
+  "num_attention_heads": 6,
+  "num_hidden_layers": 8,
+  "num_key_value_heads": 2,
+  "pad_token_id": 248044,
+  "pretraining_tp": 1,
+  "rms_norm_eps": 1e-06,
+  "rope_parameters": {
+    "rope_theta": 10000.0,
+    "rope_type": "default"
+  },
+  "tie_word_embeddings": true,
+  "transformers_version": "5.16.1",
+  "use_cache": true,
+  "vocab_size": 248320
+}

eval_results.json

@@ -0,0 +1,21 @@
+{
+  "target": "Qwen/Qwen3.5-4B (served by Ollama, 4-bit)",
+  "method": "greedy speculative decoding replayed over held-out target answers; answer tokens only",
+  "sequences": 200,
+  "rounds": 90141,
+  "proposals_per_round": 4,
+  "p_accept_given_prefix": [
+    0.3862,
+    0.3283,
+    0.3917,
+    0.4582
+  ],
+  "accepted_guesses_per_round": 0.5838,
+  "tokens_per_round": 1.5807,
+  "draft_round_ms_cuda_graph": 1.772,
+  "draft_round_ms_eager": 9.352,
+  "target_step_ms": 16.602,
+  "projected_speedup": 1.4283,
+  "run1_tokens_per_round": 1.5222,
+  "note": "speedup = tokens_per_round / (1 + draft_round / target_step); projected, not measured end to end"
+}

generation_config.json

@@ -0,0 +1,12 @@
+{
+  "_from_model_config": true,
+  "eos_token_id": [
+    248046,
+    248044
+  ],
+  "output_attentions": false,
+  "output_hidden_states": false,
+  "pad_token_id": 248044,
+  "transformers_version": "5.16.1",
+  "use_cache": true
+}

lookahead_heads.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:b437d2d4086d39f2e5f9ef07a35bb264d0618f1e9fbf1ee99c8b1b76a85efb55
+size 1769728

model.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:7dc887211a536e8ab04937212b60f7a2e240e773611ec3c2a29d396096b961ac
+size 434144792

tokenizer.json

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:06b9509352d2af50381ab2247e083b80d32d5c0aba91c272ca9ff729b6a0e523
+size 19989325

tokenizer_config.json

@@ -0,0 +1,32 @@
+{
+  "add_prefix_space": false,
+  "audio_bos_token": "<|audio_start|>",
+  "audio_eos_token": "<|audio_end|>",
+  "audio_token": "<|audio_pad|>",
+  "backend": "tokenizers",
+  "bos_token": null,
+  "clean_up_tokenization_spaces": false,
+  "eos_token": "<|im_end|>",
+  "errors": "replace",
+  "image_token": "<|image_pad|>",
+  "is_local": false,
+  "local_files_only": false,
+  "model_max_length": 262144,
+  "model_specific_special_tokens": {
+    "audio_bos_token": "<|audio_start|>",
+    "audio_eos_token": "<|audio_end|>",
+    "audio_token": "<|audio_pad|>",
+    "image_token": "<|image_pad|>",
+    "video_token": "<|video_pad|>",
+    "vision_bos_token": "<|vision_start|>",
+    "vision_eos_token": "<|vision_end|>"
+  },
+  "pad_token": "<|endoftext|>",
+  "pretokenize_regex": "(?i:'s|'t|'re|'ve|'m|'ll|'d)|[^\\r\\n\\p{L}\\p{N}]?[\\p{L}\\p{M}]+|\\p{N}| ?[^\\s\\p{L}\\p{M}\\p{N}]+[\\r\\n]*|\\s*[\\r\\n]+|\\s+(?!\\S)|\\s+",
+  "split_special_tokens": false,
+  "tokenizer_class": "Qwen2Tokenizer",
+  "unk_token": null,
+  "video_token": "<|video_pad|>",
+  "vision_bos_token": "<|vision_start|>",
+  "vision_eos_token": "<|vision_end|>"
+}