PAPER.md

# A CoLaR-Style Reasoning Compression Head for DeepSeek-V4-Flash

**Preliminary technical report**

## Abstract

We attach a lightweight **CoLaR-style reasoning compression head** to the frozen
`nvidia/DeepSeek-V4-Flash-NVFP4` backbone. The head predicts a compact 1024-d
latent of the model's next reasoning step from a mid-network hidden state, and a
symmetric decoder maps that latent back into the backbone's hidden space so the
model can *reason from a compressed latent* instead of emitting a long
token-level chain of thought. Training is two-stage — supervised regression onto
a frozen latent target (SFT), then a backbone-free GRPO stage over the head's
Gaussian latent policy. Serving the head on DeepSeek-V4's hash-routed MoE
required a compile-safe capture-and-inject path that keeps `input_ids` flowing;
the result runs on the cudagraph fast path at **~65 tok/s**. On GSM8K the head
reduces median "thinking" tokens by **~40%** and the worst-case trace length by
**~4×**. Task-accuracy is evaluated only at small scale (n=30) and is reported as
preliminary; the length/latency reduction is the robust finding. This report
documents the method and the serving process rather than a benchmark result.

## 1. Background

**CoLaR** (Compressed Latent Reasoning, arXiv:2505.16552) proposes to fold
several reasoning tokens into a single latent step, so a model reasons in a
compressed latent space rather than token by token. We apply this idea as a
*bolt-on adapter* to an existing frozen backbone rather than retraining the base
model.

**DeepSeek-V4-Flash** is a 284B-parameter (13B active) mixture-of-experts model
with 43 layers, hidden size 4096, and a 1M-token context. It exposes an explicit
`<think>` reasoning phase. Critically for serving, its MoE layers use
**hash-based expert routing keyed on `input_ids`**, which constrains how latents
can be injected (Section 4).

## 2. Method

### 2.1 Compression head

The reasoning head is a 3-layer MLP with SiLU activations:

```
h_src (4096)  →  Linear 4096→2048 → SiLU → Linear 2048→2048 → SiLU → Linear 2048→2048
             →  [mu, log_sigma],  each 1024-d
```

It reads the hidden state at **source layer 35** (`h_src`) and outputs the
parameters of a diagonal Gaussian over a 1024-d latent, with `log_sigma` clamped
to `[-10, 2]` for stability. Reasoning tokens are grouped into chunks of
**compression_factor = 4**; from the source hidden at the *last* token of group
*k*, the head predicts the compressed latent of group *k+1*.

### 2.2 Frozen target projection

The regression target is a stable readout of the **target layer 42** hidden
state, produced by a **frozen** random projection `target_proj: Linear(4096,
1024, bias=False)`. Freezing it is deliberate: a trainable target would let the
objective collapse to zero. Both the source hidden and the projected target are
**LayerNorm-standardized** — raw DeepSeek-V4 hidden states span roughly ±16, and
an unnormalized summed-MSE would sit at 1e5–1e6 and never train.

### 2.3 SFT objective

The head is trained with a **soft-MSE + entropy** loss on the predicted mean:

```
L = 0.5 · Σ_d (target_d − mu_d)²  +  0.5 · β · Σ_d log_sigma_d ,   β = 0.01
```

The MSE term regresses `mu` onto the (detached, normalized) next-group latent;
the entropy term prevents the variance from collapsing. Hyperparameters: AdamW,
lr 1e-4, warmup 100 steps of a 1000-step linear schedule, batch size 1 ×
grad-accum 4, sequence length 4096. Training data is a corpus of agentic
reasoning traces (`Glint-Research/Fable-5-traces`). The backbone is frozen
throughout; only the head (and, when reconstruction is enabled, the decoder) is
updated. The SFT loss converges from ≈507 to the single digits over ~550 steps.

### 2.4 Latent decoder

On its own the latent has no path back into the model. A **LatentDecoder** —
symmetric to the head (`1024→2048→2048→4096`) — maps the latent back to a
hidden-space vector that can be injected as an input embedding, trained with a
reconstruction MSE so the latent stays invertible.

### 2.5 GRPO over the latent policy

