A Model with No Head and Many Thoughts

AI summary

Overview

Research area: Efficient reasoning and reinforcement learning for large language models — specifically, methods that let a model "think" in continuous embedding space instead of emitting discrete tokens.

Technical level: Intermediate. The paper assumes familiarity with transformer decoding, the LM head / vocabulary projection, softmax and Gumbel-Softmax sampling, and policy-gradient RL (GRPO/PPO).

Scope: The paper introduces Soft Latent Thinking (SLT), a method that replaces the full vocabulary projection during chain-of-thought reasoning with a small learned projector, and evaluates it on mathematical and out-of-domain reasoning benchmarks against discrete-token and soft-thinking baselines.

What This Paper Is About

Large language models normally decode one token at a time by projecting a hidden state through a large vocabulary head, and every intermediate chain-of-thought step must be expressed as a discrete token. Existing "soft thinking" methods avoid discrete tokens by mixing token embeddings, but they still compute a full distribution over the entire vocabulary (roughly 150k entries) at each reasoning step, which is expensive and ties latent reasoning states to discrete token semantics. This paper asks whether reasoning can instead happen in a compressed continuous latent space, using a lightweight projector trained with reinforcement learning, so that per-step computation drops while multi-sample accuracy improves.

Key Contributions

1. A latent projector that replaces the LM head during reasoning. SLT maps a hidden state to coefficients over a learned latent basis of size K ≪ V and directly synthesizes the next continuous reasoning state, avoiding the V-way vocabulary projection entirely. Standard token decoding is kept unchanged for the final answer.

2. A trainable formulation under RL. The authors define a per-step likelihood for the latent reasoning action using the sampled Gumbel variables, producing a well-defined policy for policy-gradient updates even though intermediate states are continuous. Training only needs the projector (and basis) updated, optionally combined with LoRA adapters on backbone weights.

3. Empirical gains in the multi-sample regime with token savings. On DeepSeek-R1-Distill-Qwen-1.5B and LLaMA-3.2-3B-Instruct across five math benchmarks (AIME2024, AIME2025, AMC23, MATH-500, GSM8K), SLT reaches the highest average pass@32 among the soft-thinking approaches tested while using fewer reasoning tokens on average than SofT-GRPO.

4. A modular deployment story. Because the projector and LoRA adapters can be separated, the projector can be disabled at inference for out-of-domain tasks and the model reverts to full-vocabulary soft thinking with no penalty relative to the strongest baseline, according to the authors' out-of-domain results.

Main Findings

• Best average pass@32 among soft-thinking methods on the tested models. On DeepSeek-R1-Distill-Qwen-1.5B, SLT achieves 86.22 average pass@32 versus 83.23 for the base model and 85.18 for SofT-GRPO. On LLaMA-3.2-3B-Instruct, it achieves 60.70 versus 56.26 for the base model and 57.06 for SofT-GRPO.

• Gains concentrate at higher k, not uniformly. The authors report that SLT does not dominate at every sampling budget; individual samples can be slightly less precise (for example, on DeepSeek 57.32 average pass@1 versus 61.39 for SofT-GRPO), but increased diversity across rollouts drives improved coverage at pass@32.

• Fewer reasoning tokens on average than SofT-GRPO. On DeepSeek-R1-Distill-Qwen-1.5B, average #Token is 6073.0 for SLT versus 6517.3 for SofT-GRPO, and average #Token_c (tokens for correctly solved queries) is 4643.6 versus 5831.9.

• Lower per-step projection cost. The LM head costs 2 × d × V FLOPs, whereas the projector costs 4dK. For d = 1536, V = 150k and K = 16k, the paper reports a reduction of V/(2K) ≈ 5× on the vocabulary projection step. A mini-SGLang prototype reports an approximately 1.05× graph-level decode speedup, described as the cleaner measurement of the algorithmic gain.

• Serving throughput improves modestly in the prototype. With projector scale 8 (the largest and least favorable tested setting), TPS rises from 675.4 to 723.3 (1.071×) at batch 1 for Llama-3.2-1B, from 299.3 to 314.7 (1.051×) at batch 1 for Llama-3.2-3B, and from 519.9 to 652.9 (1.256×) at batch 4 for Llama-3.1-8B, among other listed configurations.

