nmitchko/deepseek-v4-Flash-CoLaR

🧠 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 backbone.

---

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

⚠️ This is an adapter

These weights are useless on their own. You must load the 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)

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

---

⚠️ 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:

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