# Latent Reasoning with Supervised Thinking States
**Authors**: Ido Amos, Avi Caciularu, Mor Geva, Amir Globerson, Jonathan Herzig, Lior Shani, Idan Szpektor
## Abstract
Reasoning with a chain-of-thought (CoT) enables Large Language Models (LLMs) to solve complex tasks but incurs significant inference costs due to the generation of long rationales. We propose Thinking States, a method that performs reasoning while the input is processing. Specifically, Thinking States generates sequences of thinking tokens every few input tokens, transforms the thoughts back into embedding space, and adds them to the following input tokens. This has two key advantages. First, it captures the recurrent nature of CoT, but where the thought tokens are generated as input is processing. Second, since the thoughts are represented as tokens, they can be learned from natural language supervision, and using teacher-forcing, which is parallelizable. Empirically, Thinking States outperforms other latent reasoning methods on multiple reasoning tasks, narrowing the gap to CoT on math problems, and matching its performance on 2-Hop QA with improved latency. On state-tracking tasks, we show Thinking States leads to stronger reasoning behavior than CoT, successfully extrapolating to longer sequences than seen during training.
Machine Learning, ICML
## 1 Introduction
Generating intermediate reasoning steps prior to answering enables large language models (LLMs) to tackle challenging problems with considerably greater accuracy (Wei et al., 2022). Such chain-of-thought (CoT) reasoning yields substantial performance gains, but may incur significant inference-time cost because of generating additional tokens. Thus, the challenge of reducing computational cost of reasoning while preserving accuracy has attracted considerable attention recently (Xia et al., 2025; Kang et al., 2024; Hassid et al., 2025; Hao et al., 2024).
<details>
<summary>x1.png Details</summary>
### Visual Description
## Diagram: Comparison of "Thinking States" vs. "Chain-of-Thought" in an LLM
### Overview
The image is a technical diagram comparing two different internal processing methods for a Large Language Model (LLM) when solving a simple logic problem. It is split into two distinct panels, each illustrating a different approach: "Thinking States" on the left and "Chain-of-Thought" on the right. Both panels use the same input prompt and arrive at the same final answer, but the internal mechanisms depicted are fundamentally different.
### Components/Axes
The diagram is not a chart with axes but a schematic with labeled components.
**Legend (Top-Left of Left Panel):**
* **Thinking States (S):** Represented by an orange square with a white "S".
* **Language Tokens (T):** Represented by a blue square with a white "T".
* **Compression Block (C):** Represented by a green square with a white "C".
**Left Panel: "Thinking States"**
* **Title:** "Thinking States" (bottom center).
* **Main Component:** A large, rounded rectangle labeled **"LLM"** in the center-left.
* **Internal Structure:** The LLM contains a grid of smaller, empty rounded rectangles, representing layers or processing units.
* **Data Flow:**
1. **Input Text (Bottom):** "A coin's state is heads. Alice flips, then Bob flips. What's the state? A: heads."
2. **Processing:** Arrows show data flowing from the input text into the LLM. Inside, orange "S" (Thinking State) blocks are connected to green "C" (Compression Block) blocks.
3. **Output Generation:** Arrows lead from the internal "C" blocks to blue "T" (Language Token) blocks at the top. These tokens form the output sequence: `tails`, `<EOS>`, `tails`, `<EOS>`, and finally `A: heads.`.
* **Spatial Layout:** The legend is in the top-left corner. The LLM block dominates the center. Input text is at the bottom, and output tokens are at the top.
**Right Panel: "Chain-of-Thought"**
* **Title:** "Chain-of-Thought" (bottom center).
* **Main Component:** A large, rounded rectangle (similar to the LLM block but unlabeled) in the center.
* **Internal Structure:** A grid of empty rounded rectangles, similar to the left panel but without any internal "S" or "C" blocks.
* **Data Flow:**
1. **Input Text (Bottom):** Identical to the left panel: "A coin's state is heads. Alice flips, then Bob flips. What's the state? **heads, tails, heads.** A: heads." (Note: The intermediate reasoning steps are now part of the input text in bold).
2. **Processing:** No internal state or compression blocks are shown. The flow is direct.
3. **Output Generation:** A single sequence of blue "T" (Language Token) blocks is generated at the top: `heads`, `tails`, `heads`, `A: heads.`.
* **Spatial Layout:** Mirrors the left panel. The LLM-like block is central, input text at the bottom, and output tokens are at the top.
### Detailed Analysis
The diagram contrasts two paradigms for LLM reasoning:
1. **Thinking States (Left):**
* **Mechanism:** The model uses dedicated internal componentsâ**Thinking States (S)** and **Compression Blocks (C)**âto process the problem. These are separate from the final language token generation pathway.
* **Process:** The input is processed through these internal states, which likely perform abstract reasoning or intermediate computation. The results are then "compressed" and finally translated into output language tokens.
* **Output:** The final output is a direct answer (`A: heads.`). The intermediate reasoning steps (`tails`, `<EOS>`) are generated as separate, internal tokens that are not presented as part of the final coherent answer to the user.
2. **Chain-of-Thought (Right):**
* **Mechanism:** The model relies on generating a sequence of explicit language tokens that form a reasoning trace. There are no depicted specialized internal state blocks.
* **Process:** The reasoning steps (`heads, tails, heads.`) are generated as part of the language token stream itself, appended to the original question in the input context for the final answer generation.
* **Output:** The final output includes the explicit reasoning chain followed by the answer (`A: heads.`). The "thought" process is visible in the token sequence.
### Key Observations
* **Identical Problem & Answer:** Both panels use the exact same logic puzzle (two coin flips starting from heads) and arrive at the correct answer: "heads."
* **Divergent Internal Representation:** The core difference is architectural. "Thinking States" implies a hidden, non-linguistic reasoning layer, while "Chain-of-Thought" uses the language modeling process itself to reason step-by-step.
* **Token Usage:** The "Thinking States" panel shows the generation of `<EOS>` (End of Sequence) tokens internally, suggesting segmented processing. The "Chain-of-Thought" panel shows a continuous stream of meaningful tokens.
* **Input Context:** In the Chain-of-Thought panel, the reasoning steps are explicitly included in the input text (shown in bold), indicating they are part of the prompt for the final answer step.
### Interpretation
This diagram illustrates a fundamental design choice in AI reasoning systems.
* **What it demonstrates:** It contrasts **internal, abstract computation** ("Thinking States") with **externalized, linguistic reasoning** ("Chain-of-Thought"). The former suggests a model that "thinks" in a proprietary, compressed format before speaking. The latter suggests a model that "thinks out loud" using the same language it uses to communicate.
* **Relationship between elements:** The "Thinking States" components (S, C) are positioned as intermediaries between raw input and final output, acting as a black-box reasoning engine. In "Chain-of-Thought," the reasoning is not an intermediate component but a direct, linear part of the input-output sequence.
* **Implications:** The "Thinking States" approach could be more efficient (processing compressed states) but less interpretable. The "Chain-of-Thought" approach is highly interpretable (you can see the reasoning steps) but may be less efficient as it requires generating many more tokens. The diagram suggests that both pathways can solve the same problem, highlighting that the *process* of reasoning in LLMs can be architecturally distinct from the *presentation* of reasoning. The presence of the same LLM grid in both panels implies the same base model could potentially support both modes of operation.
</details>
Figure 1: Reasoning with a Thinking State compared to Chain-of-Thought: An LLM is trained to generate task-relevant thinking sequences while processing tokens. Each generated thought sequence is transformed into a fixed-size state, denoted by $S$ , added to the following token representations.
One approach to reducing reasoning tokens is to not require thinking tokens to be in natural language, but rather be continuous embeddings (Hao et al., 2024; Shen et al., 2025). Indeed, language is highly redundant, and one can therefore expect shorter reasoning traces in this case. On the other hand, it is then harder to supervise these âcontinuous latentâ thoughts, as target values for the latents are unknown, which requires backpropagation through time (BPTT) from the desired outcome, making training considerably more challenging and limiting the number of latents that can be used in practice. An alternative approach has been to distill the thinking pattern directly to the representations of the query, rather than generate any extra tokens (Deng et al., 2024). While supporting parallel training, the resulting architecture cannot effectively condition on thoughts. As queries are processed in parallel, sequential thinking can only be realized through the LLMâs depth, which is typically fixed, resulting in reduced expressivity compared to CoT models. Put differently, CoT can be viewed as providing the model with additional depth, which Deng et al. (2024) does not benefit from.
Motivated by these previous works we present an architecture that allows conditioning on thinking as in CoT, yet does not append thinking tokens to the context or suffer from optimization with BPTT. Our Thinking States model builds on several key ideas, illustrated in Figure 1. Compute Sharing: thoughts are generated using representations of existing tokens. Recurrence: these thoughts are fed back as input (as in regular CoT) along with future tokens, not increasing the context length. Supervision with Natural Language: the thought tokens can be supervised via natural language, and thus benefit from available supervision used by prior CoT methods, and maintain interpretability.
We evaluate Thinking States on synthetic and real-world datasets, demonstrating that it is more accurate than baselines that reduce CoT lengths, and provides a significant latency advantage compared to CoT. Moreover, we demonstrate that Thinking States can be trained significantly faster than models with latent thoughts that require BPTT.
Our contributions are as follows:
1. We introduce Thinking States, a recurrent reasoning mechanism that generates natural-language thoughts while processing the input, without increasing the context length.
1. Because thoughts are represented in natural language, Thinking States can be trained with teacher-forcing the thoughts, enabling fully parallel training and avoiding the computational overhead of backpropagation through time.
1. We demonstrate that Thinking States consistently outperforms existing latent reasoning methods on various benchmarks, matches or exceeds CoT accuracy on multi-hop QA and state-tracking tasks, and achieves significant speedups compared to CoT on all tasks.
<details>
<summary>x2.png Details</summary>
### Visual Description
## Diagram: Transformer Model with KV Cache for Sequence Processing
### Overview
This image is a technical diagram illustrating the architecture and data flow of a transformer-based model processing a sequential reasoning task. It specifically highlights the use of a Key-Value (KV) Cache for efficient inference. The diagram shows how an input text sequence is processed through multiple layers, with intermediate states and cached values, to generate output tokens.
### Components/Axes
The diagram is organized into several distinct layers and components, flowing from bottom to top:
1. **Input Sequence (Bottom Layer):** A horizontal sequence of text tokens forming a complete question and answer.
* **Full Transcription:** "A coin's state is heads. Alice flips, then Bob flips. What's the state? A: heads."
* The final token "heads." is highlighted in blue, indicating it is the generated answer.
2. **KV Cache (Left Side):** A large, rounded rectangular block labeled "KV Cache". It contains two rows of empty rectangular boxes (5 in the top row, 5 in the bottom row), representing stored Key and Value vectors from previous processing steps. An arrow originates from this cache and points towards the processing chain on the right.
3. **Hidden State / Embedding Layer (L<sup>in</sup>):** A series of brown, rounded rectangular boxes labeled sequentially: `S_i`, `S_{i+1}`, `S_{i+2}`, `S_{i+3}`. These are connected by arrows and plus signs (`+`), indicating a sequential or residual connection. This layer is labeled `L^{in}` on the right side.
4. **Processing Blocks (Middle Layer):** Above each `S` box, there is a pair of green, rounded rectangular boxes labeled `T` and `C`. These likely represent Transformer blocks or specific operations (e.g., Attention and Feed-Forward networks). Arrows connect the `S` boxes to these `T`/`C` pairs.
5. **Output / Prediction Layer (L<sup>out</sup>):** At the top, there are three larger, rounded rectangular blocks labeled `Z_i`, `Z_{i+2}`, and `Z_{i+3}`. Each contains two smaller blue boxes representing output tokens:
* `Z_i`: Contains "tails" and "<EOS>".
* `Z_{i+2}`: Contains "<EOS>".
* `Z_{i+3}`: Contains "heads" and "<EOS>".
* This layer is labeled `L^{out}` on the right side.
6. **Data Flow Arrows:** Arrows indicate the direction of data propagation:
* From the KV Cache to the `S` boxes.
* Sequentially between the `S` boxes.
* From the `S` boxes up to the `T`/`C` blocks.
* From the `T`/`C` blocks up to the `Z` output blocks.
### Detailed Analysis
* **Sequence Processing:** The diagram models the processing of the 13-token input sequence. The indices `i`, `i+1`, `i+2`, `i+3` suggest the model is operating on specific, non-consecutive positions within this sequence, likely focusing on the reasoning steps ("flips") and the final answer.
* **KV Cache Role:** The KV Cache stores intermediate representations (Keys and Values) from earlier tokens. The arrow from the cache to the `S` boxes indicates that processing for later tokens (`S_{i+1}`, etc.) reuses this cached information, avoiding redundant computationâa key optimization for autoregressive generation.
* **Output Tokens:** The `Z` blocks show the model's predictions at different steps. `Z_i` predicts "tails" (an incorrect intermediate state) and an end-of-sequence token. `Z_{i+3}` correctly predicts "heads" as the final answer, followed by "<EOS>". The presence of multiple `<EOS>` tokens suggests the model can generate sequence terminators at various points.
* **Layer Labels:** `L^{in}` and `L^{out}` clearly demarcate the input embedding/hidden state layer and the final output/logit layer, respectively.
### Key Observations
1. **Non-Sequential Indexing:** The diagram skips from `S_i` to `S_{i+1}`, `S_{i+2}`, and `S_{i+3}`, implying it is highlighting specific, important steps in the reasoning chain rather than every single token.
2. **Intermediate Prediction:** The output at `Z_i` includes "tails," which is not part of the final correct answer. This suggests the model may generate or consider intermediate hypotheses during its reasoning process.
3. **Architectural Clarity:** The separation into distinct `T` and `C` blocks within the transformer layer is a common pedagogical simplification to represent the multi-head attention (`T` for "Transformer" or "Attention") and position-wise feed-forward network (`C` for "Context" or "Convolution/MLP") sub-layers.
4. **Visual Emphasis:** The final answer token "heads." in the input sequence is highlighted in blue, matching the color of the output tokens in the `Z` blocks, creating a visual link between the generated output and its place in the sequence.
### Interpretation
This diagram serves as an explanatory model for how a large language model (LLM) performs multi-step reasoning while maintaining efficiency. The core message is the **integration of a KV Cache within the autoregressive generation loop**.
* **What it demonstrates:** It shows that for a complex query requiring state tracking ("coin flips"), the model doesn't process the entire text from scratch for each new token. Instead, it leverages cached computations (KV Cache) from previous steps (`L^{in}` states `S_i`, etc.) to efficiently compute new hidden states and generate subsequent tokens (`Z_{i+2}`, `Z_{i+3}`).
* **Relationships:** The flow is cyclical and layered: Input tokens â Hidden States (`S`) â Transformer Operations (`T/C`) â Output Predictions (`Z`). The KV Cache acts as a memory bank that feeds back into this cycle, enabling context-aware generation without quadratic recomputation cost.
* **Notable Insight:** The inclusion of the incorrect prediction "tails" at `Z_i` is particularly insightful. It suggests the model's internal reasoning may involve exploring or verbalizing potential states before converging on the correct answer ("heads" at `Z_{i+3}`). This aligns with observed "chain-of-thought" behaviors in LLMs, where the model's intermediate outputs can reflect its problem-solving process, even if those specific tokens are not part of the final desired response. The diagram thus illustrates not just the architecture, but a plausible mechanism for **step-by-step reasoning** within a transformer.
</details>
(a) Inference
<details>
<summary>x3.png Details</summary>
### Visual Description
## Diagram: Sequence Prediction Model with Loss Calculation
### Overview
This image is a technical diagram illustrating a sequence-to-sequence prediction model, likely a transformer or recurrent neural network, processing a reasoning task about coin flips. The diagram shows the flow from input tokens through model components to predictions, which are then compared to ground truth labels to compute a loss. The primary language is English.
### Components/Axes
The diagram is structured vertically into three main sections:
1. **Top Section (Output & Loss):**
* **Labels (Top Row):** Two sets of target sequences.
* Left: `["tails", "<EOS>"]`
* Right: `["heads", "<EOS>"]`
* **Predictions (Second Row):** The model's predicted sequences.
* Left: `["tails", "heads"]`
* Right: `["heads", "<EOS>"]`
* **Loss Calculation:** Arrows connect corresponding positions between Labels and Predictions. Green checkmarks (â) indicate correct predictions, while a red cross (â) indicates an incorrect prediction. The word "Loss" is written between the two comparison sets, signifying the calculation of error.
2. **Middle Section (Model Architecture):**
* **T Blocks:** Two green rectangular blocks labeled "T" (likely representing Transformer decoder layers or a similar processing unit). They receive input from below and output to the prediction layer above.
* **Intermediate Representations:** Between the T blocks and the input sequence, there are multiple rows of paired, empty rectangular boxes (e.g., `[] []`), connected by vertical ellipses (`âŽ`). These represent hidden states or intermediate token representations across multiple layers or time steps.
* **Input Sequence:** A horizontal row of text tokens forming a question and answer:
`A coin's state is heads. Alice flips, then Bob flips. What's the state? A: heads.`
The final token "heads." is highlighted in blue.
3. **Bottom Section (Input Processing & Auxiliary Labels):**
* **S Blocks:** Four brown rectangular blocks labeled `S_i`, `S_{i+1}`, `S_{i+2}`, `S_{i+3}`. These are positioned below specific tokens in the input sequence ("heads.", "flips,", "flips.", "state?") and likely represent state encodings or special token embeddings.
* **C Blocks:** Four green rectangular blocks labeled "C" (likely representing encoder layers or context processors). Each `C` block feeds into a corresponding `S` block above it.
* **Labels (Bottom Row):** Another set of target sequences, aligned with the `C` blocks.
* Under first `C`: `["<EOS>"]`
* Under second `C`: `["tails", "<EOS>"]`
* Under third `C`: `["<EOS>"]`
* Under fourth `C`: `["heads", "<EOS>"]`
### Detailed Analysis
* **Task:** The model is solving a multi-step reasoning problem: Given an initial state ("heads"), two actions ("Alice flips", "Bob flips"), it must predict the final state. The correct answer is "heads" (since two flips return the coin to its original state).
* **Prediction Flow:**
1. The input text sequence is processed by the "C" (context/encoder) blocks.
2. State information (`S_i` to `S_{i+3}`) is derived and fed into the intermediate representation layers.
3. The "T" (transformer/decoder) blocks generate output predictions autoregressively.
* **Loss Calculation Detail:**
* **Left Comparison:** The model's first predicted sequence is `["tails", "heads"]`. The target label is `["tails", "<EOS>"]`.
* First token: Prediction "tails" matches label "tails" (â).
* Second token: Prediction "heads" does **not** match label "<EOS>" (â). This is an error; the model predicted an extra token instead of ending the sequence.
* **Right Comparison:** The model's second predicted sequence is `["heads", "<EOS>"]`. The target label is `["heads", "<EOS>"]`.
* First token: Prediction "heads" matches label "heads" (â).
* Second token: Prediction "<EOS>" matches label "<EOS>" (â). This is a correct prediction.
* **Text Transcription (All Visible Text):**
* `Labels`, `Predictions`, `Loss`
* `tails`, `<EOS>`, `heads`
* `T`
* `A coin's state is heads. Alice flips, then Bob flips. What's the state? A: heads.`
* `S_i`, `S_{i+1}`, `S_{i+2}`, `S_{i+3}`
* `C`
### Key Observations
1. **Error Highlighted:** The diagram explicitly shows a prediction error where the model fails to generate the `<EOS>` (End of Sequence) token at the correct time, instead generating a content token ("heads").
2. **Dual Label Sets:** There are two distinct sets of "Labels" â one at the top (for the final output sequence) and one at the bottom (likely for auxiliary training signals or intermediate state prediction).
3. **Architectural Abstraction:** The diagram uses abstract blocks ("C", "S", "T") rather than specific layer names, indicating it's a conceptual model of a sequence processing architecture.
4. **Task-Specific Example:** The coin flip problem is a classic test for logical reasoning and state tracking in AI models.
### Interpretation
This diagram serves as a pedagogical or technical illustration of how a sequence prediction model is trained using supervised learning. It demonstrates the core mechanism of comparing generated sequences against ground truth to compute a loss, which is then used to update model parameters.
The specific example highlights a common challenge in sequence generation: **proper sequence termination**. The model correctly identifies the final state ("heads") but incorrectly continues generation, failing to predict the `<EOS>` token. This suggests the model may have learned the content of the answer but not the structural rule for when to stop generating.
The presence of two label sets implies a complex training objective, possibly involving:
* **Teacher Forcing:** The bottom labels might be used to train the encoder ("C" blocks) or intermediate states.
* **Multi-Task Learning:** The model could be trained to predict both the final answer and auxiliary information.
Overall, the image encapsulates the flow of data, the structure of a modern neural network for reasoning tasks, and the fundamental process of error calculation that drives model learning. It visually answers the question: "How does a model learn to answer a question, and how do we know when it's wrong?"
</details>
(b) Training
Figure 2: Thinking States model at inference and training time. At inference time (a), Thinking States reasons by iteratively processing token chunks. At each iteration $i$ , reasoning information encoded up to layer $L^out$ is decoded by a thinking block $T$ and compressed by a compression block $C$ to a fixed-size thinking state $S_i+1$ . The state is injected to the next chunk at a shallow layer to effectively influence future computations. Chunk representations can access the history via attention layers and the KV-cache. At training time (b), predictions by $T$ are supervised by explicit annotations paired with each chunk. The same annotations are used to condition the base LLMâs reasoning, allowing parallel training.
## 2 Related Work
Latent Reasoning.
A prominent line of research focuses on adapting pretrained LLMs to internalize chain-of-thought (CoT) reasoning into compressed representations. Initial efforts distilled knowledge from a CoT-capable teacher model directly into the query representations of a student model (Deng et al., 2023). A related approach employs an iterative training curriculum that progressively removes explicit reasoning steps from the supervision signal, compelling the model to internalize the reasoning process (Deng et al., 2024). Hao et al. (2024) extended this idea by introducing recurrent âcontinuous thought tokensâ that are iteratively refined after processing the query. While promising, this recurrent formulation necessitates backpropagation through time (BPTT), which is computationally intensive and limits scalability. More recently, Shen et al. (2025) proposed adding a distillation loss that aligns the studentâs latent representations with those of a CoT-based teacher, removing reliance on the iterative curriculum. However, this approach still requires BPTT for optimization and incurs the additional computational cost of a separate teacher model.
Emergent Latent Reasoning in Standard LLMs.
Recent studies have revealed that standard LLMs can exhibit latent reasoning capabilities without explicit training. Biran et al. (2024) demonstrated this by successfully decoding intermediate reasoning steps from the modelâs hidden states, while Yang et al. (2024) showed that LLMs can arrive at correct answers on queries that inherently require multi-hop reasoning. These findings suggest that some degree of latent reasoning emerges naturally during pretraining, making multi-hop question answering, which we include in our evaluation, an important testbed for methods designed to enhance this capability.
Alternative Approaches to Latent Computation.
A distinct line of work aims to improve reasoning by providing models with additional computation time, rather than by internalizing a specific reasoning trace. Goyal et al. (2023) introduced âpause tokensâ during pretrainingâtokens without semantic meaning that give the model extra computational steps before producing output, improving performance on downstream reasoning tasks. Geiping et al. (2025) proposed dynamic depth, where the same block of layers is applied repeatedly for a variable number of iterations, effectively increasing computation per token at inference time. Finally, Zelikman et al. (2024) enabled token-level rationale generation by decoding auxiliary reasoning tokens before each prediction, training the model with reinforcement learning to produce helpful rationales. Unlike our approach, these auxiliary tokens influence only the current prediction and are discarded afterward, avoiding recurrent dependencies but also forgoing the ability to carry reasoning state across positions.
## 3 Thinking State Architecture and Supervision
We introduce the Thinking States model, illustrated in Figure 2(a). The core element of our architecture is a chunk-recurrent process of decoding thoughts from token representations, compressing thoughts into a fixed-size state, and injecting the state into representations of subsequent tokens at a shallow layer (e.g., the first layer).
Because thoughts are decoded from representations in a deep layer (e.g., second to last), we leverage the backbone LLM to perform the bulk of the computation. This allows the decoding âThinking Blockâ to remain lightweight and efficient, as it operates on a rich feature space. By injecting the compressed state into representations of future tokens at a shallow layer, we effectively condition on the reasoning process without extending the context. Sections 3.2, 3.3, provide details on how this recurrent process is executed efficiently at training and inference time.
Setting. Let an input sequence of $L$ tokens be partitioned into $K$ non-overlapping chunks $X_1,\dots,X_K$ . In the case $L\bmod câ 0$ , input tokens are grouped into chunks along with generated tokens, once those are available. Each $X_iââ^cĂ d$ denotes the input embeddings of $c$ tokens (where $c=L/K$ ) with hidden dimension $d$ , extracted after a designated shallow layer $L^in$ of the LLM.
Our architecture augments a backbone LLM $M_θ$ with two additional modules:
1. Thinking Block $T$ : A lightweight Transformer decoder (one layer in our experiments) that autoregressively generates a reasoning sequence $Z=(z_1,\dots,z_n)$ in natural language.
1. Compression Block $C$ : A Transformer encoder with a pooling layer that maps any variable-length reasoning sequence into a fixed-size state $Sââ^cĂ d$ .
Further architectural details for $T$ and $C$ are provided in Appendix A.1.
### 3.1 Architecture
Thinking States processes token chunks iteratively while maintaining a state $S_iââ^cĂ d$ . At each step $i$ , we first inject the current state into the input representations at the shallow layer, then forward through $M_θ$ :
$$
\tilde{X}_i=X_i+S_i \tag{1}
$$
$$
H_i^out=M_θ(\tilde{X}_i\mid\tilde{X}_<i) \tag{2}
$$
Here, $H_i^out$ denotes the chunk representations at a predefined deep layer $L^out$ , and past representations $\tilde{X}_<i$ are accessed through the KV-cache.
Next, $H_i^out$ is used to generate a sequence of natural-language thought tokens via the Thinking Block:
$$
Z_i+1={T}(H_i^out) \tag{3}
$$
The output $Z_i+1$ is a variable-length token sequence ending with an <EOS> token. When no reasoning is required, $Z_i+1$ may consist solely of <EOS>.
Finally, the Compression Block transforms the variable-length thought sequence into a fixed-size state for the next iteration:
$$
S_i+1={C}(Z_i+1)ââ^cĂ d \tag{4}
$$
For the first chunk, we initialize $S_1=0$ .
This design serves two key ideas: compute sharing and recurrence. As inputs are a sum of tokens and states, eq. (1 - 2), the model effectively performs next token prediction while processing and generating states, executing both computations in a single forward pass. By injecting states at a shallow layer, we allow thoughts to effectively condition future thoughts through most of the LLMâs hidden layers. Importantly, since the thought tokens are never appended to the backboneâs context window, the context length remains fixed.
### 3.2 Training with Teacher-Forced Reasoning
To supervise the generated thinking sequences $Z_i$ , we rely on explicit natural-language annotations. In our formulation, each token chunk in the training data is paired with a target reasoning sequence $Z_i^*$ of variable length. For chunks that require no intermediate reasoning, the target is simply the <EOS> token.
Access to these ground-truth sequences $Z_i^*$ serves a dual purpose. As illustrated in Figure 2(b), in addition to providing supervision for the outputs of $T$ , we use $Z_i^*$ to teacher-force the states injected into the backbone LLM. Specifically, we compute the target state as $S_i^*={C}(Z_i^*)$ and inject it instead of the model-generated state:
$$
\tilde{X}_i=X_i+S_i^* \tag{5}
$$
Since all target sequences $Z_i^*$ are available upfront, we can compute all chunk representations $H_i^out$ in a single parallel forward pass through $M_θ$ . This eliminates the need for backpropagation through time (BPTT), dramatically reducing the computational cost of training.
After this forward pass, the Thinking Block $T$ is trained to predict $Z_i^*$ via standard next-token prediction, in parallel over all chunks. Here, each $H_i^out$ is computed under the gold state history $S_1^*,\dots,S_i^*$ , so predicting $Z_i+1$ is implicitly conditioned on all prior gold reasoning steps.
We optimize the joint objective:
$$
L=L_LM+â_i=1^KL_T(Z_i,Z^*_i) \tag{6}
$$
where $L_LM$ is the standard language modeling loss and $L_T$ is cross-entropy loss over the thinking sequences. This formulation enables fully parallel training while ensuring that $T$ learns to generate thoughts consistent with a valid reasoning history.
### 3.3 Fast Prefill with Speculative Thinking
While training is fully parallelizable, naĂŻve inference remains sequential across chunks, as each $\tilde{X}_i$ depends on the previous state $S_i$ . This can significantly increase latency during query prefill.
To mitigate this overhead, we exploit a key observation: for most chunks, the generated state is trivial, that is, the Thinking Block produces only an <EOS> token. We leverage this sparsity to design an iterative yet exact prefill algorithm.
The algorithm proceeds as follows:
1. Perform a parallel forward pass over all chunks, speculating that all states are trivial.
1. For each chunk, generate thinking states using $T$ and $C$ .
1. Identify the earliest chunk $i_1$ that produces a non-trivial state. Since all chunks before $i_1$ truly have trivial states, the computation up to chunk $i_1$ is correct.
1. Cache all positions up to $i_1$ and repeat from step 1 for the remaining chunks, now conditioning on the correct state $S_i_1$ and speculating a trivial state for following chunks.
The procedure terminates when no new non-trivial states are generated, yielding exactly the same result as full sequential prefill. If $|R|$ denotes the (unknown a priori) number of chunks with non-trivial states, the algorithm completes in exactly $|R|+1$ rounds. In typical regimes where $|R|\ll K$ , this substantially reduces prefill latency. A complete description of the algorithm is provided in Appendix A.2.
### 3.4 Constructing Chunk-Level Supervision
Training the Thinking Block requires supervision that specifies both where intermediate reasoning should occur and what reasoning content to produce. We construct this supervision by aligning CoT reasoning steps to specific positions in the input query.
Consider the following example query and its CoT trajectory:
âA coinâs state is heads. Alice flips, then Bob flips. Whatâs the state?â CoT: âheads $â$ tails $â$ headsâ
Our goal is to map each reasoning step (which may span multiple tokens) to a specific position in the query where that step becomes inferable.
Step-to-Token Alignment.
We first align each reasoning step to the earliest query position where it can be logically inferred. This alignment is obtained either via a strong teacher model, e.g., Gemini 2.5-Flash (Comanici et al., 2025), or through programmatic rules when the task structure permits. The alignment is represented by inserting special indicator tokens <T> into the text:
âA coinâs state is heads. <T> Alice flips, <T> then Bob flips. <T> Whatâs the state?â
Each <T> marker indicates that the preceding context is sufficient to infer the corresponding reasoning step. We then tokenize the query and construct a corresponding reasoning array with the same length, where each position in the array is either empty or contains the reasoning step associated with the <T> token. The indicator tokens are removed after this alignment, ensuring they are never seen by the model during training.
Token-to-Chunk Alignment.
Given the aligned token sequence and reasoning array, we partition both into non-overlapping chunks of size $c$ . The target reasoning sequence for each chunk is the concatenation of all reasoning steps assigned to tokens within that chunk. For chunks with no associated reasoning steps, the target is simply the <EOS> token.
The entire process is applied only to the training data of each task we consider. Additional details can be found in Appendix A.3.
## 4 Experiments
We evaluate Thinking States on a diverse set of tasks that probe both reasoning accuracy and computational efficiency. Section 4.2 examines state-tracking, where Thinking States exhibits stronger length generalization than CoT. Section 4.3 evaluates general reasoning on Multi-Hop QA and GSM8K, showing that Thinking States surpasses latent reasoning baselines and approaches CoT performance with significant speedups. Section 4.4 presents ablation studies exploring the performanceâefficiency tradeoff as a function of recurrence depth and chunk size. Finally, Section 4.5 provides an interpretability-driven error analysis, identifying distinct failure modes and cases where Thinking States outperforms CoT.
### 4.1 Experimental Setup
All experiments use Qwen2.5-Base models (Qwen et al., 2024) (0.5B and 1.5B parameters), fine-tuned on task-specific training data.
We compare against four baselines:
1. CoT: fine-tuning on explicit chain-of-thought traces;
1. No CoT: standard fine-tuning to predict the final answer directly;
1. iCoT (Deng et al., 2024): a latent reasoning method that distills CoT into query representations via curriculum learning;
1. Coconut (Hao et al., 2024): a state-of-the-art latent reasoning method using continuous thought tokens.
We reproduce iCoT and Coconut results using the official codebase of Hao et al. (2024) with default configurations unless otherwise noted.
Metrics.
We report accuracy on gold answers across all tasks. For efficiency, we measure wall-clock speedup relative to CoT on a single A100-80GB GPU. Unlike prior work (Hao et al., 2024; Deng et al., 2024), which reports efficiency in terms of generated token count, we use wall-time because token generation with the lightweight Thinking Block $T$ is substantially faster than standard autoregressive decoding.
### 4.2 State Tracking Tasks
State tracking tasks require monitoring quantities that evolve throughout the input (Merrill et al., 2024; Kim and Schuster, 2023; Li et al., 2025). These tasks are a natural testbed for latent reasoning methods because they explicitly require sequential composition âeach step depends on the result of previous steps. We evaluate on two synthetic tasks adapted from Anil et al. (2022):
- Parity: tracking a single binary state through a sequence of flip operations.
- Variable Assignment (Vars): tracking multiple interacting integer variables, yielding more complex state dynamics.
Both tasks are expressed in natural language; formal definitions and examples are provided in Appendix A.4.
Table 1: Performance (accuracy) on State Tracking tasks, trained on samples with up to $N$ state updates and evaluated on $[N,100]$ .
| No CoT CoT Thinking States | 54.67 12.35 98.37 | 57.50 38.12 99.02 | 59.60 64.38 100.00 | 02.15 06.78 33.76 | 02.17 35.45 87.23 | 02.19 87.75 97.71 |
| --- | --- | --- | --- | --- | --- | --- |
To assess whether the underlying base model truly learns to reason via the recurrent thinking state mechanism, we employ a length generalization setting, drawing on the known extrapolation capabilities of recurrent models (delĂŠtang2023neuralnetworkschomskyhierarchy). We fine-tune all methods, with Qwen2.5-Base-0.5B as the base model, on sequences of up to $N$ operations ( $Nâ\{10,20,40\}$ ) and evaluate on lengths in $[N,100]$ , in both cases, sampled uniformly over the range. All models are trained to $100\$ in-distribution accuracy to strictly isolate extrapolation capabilities from optimization effects.
Table 1 reports out-of-distribution (OOD) accuracy for each training regime. The results demonstrate how a base LLM trained with Thinking States effectively adapts to reasoning with a recurrent mechanism, leading to stronger length generalization than that of a model reasoning via CoT. Critically, this indicates that reasoning via the recurrent mechanism offers a more effective solution than CoT, requiring a simpler training distribution to achieve robust performance.
Furthermore, while CoT generates intermediate tokens proportional to the number of state operations, considerably extending the context length and thus expressivity, Thinking States maintains a static context size. By refining the query representation through state updates, Thinking States achieves superior length generalization, suggesting the model has learned to reason through the recurrent mechanism.
### 4.3 General Reasoning Capabilities
Table 2: Performance and speedup (over CoT) on reasoning tasks for different latent reasoning baselines and CoT.
| No CoT Thinking States Coconut | 34.11 42.22 32.65 | $Ă$ 5.59 $Ă$ 2.66 $Ă$ 3.14 | 33.47 54.91 33.71 | $Ă$ 1.89 $Ă$ 1.19 $Ă$ 1.14 | 31.92 43.05 32.60 | $Ă$ 2.03 $Ă$ 1.23 $Ă$ 1.21 |
| --- | --- | --- | --- | --- | --- | --- |
| iCoT | 34.00 | $Ă$ 5.71 | 28.84 | $Ă$ 1.59 | 36.31 | $Ă$ 1.80 |
We evaluate general reasoning capabilities on two established benchmarks. First, a 2-hop QA dataset, where prior work by Biran et al. (2024); Yang et al. (2024) demonstrated standard LLMs exhibit latent reasoning to some extent, providing a useful test case for methods designed to improve latent reasoning. Second, a GSM8K-style dataset, consisting of $âź$ 400K simple math word problems and parsed CoT steps, previously used by Deng et al. (2023); Hao et al. (2024). Methods in this section use Qwen2.5-Base-1.5B as the base model.
For the 2-hop QA task, we consider two evaluation regimes. In Full Knowledge (FK), the required factual knowledge appears in the fine-tuning data, testing how well models acquire and manipulate knowledge with each method. In Parametric Knowledge (PK), examples are filtered to reflect the knowledge of the base LLM, isolating whether a method improves retrieval and manipulation of existing knowledge in the base model.
<details>
<summary>x4.png Details</summary>
### Visual Description
## Line Chart: Speedup over CoT vs. Number of Recurrence Layers
### Overview
This is a line chart plotting model accuracy against the number of layers used in a recurrence mechanism. It includes a secondary top axis showing the corresponding "Speedup over CoT" (Chain-of-Thought) for each data point. A horizontal baseline indicates the accuracy achieved without using CoT.
### Components/Axes
* **Main Title:** "Speedup over CoT" (centered at the top, with a left-pointing arrow).
* **X-Axis (Bottom):** Label: "Num Layers in Recurrence" (with a right-pointing arrow). Scale: Discrete values at 4, 8, 12, 16, 20, 24, and 26.
* **Y-Axis (Left):** Label: "Accuracy (%)". Scale: Linear from 20 to 48, with major ticks every 4 units (20, 24, 28, 32, 36, 40, 44, 48).
* **Secondary X-Axis (Top):** Displays "Speedup over CoT" multipliers corresponding to each data point below. Values: 3.12x, 2.98x, 2.97x, 2.86x, 2.78x, 2.72x, 2.66x.
* **Data Series:** A single series represented by a blue dashed line with circular markers.
* **Baseline Reference:** A horizontal red dotted line across the chart at y=34.11%, labeled "No CoT Baseline: 34.11%" in red text, positioned just above the line near the center.
### Detailed Analysis
**Data Points (Approximate Values):**
The blue line shows a clear, monotonically increasing trend in accuracy as the number of recurrence layers increases.
1. **4 Layers:** Accuracy â 22%. Corresponding Speedup = 3.12x.
2. **8 Layers:** Accuracy â 27%. Corresponding Speedup = 2.98x.
3. **12 Layers:** Accuracy â 30%. Corresponding Speedup = 2.97x.
4. **16 Layers:** Accuracy â 35%. Corresponding Speedup = 2.86x. *This point is just above the "No CoT Baseline" of 34.11%.*
5. **20 Layers:** Accuracy â 37%. Corresponding Speedup = 2.78x.
6. **24 Layers:** Accuracy â 38%. Corresponding Speedup = 2.72x.
7. **26 Layers:** Accuracy â 42%. Corresponding Speedup = 2.66x.
**Trend Verification:**
* **Accuracy Trend (Blue Line):** Slopes consistently upward from left to right. The increase is relatively steady, with a slightly steeper rise between 24 and 26 layers.
* **Speedup Trend (Top Axis Values):** Decreases monotonically from left to right (from 3.12x down to 2.66x) as the number of layers increases.
### Key Observations
1. **Trade-off Revealed:** There is a clear inverse relationship between accuracy and speedup. As the model's accuracy improves with more recurrence layers, its computational speedup over a CoT baseline diminishes.
2. **Baseline Crossover:** The model surpasses the "No CoT Baseline" accuracy (34.11%) at approximately 16 recurrence layers.
3. **Diminishing Returns on Speedup:** The largest drop in speedup (from 3.12x to 2.98x) occurs between 4 and 8 layers, while the accuracy gain is moderate. The smallest speedup decrease (2.72x to 2.66x) occurs between 24 and 26 layers, where the accuracy gain is the most pronounced.
4. **Spatial Layout:** The "No CoT Baseline" label is placed centrally for immediate reference. The speedup values are aligned directly above their corresponding data points for easy correlation.
### Interpretation
The chart demonstrates a fundamental engineering trade-off in the design of this recurrent model architecture. **Increasing the depth of recurrence (more layers) is an effective strategy for boosting task accuracy**, allowing the model to eventually outperform a standard non-CoT approach.
However, this accuracy gain comes at a direct **computational cost, measured as a reduction in inference speedup** relative to a CoT method. The "Speedup over CoT" metric suggests that while the recurrent model is faster than CoT in all shown configurations (all multipliers >1), its speed advantage erodes as it becomes more accurate.
The data suggests an optimization problem: selecting the number of recurrence layers involves balancing the desired accuracy against acceptable inference latency. The point at 16 layers represents a potential "sweet spot" where the model first exceeds the baseline accuracy while still maintaining a significant speedup (2.86x). The sharp accuracy increase at 26 layers might indicate a threshold where the added complexity yields disproportionate performance benefits, albeit with the lowest speedup in the series.
</details>
(a)
<details>
<summary>x5.png Details</summary>
### Visual Description
## Line Chart: Speedup over CoT vs. Chunk Size
### Overview
This image is a line chart titled "Speedup over CoT" that plots the accuracy (in percentage) of a method against different "Chunk Size" values. It includes a baseline comparison and annotations showing the speedup factor for each chunk size relative to a Chain-of-Thought (CoT) baseline.
### Components/Axes
* **Title:** "Speedup over CoT" (centered at the top).
* **X-Axis:** Labeled "Chunk Size" with an arrow pointing right. The axis has categorical tick marks at values: **2, 4, 8, 16, 32, 48**.
* **Y-Axis:** Labeled "Accuracy (%)" with a scale from **20 to 48**, in increments of 4 (20, 24, 28, 32, 36, 40, 44, 48).
* **Data Series:** A single blue dashed line with circular markers representing the accuracy for each chunk size.
* **Baseline:** A red dotted horizontal line at **34.11%** on the y-axis, labeled "No CoT Baseline: 34.11%" in red text, positioned just above the line in the center of the chart.
* **Annotations:** Above each data point on the blue line, a speedup factor is written at an angle (e.g., "2.42x"). These values increase from left to right.
### Detailed Analysis
The chart displays the following data points (approximate values read from the y-axis):
| Chunk Size | Approximate Accuracy (%) | Speedup Factor (over CoT) |
| :--- | :--- | :--- |
| 2 | ~35.0 | 2.42x |
| 4 | ~41.5 | 2.57x |
| 8 | ~42.0 | 2.66x |
| 16 | ~39.0 | 2.74x |
| 32 | ~37.5 | 2.83x |
| 48 | ~37.8 | 2.86x |
**Trend Verification:**
* **Accuracy Trend (Blue Line):** The line starts at ~35% for chunk size 2, rises sharply to a peak of ~42% at chunk size 8, then gradually declines to ~37.8% at chunk size 48. The overall trend is an initial increase followed by a gradual decrease.
* **Speedup Trend (Annotations):** The speedup factor shows a consistent, monotonic increase from 2.42x to 2.86x as chunk size increases from 2 to 48.
### Key Observations
1. **Peak Performance:** The highest accuracy (~42%) is achieved at a chunk size of 8.
2. **Baseline Comparison:** All tested chunk sizes yield accuracy values above the "No CoT Baseline" of 34.11%.
3. **Diverging Trends:** Accuracy and speedup do not correlate perfectly. While speedup improves steadily with larger chunk sizes, accuracy peaks at an intermediate size (8) and then diminishes.
4. **Diminishing Returns:** After chunk size 8, increasing the chunk size further leads to a decrease in accuracy, despite continued gains in speedup.
### Interpretation
The data demonstrates a trade-off between computational efficiency (speedup) and model performance (accuracy) when using chunked processing compared to a standard Chain-of-Thought (CoT) approach.
* **What it suggests:** The method being tested is consistently faster than the CoT baseline (all speedups >2.4x) and also more accurate (all points above the red baseline). This indicates a successful optimization.
* **Relationship between elements:** The "Chunk Size" parameter controls this trade-off. Smaller chunks (2) offer the least speedup but still beat the baseline. An optimal balance appears at chunk size 8, maximizing accuracy. Larger chunks (16-48) prioritize speed, yielding the highest speedups but at the cost of some accuracy loss from the peak.
* **Notable anomaly:** The accuracy does not plateau but actively decreases after the peak at chunk size 8. This suggests that beyond a certain point, the chunking strategy may introduce errors or lose context, negatively impacting performance despite faster processing.
* **Practical implication:** For applications where accuracy is paramount, a chunk size around 8 is ideal. For latency-sensitive tasks where a small accuracy drop is acceptable, larger chunk sizes (32 or 48) provide the greatest speed advantage.
</details>
(b)
Figure 3: Ablation studies over deep to shallow recurrence (a) and effects of the chunk size (b), with measured speedup over CoT (top) at each point. In (a), performance increases with the number of layers used to process and generate states, pointing to the importance of the deep to shallow design. In (b), increasing the chunk size leads to improved latency with peak performance for medium chunks, highlighting the tradeoff between per-state capacity (large chunks) and reasoning frequency (small chunks).
Table 2 summarizes accuracy and speedups for each method. Among latent-reasoning baselines, Thinking States consistently achieves the highest accuracy across all tasks while maintaining significant speedups compared to CoT. On both variants of the 2-Hop task, Thinking States reaches similar performance to CoT with improved efficiency, demonstrating a dramatic advantage over other non-CoT baselines, with more than a 20% and 7% increase in accuracy on the FK and PK variants respectively.
On the GSM task, Thinking States demonstrates considerably better performance compared to Coconut, No CoT and iCoT, improving accuracy by $âź 8\$ and maintaining a considerable $2.66Ă$ speedup. The slight latency advantage showcased by Coconut on this task leads us to evaluate if performance can be improved at the cost of higher latency. As reported in Appendix A.5.1, we observe that increasing the number of latent steps leads to a modest improvement in performance, even as latency approaches that of the CoT model. In the next section, we provide several ablation studies over possible compute-efficiency tradeoffs with Thinking States.
Finally, we analyze the training costs of Thinking States compared to training with explicit recurrence and BPTT as in Coconut. Figure 6, in Appendix A.5.3, plots the wall-clock time per forward and backward pass as a function of the number of recurrent steps/number of chunks in the input sequence. The results demonstrate that while BPTT incurs linearly increasing costs, Thinking States maintains close to constant training time due to teacher forcing. Consequently, while BPTT becomes prohibitively expensive to scale, incurring a $âź 10Ă$ cost penalty at just 10 steps, Thinking States enables efficient training regardless of reasoning depth.
### 4.4 Ablation Study
To understand the architectural factors governing Thinking States, we conduct two ablation studies over the GSM task. We ablate the impact of recurrence with a deep to shallow connection, by varying the layer to extract representations passed to $T$ , and the impact of reasoning frequency, by varying the chunk size.
Impact of Deep to Shallow Connections: In Figure 3(a), we investigate the influence of deep to shallow recurrence by varying the extraction layer for $T$ , keeping injection fixed at the first layer (as in previous experiments). Our purpose is to show the importance of the deep to shallow recurrent principle, demonstrating that allocating more of the modelâs capacity when processing and generating states leads to improved accuracy. Furthermore, this experiment highlights the inherent accuracy-latency tradeoff: utilizing fewer layers in the recurrent loop allows the remaining layers to be executed only during query prefill, resulting in higher speedups. For this analysis, we use Qwen2.5-Base-1.5B, which contains 28 hidden layers.
As shown in Figure 3(a), performance increases monotonically with the number of layers included in the recurrence, with a substantial accuracy gap of nearly 20% between the shallowest and deepest configurations. This trend confirms that maximizing the computational capacity available for processing and generating thinking states is critical for performance, effectively validating the deep-to-shallow design choice.
Impact of Chunk Size: In Figure 3(b), we study the effect of varying the chunk size. The chunk size controls two important factors: (1) the computational capacity available for state processing and generation (2) the number of expected reasoning steps generated in each call to $T$ . A small chunk size limits available capacity for each state, and is expected to generate less steps at each chunk. Conversely, larger chunks provide larger capacity, but may lead to encoding multiple sequential steps into single state. This latter scenario undermines the deep-to-shallow recurrence, as consecutive steps generated within the same chunk cannot access the full recurrent loop, relying instead on the LLMs depth and limited capacity of the lightweight Thinking Block.
The tradeoff discussed above is evident in Figure 3(b), when the chunk size is small both performance and latency degrade, which can be explained by lack of computational capacity and increased iterations during query processing. As the chunk size increases, latency improves due to fewer iterations during query processing. Increasing the chunk size beyond a certain point ( $c=8$ ) eventually leads to reduced performance, which can be explained by the model compressing too many reasoning steps into single updates. Although performance decreases, we note that Thinking States still maintains a notable gap in performance compared to other latent reasoning methods even when approaching similar speedups.
### 4.5 Error Analysis
We analyze the performance gap between Thinking States and CoT on GSM-style problems. A key advantage of Thinking States over other latent reasoning methods is interpretability: because states are compressed from natural-language thoughts, the modelâs reasoning remains inspectable.
Where Thinking States succeeds and CoT fails.
Despite the overall accuracy gap, approximately 12% of queries correctly solved by Thinking States are not solved by CoT, indicating that Thinking States mitigates certain CoT errors. Figure 4 shows two representative examples: (1) CoT hallucinates an additional reasoning step, and (2) CoT attempts an overly complex intermediate computation. In both cases, Thinking States produces a correct internal trajectory and answer. We note that both models are fine-tuned from the same base model and on identical data.
Example 1: Hallucinated Step Derrick has a bakery that makes 10 dozen doughnuts every day, [T: 10*12=120] selling at $2 per doughnut. [T: 120*2=240] How much money does Derrick make in June if he sells all the doughnuts? [T: 240*30=7200] The answer is 7200 CoT Steps: 10*12=120 120*2=240 240*30=7200 7200*6=43200 Example 2: Over-Complex Computation Jess is trying to guess the number of blue jellybeans in a jar. She can see that there are 17 green jelly beans and twice as many red jelly beans. The [T: 17*2=34] rest of the jellybeans are blue jelly beans. If there are a total of 60 jelly beans in total, [T: 17+34=51] how many blue jellybeans are there? [T: 60-51=9] The answer is 9 CoT Steps: 17*2=34 60-17-34=8
Figure 4: Examples where Thinking States succeeds but CoT fails. Thinking States reasoning is shown in green. (1) CoT hallucinates an extra step. (2) CoT attempts multiple operations in one step and errs.
Where CoT succeeds and Thinking States fails: State Ambiguity.
We identify a distinct error pattern we term state ambiguity: in many GSM queries, the quantity of interest is specified only in the final clause and may be difficult to anticipate earlier. As illustrated in Figure 5, this leads Thinking States to reason about the wrong intermediate quantity. Although the correct answer is often recoverable via a simple transformation, such recovery requires exposure to similar patterns during training.
(a) Original Richard lives in a building with 15 floors. Each floor contains 8 units. [T: 15 $Ă$ 8=120.] Three quarters of the units are occupied. [T: (3/4) $Ă$ 120=90.] Whatâs the number of unoccupied units on each floor? [T: 120-90=30.] Answer:30 (b) Disambiguated Whatâs the number of unoccupied units on each floor? Richard lives in a building with 15 floors. Each floor contains 8 units. Three quarters of the units are occupied. [T: 0.75 $Ă$ 8=6] Whatâs the number of unoccupied units on each floor? Answer:2
Figure 5: Illustration of state ambiguity. In (a), Thinking States generates valid thoughts that are not aligned with the final clause in the query, leading to an error. By adding the information to the begining of the query, as in (b), the same model produces a correct answer, without training on the new format.
We validate this hypothesis by prepending the final question to the start of the prompt for all error cases. This zero-shot intervention improves accuracy from 42.22 to 48.65, despite the augmented queries being out-of-distribution. Notably, this limitation stems from the combination of Thinking States with a causal autoregressive backbone rather than the thinking-state mechanism itself; bidirectional query processing could, in principle, identify the relevant quantity before committing to intermediate state updates.
## 5 Conclusions
In this work, we introduce Thinking States, a method that enables reasoning while processing input tokens without extending the modelâs context. By relying on a recurrent state representation of the thinking trajectory, the model effectively conditions its token processing on intermediate reasoning steps Empirically, this leads to improved performance compared to other latent reasoning baselines while maintaining a significant efficiency advantage over Chain-of-Thought (CoT).
Furthermore, a key contribution of our approach is the specialized data construction process that maps reasoning steps to specific input chunks. This formulation ensures that the recurrent process does not optimize with BPTT during training. Instead, training is fully parallelizable via teacher forcing, leading to a substantial reduction in computational cost compared to prior recurrent methods.
Our framework opens several avenues for future research. First, while we currently target query prefill, extending Thinking States to the decoding phase could enable allocating dynamic compute during token generation. Second, while our current approach relies on Supervised Fine-Tuning over existing CoT data, this model can serve as an effective âwarm-startâ initialization for Reinforcement Learning. Starting from a model that has already learned to compress reasoning into latent states could stabilize RL training, allowing the model to subsequently optimize its internal thinking process beyond the constraints of human-generated traces.
## Impact Statement
This paper presents a method for improving the efficiency of reasoning in large language models. By reducing computational costs relative to chain-of-thought approaches, our work may contribute to lowering the energy consumption and environmental footprint of deploying reasoning-capable language models. Additionally, because our method generates interpretable natural-language thoughts, it may support efforts toward more transparent and auditable AI systems. We do not foresee unique ethical concerns beyond those broadly associated with advances in language model capabilities.
## References
- C. Anil, Y. Wu, A. Andreassen, A. Lewkowycz, V. Misra, V. Ramasesh, A. Slone, G. Gur-Ari, E. Dyer, and B. Neyshabur (2022) Exploring length generalization in large language models. ArXiv preprint abs/2207.04901. External Links: Link Cited by: §A.4, §A.4, §4.2.
- E. Biran, D. Gottesman, S. Yang, M. Geva, and A. Globerson (2024) Hopping too late: exploring the limitations of large language models on multi-hop queries. ArXiv preprint abs/2406.12775. External Links: Link Cited by: §2, §4.3.
- G. Comanici, E. Bieber, M. Schaekermann, I. Pasupat, N. Sachdeva, I. Dhillon, M. Blistein, O. Ram, D. Zhang, E. Rosen, et al. (2025) Gemini 2.5: pushing the frontier with advanced reasoning, multimodality, long context, and next generation agentic capabilities. ArXiv preprint abs/2507.06261. External Links: Link Cited by: §A.5.2, §3.4.
- M. Dehghani, S. Gouws, O. Vinyals, J. Uszkoreit, and L. Kaiser (2019) Universal transformers. In 7th International Conference on Learning Representations, ICLR 2019, New Orleans, LA, USA, May 6-9, 2019, External Links: Link Cited by: §A.4.
- Y. Deng, Y. Choi, and S. Shieber (2024) From explicit cot to implicit cot: learning to internalize cot step by step. ArXiv preprint abs/2405.14838. External Links: Link Cited by: §1, §2, item (iii), §4.1, footnote 1.
- Y. Deng, K. Prasad, R. Fernandez, P. Smolensky, V. Chaudhary, and S. Shieber (2023) Implicit chain of thought reasoning via knowledge distillation. ArXiv preprint abs/2311.01460. External Links: Link Cited by: §2, §4.3.
- J. Geiping, S. McLeish, N. Jain, J. Kirchenbauer, S. Singh, B. R. Bartoldson, B. Kailkhura, A. Bhatele, and T. Goldstein (2025) Scaling up test-time compute with latent reasoning: a recurrent depth approach. ArXiv preprint abs/2502.05171. External Links: Link Cited by: §2.
- S. Goyal, Z. Ji, A. S. Rawat, A. K. Menon, S. Kumar, and V. Nagarajan (2023) Think before you speak: training language models with pause tokens. ArXiv preprint abs/2310.02226. External Links: Link Cited by: §2.
- S. Hao, S. Sukhbaatar, D. Su, X. Li, Z. Hu, J. Weston, and Y. Tian (2024) Training large language models to reason in a continuous latent space. ArXiv preprint abs/2412.06769. External Links: Link Cited by: Table 3, Table 3, §1, §1, §2, item (iv), §4.1, §4.1, §4.3.
- M. Hassid, G. Synnaeve, Y. Adi, and R. Schwartz (2025) Donât overthink it. preferring shorter thinking chains for improved llm reasoning. ArXiv preprint abs/2505.17813. External Links: Link Cited by: §1.
- Y. Kang, X. Sun, L. Chen, and W. Zou (2024) C3oT: generating shorter chain-of-thought without compromising effectiveness. ArXiv preprint abs/2412.11664. External Links: Link Cited by: §1.
- G. Kaplan, M. Oren, Y. Reif, and R. Schwartz (2024) From tokens to words: on the inner lexicon of llms. ArXiv preprint abs/2410.05864. External Links: Link Cited by: §A.1.
- N. Kim and S. Schuster (2023) Entity tracking in language models. ArXiv preprint abs/2305.02363. External Links: Link Cited by: §4.2.
- B. Z. Li, Z. C. Guo, and J. Andreas (2025) (How) do language models track state?. ArXiv preprint abs/2503.02854. External Links: Link Cited by: §4.2.
- W. Merrill, J. Petty, and A. Sabharwal (2024) The illusion of state in state-space models. In Proceedings of the 41st International Conference on Machine Learning, ICMLâ24. Cited by: §4.2.
- Qwen, :, A. Yang, B. Yang, B. Zhang, B. Hui, B. Zheng, B. Yu, C. Li, D. Liu, F. Huang, H. Wei, H. Lin, J. Yang, J. Tu, J. Zhang, J. Yang, J. Yang, J. Zhou, J. Lin, K. Dang, K. Lu, K. Bao, K. Yang, L. Yu, M. Li, M. Xue, P. Zhang, Q. Zhu, R. Men, R. Lin, T. Li, T. Tang, T. Xia, X. Ren, X. Ren, Y. Fan, Y. Su, Y. Zhang, Y. Wan, Y. Liu, Z. Cui, Z. Zhang, and Z. Qiu (2024) Qwen2.5 technical report. ArXiv preprint abs/2412.15115. External Links: Link Cited by: §4.1.
- Z. Shen, H. Yan, L. Zhang, Z. Hu, Y. Du, and Y. He (2025) CODI: compressing chain-of-thought into continuous space via self-distillation. ArXiv preprint abs/2502.21074. External Links: Link Cited by: §1, §2.
- J. Wei, X. Wang, D. Schuurmans, M. Bosma, B. Ichter, F. Xia, E. Chi, Q. Le, and D. Zhou (2022) Chain-of-thought prompting elicits reasoning in large language models. ArXiv preprint abs/2201.11903. External Links: Link Cited by: §1.
- H. Xia, C. T. Leong, W. Wang, Y. Li, and W. Li (2025) TokenSkip: controllable chain-of-thought compression in llms. ArXiv preprint abs/2502.12067. External Links: Link Cited by: §1.
- S. Yang, N. Kassner, E. Gribovskaya, S. Riedel, and M. Geva (2024) Do large language models perform latent multi-hop reasoning without exploiting shortcuts?. ArXiv preprint abs/2411.16679. External Links: Link Cited by: §2, §4.3.
- E. Zelikman, G. Harik, Y. Shao, V. Jayasiri, N. Haber, and N. D. Goodman (2024) Quiet-star: language models can teach themselves to think before speaking. ArXiv preprint abs/2403.09629. External Links: Link Cited by: §2.
## Appendix A Method - Additional Details
### A.1 Architecture
We provide additional details on the thinking and compression blocks, $T$ and $C$ and their integration with the base LLM, $M_θ$ . In all reported experiments, except when explicitly stated in section 4.4, we set the injection and extraction layer, $L^in,L^out$ , to be after the first layer (excluding the token embedding layer) and after the second-to-last layer respectively.
Both $T$ and $C$ share the same underlying core module, a causal Transformer block, sharing configurations and hyper-parameters (i.e., hidden size, number of attention heads, sliding/dense attention etc.) as the layers in $M_θ$ . In our experiments, $M_θ$ follows the architecture of Qwen2.5 models, with model sizes listed in the main text. As the thinking head, $T$ , is used for auto-regressive generation of language tokens, and is expected to return token embeddings, it additionally includes an unembedding and embedding layer.
This architecture facilitates the function of the thinking head and compression head as follows. At inference time, $T$ produces thoughts in natural-language through autoregressive generation, initially conditioned on the chunk of hidden representations $H_i^outââ^cĂ d$ , where $c,d$ denote the chunk size and the hidden size of $M_θ$ . Each token predicted by $T$ is passed through the embedding layer, and the resulting embedding is appended to the context of $T$ . This results in a distribution mismatch, as $H_i^out$ (extracted from a deep layer) share the context with standard token embeddings. However, we observe that this does not lead to a significant limitations when the block is finetuned for next-token-prediction in the same setup, even with relatively little data as in 2-Hop PK
For $C$ , to use the causal Transformer as a compression module, we apply a single parallel forward pass to all input token embeddings using a causal attention mask (though bidirectional attention is also valid) to obtain a contextualize representations. We then extract the last $c$ representations to serve as the state $S_i$ injected to the next chunk, appending padding tokens to the input sequence if it is shorter than $c$ . Through finetuning, $C$ learns to compress relevant information into future positions. We also experimented with appending a set of learnable queries appended to the query to make the compression explicit, however, this approach performed worse in practice.
To initialize the parameters of the two modules, we utilize existing components from $M_θ$ . Namely, the Transformer block in $T$ is initialized as an independent copy of the last layer of $M_θ$ , while $C$ is initialized as an independent copy of the first layer (excluding the token embedding layer). The unembedding layer in $T$ is initialized as an independent copy of the unembedding layer of $M_θ$ while the embedding layer is shared with $M_θ$ (i.e., is not an independent copy).
For $T$ , the motivation is to leverage parameters already trained to convert hidden representations into token predictions. For $C$ , the motivation is to leverage parameters trained to contextualize token embeddings (Kaplan et al., 2024) and generate a state sharing a latent space with the injection layer, which immediately follows the first layer.
### A.2 Fast Prefill with Speculated Thinking
We provide a more detailed description of the fast prefill algorithm detailed in Section 3.3 in algorithm 1 Full version can be reviewed in our codebase, which we plan to fully release. The idea is using the sparsity of non-trivial states, generated by more than an <EOS>, to a few allow parallel execution instead of iteratively processing all chunks in the query. This is achieved by always speculating future states are trivial, applying a parallel pass conditioned on the speculation, then correcting for the first location where the assumption breaks and repeating.
This process is guaranteed to terminate in exactly $|R|+1$ iterations, where $|R|$ is the number of non-trivial states generated, which is unknown in advance. As each iteration after the first implies exactly one non-trivial state was generated, this leads to exactly $|R|+1$ rounds.
In algorithm 1, we explicitly denote that $M_θ$ returns a set of hidden representations $H$ from all layers, only then extracting the target representation $H_i^out$ to emphasize all representations are stored in the cache. Additionally, for ease of presentation, the algorithm omits the fact that when finished, the representation of the last position in the input is used to generate the next token, as is in standard decoding.
Algorithm 1 Fast Speculative Prefill
Input: Query chunks $X_1\dots K$ , Backbone $M_θ$ , Thinking Block $T_Ď$ , Compression $C_Ď$ , Chunk Size $c$
$t=0$
$kv=init\_cache()$
$S_i=\hat{0}ââ^cĂ d iâ[0,\dots,K-1]$ {Speculate thoughts}
$S=concat([S_1,\dots,S_k])$
$X=concat([X_1,\dots,X_k])$
$finished=false$
while $not finished$ do
$\tilde{X}=S+X$
$H=M_θ(\tilde{X},past\_kv=kv)$
$H^out=extract(H,L^out)$
$H_i^out=chunk(H^out,c) iâ[t,\dots,K]$
$S_i+1=C_Ď_2(T_Ď_1(H_i^out))$
$t=find\_first\_true(S_iâ \hat{0})$
if $t$ is $None$ then
$finished=true$
$kv.cache(H,max\_index=None)$
else
$kv.cache(H,max\_index=t)$
$S=concat([S_t,\hat{0},\dots,\hat{0}])$
$X=concat([X_t,X_t+1,\dots,X_k])$
end if
end while
return $kv,S_-1$
### A.3 Constructing Chunk-Level Supervision
As mentioned in the main text, Section 3.4, chunk level supervision for the Thinking Block is obtained by mapping reasoning steps to specific positions in the input query. The mapping is realized by placing special indicator tokens in the query text.
Once the query is tokenized, the special tokens allow us to construct a matching reasoning array sharing the same length. At every position in the query containing an indicator token, we insert the matching reasoning sequence to the same position in the reasoning array. In the vast majority of cases, the desired supervision is obtained by shifting the reasoning array one position to the left, then removing all entries containing a special token in the query sequence, from the query tokens and the reasoning array.
A small number of cases contain multiple consecutive indicator tokens that are placed in succession. In these cases, we assign the matching steps in order, starting from the first token position preceding the first indicator. In case the reasoning steps âoverflowâ the token sequence length, these steps are concatenated together with a special separator token.
### A.4 State Tracking Tasks
We provide the definitions for the state tracking tasks used for evaluation in Section 4.2, âParityâ and âVariable Assignmentâ (abbreviated to âVarsâ).
Parity: The parity task acts as a natural language version of the standard parity prediction problem over binary sequences, which is commonly used when evaluating extrapolation capabilities in various architectures (Dehghani et al., 2019; Anil et al., 2022; delĂŠtang2023neuralnetworkschomskyhierarchy). Over binary sequences, e.g., $s=00101011$ , the task is predicting $sum(s)_\bmod 2$ . To represent the task in natural language, we transform the binary sequence into a query describing operations applied to a coin by two entities. When providing supervision to $T$ , the indicator token mentioned in Section 3.4 is placed after each operation, paired with the current state.
An abbreviated query is provided in Figure 2(a) and a full query, equivalent to the binary sequence $s=1011$ , is given here, along with indicator tokens and the matching labels:
âThe coin starts at state heads. <T> Alice doesnât flip the coin. <T> Bob flips the coin. <T> Alice flips the coin. <T> The final state of the coin is heads. â
Reasoning Targets: [âheads <eos> â, âheads <eos> â, âtails <eos> â, âheads <eos> â]
Where <T> denotes the indicator token and blue denotes text that converted into target tokens for predictions. In this formulation, increasing the number of state operations directly translates to longer binary sequences, i.e., more operations applied to the coin.
Variable Assignment (Vars): This task extends the state tracking setting to multiple integer-valued variables whose updates may depend on one another, requiring complex state updates (Anil et al., 2022). Unlike Parity, which tracks a single binary state, Vars requires maintaining the values of multiple variables (e.g., $a,b,c$ ) and executing arithmetic operations (e.g., $b=b+a$ ) that reference the current state of other variables.
The task is represented in natural language as a sequence of assignment operations. To generate supervision for $T$ , similar to the Parity task, we insert the indicator token <T> after each operation paired with the value of the updated variable.
An example query, equivalent to the provided raw input, is formatted below with indicator tokens and the corresponding reasoning targets (derived using modular arithmetic consistent with the final values):
âTrack the variables values: a=1; b=2 a=a+b <T> b=b+a <T> b=b+3 <T> Final values: a=3 b=8 â
Reasoning Targets: [âa=3 <eos> â, âb=5 <eos> â, âb=8 <eos> â]
### A.5 General Reasoning Capabilities
#### A.5.1 Coconut - Additional Results
As mentioned in the main paper, we scale computation with Coconut by increasing the number of latent thinking tokens. This is achieved by increasing the number of latent tokens introduced in each round of the curriculum, out of 3 rounds in total, excluding the initial phase training the CoT checkpoint used for initialization. As is evident in Table 3 The experiment in the main paper uses $2$ thinking tokens per round, $6$ in total.
Table 3: Coconut (Hao et al., 2024) performance and speedup as computation increases, by increasing the number of latent steps.
| 6 9 12 | 32.65 33.48 31.29 | 3.14 $Ă$ 2.53 $Ă$ 2.08 $Ă$ |
| --- | --- | --- |
| 15 | 32.35 | 1.77 $Ă$ |
| 21 | 32.65 | 1.38 $Ă$ |
#### A.5.2 Data Construction for GSM
As mentioned in the introduction, for the GSM task, we rely on a strong teacher model (Gemini-2.5 Flash (Comanici et al., 2025)) to transform existing parsed CoT data into the mapping used to supervise the thinking head. We achieve this by prompting the teacher model with a prompt similar to the one presented in Figure 7, removing two of the in context examples for readability. Using the structured output, we are able to extract the transformed query and matching thinking steps. To validate the process, we only include samples whose queries and thinking steps provide an exact string match to the source query and thinking steps, once the indicator tokens <THINK> are removed, discarding any samples that result in a mismatch. After validation we are left with $375,101/385,620$ samples, or $97.2\$ .
<details>
<summary>x6.png Details</summary>
### Visual Description
## Line Chart: Performance Comparison of "Coconut" vs. "Thinking States" Methods
### Overview
The image is a line chart comparing the computational time (in seconds) of two methods, "Coconut" and "Thinking States," as a function of problem size ("Number of Chunks / Latents"). Each method is further broken down into "Forward" and "Backward" passes. The chart uses a logarithmic scale for the time axis.
### Components/Axes
* **Chart Type:** Line chart with error bars.
* **X-Axis:**
* **Label:** `Number of Chunks / Latents`
* **Scale:** Linear, with major tick marks at 2, 4, 6, 8, 10, 12, 14, and 16.
* **Y-Axis:**
* **Label:** `Time (s)`
* **Scale:** Logarithmic (base 10), ranging from `10^-1` (0.1 seconds) to `10^0` (1.0 second).
* **Legend (Positioned at the top center of the chart area):**
* **Coconut:** Represented by an orange line.
* **Thinking States:** Represented by a blue line.
* **Forward:** Represented by a solid line with circular markers (â).
* **Backward:** Represented by a dashed line with square markers (â ).
* **Combined Interpretation:** The legend defines four distinct data series:
1. **Coconut Forward:** Orange solid line with circle markers.
2. **Coconut Backward:** Orange dashed line with square markers.
3. **Thinking States Forward:** Blue solid line with circle markers.
4. **Thinking States Backward:** Blue dashed line with square markers.
### Detailed Analysis
**Data Series & Trends:**
1. **Coconut Forward (Orange, Solid, Circles):**
* **Trend:** Shows a strong, approximately linear upward trend on the logarithmic scale, indicating an exponential increase in time with the number of chunks/latents.
* **Approximate Data Points (with uncertainty from visual estimation):**
* At 1 chunk: ~0.08 s (with a large error bar extending below 0.05 s).
* At 2 chunks: ~0.15 s.
* At 4 chunks: ~0.25 s.
* At 6 chunks: ~0.35 s.
* At 8 chunks: ~0.45 s.
* At 10 chunks: ~0.55 s.
* At 12 chunks: ~0.65 s.
* At 14 chunks: ~0.75 s.
* At 16 chunks: ~0.85 s.
2. **Coconut Backward (Orange, Dashed, Squares):**
* **Trend:** Also shows a strong, approximately linear upward trend on the logarithmic scale, parallel to but consistently higher than the Coconut Forward series.
* **Approximate Data Points:**
* At 1 chunk: ~0.1 s.
* At 2 chunks: ~0.2 s.
* At 4 chunks: ~0.35 s.
* At 6 chunks: ~0.5 s.
* At 8 chunks: ~0.65 s.
* At 10 chunks: ~0.8 s.
* At 12 chunks: ~1.0 s.
* At 14 chunks: ~1.1 s.
* At 16 chunks: ~1.2 s.
3. **Thinking States Forward (Blue, Solid, Circles):**
* **Trend:** Appears nearly flat, showing a very slight, almost negligible increase in time as the number of chunks/latents increases.
* **Approximate Data Points:** Consistently clustered around ~0.09 s across all x-axis values (from 1 to 16 chunks). Error bars are present but relatively small.
4. **Thinking States Backward (Blue, Dashed, Squares):**
* **Trend:** Appears completely flat, showing no discernible increase in time with problem size.
* **Approximate Data Points:** Consistently clustered around ~0.07 s across all x-axis values. This series is consistently the lowest in time.
### Key Observations
1. **Divergent Scaling:** The most prominent feature is the stark difference in scaling behavior. The "Coconut" method's time cost grows exponentially with problem size (linear on log plot), while the "Thinking States" method's time cost is nearly constant.
2. **Forward vs. Backward Cost:** For the "Coconut" method, the Backward pass is consistently more time-consuming than the Forward pass. For "Thinking States," the difference between Forward and Backward passes is minimal, with the Backward pass being slightly faster.
3. **Performance Crossover:** At the smallest problem size (1 chunk), the times for all methods are relatively close (within the 0.05-0.1 s range). As the number of chunks increases, the performance gap widens dramatically, with "Coconut" becoming orders of magnitude slower than "Thinking States" at 16 chunks.
4. **Error Bars:** The "Coconut" series, particularly at lower chunk counts, show larger error bars, suggesting higher variance in its execution time compared to the more stable "Thinking States" method.
### Interpretation
This chart demonstrates a fundamental difference in the algorithmic complexity or implementation efficiency between the "Coconut" and "Thinking States" approaches for the measured task.
* **"Coconut"** exhibits **poor scalability**. Its computational time increases exponentially with the problem size (number of chunks/latents). This suggests its algorithm has a high time complexity (e.g., O(n^k) where k>1), making it increasingly impractical for larger problems. The higher cost of the Backward pass might indicate a more complex gradient computation or optimization step.
* **"Thinking States"** exhibits **excellent, near-constant-time scalability**. Its performance is largely independent of the problem size within the tested range. This is characteristic of algorithms with O(1) or very low O(log n) complexity for this parameter. This property is highly desirable, as it means the method can handle larger inputs without a significant performance penalty. The minimal difference between Forward and Backward passes suggests a symmetric and efficient computation.
**Conclusion:** The data strongly suggests that for tasks involving a variable number of "Chunks / Latents," the "Thinking States" method is vastly more efficient and scalable than the "Coconut" method. The choice between them would heavily depend on the expected problem size; for small, fixed-size problems, the difference may be negligible, but for any application where the input size can grow, "Thinking States" is the clearly superior choice from a performance perspective.
</details>
Figure 6: Training Cost Analysis Comparison of wall-clock time for forward and backward passes on a fixed-length input ( $L=128$ ) as a function of the number of recurrent steps (chunks). Methods relying on BPTT demonstrate linear scaling in time, circumvented by teacher forcing in Thinking States despite the added dimensionality and computational overhead of $T$ .
#### A.5.3 Training Compute
We quantify the computational advantage of training Thinking States with teacher forcing compared to latent reasoning methods that rely on Backpropagation Through Time (BPTT), such as Coconut. Figure 6 illustrates the wall-clock time required for a single forward and backward pass on a fixed-length sequence, as a function of the number of recurrent steps (Coconut) or number of chunks Here the chunk size accordingly, we note that $k$ chunks imply $k$ recurrent steps at inference time. (Thinking States), by controlling the chunk size relative to the sequence length. We use a batch size of 1 and a sequence length of 128 tokens.
As shown in Figure 6, training with BPTT exhibits a linear increase in wall-time for both forward and backward passes as the number of latent steps grows. In contrast, Thinking States maintains close to constant training time regardless of the number of chunks, as the availability of ground-truth states allows for fully parallelized computation. The small variation in wall-time for Thinking States is explained by the additional dimensionality of decoding reasoning sequence from fixed-size token chunks. Yet, as visible in the figure, these variations are small compared to the increase in cost due to BPTT.
Prompt for GSM Data Generation
You are an expert in computational linguistics. Your task is to augment a given query with thinking markers (<THINK>) based on a provided reasoning trace. This involves identifying the precise locations in the text where specific calculations can be performed. Task Instructions: 1.
Analyze Inputs: You will receive two inputs: â˘
A Query: A question or word problem. â˘
A Reasoning Trace: An ordered list of strings, where each string is a distinct mathematical calculation required to solve the query. 2.
Produce Two Outputs: â˘
query: This is the original Query text, but with the special token <THINK> inserted at specific locations. â˘
thinking: This is the ordered list of calculations. 3.
Rules for <THINK> Placement: â˘
The number of <THINK> tokens inserted into the query must be exactly equal to the number of calculations in the Reasoning Trace. â˘
The order of the <THINK> tokens in the query must correspond one-to-one with the order of the calculations in the Reasoning Trace. â˘
For each calculation, you must insert its corresponding <THINK> token at the earliest possible location in the query. This location is defined as the point immediately after the word or phrase that provides the final piece of information needed to perform that specific calculation. 4.
Rule for the thinking List: â˘
The thinking output list is simply a direct copy of the ordered list of calculations provided in the Input Reasoning Trace. 5.
Output Format: â˘
The output should be provided as a python dictionary with two keys: query and thinking. â˘
The values should be wrapped in <<>> to support easy parsing. â˘
Example: { "query": "<<Your modified query here>>", "thinking": <<[â<<calculation1>>â, â<<calculation2>>â, ...]>> }
Examples Example 1: â˘
Input Query: âHannah has three dogs. The first dog eats 1.5 cups of dog food a day. The second dog eats twice as much while the third dog eats 2.5 cups more than the second dog. How many cups of dog food should Hannah prepare in a day for her three dogs?â â˘
Input Reasoning Trace: [â<<1.5*2=3>>â, â<<3+2.5=5.5>>â, â<<1.5+3+5.5=10>>â] Required Output: { "query": <<âHannah has three dogs. The first dog eats 1.5 cups of dog food a day. The second dog eats twice as much<THINK> while the third dog eats 2.5 cups more than the second dog.<THINK> How many cups of dog food should Hannah prepare in a day for her three dogs?<THINK> â>>, "thinking": <<[â<<1.5*2=3>>â, â<<3+2.5=5.5>>â, â<<1.5+3+5.5=10>>â]>> }
Task Now, perform this transformation for the following input and return your output as a python dictionary. Query: <<QUERY>> Reasoning Trace: <<REASONING_TRACE>> Produce your output as a python dictionary in the specified format below: { "query": <<ADD YOUR OUTPUT HERE>>, "thinking": <<ADD YOUR OUTPUT HERE>> }
Figure 7: The abbreviated version of few-shot prompt used to generate the alignment data. The prompt instructs the model to act as a computational linguist and inject <THINK> tokens into the query corresponding to specific reasoning steps. The full prompt used is identical except for containing 2 additional in-context examples, similar to the example provided.