PaperlayerView exact revision on arXiv

Source-linked AI summary

A Model with No Head and Many Thoughts

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

arXiv:2608.31069v1cs.LGcs.CL

TL;DR

Large language models incur costly full-vocabulary projections during reasoning, and existing soft-thinking methods keep intermediate states tied to token-embedding mixtures. Soft Latent Thinking replaces that projection with a lightweight latent projector for continuous reasoning, while retaining standard decoding for final responses. Across evaluations, it reduces per-step compute and improves higher-k performance, with modular deployment but limited out-of-domain transfer from domain-specific initialization.

  • Problem

    Existing soft-thinking approaches retain the full vocabulary head, tying latent reasoning states to token semantics and creating a computational bottleneck for long reasoning traces.

  • Method

    Soft Latent Thinking replaces the vocabulary projection only during reasoning with a compact latent projector that maps hidden states to continuous embeddings over a learned basis, trained with a Gumbel-based per-step likelihood.

  • Results

    Soft Latent Thinking outperforms SofT-GRPO and other baselines at higher pass@k while reducing reasoning-token use and per-step computation.

  • Takeaways & Limitations

    SLT supports continuous reasoning without discrete intermediate token generation and enables modular projector selection while sharing a frozen backbone.

  • Takeaways & Limitations

    A math-initialized projector limits out-of-domain transfer: performance degrades on HumanEval, although disabling the projector preserves performance with standard soft thinking and LoRA.

Abstract

from arXiv · show

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.

1 Introduction

Existing soft-thinking methods retain the full vocabulary head, tying reasoning to token embeddings and making long traces expensive. Soft Latent Thinking replaces that operation during reasoning with a compact latent projector, supports RL training, and improves higher-k performance while using fewer reasoning tokens.

  • Motivation: Existing soft-thinking methods form a V-way vocabulary distribution at every reasoning step, creating computational and representational limitations.Their intermediate states remain tied to the span of token embeddings, while vocabulary projection and normalization become bottlenecks for long traces.
  • Method: Soft Latent Thinking replaces the vocabulary projection during reasoning with a compact projector over a learned latent basis of size K ≪ V.The projector directly synthesizes the next continuous reasoning state instead of producing a distribution over tokens.
  • Method: SLT defines a tractable per-step likelihood through sampled Gumbel variables, enabling policy-gradient RL for continuous latent reasoning.Training can update the projector and basis alone or combine them with LoRA without full model retraining.
  • Method: Figure 1 contrasts discrete-token decoding, vocabulary-based soft thinking, and SLT's direct soft-embedding generation without full vocabulary projection.SLT uses a lightweight encoder-decoder projector during reasoning.
  • Results: Compared with SofT-GRPO, SLT improves average pass@32 across five mathematical reasoning benchmarks while using fewer reasoning tokens and a cheaper per-step operator.The analysis attributes higher-k gains primarily to greater rollout diversity and coverage, despite slightly lower precision for individual samples.

2 Related Work

Related work extends chain-of-thought reasoning into continuous embedding space and adds stochastic sampling or reinforcement learning to improve exploration and optimization. These approaches address limitations of discrete intermediate tokens but introduce distinct challenges around diversity and policy-gradient likelihoods.

  • Chain-of-thought: Chain-of-thought improves complex-task performance by allocating extra computation to intermediate step-by-step reasoning.The related work frames CoT as often outperforming direct prediction.
  • Continuous reasoning: Continuous-space methods maintain latent reasoning trajectories without committing to discrete surface forms at intermediate steps.Examples include recurrent hidden-state feedback in Coconut and continuous denoising in Diffusion-of-Thought.
  • Soft thinking: Soft thinking constructs intermediate steps as weighted mixtures of token embeddings, enabling smooth transitions across token semantics but potentially causing deterministic behavior.Gumbel-Softmax sampling was introduced to promote diversity in reasoning paths.
  • Reinforcement learning: RLVR methods such as GRPO sample trajectory groups and favor higher-reward samples, while early soft-thinking integrations underperform discrete-token counterparts because soft tokens limit exploration.This limitation motivates stochastic treatment of continuous reasoning paths.

3 Background