• Training the projector matters. Replacing trained projector weights with initial ones at fixed LoRA gives 28.3/70.1/76.7 at @1/@16/@32 versus 28.7/74.3/80.0 for the trained projector on AIME2024 with DeepSeek-R1-Distill-Qwen-1.5B, and increases token usage from 10280 to 11383. Full-vocabulary soft thinking with the same LoRA weights reaches 29.5/72.3/80.0 with 11242 tokens.

• Naive vocabulary compression at inference alone hurts. With no fine-tuning, an initialized projector scores 11.3/56.3/63.3 versus 27.3/66.7/70.0 for full-vocabulary soft thinking. Applied to a SofT-GRPO checkpoint, the pruned projector scores 15.8/60.6/70.0 versus 32.6/76.7/80.0.

• Projector size trades single-attempt focus against multi-sample diversity. K ≈ 1.5k gives 20.4/58.2/66.7, K ≈ 4k gives 29.7/62.0/66.7, and K ≈ 8k gives 28.7/74.3/80.0 on AIME2024. The text states that K ≈ 6k achieves the best pass@1 and K ≈ 12k the best pass@16 and pass@32, and that K ≈ 12k is optimal overall.

• Joint training of projector and backbone is necessary. Training only the projector with a frozen backbone yields 26.0/65.4/70.0, weaker than the base model, while joint projector plus LoRA training gives 28.7/74.3/80.0 with 10280 tokens.

• Temperature is model-dependent for training, shared for inference. τ_g = 0.1 worked best for LLaMA-3.2-3B-Instruct and τ_g = 0.5 for DeepSeek-R1-Distill-Qwen-1.5B during training; both models performed best with τ_g = 0.5 at inference.

• A preliminary 9B check reproduced the qualitative pattern. On AIME2024 with Qwen3.5-9B, the projector model scored 68.4/92.8/93.3 pass@1/@16/@32 with 9288.1 average tokens, versus 76.4/93.2/93.3 with 23292.1 average tokens for the base model. The authors explicitly describe this as an untuned sanity check, not a scaling study.

• Out-of-domain behavior is mixed but has a workaround. On GPQA Diamond the projector preserves accuracy; on HumanEval it degrades (SLT 54.0/77.9/87.2 versus SofT-GRPO 71.2/88.4/94.5 at @1/@8/@32). Disabling the projector gives 35.4/83.0/97.0 on GPQA and 70.0/88.0/92.7 on HumanEval.

Methodology in Plain English

Normally, when a model "thinks," every thinking step is a token: the model takes its hidden state, multiplies it by a big matrix to score every word in the vocabulary, samples one, and feeds that word's embedding back in. Soft-thinking variants skip the sampling and feed back a weighted blend of every word embedding instead — but they still pay for scoring all ~150k vocabulary entries each step, and the blend is still confined to being a mix of real words.

This paper keeps the blending idea but changes what the model blends. A small encoder compresses the hidden state into just K numbers — the weights over a learned set of latent directions rather than over words. Those weights are sampled with Gumbel-Softmax, which injects controlled randomness so the model explores different reasoning paths. A small decoder then turns those weights back into a vector in the model's embedding space, which is fed in as the next reasoning step. Because this whole operation lives on K ≈ 12k–24k dimensions instead of V ≈ 150k, it is cheaper, and because the latent directions are learned rather than copied from the vocabulary, reasoning states are not forced to look like blends of words.

For training, the authors reuse the trick from SofT-GRPO: instead of trying to assign a probability to a continuous vector, they assign likelihood to the random Gumbel noise that produced the sampled mixture. That noise is fixed for a given rollout, so the only thing depending on the current policy parameters is the K-way probability distribution — which makes a standard GRPO importance ratio well defined. The final answer is still decoded normally with the LM head, and the reward comes from whether the answer is correct. Reasoning stops when the soft embedding becomes sufficiently cosine-similar to the </think> token embedding (or the boxed token embedding for models without explicit thinking boundaries). The projector is initialized by copying rows from the pretrained LM head and embedding table for the K most frequent tokens in the target domain; the frequency distribution was computed on the math subset of OpenThoughts-114k, where the top 5000 tokens cover almost 99% of occurrences. Training used the DeepScaleR dataset with the backbone frozen and LoRA adapters of rank 64 on all attention and MLP modules trained jointly with the projector.

