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.