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,196 @@
+---
+language:
+- en
+license: apache-2.0
+base_model: Qwen/Qwen3-1.7B
+tags:
+- continuous-thought
+- latent-reasoning
+- distillation
+- gsm8k
+- codi
+datasets:
+- openai/gsm8k
+metrics:
+- accuracy
+model-index:
+- name: r11_rw_perstep_g1_ans256
+  results:
+  - task:
+      type: math-reasoning
+      name: Grade School Math
+    dataset:
+      name: GSM8k
+      type: openai/gsm8k
+      split: test
+    metrics:
+    - type: accuracy
+      value: 82.7
+      name: Accuracy
+---
+
+# r11_rw_perstep_g1_ans256
+
+**Best reward-weighted — correct-only per-step distillation (82.7%)**
+
+- **Best reward-weighted result**: 82.7% on GSM8k
+- Uses reward-weighted teacher (average over CORRECT rollouts only)
+- Per-step distillation at every latent step with γ=1.0
+- Trained with max_answer_len=256
+
+## Overview
+
+This model implements **Continuous Thought Distillation (CODI)** — an autoregressive latent
+reasoning loop that processes K=6 continuous thought steps before generating a text answer.
+Teacher hidden states are extracted from multiple chain-of-thought rollouts generated by the
+base model and distilled into the latent representations.
+
+**Key idea**: Instead of distilling from a single reasoning trace, we aggregate hidden states
+from multiple rollouts (16 per problem). The teacher signal is the reward-weighted average of CORRECT rollouts only.
+
+## Architecture
+
+- **Base model**: Qwen3-1.7B (1.7B parameters)
+- **Fine-tuning**: LoRA (rank=32, alpha=16) on q/k/v/o/gate/up/down_proj
+- **Projection head**: Linear(2048, 2048) → GELU → Linear(2048, 2048) → LayerNorm(2048)
+- **Latent steps**: K=6 autoregressive continuous thought steps
+- **Inference**: Process prompt → K latent steps via ProjectionHead + KV cache → greedy text generation
+
+## Training Details
+
+| Parameter | Value |
+|---|---|
+| Mode | `codi_reward_weighted` |
+| Per-step distillation | `True` |
+| Distillation γ | 1.0 |
+| Learning rate | 0.0002 |
+| Epochs | 3 |
+| Batch size (per GPU) | 1 |
+| Gradient accumulation | 16 |
+| Effective batch size | 128 (across 8 GPUs) |
+| Max answer length | 256 |
+| Latent steps (K) | 6 |
+| Task | GSM8k (7,473 training problems) |
+| Rollouts per problem | 16 |
+| **GSM8k test accuracy** | **82.7%** |
+
+## How to Use
+
+### Requirements
+
+```bash
+pip install torch transformers
+```
+
+### Inference Code
+
+```python
+import torch
+import torch.nn as nn
+from transformers import AutoModelForCausalLM, AutoTokenizer
+
+# Load model
+model_name = "LakshyAAAgrawal/continuous-thought-r11_rw_perstep_g1_ans256"
+tokenizer = AutoTokenizer.from_pretrained(model_name)
+model = AutoModelForCausalLM.from_pretrained(
+    model_name, torch_dtype=torch.bfloat16, device_map="auto"
+)
+model.eval()
+
+# Load projection head
+class ProjectionHead(nn.Module):
+    def __init__(self, hidden_size):
+        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)
+
+proj = ProjectionHead(model.config.hidden_size)
+proj.load_state_dict(torch.load(
+    hf_hub_download(model_name, "projection_head.pt"), map_location="cpu"
+))
+proj = proj.to(model.dtype).to(model.device).eval()
+
+# Generate with latent reasoning
+question = "Natalia sold clips to 48 of her friends in April, and then she sold half as many clips in May. How many clips did Natalia sell altogether in April and May?"
+messages = [{"role": "user", "content": f"Solve the following math problem step by step. Show your work and put your final numerical answer after ####.\n\n{question}"}]
+prompt = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
+inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
+
+# Step 1: Process prompt
+with torch.no_grad():
+    out = model(**inputs, output_hidden_states=True, use_cache=True)
+    past_kv = out.past_key_values
+    latent = out.hidden_states[-1][:, -1, :]  # last token hidden state
+
+    # Step 2: K latent reasoning steps
+    mask = inputs["attention_mask"].clone()
+    for k in range(6):
+        latent = proj(latent)
+        mask = torch.cat([mask, torch.ones(1, 1, device=mask.device, dtype=mask.dtype)], 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 text generation
+    next_token = out.logits[:, -1, :].argmax(dim=-1)
+    generated = [next_token]
+    for _ in range(2047):
+        if next_token.item() == tokenizer.eos_token_id:
+            break
+        mask = torch.cat([mask, torch.ones(1, 1, device=mask.device, dtype=mask.dtype)], 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)
+        generated.append(next_token)
+
+    response = tokenizer.decode(torch.cat(generated), skip_special_tokens=True)
+    print(response)
+```
+
+### Evaluation
+
+To evaluate on the full GSM8k test set, use the evaluation script from our repository:
+
+```bash
+python evaluate.py \
+    --model_dir LakshyAAAgrawal/continuous-thought-r11_rw_perstep_g1_ans256 \
+    --mode codi_reward_weighted \
+    --output results/r11_rw_perstep_g1_ans256.json \
+    --max_new_tokens 2048 \
+    --num_latent 6
+```
+
+**Important**: Use `max_new_tokens=2048` for evaluation. The model generates verbose
+chain-of-thought text after the latent steps, requiring more tokens than standard models.
+
+## Results Comparison
+
+| Model | Mode | Per-step | γ | ans_len | GSM8k Accuracy |
+|---|---|---|---|---|---|
+| **This model** | **codi_reward_weighted** | **True** | **1.0** | **256** | **82.7%** |
+| Qwen3-1.7B (base) | — | — | — | — | 77.3% |
+| Discrete SFT | sft | — | — | — | 80.7% |
+| CODI RW final-step | rw | no | 1.0 | 128 | 81.0% |
+| CODI Uniform per-step ans256 | uniform | yes | 2.0 | 256 | 83.2% |
+| CODI RW per-step ans256 | rw | yes | 1.0 | 256 | 82.7% |
+
+## Citation
+
+If you use this model, please cite:
+
+```bibtex
+@misc{continuous-thought-2025,
+  title={Continuous Thought Distillation: Reward-Weighted Multi-Trace Reasoning},
+  author={Lakshya Agrawal},
+  year={2025},
+  url={https://huggingface.co/LakshyAAAgrawal/continuous-thought-r11_rw_perstep_g1_ans256}
+}
+```

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:b59c1ed931b3aef61396f2b9f7176ae5891cb98695607983f1902b6ed32a4098
+size 4063515640

projection_head.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:88ecc9fa08f454c7c8909d81fdafd299b85441540c5e2d6f576b2e2d75c06aee
+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_reward_weighted",
+  "model": "Qwen/Qwen3-1.7B",
+  "rollouts": "data/rollouts.jsonl",
+  "teacher_states": "data/teacher_states_v3.pt",
+  "output_dir": "checkpoints/r11_rw_perstep_g1_ans256",
+  "epochs": 3,
+  "batch_size": 1,
+  "grad_accum": 16,
+  "lr": 0.0002,
+  "gamma": 1.0,
+  "warmup_ratio": 0.03,
+  "num_latent": 6,
+  "max_prompt_len": 512,
+  "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": 0.5783941572043985,
+  "total_time_s": 2959.5815694332123
+}