README.md

---
license: other
license_name: deepseek
license_link: https://huggingface.co/nvidia/DeepSeek-V4-Flash-NVFP4
base_model: nvidia/DeepSeek-V4-Flash-NVFP4
base_model_relation: adapter
library_name: pytorch
pipeline_tag: text-generation
tags:
  - colar
  - reasoning-compression
  - latent-reasoning
  - deepseek-v4
  - chain-of-thought
---

<div align="center">

# 🧠 DeepSeek-V4-Flash · CoLaR Reasoning Compression Head

**Reason in a compact latent space — not a long token-by-token trace.**

*A lightweight adapter for the frozen [`nvidia/DeepSeek-V4-Flash-NVFP4`](https://huggingface.co/nvidia/DeepSeek-V4-Flash-NVFP4) backbone.*

[![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.

---

## 📦 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

`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.

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.

> **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

Three components, bundled in every checkpoint (`sft-c16` adds a fourth, the stop head):

```
                   layer 35 hidden (4096-d)
                            │
                            ▼  LayerNorm
        ┌─────────── ReasoningCompressionHead ───────────┐
        │  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
        ┌─────────────── LatentDecoder ──────────────────┐
        │  Linear 1024→2048 · SiLU                        │
        │  Linear 2048→2048 · SiLU                        │
        │  Linear 2048→4096  →  hidden vector (4096-d)    │
        └─────────────────────────────────────────────────┘
                            │
                            ▼  injected back into the residual stream
                        DeepSeek-V4-Flash backbone (frozen)
```

- **`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

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]
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 mirrored in the safetensors metadata (`config`, `format_version`,
`subdicts`) and in the sibling `config.json`.

---

## 🚀 Load a head (self-contained — no repo import)

```python
import json
import torch
import torch.nn.functional as F
from torch import nn
from safetensors.torch import load_file
from safetensors import safe_open

CKPT = "sft-c16/colar_head_sft.safetensors"   # or grpo/…, sft/…

class ReasoningCompressionHead(nn.Module):
    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(
            nn.Linear(hidden_size, mlp_dim), nn.SiLU(),
            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__()
        mlp_dim = mlp_dim or hidden_size // 2
        self.net = nn.Sequential(
            nn.Linear(latent_dim, mlp_dim), nn.SiLU(),
            nn.Linear(mlp_dim, mlp_dim), nn.SiLU(),
            nn.Linear(mlp_dim, hidden_size),
        )
    def forward(self, z):
        return self.net(z)

# read geometry from the safetensors metadata
with safe_open(CKPT, framework="pt") as f:
    cfg = json.loads(f.metadata()["config"])
hs, ld = cfg["hidden_size"], cfg["latent_dim"]

flat = load_file(CKPT)
def sub(prefix):
    return {k[len(prefix):]: v for k, v in flat.items() if k.startswith(prefix)}

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."))

# 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 (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 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)

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.

| Metric | Base (no injection) | +Head (2k budget) | +Head (14.5k budget) |
|---|---|---|---|
| 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>`% 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

- Evaluation is **small-n (30 docs, GSM8K only)**; treat accuracy as directional.
- No full-benchmark or multi-task evaluation yet.
- `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

Method after **CoLaR — Compressed Latent Reasoning**:

```bibtex
@article{colar2025,
  title   = {CoLaR: Compressed Latent Reasoning},
  journal = {arXiv preprint arXiv:2505.16552},
  year    = {2025}
}
```