Because the vLLM backbone is frozen and outside the autograd graph, the textbook
"log-prob of generated tokens" GRPO is unavailable. We instead treat the head as
a **stochastic policy over the injected latent**:

```
mu, log_sigma = head(h_src);   z ~ N(mu, diag(sigma²));   decode(z) → inject → generate → reward
```

REINFORCE then gives a gradient into the head alone,
`grad = E[ −log N(z | mu, sigma) · A ]`, with `A` the group-relative advantage
(the "G" in GRPO). The decoder and backbone are the non-differentiable
environment. GRPO uses AdamW at lr 1e-5, KL coefficient 0.1, 500 steps, with a
token-overlap / exact-match reward.

## 3. Checkpoints

Each checkpoint is a self-describing **v2 bundle**: the head, decoder, and
`target_proj` state dicts plus a `config` block (`hidden_size`, `latent_dim`,
`source_layer`, `target_layer`, `compression_factor`). We release two stages —
`sft` (post-supervised) and `grpo` (post-RL) — as flattened `.safetensors`
(prefixed keys `reasoning_head.*`, `decoder.*`, `target_proj.*`) alongside the
original `.pt` bundles.

## 4. Serving process (engineering contribution)

Hosting the head for fast inference was the main practical hurdle.

**Injection.** vLLM's native `prompt_embeds` path nulls `input_ids` when
embeddings are supplied — but DeepSeek-V4's **hash-based MoE routing needs
`input_ids`**, so `--enable-prompt-embeds` crashes the engine at startup
(`DeepSeek V4 hash MoE routing requires input_ids`). We instead inject through an
**`embed_tokens` forward hook**: the hook overwrites the embedding at chosen
positions with the decoded latent while token ids continue to flow, so hash-MoE
routing still works. The decoded latent is injected at the final `<think>` prompt
position.

**Compile-safe capture.** Naively capturing a layer-35 hidden state forces
`enforce_eager` and forfeits cudagraphs. We instead write the capture into a 2-D
`(tokens, hidden)` buffer allocated **outside** the cudagraph memory pool, filled
by an in-loop `copy_`, so capture is compatible with compiled execution. A
byte-exact retrain gate confirmed the compiled-buffer capture equals the eager
forward-hook capture on the same prefill (`max_abs_diff = 0.0`), so no head
retraining was needed to move onto the fast path.

**Result.** The head serves on the compiled / cudagraph path at **~65 tok/s**,
with a one-time ~5-minute startup (weight load + cudagraph capture).

## 5. Results (preliminary)

GSM8K, 8-shot, temperature 0, GRPO head, **n = 30 documents**. Accuracy at this
scale is noisy; the **robust signal is reasoning-length reduction**.

| 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% reduction in median thinking tokens**; the longest trace shrinks from
  1740–2048 to ~415 tokens (**~4×**, ~3% of the available budget).
- The head **never skips** reasoning (skip% = 0) — it compresses rather than
  bypasses it.
- The drop in closed-`</think>`% is a **formatting artifact**, not truncation: a
  78× larger token budget did not increase closure, and the token-F1 GRPO reward
  never rewarded emitting the closing tag. Adding a `</think>`-closure bonus to
  the reward is the clear next step.
- The published DeepSeek-V4-Flash-Base GSM8K score (90.8) is quoted for context
  only; it is not a baseline reproduced by this harness.

## 6. Limitations and future work

- Evaluation is **small-n (30 docs, single task)**; a full multi-task benchmark
  is needed before any accuracy claim.
- **Closure reward:** add a `</think>`-closure bonus and retrain GRPO.
- **Closed-loop latent decoding** (emitting reasoning steps with no token id,
  bypassing the sampler) requires scheduler / KV-slot changes and is future work;
  current injection is prompt-position only.
- **Hardware ceiling:** on 2× RTX PRO 6000 (96 GiB), after weights and cudagraphs
  only ~6.4 GiB of KV cache remains, capping usable context at ≈15.8k tokens.

## References

- CoLaR: Compressed Latent Reasoning. arXiv:2505.16552, 2025.
- DeepSeek-V4-Flash — `deepseek-ai/DeepSeek-V4-Flash`; NVFP4 build
  `nvidia/DeepSeek-V4-Flash-NVFP4`.