---
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
---
# ๐ง 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.*
[](https://arxiv.org/abs/2505.16552)
[](https://huggingface.co/nvidia/DeepSeek-V4-Flash-NVFP4)
[](#-this-is-an-adapter)
[](#-checkpoint-layout)
---
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 `` % | 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-``% 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/`` 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}
}
```