Refresh README; document sft-c16 (latent-only reasoning, c=16, learned stop)

README.md

@@ -14,51 +14,71 @@ tags:
   - chain-of-thought
 ---
 
-# DeepSeek-V4-Flash CoLaR — Reasoning Compression Head
+<div align="center">
 
-A **CoLaR-style reasoning compression head** that attaches to the frozen
-[`nvidia/DeepSeek-V4-Flash-NVFP4`](https://huggingface.co/nvidia/DeepSeek-V4-Flash-NVFP4)
-backbone. It learns a compact latent representation of the model's chain-of-thought
-and injects it back into the backbone so the model **reasons from a compressed
-latent** rather than emitting a long token-by-token trace.
+# 🧠 DeepSeek-V4-Flash · CoLaR Reasoning Compression Head
 
-> **This is an adapter, not a standalone model.** These weights (~136 MB per
-> checkpoint) are useless without the DeepSeek-V4-Flash backbone. You must load
-> the backbone separately.
+**Reason in a compact latent space — not a long token-by-token trace.**
 
-**Headline result (preliminary):** on GSM8K the head cuts **median "thinking"
-tokens by ~40%** and shrinks the **longest trace by ~4×** while serving at
-**~65 tok/s** on the compile-safe fast path — with no separate reasoning model.
-Task-accuracy numbers are small-n (n=30) and reported below with caveats.
+*A lightweight adapter for the frozen [`nvidia/DeepSeek-V4-Flash-NVFP4`](https://huggingface.co/nvidia/DeepSeek-V4-Flash-NVFP4) backbone.*
 
-The method follows **CoLaR** (Compressed Latent Reasoning, [arXiv:2505.16552](https://arxiv.org/abs/2505.16552)).
+[![Method: CoLaR](https://img.shields.io/badge/method-CoLaR-6c5ce7)](https://arxiv.org/abs/2505.16552)
+[![Base: DeepSeek--V4--Flash](https://img.shields.io/badge/base-DeepSeek--V4--Flash--NVFP4-0984e3)](https://huggingface.co/nvidia/DeepSeek-V4-Flash-NVFP4)
+[![Type: adapter](https://img.shields.io/badge/type-adapter-00b894)](#-this-is-an-adapter)
+[![Format: safetensors](https://img.shields.io/badge/format-safetensors-fdcb6e)](#-checkpoint-layout)
+
+</div>
+
+---
+
+The **CoLaR reasoning compression head** learns a compact latent representation of the
+model's chain-of-thought and feeds it back into the backbone, so the model **reasons from
+a compressed latent** instead of spelling out every reasoning token. The backbone stays
+frozen; only this small head (~136 MB) is trained.
+
+Method: **CoLaR — Compressed Latent Reasoning** ([arXiv:2505.16552](https://arxiv.org/abs/2505.16552)).
+
+> ### ⚠️ This is an adapter
+> These weights are **useless on their own**. You must load the
+> [`nvidia/DeepSeek-V4-Flash-NVFP4`](https://huggingface.co/nvidia/DeepSeek-V4-Flash-NVFP4)
+> backbone separately and attach this head to it.
 
 ---
 
-## What's in this repo
+## 📦 Checkpoints
+
+| Variant | Path | Stage | Reasoning mode |
+|---|---|---|---|
+| **`sft-c16`** ✨ *latest* | `sft-c16/` | closed-loop SFT, `c = 16` | **Latent-only reasoning** — see below |
+| `grpo` | `grpo/` | SFT → GRPO (latent-policy RL) | Prompt-position latent injection |
+| `sft` | `sft/` | SFT (soft-MSE latent regression) | Prompt-position latent injection |
+
+Each variant folder holds `colar_head_sft.{safetensors,pt}` (identical weights, two formats)
+and a `config.json` with its geometry.
+
+### ✨ `sft-c16` — latent-only reasoning
 
-| Path | Description |
-|---|---|
-| `sft/colar_head_sft.safetensors` | Head after **SFT** (soft-MSE latent regression). |
-| `grpo/colar_head_grpo.safetensors` | Head after **SFT → GRPO** (latent-policy RL). **Recommended for serving.** |
-| `sft/colar_head_sft.pt`, `grpo/colar_head_grpo.pt` | The original v2 bundle `.pt` files (same weights; pickle format). |
-| `*/config.json` | Per-checkpoint geometry + which sub-modules are present. |
-| `config.json` | Top-level architecture description. |
-| `PAPER.md` | Short technical report (method, serving, results). |
-| `results/` | Raw eval JSON + HTML bar charts backing the results table. |
+`sft-c16` is the newest head and works differently from the earlier checkpoints. It is a
+**latent-only reasoning** head: the model performs its *entire* reasoning phase inside the
+compressed latent space — one latent step stands in for **`c = 16`** reasoning tokens — and
+a **learned stop head** decides when the reasoning is complete and the model should begin
+emitting its answer. There is no accompanying token-by-token thinking trace to read; the
+reasoning happens in the latents, and the model surfaces only the final answer.
 
-**Which to use:** `grpo` is the checkpoint used for the results below and for
-serving. `sft` is the pre-RL stage, useful for comparison/ablation.
+This is what the extra `stop_head` sub-module (present only in `sft-c16`) provides: a small
+classifier over the running latent state that fires a learned "end-of-reasoning" signal, so
+the latent phase self-terminates at a variable, content-dependent depth rather than running
+a fixed number of steps.
 
-> **Code is not released yet.** The training scripts and the vLLM serving addon
-> will be published separately later. This repo ships weights + docs only; the
-> loader snippet below is fully self-contained and needs no first-party code.
+> **Serving parameters (temperature, stop threshold, latent budget, etc.) are intentionally
+> not prescribed here.** They interact with your prompts and decoding setup — find the
+> settings that work best for your use case.
 
 ---
 
-## Architecture
+## 🏗️ Architecture
 
-Three components, all bundled in each checkpoint:
+Three components, bundled in every checkpoint (`sft-c16` adds a fourth, the stop head):
 
 ```
                    layer 35 hidden (4096-d)
@@ -68,6 +88,9 @@ Three components, all bundled in each checkpoint:
         │  Linear 4096→2048 · SiLU                        │
         │  Linear 2048→2048 · SiLU                        │
         │  Linear 2048→2048  →  [mu, log_sigma] (1024-d)  │
+        │                                                 │
+        │  stop_head (sft-c16 only):                      │
+        │     Linear 4096→1024 · SiLU · Linear 1024→1     │  → learned end-of-reasoning
         └─────────────────────────────────────────────────┘
                             │ mu (1024-d latent)
                             ▼  LayerNorm
@@ -77,67 +100,53 @@ Three components, all bundled in each checkpoint:
         │  Linear 2048→4096  →  hidden vector (4096-d)    │
         └─────────────────────────────────────────────────┘
                             │
-                            ▼  injected at the final <think> prompt position
+                            ▼  injected back into the residual stream
                         DeepSeek-V4-Flash backbone (frozen)
 ```
 
-- **`target_proj`** — a frozen `Linear(4096, 1024, bias=False)` that defines the
-  SFT regression target (a stable readout of the layer-42 hidden state). Included
-  for reproducibility; not needed at inference.
-- **Source layer 35 → target layer 42**, **`compression_factor = 4`** (4
-  consecutive reasoning tokens fold into one latent step), **`latent_dim = 1024`**,
-  **`hidden_size = 4096`**.
-
-| Config | Value |
-|---|---|
-| base model | `nvidia/DeepSeek-V4-Flash-NVFP4` |
-| hidden_size | 4096 |
-| latent_dim | 1024 |
-| mlp_dim | 2048 |
-| source_layer / target_layer | 35 / 42 |
-| compression_factor | 4 |
-| activation | SiLU |
-| checkpoint format | v2 bundle (`reasoning_head` + `decoder` + `target_proj` + `config`) |
+- **`target_proj`** — a frozen `Linear(4096, 1024, bias=False)` defining the SFT regression
+  target (a stable readout of the layer-42 hidden state). Kept for reproducibility; not
+  needed at inference.
+- **`stop_head`** — *(`sft-c16` only)* the learned end-of-reasoning classifier that enables
+  latent-only reasoning.
+
+| Config | `sft` / `grpo` | `sft-c16` |
+|---|---|---|
+| base model | `nvidia/DeepSeek-V4-Flash-NVFP4` | ← same |
+| hidden_size | 4096 | 4096 |
+| latent_dim | 1024 | 1024 |
+| mlp_dim | 2048 | 2048 |
+| source_layer / target_layer | 35 / 42 | 35 / 42 |
+| **compression_factor** | **4** | **16** |
+| learned stop head | — | ✅ |
+| activation | SiLU | SiLU |
+| checkpoint format | v2 bundle | v3 bundle (adds `stop_head`) |
 
 ---
 
-## Checkpoint layout
+## 🗂️ Checkpoint layout
 
-The `.safetensors` files hold a single flat tensor dict; the original three
-sub-modules are distinguished by key prefix:
+Each `.safetensors` file holds a single flat tensor dict; the sub-modules are distinguished
+by key prefix. For `sft-c16`:
 
 ```
-reasoning_head.net.0.weight  [2048, 4096]   reasoning_head.net.0.bias  [2048]
-reasoning_head.net.2.weight  [2048, 2048]   reasoning_head.net.2.bias  [2048]
-reasoning_head.net.4.weight  [2048, 2048]   reasoning_head.net.4.bias  [2048]
-decoder.net.0.weight         [2048, 1024]   decoder.net.0.bias         [2048]
-decoder.net.2.weight         [2048, 2048]   decoder.net.2.bias         [2048]
-decoder.net.4.weight         [4096, 2048]   decoder.net.4.bias         [4096]
-target_proj.weight           [1024, 4096]
+reasoning_head.net.0.weight        [2048, 4096]   reasoning_head.net.0.bias        [2048]
+reasoning_head.net.2.weight        [2048, 2048]   reasoning_head.net.2.bias        [2048]
+reasoning_head.net.4.weight        [2048, 2048]   reasoning_head.net.4.bias        [2048]
+reasoning_head.stop_head.0.weight  [1024, 4096]   reasoning_head.stop_head.0.bias  [1024]   ← sft-c16 only
+reasoning_head.stop_head.2.weight  [1, 1024]      reasoning_head.stop_head.2.bias  [1]      ← sft-c16 only
+decoder.net.0.weight               [2048, 1024]   decoder.net.0.bias               [2048]
+decoder.net.2.weight               [2048, 2048]   decoder.net.2.bias               [2048]
+decoder.net.4.weight               [4096, 2048]   decoder.net.4.bias               [4096]
+target_proj.weight                 [1024, 4096]
 ```
 
-The geometry is also stored in the safetensors metadata (`config`,
-`format_version`, `subdicts`) and in the sibling `config.json`.
+The geometry is mirrored in the safetensors metadata (`config`, `format_version`,
+`subdicts`) and in the sibling `config.json`.
 
 ---
-**FAQ**:
 
-
-| Question | Answer |
-|---|---|
-| What does this repo provide? | A **CoLaR-style reasoning compression head** that attaches to the frozen [`nvidia/DeepSeek-V4-Flash-NVFP4`](https://huggingface.co/nvidia/DeepSeek-V4-Flash-NVFP4) backbone. It learns a compact latent representation of the model's chain-of-thought and injects it back into the backbone so the model **reasons from a compressed latent** rather than emitting a long token-by-token trace. |
-| Which checkpoint is recommended for production use? | Use the `grpo` checkpoint (e.g., `grpo/colar_head_grpo.safetensors`). The `sft` checkpoint is the pre-RL stage, useful for comparison/ablation. |
-| What license applies? | Other — see [`LICENSE`](./LICENSE) file. |
-| Which backbone model does CoLaR require? | The frozen DeepSeek-V4-Flash-NVFP4 backbone (~136 MB per checkpoint); these weights are useless without the backbone, which must be loaded separately. |
-| How do I load the head weights without importing the repository? | Load using `safe_open` and `load_file` (see the self-contained loader snippet in README); the code requires only PyTorch, safetensors, and torch.nn.functional — no first-party imports. |
-| Is this compatible with vLLM or other inference engines? | Not yet — vLLM support is planned in a fork; full details are in [`PAPER.md`](./PAPER.md). |
-| What improvements does CoLaR bring compared to the base model? | Median "thinking" tokens drop by **~40%** and the longest trace shrinks **~4×**, with service speed **~65 tok/s** on the compile-safe fast path; preliminary lm_eval shows **10–20% better results** (GSM8K n=30). |
-| What are the known limitations or caveats of using CoLaR? | Evaluation is **small-n (30 docs, GSM8K only)**; no full-benchmark or multi-task evaluation yet; closed-loop latent decoding (generating a step with no token id) is future work; current injection is at the prompt/`` position only. |
-
-
----
-
-## Load the head (self-contained, no repo import)
+## 🚀 Load a head (self-contained — no repo import)
 
 ```python
 import json
@@ -147,10 +156,10 @@ from torch import nn
 from safetensors.torch import load_file
 from safetensors import safe_open
 
-CKPT = "grpo/colar_head_grpo.safetensors"  # or sft/colar_head_sft.safetensors
+CKPT = "sft-c16/colar_head_sft.safetensors"   # or grpo/…, sft/…
 
 class ReasoningCompressionHead(nn.Module):
-    def __init__(self, hidden_size, latent_dim, mlp_dim=None):
+    def __init__(self, hidden_size, latent_dim, mlp_dim=None, stop_head=False):
         super().__init__()
         mlp_dim = mlp_dim or hidden_size // 2
         self.net = nn.Sequential(
@@ -158,10 +167,19 @@ class ReasoningCompressionHead(nn.Module):
             nn.Linear(mlp_dim, mlp_dim), nn.SiLU(),
             nn.Linear(mlp_dim, 2 * latent_dim),
         )
+        # sft-c16 (v3): learned end-of-reasoning classifier over the latent state
+        self.stop_head = nn.Sequential(
+            nn.Linear(hidden_size, latent_dim), nn.SiLU(),
+            nn.Linear(latent_dim, 1),
+        ) if stop_head else None
+
     def forward(self, h):
         mu, log_sigma = self.net(h).chunk(2, dim=-1)
         return mu, log_sigma.clamp(-10.0, 2.0)
 
+    def stop_logit(self, h):
+        return self.stop_head(h)   # sigmoid(stop_logit) > threshold ⇒ end reasoning
+
 class LatentDecoder(nn.Module):
     def __init__(self, hidden_size, latent_dim, mlp_dim=None):
         super().__init__()
@@ -183,73 +201,79 @@ flat = load_file(CKPT)
 def sub(prefix):
     return {k[len(prefix):]: v for k, v in flat.items() if k.startswith(prefix)}
 
-head = ReasoningCompressionHead(hs, ld); head.load_state_dict(sub("reasoning_head.")); head.eval()
-decoder = LatentDecoder(hs, ld);        decoder.load_state_dict(sub("decoder."));        decoder.eval()
+has_stop = any(k.startswith("reasoning_head.stop_head.") for k in flat)
+head = ReasoningCompressionHead(hs, ld, stop_head=has_stop)
+head.load_state_dict(sub("reasoning_head.")); head.eval()
+decoder = LatentDecoder(hs, ld); decoder.load_state_dict(sub("decoder.")); decoder.eval()
 target_proj = nn.Linear(hs, ld, bias=False); target_proj.load_state_dict(sub("target_proj."))
 
-# usage: h35 is the layer-35 hidden state at the final <think> position, shape (B, 4096)
-# mu, _ = head(F.layer_norm(h35, (hs,)))
-# inject = decoder(F.layer_norm(mu, (ld,)))   # (B, 4096) -> overwrite the embedding at that position
+# One latent step, given h35 = layer-35 hidden at the current position, shape (B, 4096):
+#   mu, _  = head(F.layer_norm(h35, (hs,)))
+#   inject = decoder(F.layer_norm(mu, (ld,)))          # (B, 4096) → back into the residual stream
+#   p_stop = head.stop_logit(F.layer_norm(h35, (hs,))).sigmoid()   # sft-c16: end reasoning when high
 ```
 
 ---
 
-## How it's served
+## 🔬 How it's served (design note)
 
-DeepSeek-V4 uses **hash-based MoE expert routing keyed on `input_ids`**, so
-vLLM's native `prompt_embeds` injection path crashes the engine
-(`hash MoE routing requires input_ids`). Injection therefore uses an
-**`embed_tokens` forward hook** that overwrites the embedding at chosen positions
-with the decoded latent while token ids keep flowing for routing. Combined with a
-compile-safe capture buffer, this runs on the **cudagraph fast path (~65 tok/s)**
-instead of `enforce_eager`. Full details in [`PAPER.md`](./PAPER.md). The serving
-addon itself is not released yet.
+DeepSeek-V4 uses **hash-based MoE expert routing keyed on `input_ids`**, so vLLM's native
+`prompt_embeds` injection path crashes the engine (`hash MoE routing requires input_ids`).
+Injection instead uses an **`embed_tokens` forward hook** that overwrites the embedding at
+the chosen position with the decoded latent while token ids keep flowing for routing.
+Combined with a compile-safe capture buffer, this runs on the **cudagraph fast path** rather
+than `enforce_eager`. Full details in [`PAPER.md`](./PAPER.md). *The serving addon is not
+released yet.*
 
 ---
 
-## Results (preliminary)
+## 📊 Results (preliminary)
 
 GSM8K, 8-shot, temperature 0, DeepSeek-V4-Flash-NVFP4 backbone, GRPO head.
-**n = 30 documents** — accuracy is small-n and noisy; the **robust signal is the
-reduction in reasoning length**, not the exact-match score.
+**n = 30 documents** — accuracy is small-n and noisy; the **robust signal is the reduction
+in reasoning length**, not the exact-match score.
 
 | Metric | Base (no injection) | +Head (2k budget) | +Head (14.5k budget) |
 |---|---|---|---|
-| exact-match | 40.0 / 33.3 | 40.0 | 46.7 |
+| exact-match | 40.0 / 33.3 | 40.0 | **46.7** |
 | closed `</think>` % | 97–100 | 30 | 40 |
 | median think tokens | 146–187 | 106 | **103** |
 | max think tokens | 1740–2048 | 472 | **415** |
 
 - **~40% fewer median thinking tokens** and a **~4× shorter worst-case trace.**
-- The head **never skips** reasoning (skip% = 0) — it compresses it.
-- The low closed-`</think>`% is a **formatting artifact**: the GRPO reward
-  (token-F1) never rewarded emitting the closing tag. A larger 78× budget did not
-  make traces close, confirming it is not truncation. Adding a closure bonus to
-  the reward is the obvious next step.
-- For context only, the published DeepSeek-V4-Flash-Base scores **90.8** on GSM8K
-  under its own full harness; that is **not** a baseline we reproduce here (our
-  base column is the same backbone under our small-n thinking-mode harness).
+- The head **never skips** reasoning (skip% = 0) — it *compresses* it.
+- The low closed-`</think>`% for `grpo` is a **formatting artifact**: its token-F1 reward
+  never rewarded emitting the closing tag (a larger 78× budget did not make traces close,
+  confirming it is not truncation). The `sft-c16` head's learned **stop** addresses exactly
+  this by terminating the latent phase on a learned signal.
+- For context only, published DeepSeek-V4-Flash-Base scores **90.8** on GSM8K under its own
+  full harness — **not** a baseline reproduced here.
 
 Raw numbers and bar charts are in [`results/`](./results).
 
 ---
 
-## Limitations
+## ⚠️ Limitations
 
 - Evaluation is **small-n (30 docs, GSM8K only)**; treat accuracy as directional.
 - No full-benchmark or multi-task evaluation yet.
-- Closed-loop latent decoding (generating a step with no token id) is future work;
-  current injection is at the prompt/`<think>` position only.
+- `sft` / `grpo` inject at the prompt/`<think>` position only. `sft-c16` runs the full
+  latent-only closed loop with a learned stop; broader eval of it is ongoing.
 - Measured on 2× RTX PRO 6000 (96 GiB); usable context ≈ 15.8k tokens on that box.
+- **Code is not released yet.** Training scripts and the vLLM serving addon will be
+  published separately; this repo ships weights + docs, and the loader above needs no
+  first-party code.
+
+---
 
-## Citation
+## 📄 Citation
 
-Method after CoLaR — Compressed Latent Reasoning:
+Method after **CoLaR — Compressed Latent Reasoning**:
 
 ```bibtex
 @article{colar2025,
-  title  = {CoLaR: Compressed Latent Reasoning},
+  title   = {CoLaR: Compressed Latent Reasoning},
   journal = {arXiv preprint arXiv:2505.16552},
-  year   = {2025}
+  year    = {2025}
 }
 ```