Upload folder using huggingface_hub

.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,184 @@
+---
+language:
+- en
+license: apache-2.0
+library_name: transformers
+base_model: Qwen/Qwen3-1.7B
+tags:
+- qthink
+- latent-reasoning
+- distillation
+- qwen3
+- lora
+- tooluse
+datasets:
+- tooluse
+pipeline_tag: text-generation
+model-index:
+- name: QThink-Qwen3-1.7B-Tooluse
+  results:
+  - task:
+      type: text-generation
+      name: Tooluse
+    dataset:
+      name: Tooluse
+      type: tooluse
+      split: test
+    metrics:
+    - type: accuracy
+      value: 48.5% EM
+      name: Accuracy
+---
+
+# QThink-Qwen3-1.7B-Tooluse
+
+**QThink: Parallel Latent Reasoning via Per-Step Distillation of Multiple Rollouts**
+
+This model replaces explicit chain-of-thought (`<think>...</think>`) with **6 latent forward passes** through a learned projection head, achieving **48.5% EM** on Tooluse.
+
+## How QThink Works
+
+Instead of generating thousands of reasoning tokens, QThink:
+
+1. **Processes the prompt** through the base model
+2. **Runs K=6 latent steps**: each step applies a ProjectionHead (`Linear(2048,2048) → GELU → Linear(2048,2048) → LayerNorm(2048)`) to the hidden state, then feeds it back through the model via `inputs_embeds` + `past_key_values`
+3. **Generates the answer** directly from the last latent step's hidden state — no `<think>` block needed
+
+### Training: Per-Step Distillation from Multiple Rollouts
+
+1. Generate G=16 chain-of-thought rollouts per training problem using the base model
+2. For each rollout, extract hidden states at K=6 evenly-spaced positions within the `<think>` block
+3. Average these hidden states across **all rollouts** (uniform — including incorrect ones) at each step
+4. Train each latent step to match the corresponding teacher state via L1 loss, jointly with cross-entropy on the answer
+
+The key innovations:
+- **Per-step distillation**: Every latent step gets direct supervision, not just the final one
+- **Uniform multi-rollout teachers**: Averaging over ALL rollouts (correct + incorrect) outperforms using only correct rollouts
+- **ans256 training**: Training with longer answer targets (+3pp improvement)
+
+## Results
+
+| Model | Exact Match | Action Accuracy |
+|-------|------------|----------------|
+| **QThink uniform per-step (ours)** | **48.5%** | 67.6% |
+| Base Qwen3-1.7B | 47.1% | 75.0% |
+| SFT | 45.6% | 73.5% |
+| QThink RW final-step (CODI) | 42.6% | 67.6% |
+
+
+### Cross-Benchmark Results
+
+QThink uniform per-step is the **best model on all 3 benchmarks**:
+
+| Benchmark | QThink (ours) | SFT | Base | CODI |
+|-----------|--------------|-----|------|------|
+| GSM8k | **83.2%** | 80.7% | 77.3% | 80.4% |
+| MATH-500 | **43.6%** | 38.2% | 33.6% | 28.0% |
+| Tooluse | **48.5%** | 45.6% | 47.1% | 42.6% |
+
+## Model Details
+
+| Parameter | Value |
+|-----------|-------|
+| Base model | [Qwen/Qwen3-1.7B](https://huggingface.co/Qwen/Qwen3-1.7B) |
+| Fine-tuning | LoRA (rank=32, alpha=16) |
+| Distillation mode | Uniform (all rollouts) |
+| Per-step distillation | Yes (K=6 steps) |
+| Distillation weight (γ) | 2.0 |
+| Learning rate | 0.0002 |
+| Epochs | 3 |
+| Batch size × grad accum × GPUs | 1 × 16 × 8 = 128 effective |
+| Max answer length | 256 tokens |
+| Max prompt length | 1024 tokens |
+| Rollouts per problem | 16 |
+| Dataset | Tooluse (from [SDPO](https://github.com/lasgroup/SDPO)) — 4,046 train problems, 68 test problems |
+
+## Architecture
+
+The checkpoint contains:
+- **`model.safetensors`**: Full Qwen3-1.7B weights with merged LoRA adapters
+- **`projection_head.pt`**: ProjectionHead weights (PyTorch state dict)
+  - `mlp.0`: Linear(2048 → 2048) + bias
+  - `mlp.2`: Linear(2048 → 2048) + bias
+  - `mlp.3`: LayerNorm(2048)
+
+## Usage
+
+```python
+import torch
+import torch.nn as nn
+from transformers import AutoModelForCausalLM, AutoTokenizer
+from huggingface_hub import hf_hub_download
+
+class ProjectionHead(nn.Module):
+    def __init__(self, hidden_size=2048):
+        super().__init__()
+        self.mlp = nn.Sequential(
+            nn.Linear(hidden_size, hidden_size),
+            nn.GELU(),
+            nn.Linear(hidden_size, hidden_size),
+            nn.LayerNorm(hidden_size),
+        )
+    def forward(self, x):
+        return self.mlp(x)
+
+# Load model and projection head
+repo_id = "LakshyAAAgrawal/QThink-Qwen3-1.7B-Tooluse"
+model = AutoModelForCausalLM.from_pretrained(repo_id, torch_dtype=torch.bfloat16, device_map="auto")
+tokenizer = AutoTokenizer.from_pretrained(repo_id)
+proj = ProjectionHead(2048).to(model.device).to(torch.bfloat16)
+proj.load_state_dict(torch.load(
+    hf_hub_download(repo_id, "projection_head.pt"), map_location=model.device
+))
+proj.eval()
+model.eval()
+
+# Prepare input
+question = "What's the weather like in San Francisco?"
+messages = [{"role": "user", "content": question}]
+text = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True, enable_thinking=False)
+inputs = tokenizer(text, return_tensors="pt").to(model.device)
+
+with torch.no_grad():
+    # Step 1: Process prompt
+    out = model(**inputs, output_hidden_states=True, use_cache=True)
+    past_kv = out.past_key_values
+    latent = out.hidden_states[-1][:, -1, :]
+
+    # Step 2: K=6 latent reasoning steps
+    mask = inputs["attention_mask"].clone()
+    for k in range(6):
+        latent = proj(latent)
+        mask = torch.cat([mask, mask.new_ones(1, 1)], dim=1)
+        out = model(inputs_embeds=latent.unsqueeze(1), attention_mask=mask,
+                    past_key_values=past_kv, output_hidden_states=True, use_cache=True)
+        past_kv = out.past_key_values
+        latent = out.hidden_states[-1][:, -1, :]
+
+    # Step 3: Greedy decode answer
+    next_token = out.logits[:, -1, :].argmax(dim=-1)
+    tokens = [next_token]
+    eos_id = tokenizer.eos_token_id
+    for _ in range(2047):
+        if next_token.item() == eos_id:
+            break
+        mask = torch.cat([mask, mask.new_ones(1, 1)], dim=1)
+        out = model(input_ids=next_token.unsqueeze(0), attention_mask=mask,
+                    past_key_values=past_kv, use_cache=True)
+        past_kv = out.past_key_values
+        next_token = out.logits[:, -1, :].argmax(dim=-1)
+        tokens.append(next_token)
+
+print(tokenizer.decode(torch.cat(tokens), skip_special_tokens=True))
+```
+
+## Citation
+
+```bibtex
+@misc{qthink2025,
+  title={QThink: Parallel Latent Reasoning via Per-Step Distillation of Multiple Rollouts},
+  author={Lakshya Agrawal},
+  year={2025},
+  url={https://huggingface.co/LakshyAAAgrawal/QThink-Qwen3-1.7B-Tooluse}
+}
+```

chat_template.jinja

@@ -0,0 +1,89 @@
+{%- if tools %}
+    {{- '<|im_start|>system\n' }}
+    {%- if messages[0].role == 'system' %}
+        {{- messages[0].content + '\n\n' }}
+    {%- endif %}
+    {{- "# Tools\n\nYou may call one or more functions to assist with the user query.\n\nYou are provided with function signatures within <tools></tools> XML tags:\n<tools>" }}
+    {%- for tool in tools %}
+        {{- "\n" }}
+        {{- tool | tojson }}
+    {%- endfor %}
+    {{- "\n</tools>\n\nFor each function call, return a json object with function name and arguments within <tool_call></tool_call> XML tags:\n<tool_call>\n{\"name\": <function-name>, \"arguments\": <args-json-object>}\n</tool_call><|im_end|>\n" }}
+{%- else %}
+    {%- if messages[0].role == 'system' %}
+        {{- '<|im_start|>system\n' + messages[0].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" and message.content is string and not(message.content.startswith('<tool_response>') and message.content.endswith('</tool_response>')) %}
+        {%- set ns.multi_step_tool = false %}
+        {%- set ns.last_query_index = index %}
+    {%- endif %}
+{%- endfor %}
+{%- for message in messages %}
+    {%- if message.content is string %}
+        {%- set content = message.content %}
+    {%- else %}
+        {%- set content = '' %}
+    {%- endif %}
+    {%- if (message.role == "user") or (message.role == "system" and not loop.first) %}
+        {{- '<|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 %}
+        {%- if loop.index0 > ns.last_query_index %}
+            {%- if loop.last or (not loop.last and reasoning_content) %}
+                {{- '<|im_start|>' + message.role + '\n<think>\n' + reasoning_content.strip('\n') + '\n</think>\n\n' + content.lstrip('\n') }}
+            {%- else %}
+                {{- '<|im_start|>' + message.role + '\n' + content }}
+            {%- endif %}
+        {%- else %}
+            {{- '<|im_start|>' + message.role + '\n' + content }}
+        {%- endif %}
+        {%- if message.tool_calls %}
+            {%- for tool_call in message.tool_calls %}
+                {%- if (loop.first and content) or (not loop.first) %}
+                    {{- '\n' }}
+                {%- endif %}
+                {%- if tool_call.function %}
+                    {%- set tool_call = tool_call.function %}
+                {%- endif %}
+                {{- '<tool_call>\n{"name": "' }}
+                {{- tool_call.name }}
+                {{- '", "arguments": ' }}
+                {%- if tool_call.arguments is string %}
+                    {{- tool_call.arguments }}
+                {%- else %}
+                    {{- tool_call.arguments | tojson }}
+                {%- endif %}
+                {{- '}\n</tool_call>' }}
+            {%- endfor %}
+        {%- endif %}
+        {{- '<|im_end|>\n' }}
+    {%- elif message.role == "tool" %}
+        {%- if loop.first or (messages[loop.index0 - 1].role != "tool") %}
+            {{- '<|im_start|>user' }}
+        {%- endif %}
+        {{- '\n<tool_response>\n' }}
+        {{- content }}
+        {{- '\n</tool_response>' }}
+        {%- if loop.last or (messages[loop.index0 + 1].role != "tool") %}
+            {{- '<|im_end|>\n' }}
+        {%- endif %}
+    {%- 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' }}
+    {%- endif %}
+{%- endif %}

config.json

@@ -0,0 +1,63 @@
+{
+  "architectures": [
+    "Qwen3ForCausalLM"
+  ],
+  "attention_bias": false,
+  "attention_dropout": 0.0,
+  "bos_token_id": 151643,
+  "dtype": "bfloat16",
+  "eos_token_id": 151645,
+  "head_dim": 128,
+  "hidden_act": "silu",
+  "hidden_size": 2048,
+  "initializer_range": 0.02,
+  "intermediate_size": 6144,
+  "layer_types": [
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention"
+  ],
+  "max_position_embeddings": 40960,
+  "max_window_layers": 28,
+  "model_type": "qwen3",
+  "num_attention_heads": 16,
+  "num_hidden_layers": 28,
+  "num_key_value_heads": 8,
+  "pad_token_id": null,
+  "rms_norm_eps": 1e-06,
+  "rope_parameters": {
+    "rope_theta": 1000000,
+    "rope_type": "default"
+  },
+  "sliding_window": null,
+  "tie_word_embeddings": true,
+  "transformers_version": "5.3.0",
+  "use_cache": true,
+  "use_sliding_window": false,
+  "vocab_size": 151936
+}

generation_config.json

@@ -0,0 +1,13 @@
+{
+  "bos_token_id": 151643,
+  "do_sample": true,
+  "eos_token_id": [
+    151645,
+    151643
+  ],
+  "pad_token_id": 151643,
+  "temperature": 0.6,
+  "top_k": 20,
+  "top_p": 0.95,
+  "transformers_version": "5.3.0"
+}

model.safetensors

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

projection_head.pt

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

tokenizer.json

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

tokenizer_config.json

@@ -0,0 +1,29 @@
+{
+  "add_prefix_space": false,
+  "backend": "tokenizers",
+  "bos_token": null,
+  "clean_up_tokenization_spaces": false,
+  "eos_token": "<|im_end|>",
+  "errors": "replace",
+  "extra_special_tokens": [
+    "<|im_start|>",
+    "<|im_end|>",
+    "<|object_ref_start|>",
+    "<|object_ref_end|>",
+    "<|box_start|>",
+    "<|box_end|>",
+    "<|quad_start|>",
+    "<|quad_end|>",
+    "<|vision_start|>",
+    "<|vision_end|>",
+    "<|vision_pad|>",
+    "<|image_pad|>",
+    "<|video_pad|>"
+  ],
+  "is_local": false,
+  "model_max_length": 131072,
+  "pad_token": "<|endoftext|>",
+  "split_special_tokens": false,
+  "tokenizer_class": "Qwen2Tokenizer",
+  "unk_token": null
+}

training_config.json

@@ -0,0 +1,25 @@
+{
+  "mode": "codi_uniform",
+  "model": "Qwen/Qwen3-1.7B",
+  "rollouts": "data/tooluse_rollouts.jsonl",
+  "teacher_states": "data/tooluse_teacher_states.pt",
+  "output_dir": "checkpoints/r14_tooluse_qthink_uniform_perstep_g2_ans256",
+  "epochs": 3,
+  "batch_size": 1,
+  "grad_accum": 16,
+  "lr": 0.0002,
+  "gamma": 2.0,
+  "warmup_ratio": 0.03,
+  "num_latent": 6,
+  "max_prompt_len": 1024,
+  "max_answer_len": 256,
+  "max_len": 1024,
+  "answer_format": "full",
+  "per_step_distill": true,
+  "log_every": 10,
+  "seed": 42,
+  "lora_rank": 32,
+  "lora_alpha": 16,
+  "best_loss": 1.1106567853995464,
+  "total_time_s": 1602.5439190864563
+}