2602.08332

# 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 \n ## Diagram: LLM Thinking States vs. Chain-of-Thought ### Overview The image presents a comparative diagram illustrating two approaches within a Large Language Model (LLM): "Thinking States" and "Chain-of-Thought." Both diagrams depict the internal processing of a simple coin flip scenario ("A coin state is heads. Alice flips, then Bob flips. What's the state? A: heads."). The diagrams use a block-based representation to visualize the flow of information through the LLM. ### Components/Axes The diagram is divided into two main sections, each representing a different approach. Each section contains: * **LLM Block:** A large rectangular block labeled "LLM" representing the language model itself. * **Thinking States/Chain-of-Thought Labels:** Labels indicating the specific approach being visualized. * **Thinking States/Language Tokens/Compression Block:** Labels for the different layers of processing within the LLM. * **Input Text:** The prompt "A coin state is heads. Alice flips, then Bob flips. What's the state? A: heads." is present below each diagram. * **Blocks:** Smaller rectangular blocks within the LLM block, representing internal states or computations. * **Arrows:** Arrows indicating the flow of information between blocks. * **Labels within Blocks:** Text labels within the blocks, such as "tails", "heads", "<EOS>", "s", "c", and "A". ### Detailed Analysis or Content Details **Thinking States Diagram (Left Side):** * **Thinking States Layer:** Contains blocks labeled "s" (approximately 6x4 grid). Two blocks are highlighted in orange. * **Language Tokens Layer:** Contains blocks labeled "t" (approximately 6x4 grid). * **Compression Block Layer:** Contains blocks labeled "c" (approximately 6x4 grid). * **Flow:** The flow starts with "tails" and "<EOS>" entering the Thinking States layer. These values propagate through the Compression Block layer, then to the Language Tokens layer. The final output is "A: heads." * **Input:** "A coin state is heads. Alice flips, then Bob flips. What's the state? A: heads." **Chain-of-Thought Diagram (Right Side):** * **LLM Block:** Contains a grid of blocks (approximately 6x8). * **Flow:** The flow starts with "heads", "tails", "heads", "heads" entering the LLM block. These values propagate through the block grid. The final output is "A: heads." * **Input:** "A coin state is heads. Alice flips, then Bob flips. What's the state? A: heads." ### Key Observations * **Thinking States:** Appears to compress information into a smaller "Thinking States" representation before generating the final answer. The orange highlighted blocks suggest a focus on specific states. * **Chain-of-Thought:** Shows a more direct flow of information, with intermediate states ("heads", "tails", "heads", "heads") being explicitly represented within the LLM. * **Output Consistency:** Both approaches arrive at the same answer ("A: heads."). * **Visual Complexity:** The Chain-of-Thought diagram appears more visually complex, with more blocks and a longer flow path. ### Interpretation The diagram illustrates two different strategies for how an LLM might approach a reasoning task. "Thinking States" suggests a compression-based approach, where the LLM first reduces the input into a more compact internal representation before generating the output. "Chain-of-Thought" demonstrates a more explicit reasoning process, where the LLM maintains a record of intermediate states and steps. The fact that both approaches yield the same result suggests that both are capable of solving the problem, but they do so in different ways. The Chain-of-Thought approach might be more interpretable, as it explicitly shows the reasoning steps, while the Thinking States approach might be more efficient in terms of computational resources. The diagram highlights the internal workings of LLMs and how they can be designed to perform reasoning tasks. It suggests that there are multiple ways to achieve the same result, and the choice of approach can have implications for efficiency, interpretability, and other factors. The diagram is a conceptual illustration and does not provide specific numerical data or performance metrics. It is a qualitative comparison of two different architectural approaches. </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: Recurrent Neural Network with KV Cache Illustration ### Overview The image depicts a diagram illustrating a recurrent neural network architecture incorporating a Key-Value (KV) cache. The diagram shows the flow of information through the network, with a focus on how past states are stored and retrieved using the KV cache. The diagram also includes a textual example of a coin flip scenario. ### Components/Axes The diagram consists of several key components: * **L_in**: Represents the input layer. * **L_out**: Represents the output layer. * **S_i, S_i+1, S_i+2, S_i+3**: Represent the hidden states at different time steps. The "+" symbol indicates an addition operation. * **KV Cache**: A storage mechanism for past key-value pairs. * **T**: Represents a "Transformer" block within the KV cache. * **C**: Represents a "Cache" block within the KV cache. * **Z_i, Z_i+1, Z_i+2, Z_i+3**: Represent the output tokens at different time steps. * **Textual Example**: "A coin's state is heads. Alice flips, then Bob flips. What's the state? A: heads." ### Detailed Analysis or Content Details The diagram shows a sequential flow of information. 1. **Input Layer (L_in)**: The input layer consists of a series of unlabeled boxes, representing the initial input sequence. 2. **Hidden States**: The input is processed through a series of hidden states (S_i to S_i+3). Each hidden state is calculated by adding the previous hidden state to some input. 3. **KV Cache**: The KV cache stores key-value pairs derived from the hidden states. The cache is structured with "Transformer" (T) and "Cache" (C) blocks. The arrows indicate that the hidden states are used to populate the KV cache. 4. **Output Layer (L_out)**: The output layer receives information from both the current hidden state and the KV cache. The KV cache provides context from previous time steps. 5. **Output Tokens**: The output layer generates output tokens (Z_i to Z_i+3). * Z_i contains the text "tails" and "<EOS>". * Z_i+1 contains the text "<EOS>". * Z_i+2 contains the text "heads" and "<EOS>". The textual example at the bottom provides a simple scenario: "A coin's state is heads. Alice flips, then Bob flips. What's the state? A: heads." This example likely illustrates how the network can track state changes over time. ### Key Observations * The KV cache is crucial for maintaining context across time steps. * The diagram suggests a recurrent process where the hidden state at each time step depends on the previous hidden state and the current input. * The "<EOS>" token likely signifies the end of a sequence. * The Transformer and Cache blocks within the KV cache suggest a more complex internal structure. ### Interpretation This diagram illustrates a recurrent neural network architecture designed to handle sequential data. The KV cache is a key component, allowing the network to access and utilize information from past time steps. This is particularly important for tasks where context is crucial, such as natural language processing or time series analysis. The coin flip example demonstrates how the network can track state changes over time, suggesting its ability to model dynamic systems. The use of "Transformer" blocks within the KV cache hints at the incorporation of attention mechanisms, which can further enhance the network's ability to focus on relevant information from the past. The diagram is a simplified representation of a more complex system, but it effectively conveys the core principles of recurrent processing and the role of the KV cache in maintaining context. The diagram suggests a model capable of predicting future states based on past observations. The use of the KV cache is a technique to improve the efficiency and performance of such models. </details> (a) Inference <details> <summary>x3.png Details</summary>  ### Visual Description \n ## Diagram: Coin Flip Sequence & Prediction ### Overview The image depicts a diagram illustrating a sequence of coin flips and a prediction model attempting to predict the outcome. It shows two parallel sequences, one representing the true labels and the other representing the model's predictions. A loss function is indicated where the prediction differs from the label. The diagram also includes a question-answer scenario related to the coin flip sequence. ### Components/Axes The diagram is structured into two main sections, separated by a question. The top section shows the "Labels" and "Predictions" with associated arrows and a loss calculation. The bottom section shows a sequence of states (S_i, S_i+1, S_i+2, S_i+3) with corresponding labels. Labels: "tails", "<EOS>", "heads" Predictions: "tails", "heads", "heads", "<EOS>" Loss: "L_oss" States: S_i, S_i+1, S_i+2, S_i+3 Question: "A coin's state is heads. Alice flips, then Bob flips. What's the state? A: heads." ### Detailed Analysis or Content Details **Top Section:** * **Labels:** The labels sequence starts with "tails", followed by "<EOS>". Then "heads" followed by "<EOS>". The arrows indicate a flow from the label to the prediction. * **Predictions:** The predictions sequence starts with "tails", then "heads", then "heads", and ends with "<EOS>". * **Loss:** A red "X" symbol labeled "L_oss" connects the "tails" label to the "heads" prediction, indicating a discrepancy and thus a loss. **Bottom Section:** * **States:** Four states are shown: S_i, S_i+1, S_i+2, and S_i+3. Each state is represented by a "C" inside a square. * **Labels (Bottom):** The labels associated with the states are: "<EOS>", "tails", "<EOS>", "heads". The arrows indicate a flow from the label to the state. **Question & Answer:** * The question states: "A coin's state is heads. Alice flips, then Bob flips. What's the state? A: heads." ### Key Observations * The model correctly predicts "heads" after an initial incorrect prediction of "tails". * The "<EOS>" token appears to signify the end of a sequence. * The loss is only incurred when the prediction does not match the label. * The question-answer section provides context for the coin flip sequence. ### Interpretation This diagram illustrates a simplified model for predicting coin flip outcomes. The model attempts to learn the sequence of flips and predict the next state. The loss function quantifies the error between the predicted and actual outcomes. The "<EOS>" token suggests the model is trained on sequences of variable length. The question-answer section implies a sequential process where Alice and Bob take turns flipping the coin, and the model is trying to predict the final state. The diagram demonstrates a basic example of sequence prediction, potentially within a reinforcement learning or Markov decision process framework. The use of "C" within the state boxes is unclear without further context, but could represent the coin itself or a hidden state within the model. The diagram is a conceptual illustration rather than a presentation of specific numerical data. </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 $\mathbf{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 $\mathbf{X}_{1},\dots,\mathbf{X}_{K}$ . In the case $L\bmod c\neq 0$ , input tokens are grouped into chunks along with generated tokens, once those are available. Each $\mathbf{X}_{i}\in\mathbb{R}^{c\times 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_{\theta}$ with two additional modules: 1. Thinking Block $T$ : A lightweight Transformer decoder (one layer in our experiments) that autoregressively generates a reasoning sequence $\mathbf{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 $\mathbf{S}\in\mathbb{R}^{c\times 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 $\mathbf{S}_{i}\in\mathbb{R}^{c\times d}$ . At each step $i$ , we first inject the current state into the input representations at the shallow layer, then forward through $M_{\theta}$ : $$ \tilde{\mathbf{X}}_{i}=\mathbf{X}_{i}+\mathbf{S}_{i} \tag{1} $$ $$ \mathbf{H}_{i}^{out}=M_{\theta}(\tilde{\mathbf{X}}_{i}\mid\tilde{\mathbf{X}}_{<i}) \tag{2} $$ Here, $\mathbf{H}_{i}^{out}$ denotes the chunk representations at a predefined deep layer $L^{out}$ , and past representations $\tilde{\mathbf{X}}_{<i}$ are accessed through the KV-cache. Next, $\mathbf{H}_{i}^{out}$ is used to generate a sequence of natural-language thought tokens via the Thinking Block: $$ \mathbf{Z}_{i+1}={T}(\mathbf{H}_{i}^{out}) \tag{3} $$ The output $\mathbf{Z}_{i+1}$ is a variable-length token sequence ending with an <EOS> token. When no reasoning is required, $\mathbf{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: $$ \mathbf{S}_{i+1}={C}(\mathbf{Z}_{i+1})\in\mathbb{R}^{c\times d} \tag{4} $$ For the first chunk, we initialize $\mathbf{S}_{1}=\mathbf{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 $\mathbf{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 $\mathbf{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 $\mathbf{Z}_{i}^{*}$ serves a dual purpose. As illustrated in Figure 2(b), in addition to providing supervision for the outputs of $T$ , we use $\mathbf{Z}_{i}^{*}$ to teacher-force the states injected into the backbone LLM. Specifically, we compute the target state as $\mathbf{S}_{i}^{*}={C}(\mathbf{Z}_{i}^{*})$ and inject it instead of the model-generated state: $$ \tilde{\mathbf{X}}_{i}=\mathbf{X}_{i}+\mathbf{S}_{i}^{*} \tag{5} $$ Since all target sequences $\mathbf{Z}_{i}^{*}$ are available upfront, we can compute all chunk representations $\mathbf{H}_{i}^{out}$ in a single parallel forward pass through $M_{\theta}$ . 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 $\mathbf{Z}_{i}^{*}$ via standard next-token prediction, in parallel over all chunks. Here, each $\mathbf{H}_{i}^{out}$ is computed under the gold state history $\mathbf{S}_{1}^{*},\dots,\mathbf{S}_{i}^{*}$ , so predicting $\mathbf{Z}_{i+1}$ is implicitly conditioned on all prior gold reasoning steps. We optimize the joint objective: $$ \mathcal{L}=\mathcal{L}_{\text{LM}}+\sum_{i=1}^{K}\mathcal{L}_{\text{T}}(\mathbf{Z}_{i},\mathbf{Z}^{*}_{i}) \tag{6} $$ where $\mathcal{L}_{\text{LM}}$ is the standard language modeling loss and $\mathcal{L}_{\text{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{\mathbf{X}}_{i}$ depends on the previous state $\mathbf{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 $\mathbf{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 $\rightarrow$ tails $\rightarrow$ 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\in\{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 | $\times$ 5.59 $\times$ 2.66 $\times$ 3.14 | 33.47 54.91 33.71 | $\times$ 1.89 $\times$ 1.19 $\times$ 1.14 | 31.92 43.05 32.60 | $\times$ 2.03 $\times$ 1.23 $\times$ 1.21 | | --- | --- | --- | --- | --- | --- | --- | | iCoT | 34.00 | $\times$ 5.71 | 28.84 | $\times$ 1.59 | 36.31 | $\times$ 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 $\sim$ 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 \n ## Chart: Accuracy vs. Number of Layers in Recurrence with Speedup over CoT ### Overview The image presents a line chart illustrating the relationship between the number of layers in recurrence and the resulting accuracy, alongside speedup factors compared to a "Chain of Thought" (CoT) baseline. A horizontal dashed line indicates the CoT baseline accuracy. Speedup values are displayed above the data points. ### Components/Axes * **X-axis:** "Num Layers in Recurrence" - Scale ranges from 4 to 26, with markers at 4, 8, 12, 16, 20, and 24/26. * **Y-axis:** "Accuracy (%)" - Scale ranges from 20 to 48, with markers at 20, 24, 28, 32, 36, 40, 44, and 48. * **Data Series:** A single blue dashed line representing accuracy as a function of the number of layers. * **Baseline:** A horizontal red dashed line labeled "No CoT Baseline: 34.11%". * **Speedup Indicators:** Text labels above each data point indicating the speedup factor over the CoT baseline (e.g., "3.12x", "2.98x"). * **Legend:** No explicit legend is present, but the labels and line color clearly define the data series. ### Detailed Analysis The blue dashed line shows an increasing trend in accuracy as the number of layers in recurrence increases. * **4 Layers:** Accuracy is approximately 21%. Speedup is 3.12x. * **8 Layers:** Accuracy is approximately 25%. Speedup is 2.98x. * **12 Layers:** Accuracy is approximately 30%. Speedup is 2.97x. * **16 Layers:** Accuracy is approximately 32%. Speedup is 2.86x. * **20 Layers:** Accuracy is approximately 37%. Speedup is 2.78x. * **24/26 Layers:** Accuracy is approximately 40%. Speedup is 2.76x. The CoT baseline accuracy is fixed at 34.11%. The accuracy values surpass the baseline at approximately 16 layers. ### Key Observations * Accuracy increases rapidly with the number of layers initially, then the rate of increase slows down. * The speedup over CoT decreases slightly as the number of layers increases, but remains consistently above 2.76x. * The accuracy surpasses the CoT baseline at 16 layers. * The data suggests diminishing returns in accuracy gains with increasing layers beyond 16. ### Interpretation The chart demonstrates that increasing the number of layers in recurrence improves accuracy, but with diminishing returns. The speedup over the CoT baseline indicates that this approach is significantly faster than the baseline method, even with a relatively small number of layers. The CoT baseline represents a performance threshold. The data suggests that a model with around 16-20 layers provides a good balance between accuracy and speedup. The slight decrease in speedup with increasing layers might indicate increased computational cost outweighing the marginal accuracy gains. The chart implies that the recurrence-based approach is a viable alternative to CoT, offering both improved accuracy and efficiency. The data points are relatively sparse, so interpolation between them should be done with caution. </details> (a) <details> <summary>x5.png Details</summary>  ### Visual Description \n ## Line Chart: Accuracy vs. Chunk Size with Speedup over CoT ### Overview This line chart depicts the relationship between Chunk Size and Accuracy, along with the Speedup over Chain of Thought (CoT). The chart shows how accuracy changes as the chunk size increases, and provides a baseline accuracy for the "No CoT" method. Speedup values are displayed above each data point. ### Components/Axes * **X-axis:** Chunk Size (ranging from 2 to 48, with markers at 2, 4, 8, 16, 32, and 48). * **Y-axis:** Accuracy (%) (ranging from 20% to 48%, with markers at 24%, 28%, 32%, 36%, 40%, 44%, and 48%). * **Data Series:** A single dashed blue line representing Accuracy vs. Chunk Size. * **Baseline:** A red dotted horizontal line labeled "No CoT Baseline: 34.11%". * **Speedup Labels:** Text labels above each data point indicating the speedup factor over CoT (e.g., "2.42x", "2.57x", etc.). * **Title:** "Speedup over CoT" positioned above the chart. ### Detailed Analysis The blue line representing Accuracy generally increases with Chunk Size, reaches a peak, and then slightly decreases. * **Chunk Size = 2:** Accuracy ≈ 35%. Speedup = 2.42x. * **Chunk Size = 4:** Accuracy ≈ 41%. Speedup = 2.57x. * **Chunk Size = 8:** Accuracy ≈ 42%. Speedup = 2.66x. * **Chunk Size = 16:** Accuracy ≈ 39%. Speedup = 2.74x. * **Chunk Size = 32:** Accuracy ≈ 38%. Speedup = 2.83x. * **Chunk Size = 48:** Accuracy ≈ 37%. Speedup = 2.86x. The "No CoT Baseline" is at 34.11% accuracy. ### Key Observations * Accuracy peaks around a Chunk Size of 8, then declines slightly. * Speedup over CoT consistently increases with Chunk Size, even as accuracy plateaus or declines. * All accuracy values are above the "No CoT Baseline" of 34.11%. * The increase in speedup appears to be diminishing as chunk size increases. ### Interpretation The data suggests that increasing the chunk size initially improves accuracy, but beyond a certain point (around Chunk Size 8), further increases in chunk size do not lead to significant gains in accuracy and may even cause a slight decrease. However, the speedup over CoT continues to increase with chunk size, indicating a trade-off between accuracy and computational efficiency. The consistent accuracy above the "No CoT Baseline" suggests that the method used is generally more effective than not using CoT. The diminishing returns in speedup as chunk size increases might indicate that there are limitations to how much benefit can be gained from simply increasing the chunk size. Further investigation might be needed to determine the optimal chunk size for maximizing both accuracy and efficiency. The chart demonstrates a clear relationship between chunk size, accuracy, and computational speed, providing valuable insights for optimizing the performance of the system. </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 $\sim 8\$ and maintaining a considerable $2.66\times$ 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 $\sim 10\times$ 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 $\times$ 8=120.] Three quarters of the units are occupied. [T: (3/4) $\times$ 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 $\times$ 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_{\theta}$ . 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_{\theta}$ . In our experiments, $M_{\theta}$ 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 $\mathbf{H}_{i}^{out}\in\mathbb{R}^{c\times d}$ , where $c,d$ denote the chunk size and the hidden size of $M_{\theta}$ . 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 $\mathbf{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 $\mathbf{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_{\theta}$ . Namely, the Transformer block in $T$ is initialized as an independent copy of the last layer of $M_{\theta}$ , 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_{\theta}$ while the embedding layer is shared with $M_{\theta}$ (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_{\theta}$ returns a set of hidden representations $H$ from all layers, only then extracting the target representation $\mathbf{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 $\mathbf{X}_{1\dots K}$ , Backbone $M_{\theta}$ , Thinking Block $T_{\phi}$ , Compression $C_{\phi}$ , Chunk Size $c$ $t=0$ $kv=init\_cache()$ $\mathbf{S}_{i}=\hat{0}\in\mathbb{R}^{c\times d}\quad i\in[0,\dots,K-1]$ {Speculate thoughts} $S=concat([\mathbf{S}_{1},\dots,\mathbf{S}_{k}])$ $X=concat([\mathbf{X}_{1},\dots,\mathbf{X}_{k}])$ $finished=false$ while $not\ \ finished$ do $\tilde{X}=S+X$ $H=M_{\theta}(\tilde{X},past\_kv=kv)$ $H^{out}=extract(H,L^{out})$ $\mathbf{H}_{i}^{out}=chunk(H^{out},c)\quad i\in[t,\dots,K]$ $\mathbf{S}_{i+1}=\mathcal{C}_{\phi_{2}}(T_{\phi_{1}}(\mathbf{H}_{i}^{out}))$ $t=find\_first\_true(\mathbf{S}_{i}\neq\hat{0})$ if $t$ is $None$ then $finished=true$ $kv.cache(H,max\_index=None)$ else $kv.cache(H,max\_index=t)$ $S=concat([\mathbf{S}_{t},\hat{0},\dots,\hat{0}])$ $X=concat([\mathbf{X}_{t},\mathbf{X}_{t+1},\dots,\mathbf{X}_{k}])$ end if end while return $kv,\mathbf{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 $\times$ 2.53 $\times$ 2.08 $\times$ | | --- | --- | --- | | 15 | 32.35 | 1.77 $\times$ | | 21 | 32.65 | 1.38 $\times$ | #### 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 \n ## Line Chart: Performance vs. Chunk Size ### Overview This image presents a line chart comparing the time taken (in seconds) by different methods – Coconut, Thinking States, Forward, and Backward – as the number of chunks/latents increases from 1 to 16. The y-axis uses a logarithmic scale. ### Components/Axes * **X-axis:** "Number of Chunks / Latents" ranging from 1 to 16. The axis is linearly scaled. * **Y-axis:** "Time (s)" ranging from approximately 0.01 to 100 seconds. The axis is logarithmically scaled, with markers at 0.1 and 10. * **Legend:** Located at the top-center of the chart. * Coconut (Orange solid line) * Thinking States (Blue solid line) * Forward (Gray solid line with circle markers) * Backward (Blue dashed line with square markers) ### Detailed Analysis The chart displays four distinct lines representing the performance of each method. * **Coconut (Orange):** The line slopes upward, indicating increasing time with increasing chunks/latents. * At 1 chunk/latent: ~0.03 seconds * At 2 chunks/latents: ~0.2 seconds * At 4 chunks/latents: ~0.6 seconds * At 6 chunks/latents: ~1.2 seconds * At 8 chunks/latents: ~1.7 seconds * At 10 chunks/latents: ~2.2 seconds * At 12 chunks/latents: ~2.7 seconds * At 14 chunks/latents: ~3.2 seconds * At 16 chunks/latents: ~3.7 seconds * **Thinking States (Blue):** The line is relatively flat, indicating minimal change in time with increasing chunks/latents. * At 1 chunk/latent: ~0.04 seconds * At 2 chunks/latents: ~0.05 seconds * At 4 chunks/latents: ~0.05 seconds * At 6 chunks/latents: ~0.05 seconds * At 8 chunks/latents: ~0.05 seconds * At 10 chunks/latents: ~0.05 seconds * At 12 chunks/latents: ~0.05 seconds * At 14 chunks/latents: ~0.05 seconds * At 16 chunks/latents: ~0.05 seconds * **Forward (Gray):** The line is nearly horizontal, showing very little variation in time. * At 1 chunk/latent: ~0.06 seconds * At 2 chunks/latents: ~0.06 seconds * At 4 chunks/latents: ~0.06 seconds * At 6 chunks/latents: ~0.06 seconds * At 8 chunks/latents: ~0.06 seconds * At 10 chunks/latents: ~0.06 seconds * At 12 chunks/latents: ~0.06 seconds * At 14 chunks/latents: ~0.06 seconds * At 16 chunks/latents: ~0.06 seconds * **Backward (Blue dashed):** The line is also relatively flat, but slightly below the "Thinking States" line. * At 1 chunk/latent: ~0.03 seconds * At 2 chunks/latents: ~0.04 seconds * At 4 chunks/latents: ~0.04 seconds * At 6 chunks/latents: ~0.04 seconds * At 8 chunks/latents: ~0.04 seconds * At 10 chunks/latents: ~0.04 seconds * At 12 chunks/latents: ~0.04 seconds * At 14 chunks/latents: ~0.04 seconds * At 16 chunks/latents: ~0.04 seconds ### Key Observations * The "Coconut" method exhibits a clear positive correlation between the number of chunks/latents and the time taken. * "Thinking States", "Forward", and "Backward" methods show minimal sensitivity to the number of chunks/latents, maintaining relatively constant execution times. * The logarithmic scale on the y-axis emphasizes the significant difference in time scaling between "Coconut" and the other methods. * "Coconut" is significantly slower than the other methods, especially as the number of chunks/latents increases. ### Interpretation The data suggests that the "Coconut" method's computational cost scales with the size of the input (number of chunks/latents), while the other methods ("Thinking States", "Forward", and "Backward") are largely independent of this parameter. This could indicate that "Coconut" involves a process that requires more computation per chunk/latent, such as iterative refinement or complex calculations. The flat lines for the other methods suggest they may utilize a more efficient algorithm or have a fixed computational cost regardless of input size. The logarithmic scale highlights the exponential increase in time for "Coconut", making it less suitable for large numbers of chunks/latents. The "Backward" method is slightly faster than "Thinking States" and "Forward" but the difference is minimal. This data could be used to inform the selection of an appropriate method based on the expected input size and performance requirements. </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.