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