The background describes how standard decoding and soft thinking generate intermediate representations, then explains how stochastic soft thinking makes continuous trajectories compatible with policy-gradient optimization. SofT-GRPO assigns likelihoods to underlying Gumbel variables rather than directly to soft-token vectors.

  • Standard decoding: Standard autoregressive decoding projects each hidden state through a full vocabulary head before selecting the next token.The hidden state h_t is mapped through W_head over a vocabulary of size V.
  • Soft thinking: Soft thinking computes vocabulary probabilities with the LM head and feeds a probability-weighted embedding mixture into the next step.This enables continuous reasoning without training but still requires a full softmax over V.
  • Stochastic soft thinking: Gumbel-Softmax sampling uses vocabulary log-probabilities and a temperature parameter to encourage diverse soft reasoning paths.The resulting soft token is used as the next input.
  • SofT-GRPO: SofT-GRPO represents each reasoning step as a continuous vector and must address both limited exploration and the absence of categorical per-step probabilities.Its solution assigns likelihood to underlying Gumbel variables rather than directly to the soft token.
  • Policy update: During policy updates, rollout Gumbel variables remain fixed, so the policy-dependent term is the candidate probability p_t,i.The importance ratio therefore evaluates how the current policy changes likelihood relative to the old policy.

4 Soft Latent Thinking

Soft Latent Thinking replaces the vocabulary projection used during latent reasoning with a learned, lower-dimensional projector and continuous latent basis. It preserves standard token decoding for final answers while reducing reasoning-step computation and allowing states beyond mixtures of discrete token embeddings.

  • Latent projector: Soft Latent Thinking replaces full-vocabulary soft-token formation with an independent latent projector of size K ≪ V, then decodes the latent mixture into model embedding space.Unlike full-vocabulary soft thinking, the latent space is not constrained to discrete token semantics.
  • Latent projector: K ≈12k–24k instead of V ≈150k reduces per-step reasoning computation while removing the requirement that intermediate states be expressible as mixtures of discrete token embeddings.The final response still uses standard LM-head decoding.
  • Latent projector: The encoder maps hidden state h_t to K logits, applies Gumbel–Softmax over K categories, and produces mixture weights for the next soft embedding.The resulting soft embedding is fed into the next model step during reasoning.
  • Initialization: The latent basis is initialized from frequent domain-token embeddings, but the encoder and decoder basis are subsequently trained independently from the pretrained vocabulary head.Although the top 5000 tokens cover almost 99% of mathematical-reasoning token occurrences, ablations find K ≈12k optimal, suggesting benefits from rarer tokens.
  • Training: Gumbel-perturbed latent actions provide per-step likelihoods for GRPO policy-gradient updates, enabling training with projector-and-basis adaptation alone or with optional LoRA.The method replays sampled soft embeddings and evaluates current-policy likelihoods for the latent trajectory.
  • Inference: Latent reasoning ends when the soft embedding aligns with a thinking-boundary embedding, after which the model switches to standard autoregressive answer decoding.The boundary is </think> when available and otherwise may use the \boxed embedding.

5 Experiments

Experiments evaluate Soft Latent Thinking across mathematical, out-of-domain, efficiency, and ablation settings. SLT improves higher-budget multi-sample accuracy while reducing reasoning cost, but its benefits depend on projector training, size, temperature, and domain alignment.

  • Experimental setup: Experiments use DeepSeek-R1-Distill-Qwen-1.5B and LLaMA-3.2-3B-Instruct across five mathematical reasoning benchmarks, with comparisons against base, GRPO, Soft-Thinking, and SofT-GRPO models.Training uses DeepScaleR with Soft-GRPO rewards, frozen backbones, rank-64 LoRA adapters, and the projector.
  • Main results: 86.22 average pass@32 versus 83.23 for the base model and 85.18 for SofT-GRPO on Qwen; 60.70 versus 56.26 and 57.06 on LLaMA.The gains primarily improve the multi-sample accuracy–efficiency tradeoff rather than uniformly dominating at every sampling budget.
  • Main results: Higher-k gains are driven primarily by greater diversity across reasoning rollouts, although individual samples can be slightly less precise than full-vocabulary soft thinking.This makes SLT most useful when multiple rollouts are affordable or already required, including RL training and pass@k inference.
  • Out-of-domain evaluation: 97.0 versus 95.5 on GPQA and 92.7 versus 94.5 on HumanEval when the projector is disabled, matching or exceeding SofT-GRPO out of domain.The projector transfers well to GPQA but degrades HumanEval with math-initialized representations; bypassing it preserves LoRA performance.
  • Computational efficiency: SLT uses fewer reasoning tokens and lower per-step FLOPs than SofT-GRPO, producing computational savings during chain-of-thought generation.For d = 1536, V = 150k, and K = 16k, the vocabulary-projection step is reduced by approximately 5×; graph-level decode speedup is approximately 1.05×.
  • Ablations: Projector training and joint backbone adaptation are necessary: untrained compression degrades performance, while joint projector-and-LoRA training achieves the best results.K ≈6k gives the best pass@1, whereas K ≈12k gives the best pass@16 and pass@32; inference temperature τg = 0.5 works best for both models.

