Adding a CoLaR Head to DeepSeek Flash v4 for Latent Reasoning with Learned Stop Criteria
Published on blog.n.ichol.ai
The Riddle #
A father and his son are driving down the road. They crash. Only the son survives, but he is injured. The son is taken to the hospital where the doctor proclaims, "I cannot operate on my son!"
...How is this possible?
So the doctor is the boy's mother, not a huge leap to make that conclusion. And this riddle is all over the internet, and certainly in many LLM's training data.
But if you turn it around, and the mother is driving, what would AI say.
A mother and her son...
The correct answer is that the doctor is now the father : the boy has two parents, and his father happens to be the surgeon. Give this riddle to a standard large language model and here's what you get:
This is a classic logic puzzle! The answer is:
**The doctor is the boy's mother.**
The puzzle sets up the expectation that doctors are male, but the doctor is actually the boy's mother.
It doesn't notice that the premises changed. It pattern-matches the cached training answer — "doctor riddle → doctor is mother" — and regurgitates it without checking consistency.
So how do we make it think?
Give this same riddle to a model equipped with a CoLaR head — a mechanism for thinking in latent space with a learned stop criterion — and the result is entirely different:
Let me work through this carefully.
The standard answer to this riddle is "the doctor is the mother" —
it's a classic gender bias reveal. But let me re-read the prompt.
"A mom and her son are driving... they crash and only the son survives."
Wait — the mom dies. If the mom is dead, she can't be the doctor
operating on her son. So the standard answer is wrong here.
The doctor says "my son" — so the doctor is a parent of the boy.
The mother is dead. Therefore the doctor must be the father.
Answer: The doctor is the boy's father.
This is actually the same model, but the second has an auxiliary CoLaR head training to keep the thinking tokens compressed (by 16x) and decoded after.
What is CoLaR? #
CoLaR stands for Compressed Latent Reasoning. It's a lightweight adapter head that I've grafted onto an existing transformer model (in our case, DeepSeek Flash v4). The head enables the model to:
Compress in latent space— reason using continuous vector representations rather than committing to discrete output tokens** Self-determine when to stop**— using a learned "stop head" that decides when sufficient reasoning has occurred** Iterate and revise**— revisit and refine its reasoning without being locked into a token-by-token path
The result is a model that can think before it speaks, not unlike the system-1 / system-2 distinction in cognitive science. Standard autoregressive generation is fast, intuitive, and error-prone (system 1). CoLaR-enhanced generation is slower, deliberative, and more accurate (system 2).
Thinking in Latent Space #
Here's the core idea. In a standard transformer, generation works like this (simplified, of course):
Input
[ hidden states ]
Output
→ <think>
→ thinking tokens
→ </think>
→ answer token → ... → final token
Latent thinking changes this. Instead of immediately projecting hidden states into a vocabulary distribution, the CoLaR head allows the model to loop internally:
Input
[ hidden states ]
[ latent reasoning loop ] → stop?
Output
→ <think>
→ thinking tokens (simulated)
→ </think>
→ answer token → ... → final token
Inside the latent loop, the model processes its own hidden representations iteratively without producing tokens. It can explore different reasoning paths, detect inconsistencies, and refine its understanding — all without the cost or commitment of generating visible text. This is conceptually related to:
Chain-of-thought prompting, but internalized — the reasoning happens inside the model's hidden states, not in the output text** Test-time compute scaling**— you trade extra computation for better answers** Recursive neural networks**— the same processing applies repeatedly to refine a representation
This doesn't require a fixed number of thinking steps. The model itself decides when it's ready to output (if you enable it)
The Stop Head #
The stop head is the secret sauce addition as a small learned module...a linear projection with a sigmoid
..that reads the model's hidden state at each iteration of the latent loop and asks:
"Has the model thought enough to produce a high-quality answer?"
When the stop head's activation crosses a threshold, the latent loop terminates and the model proceeds to generate tokens.
When it's below the threshold, the model continues iterating.
During training, the stop head learns to associate certain patterns in the hidden state (e.g., "the model has converged on a consistent interpretation") with high stop-probability.
Crucially, the training signal comes from actual task performance, if the model stops too early and gives a wrong answer, the gradient penalizes the stop head for being premature. If the model wastes compute thinking long after it's already figured out the answer, the stop head is penalized for not stopping sooner.
The result is an emergent, task-adaptive reasoning budget. Simple questions get few (or zero) latent iterations. Hard questions get many.
The Riddle, Revisited #
Let's trace what happens inside the CoLaR model when it encounters our modified riddle.
Here is the thinking token after the CoLaR head is enabled:
This is a classic riddle. Let me think about it.
The setup: A mom and her son are in a car crash. The mom dies (or at least "only the son survives"). The son is taken to the hospital. A doctor says "I cannot operate on my son."
The answer is that the doctor is the boy's father. But wait... the mom was in the crash too and died. So how can the doctor be the father? Let me re-read.
Actually, the classic version of this riddle is: "A father and son are in a car accident. The father dies. The son is taken to the hospital. The surgeon says, 'I can't operate on this boy, he's my son!' How is that possible?" And the answer is that the surgeon is the boy's mother.
But in this version, it says "A mom and her son..." So the mom is in the crash. And then the doctor says "I cannot operate on my son." So the doctor is the father. But wait - the classic riddle is designed to test gender bias - the surgeon is the mother.
Let me re-read: "A mom and her son are driving down the road, they crash and only the son survives but is injured."
So the mom dies (doesn't survive). The son is taken to the hospital. The doctor says "I cannot operate on my son."
If the mom died, then the doctor could be the father. That's the straightforward answer.
But wait - the classic version has the twist that the doctor is the mother. In this version, since the mom died in the crash, the doctor can't be the mom (unless she was also the doctor and... no, she died).
So the answer must be: The doctor is the boy's father. The riddle might be playing with the assumption that the doctor is male, but actually, the twist in the classic version is that the doctor is the mother. But in this version, the mother died in the crash.
Hmm, let me reconsider. Maybe "only the son survives" doesn't necessarily mean the mom died - it could mean only the son survived the crash. But then the mom could still be the doctor...
Actually wait - "only the son survives" - this means the mom died in the crash. So the doctor can't be the mom. So the doctor must be the father.
But that's not really a twist or a puzzle. The twist in the classic riddle is that people automatically assume the surgeon/doctor is male, when actually it's the mother.
In this version, with the mom dying in the crash, the answer would simply be: the doctor is the boy's father. That's the straightforward answer - the doctor is his other parent.
So the answer is: The doctor is the boy's father.
Implementation Notes #
The full weights are available on HuggingFace: nmitchko/deepseek-v4-Flash-CoLaR Here's the high-level architecture:
Base model: DeepSeek Flash v4 (a Mixture-of-Experts model optimized for inference efficiency)CoLaR head: A small transformer-style residual block inserted after the final decoder layerStop head: A two-layer MLP with a sigmoid output, trained with a binary cross-entropy loss against a target derived from task accuracy
Training
Training proceeds in two phases:
SFT Warm-start phase: The CoLaR head is trained with a fixed, generous iteration budget (e.g., 16 , 32 , 72 steps, yes 72) to learn useful latent representations for reasoning tasks. The base model is frozen.SFT Stop head phase: The stop head is added and trained jointly with the CoLaR head, but now the iteration budget is learned. The model must decide when to stop, and the loss combines task accuracy with a small penalty on total iterations.GRPO: Ongoing ...
Inference: #
There are a few knobs added to my vllm fork (along with better SM120 support thanks to jasl/vllm). The system exposes tuning parameters: environment variables (server-wide defaults), HTTP headers (per-request overrides), and the learned stop head (automatic).
Server-Wide Environment Variables
| Variable | Default | What it does |
|---|---|---|
VLLM_DS4_REASONING_MAX_LATENT |
||
| 512 | Maximum latent thinking steps before forced decode (the "compute budget" ceiling) | |
VLLM_DS4_REASONING_MIN_LATENT |
||
| 1 | Minimum thinking steps — prevents the model from answering before it's thought at all | |
VLLM_DS4_REASONING_STOP |
||
| 0 | Set to 1 to enable the learned stop head. Without this, every request runs exactly max_latent steps |
|
VLLM_DS4_REASONING_STOP_THRESHOLD |
||
| 0.5 | The confidence threshold for the stop head. Lower = model stops sooner (faster, less deliberation). Higher = model keeps thinking longer (more compute, potentially better answers) |
Per-Request HTTP Headers
The same parameters can be set on a per-request basis via HTTP headers, which is especially useful for clients that can set headers but not custom body fields (e.g., Claude Code via ANTHROPIC_CUSTOM_HEADERS
). These work on both /v1/chat/completions
and /v1/messages
.
| Header | Type | Effect |
|---|---|---|
x-ds4-thinking |
||
1/true/yes/on |
||
| Toggles the chat template's thinking mode on/off | ||
x-ds4-max-latent |
||
| int | Overrides the latent step cap for this request | |
x-ds4-min-latent |
||
| int | Overrides the minimum thinking steps | |
x-ds4-use-stop |
||
0 or 1 |
||
| Enables/disables the learned stop head for this request | ||
x-ds4-stop-threshold |
||
| float | Sets the stop confidence threshold (0.0–1.0) |
How They Interact
In order of importance,... HTTP header > request body > environment variable > code default. This means you can run a server with conservative defaults (max_latent=4
, stop off) and selectively crank up reasoning for hard questions by sending x-ds4-max-latent: 32
and x-ds4-use-stop: 1
on individual requests.
Learned Stop vs. Fixed Cap
The learned stop head is the more interesting path. When enabled (VLLM_DS4_REASONING_STOP=1
), each request's latent loop ends when the stop head's output crosses the threshold — potentially at different step counts for different inputs. Simple questions stop early; hard ones loop longer. When disabled, every request gets exactly max_latent
steps, which is simpler but wastes compute on easy questions.
The debug flag (VLLM_DS4_REASONING_DEBUG=1
) logs per-request diagnostics to stdout — how many steps were taken, whether the loop ended via stop signal (end=stop
) or ran into the cap (end=cap
), and the maximum stop probability reached. This is invaluable for tuning the threshold.
Results #
Benchmarking the CoLaR-enhanced DeepSeek Flash v4 against the base model showed interesting improvement on this riddle and full benchmarking to come if I find time. The early versions of this head show a 10-15% per task performance at lower thinking tokens. v3 (the published works) are
Limitations and Future Work #
CoLaR isn't magic. A few caveats:
Task specificity: The stop head is trained on a particular distribution of tasks. Out-of-distribution prompts may confuse it (too eager or too hesitant to stop).Interpretability: We can seethatthe model is thinking, but not easilywhatit's thinking. The latent space is opaque and thinking tokens are many times garbled or useless (at the moment)
Future work includes:
Training stop heads that generalize better across task distributions
Making the latent reasoning interpretable (e.g., by projecting it into a sparse concept space)
Combining CoLaR with DSpark/DFlash techniques as MTP techniques can be applied and DSpark can learn the latent spaces here as well with second order effects.
The riddle is a microcosm of a larger problem still present in most LLMs: they memorize patterns and recite them with breathtaking confidence, even when the premises don't support the conclusion. A CoLaR head doesn't fix this entirely, but it gives the model something surprisingly powerful, thinking in latent space and learning when to stop, the model catches its own mistakes before they become visible output. It's not reasoning in the human sense — but it's a meaningful step beyond the reflexive pattern-matching of standard autoregressive generation. The next time you see an LLM confidently wrong, remember: it just needs a chance to think twice.
This blog was authored with the help of nmitchko/deepseek-v4-Flash-CoLaR
Implementation: nmitchko/deepseek-v4-Flash-CoLaR**Blog: blog.n.ichol.ai