Latent Reasoning with Supervised Thinking States

Paper: 2602.08332 Authors: Ido Amos, Avi Caciularu, Mor Geva, Amir Globerson, Jonathan Herzig, Lior Shani, Idan Szpektor Categories: cs.CL, cs.LG

Abstract

This paper proposes “Thinking States” that enable language models to reason while processing input, generating thinking tokens at regular intervals and converting them back to embeddings. This approach mirrors iterative chain-of-thought but operates during input processing rather than output generation. The token-based thoughts can be learned from natural language supervision using parallelizable teacher-forcing. Experimental results show the method approaches CoT quality on math tasks, matches it on 2-Hop QA with better latency, and achieves superior performance on state-tracking with better generalization.

Key Contributions

• Thinking States architecture: A novel mechanism for injecting reasoning steps during input processing
• Parallelizable training: Teacher-forcing enables efficient training unlike sequential CoT generation
• Latency advantages: Reasoning happens during input processing, reducing output generation time
• Strong generalization: Superior out-of-distribution performance on state-tracking tasks

Motivation

Standard chain-of-thought reasoning has a fundamental efficiency problem: reasoning tokens are generated sequentially during output, adding significant latency. What if the model could “think” while reading the input instead?

Thinking States address this by inserting special thinking tokens at regular intervals during input processing. These tokens serve as computational “scratch space” where the model can perform intermediate reasoning steps.

Architecture

Thinking Token Insertion

Given an input sequence $x_1, x_2, \ldots, x_n$, Thinking States inserts thinking tokens at fixed intervals $k$:

$x_1, \ldots, x_k, [\text{THINK}], x_{k+1}, \ldots, x_{2k}, [\text{THINK}], \ldots$

Each $[\text{THINK}]$ token:

1. Is processed by the model to produce a hidden state
2. The hidden state is projected to a token prediction
3. The predicted token’s embedding replaces the thinking token’s representation
4. Processing continues with this updated representation

The Think-Project-Replace Cycle

At each thinking position $t$, the model:

1. Computes hidden state $h_t$ from all previous tokens including prior thinking states
2. Projects $h_t$ through the language model head to get logits over the vocabulary
3. Selects the top token $w_t = \arg\max(\text{LM\_head}(h_t))$
4. Replaces the thinking token embedding with $\text{Embed}(w_t)$

This creates a discrete bottleneck that forces the model to compress its intermediate reasoning into interpretable tokens.

Training with Teacher-Forcing

A key advantage is that thinking tokens can be supervised with natural language targets. Given a reasoning chain from a teacher model, the thinking tokens are trained to predict the corresponding reasoning words:

$\mathcal{L} = \sum_{t \in \text{think\_positions}} -\log P(w_t^* | x_{<t}, \text{think}_{<t})$

where $w_t^*$ is the target thinking word from the teacher. Because the targets are known during training, this can be parallelized via teacher-forcing, unlike autoregressive CoT generation.

Experimental Results

Math Reasoning

On mathematical reasoning benchmarks, Thinking States approaches the quality of full chain-of-thought:

• Closes the gap between direct answering and CoT by a significant margin
• The thinking tokens learn to perform intermediate calculations
• Performance scales with the number of thinking tokens inserted

2-Hop Question Answering

On 2-Hop QA tasks requiring compositional reasoning:

• Matches CoT accuracy while providing better laten model learns to resolve the first hop during input processing
• Thinking tokens capture bridge entity identification

State Tracking

On state-tracking tasks (e.g., tracking object positions through a sequence of operations):

• Outperforms CoT with better generalization to longer sequences
• The iterative nature of thinking states naturally aligns with incremental state updates
• Superior out-of-distribution performance on sequences longer than training examples

Analysis

What Do Thinking Tokens Learn?

Inspection of the predicted thinking tokens reveals interpretable intermediate reasoning:

• On math: parutation results and carry operations
• On QA: bridge entities and intermediate facts
• On state tracking: current state representations after each operation

Latency Benefits

The key efficiency advantage comes from the parallelization of reasoning:

Method	Reasoning Phase	Output Phase	Total Latency
Direct	None	Short	Low
CoT	None	Long (sequential)	High
Thinking States	During input (parallel)	Short	Medium

Since input processing can leverage parallel computation (alns are known), thinking during input is more efficient than thinking during output.

Scaling Behavior

The number of thinking tokens per interval affects performance:

• Too few: insufficient computational budget for reasoning
• Too many: diminishing returns and increased input processing time
• Optimal interval depends on task complexity

Significance

Thinking States represent a compelling middle ground between direct answering and full chain-of-thought reasoning. By shifting reasoning computation from the sequential output phase to the parallelizable input phase, the method achieves strong reasoning performance with better latency characteristics. The ability to supervise thinking tokens with natural language makes training straightforward and the intermediate reasoning interpretable. This work opens a promising direction for efficient reasoning in language models.