6 Discussion

SLT reduces reasoning-time projection cost with a decoupled latent projector and supports modular domain adaptation. Its stochastic latent operator favors diverse reasoning paths, improving higher-k coverage despite potentially lower single-sample precision.

  • Training efficiency: ∼5–10× lower per-step projection cost follows from replacing V ≈150k with K ≈12–24k, with savings compounding across long traces and GRPO rollouts.End-to-end gains are smaller because attention and MLP layers remain.
  • Decoupled embeddings: SLT decouples latent-reasoning embeddings from token identity, allowing its decoder to learn representations optimized for reasoning.Standard soft thinking reuses the embedding table for both latent reasoning and discrete output.
  • LoRA as decoding adapter: LoRA adapts a frozen backbone to projector-generated soft embeddings, helping preserve out-of-domain knowledge when the projector is disabled.This separates projector-based reasoning from backbone adaptation without requiring full model retraining.
  • Modularity: Separate projectors can target math, code, or science while sharing one frozen backbone, and users can select or disable them at inference.This creates a plug-and-play deployment setup for domain-specific reasoning.
  • Exploration via stochasticity: Controlled Gumbel-Softmax stochasticity improves higher-k coverage by exploring diverse reasoning paths, whereas full-vocabulary soft thinking has higher pass@1 but lower pass@16.The trade-off favors diverse rollouts at larger k rather than maximum precision from an individual sample.

7 Conclusion

Soft Latent Thinking uses a lightweight compressed-vocabulary projector to perform continuous latent reasoning while reducing chain-of-thought computation. Experiments report superior higher-k performance and improved token efficiency, with stochasticity supporting diverse reasoning paths.

  • 7 Conclusion: SLT enables continuous embedding-space reasoning through a lightweight projector over a compressed vocabulary.It replaces the full vocabulary projection during chain-of-thought.
  • 7 Conclusion: Replacing the full vocabulary projection reduces per-step FLOPs during chain-of-thought while maintaining competitive or superior performance.The computational reduction comes from using a small learned projector.
  • 7 Conclusion: Experiments on DeepSeek-R1-Distill-Qwen-1.5B and LLaMA-3.2-3B-Instruct report higher pass@k than SofT-GRPO and other baselines, with improved token efficiency.The reported advantage is emphasized at higher pass@k values.
  • 7 Conclusion: The compressed vocabulary introduces beneficial stochasticity that explores diverse reasoning paths across samples.The analysis attributes this exploration to subtle adjustments in token selection by the projector.

8 Limitations

The projector’s domain-specific initialization limits transfer beyond its target domain when enabled. On HumanEval, a math-initialized projector degrades performance, while disabling it and using standard soft thinking with LoRA preserves performance.

  • 8 Limitations: Domain-specific projector initialization limits out-of-domain transfer when the projector is enabled.The limitation is linked to initialization on tokens such as mathematical vocabulary.
  • 8 Limitations: On HumanEval, code performance degrades with the math-initialized projector, but disabling it and using standard soft thinking with LoRA preserves performance.The paper suggests multi-domain projectors and domain-agnostic initialization as future directions.
Loading 2608.31069v1…