Why This Matters

Impact on research. The paper challenges an implicit assumption in soft-thinking work: that latent reasoning must live in the span of token embeddings and must pay the full vocabulary projection cost. It shows that decoupling the reasoning operator from the output head is both cheaper and, in the multi-sample regime, more accurate — and it reframes vocabulary compression as something that must be trained rather than applied post hoc. It also isolates a diversity effect: the compressed vocabulary is less precise per sample but produces better coverage at high k, which matters for RL rollouts and for evaluation protocols that report pass@k.

Real-world applications:

• Reinforcement learning training pipelines for reasoning models, where many rollouts per query are already required and a cheaper per-step reasoning operator plus shorter traces compounds over thousands of updates.
• Multi-sample inference and answer aggregation, where generating many candidate solutions and picking or voting among them is the norm — the setting where SLT's pass@32 advantage is largest.
• Modular domain deployment, where a shared frozen backbone is paired with different domain-specific projectors (math, science) and the projector is disabled for general-purpose or code-oriented traffic.
• Latency- and cost-sensitive serving of long chain-of-thought traces, where the paper's prototype reports higher throughput across the tested models and batch sizes.

Industry relevance. The method requires only LoRA adapters on a frozen backbone plus a small projector, which is a comparatively cheap post-training recipe rather than a full pretraining change. Bypassing the LM head also directly targets the dominant per-step cost in long reasoning traces, and the reported graph-level decode speedup of about 1.05× is a measurement an inference-serving team can reason about. The plug-and-play framing — separate projector for reasoning capacity, LoRA for domain adaptation — maps cleanly onto how production teams already manage adapters.

Future Directions

• Multi-domain and domain-agnostic projector initialization. The authors note the projector is initialized on domain-specific tokens, which limits out-of-domain transfer when it is enabled, and suggest exploring multi-domain projectors and better domain-agnostic initializations.

• Better initialization of the latent basis. A preliminary PCA/SVD initialization over the embedding table performed poorly, which the authors read as evidence that useful reasoning traces occupy a specialized subspace rather than a high-variance global embedding subspace — leaving the question of what a better basis would look like open.

• A full stopping-threshold sweep and scale study. Only preliminary sensitivity checks were run on the cosine-similarity stopping thresholds (intermediate similarity was observed to stay below 0.2 and rise sharply near the final answer), and the larger-model result on Qwen3.5-9B is a single untuned configuration rather than a full baseline sweep or scale-specific hyperparameter tuning.

• Understanding and controlling the diversity–accuracy tradeoff. The paper attributes pass@k gains to increased rollout diversity but does not report a mechanism for tuning that tradeoff directly, which is a natural next question for both single-attempt accuracy and RL training stability.

Target Audience

Researchers and engineers working on LLM reasoning efficiency, chain-of-thought alternatives, and reinforcement learning with verifiable rewards for reasoning models. It is also relevant to inference-serving and post-training engineers interested in adapter-based modular deployment, and to readers already familiar with soft-thinking and SofT-GRPO who want to see whether the vocabulary projection can be removed from the reasoning loop.

Authors’ abstract

Large language models decode by projecting hidden states through a large vocabulary head at every step. This operation is computationally costly and forces all reasoning to be expressed in discrete tokens. We introduce Soft Latent Thinking, a method that replaces the LM head during reasoning with a lightweight projector, enabling autoregressive rollout in embedding space where reasoning steps remain continuous rather than tokenized. Experiments on DeepSeek-Qwen-1.5B and LLaMA-3.2-3B show that Soft Latent Thinking consistently improves pass@k across all k while reducing per-step compute during chain-of-thought. Our method achieves the highest pass@32 among all soft-thinking approaches, demonstrating that effective reasoning can be carried out in continuous space without discrete token generation.