Skip the vocab head at think time: Soft Latent Thinking hits 86.22 avg pass@32

A Model with No Head and Many Thoughts

Nikita Koriagin, Yaroslav Aksenov, George Bredis, Gleb Gerasimov, Nikita Balagansky, Daniil Gavrilov

EMNLP 2026

cs.LG, cs.CL

2026-09-01

Soft Latent Thinking replaces the 150k vocab head with a ~12k projector at think time. On Qwen-1.5B, avg pass@32 rises from SofT-GRPO's 85.18 to 86.22; Mean@1 falls.

What problem this solves

Soft thinking already replaced discrete CoT tokens with a probability-weighted mix of embeddings. The Zhang et al. 2025 recipe still runs a full vocabulary softmax first: a V-way distribution, V around 150k, then mixes every token embedding. On long traces that projection is expensive, and every latent state has to live in the span of the embedding table.

SofT-GRPO added Gumbel-Softmax so GRPO can train this path. The vocab head stayed. T-Tech asks a narrower question: during thinking only, can you drop the head, swap in a much smaller projector, and keep the same Gumbel-GRPO loop.

Method

Soft Latent Thinking (SLT) changes the reasoning operator. Final answers still go through the original LM head.

Each step takes the last-layer hidden state. A linear encoder compresses it to K logits, Gumbel-Softmax mixes over those K directions, and a linear decoder maps the mixture back into embedding space as the next input. K is 8 times hidden size: 12288 for Qwen-1.5B, 24576 for LLaMA-3.2-3B. Against a 150k vocab, that projection is roughly 5 to 10 times cheaper.

Init copies the K most frequent domain tokens. From the math split of OpenThoughts-114k, matching LM-head rows go into the encoder and embedding rows into the decoder, then the copies untie and train with LoRA. The backbone stays frozen. LoRA rank 64 covers all attention and MLP modules. Data is DeepScaleR. Reward is correctness only.

Stopping is crude: when the soft embedding's cosine to </think> (or \boxed on LLaMA) clears a threshold, decoding falls back to tokens. Mid-trace that cosine usually sits below 0.2 and jumps near the answer, so the threshold is not very brittle.

Policy gradients do not put a density on the soft embedding. They put likelihood on the stored Gumbel variables from sampling. Same trick as SofT-GRPO, with V replaced by K.

Results

Five math benchmarks: AIME 2024/2025, AMC23, MATH-500, GSM8K. @1 is mean Pass@1 over 32 samples; @16 and @32 are Pass@k.

Method (Qwen-1.5B)Mean@1pass@16pass@32
Discrete CoT, no FT58.0980.5483.23
Discrete GRPO61.2880.1682.39
SofT-GRPO61.3983.5485.18
SLT57.3282.6686.22

The abstract says pass@k improves at every k. Table 1 does not. On Qwen-1.5B, SLT Mean@1 is 57.32 against SofT-GRPO 61.39, discrete GRPO 61.28, and untuned discrete CoT 58.09. The win is coverage: average pass@32 reaches 86.22 versus 85.18 and 83.23. LLaMA-3.2-3B has the same shape. Mean@1 drops from SofT-GRPO's 32.83 to 26.58; pass@32 rises from 57.06 to 60.70.

Traces are shorter. Qwen five-set average: 6073 tokens versus SofT-GRPO's 6517; on correctly solved items, 4644 versus 5832. On AIME2024 correct traces, 7006 versus 10756.

Ablations pin the cause. Compress the vocab at inference with no joint training, and AIME2024 pass@32 falls from full-vocab soft thinking's 70.0 to 63.3. Train the projector with a fully frozen backbone: pass@32 is 70.0, below the untuned discrete 73.3. K=1536 is too tight (66.7). K=6144 has the best single-shot AIME Mean@1 at 29.7, but pass@32 stays 66.7. K=12288 is the one that lifts pass@16/@32 to 74.3/80.0.

Out of domain is blunt. With the math projector on, HumanEval pass@32 falls from SofT-GRPO's 94.5 to 87.2. Turn the projector off and keep the jointly trained LoRA: 92.7. On GPQA Diamond the on/off numbers are 95.5 and 97.0; science notation overlaps math, so the hit is smaller.

Qwen3.5-9B is one untuned AIME2024 probe: projector pass@1 68.4 versus base 76.4, pass@32 tied at 93.3, average tokens 9288 versus 23292. Same direction. Not a scaling study.

The vocab projection is about 5 times cheaper; attention and MLP still dominate. A mini-SGLang prototype sees about 1.05x decode. The best cell is Llama-3.1-8B at batch 4, 1.256x.

Why it matters

This is for settings that already sample many rollouts: RL training, pass@k, best-of-n. Do not use it for a single answer. Mean@1 is a regression.

Deployment can be modular. Turn the projector on in-domain for coverage and shorter traces; turn it off out of domain and keep LoRA. No full-parameter finetune.

The paper does not show that continuous thought beats discrete tokens. It shows that, on the soft-thinking plus GRPO line, swapping the vocab head for a compressed projector trades single-sample precision for diversity and a cheaper reasoning step.

Limitations

The authors say the projector is initialized on math-frequent tokens, so leaving it on hurts transfer. The HumanEval drop is in the main text. Early PCA/SVD init over the embedding table failed, which they read as evidence that useful traces sit in a specialized subspace, not a high-variance global one.

No full sweep of the stop threshold. The 9B run is a single untuned point. There is no Coconut-style hidden-state loop baseline.

The abstract and Table 1 disagree. The abstract claims pass@k gains at every k; Mean@1 is clearly below SofT-GRPO and discrete GRPO. Later sections walk that back and credit diversity plus high-k coverage. Trust the table.

End-to-end speedup is about 5%. If the vocab head is not the bottleneck, the projector barely shows up.

Terms

Source

What people are saying

Related papers

All paper explainers