Emergence of Superposition Mechanisms: Unveiling the Training Dynamics of Chain of Thought

Whether gradient methods can naturally learn superposition in continuous chain of thought is the question this paper addresses.

• Setting: the authors say long discrete chain of thought is costly, so they keep the reasoning trajectory in a continuous latent space. Zhu et al. (2025) write a class of reasoning as directed graph reachability and give a parameter construction for a two-layer transformer that maintains superposition of multiple search trajectories with O(n) steps of continuous thought. n is the number of vertices.
• Gap: that construction does not analyze whether gradient training reaches this solution. The authors contrast Tian et al. (2023a), Nichani et al. (2024a), Nguyen and Nguyen-Tang (2025): in those analyses without continuous thought, logits grow logarithmically and are therefore not bounded above.
• Hypothesis: when each sample demonstrates only one shortest path, superposition can still emerge. The authors connect this to the index-matching logit. This quantity measures local search strength: too small and local structure is barely exploited; too large and it may commit to one trajectory based only on local features such as in-degree.
• This paper: analyzes gradient flow in the two stages of thought generation and answer prediction on a simplified two-layer transformer, and tracks the logit in experiments. The authors state that under their mild assumptions, this logit first rises and then remains bounded, giving comparable weights to multiple locally feasible trajectories and forming superposition.

Related work is divided into three groups according to Section 1.1. The experimental baselines are two continuous thought losses, and the data comes from ProsQA.

Text chain of thought

• Wei et al. (2022) use chain-of-thought for test-time scaling. Khot et al. (2022), Zhou et al. (2022) rely only on prompting; Yue et al. (2023), Yu et al. (2023), Wang et al. (2023a), Shao et al. (2024) write chain of thought into training samples.
• Liu et al. (2022), Feng et al. (2023), Merrill and Sabharwal (2023), Li et al. (2024b) discuss chain of thought and transformer expressivity. Zhu et al. (2024) use training dynamics to examine the role of chain of thought in two-hop reasoning. Wen et al. (2024), Kim and Suzuki (2024) study chain of thought in training data and sample efficiency.

Latent-space reasoning

• Goyal et al. (2023) use learnable pause tokens. The paper writes that the theoretical result of London and Kanade (2025) is that pause tokens can strictly increase transformer expressivity. Pfau et al. (2024) study filler tokens, Wang et al. (2023b) use planning tokens, Su et al. (2025) use latent-space abstract tokens.
• Hao et al. (2024) propose continuous chain of thought COCONUT. Zhu et al. (2025) analyze its advantages from expressivity and give a parameter construction on graph reachability. The paper says that construction does not analyze whether gradient training can naturally reach it.

Training dynamics

• The paper lists Jelassi et al. (2022) through Guo et al. (2024) as transformer optimization analyses without expanding on each one.
• Another line of work uses training dynamics to understand reasoning patterns; the paper names induction head (Nichani et al., 2024a), reversal curse (Zhu et al., 2024; Ma et al., 2026), chain of thought (Wen et al., 2024; Kim and Suzuki, 2024; Huang et al., 2025b), fact recall (Nichani et al., 2024b), two-hop reasoning (Guo et al., 2025; Lin et al., 2025), and out-of-context reasoning (Huang et al., 2025a).
• The paper explicitly contrasts Tian et al. (2023a), Nichani et al. (2024a), Nguyen and Nguyen-Tang (2025): in those settings attention logits diverge, whereas continuous thought training typically yields bounded logits.

Baselines and data

• The experimental baselines are the COCONUT loss Lcoco and the COCONUT-BFS-exp variant in Appendix E.2. The theory also contrasts COCONUT-BFS, which enumerates all next-layer nodes. The data is the subset of ProsQA (Hao et al., 2024) used by Zhu et al. (2025).

The authors place this paper outside text chain of thought: the thought vector stays in continuous space. Relative to the expressivity construction of Zhu et al. (2025), this paper analyzes how gradient training produces superposition.

The overall approach is to split COCONUT into thought generation and answer prediction, and track the index-matching logit on a simplified two-layer transformer. The theory uses gradient flow; the experiments use curriculum learning, with details in question 4.

Task, prompt, and continuous thought

• Task: directed graph G=(V,E) with n vertices. A fixed root r and two candidate endpoints, exactly one c* reachable from r and the other c⊥ unreachable. Also fixed is a shortest path p from r to c*.
• Prompt: a sentence-start token, m edges (source node, target node, edge token), a question token, two candidates, a reasoning token, and the root. Prompt length t0=3m+6. Position notation follows Zhu et al. (2025).
• Continuous decoding: discrete chain of thought first samples the output into tokens and then embeds them; continuous chain of thought feeds the transformer output directly back into the sequence. After C steps, an answer token is appended, and then the vocabulary entry with the largest inner product with the output is taken.
• Curriculum: stage c+1 first autoregressively generates the first c continuous thoughts without loss, then supervises the next thought. Each instance has only one shortest-path demonstration.

Second layer and local search strength

• The authors do not rederive the first-layer copy, citing existing work on the copy and its formation. The hidden state of an edge token is the sum of the source and target embeddings; the current thought is set to a normalized superposition of embeddings of nodes reachable within c steps.
• Assumption 1: Et equals E; Es, Et are orthonormal on the vocabulary and orthogonal to each other.
• The second layer has only attention. The authors use linear attention with V fixed to the identity, writing the weights as a sum of per-vertex terms μv E(v)Es(v) transposed, with μv(0)=0. This μ is the index-matching logit.
• Remark 1: More general weights contain n by n cross terms. Permuting vertex labels collapses them to two parameters; the authors then set off-diagonal terms to 0. The term from the answer token to the reasoning token is left for the prediction stage.

Two losses and one-step expansion

• The paper recounts the experiment of Zhu et al. (2025): even a 12-layer transformer struggles to directly predict a shortest path of only 3 or 4 steps. This paper therefore does not require global planning, only that the next thought carries nodes reachable within c+1 steps.
• COCONUT-BFS encourages any node in Nc+1. COCONUT only raises the logit of the demonstration node pc+1. Both average over vertex permutations. The authors say practice corresponds to the latter, since each question usually has only one demonstration.
• Under gradient flow with symmetric initialization, all μv merge into a single μ(t). Theorem 4 in Appendix B gives μ̇(t)=α/(√nK)(d_⋆−F(μ(t))). d* is the restricted in-degree of the demonstration node within the current reachable set, K is the size of that set, and F is smooth, strictly increasing, with limits 0 and the maximum restricted in-degree at the two ends.
• Theorem 1: under COCONUT-BFS, μ grows at least logarithmically and the attention logit is unbounded. Under COCONUT, if d* is less than the maximum restricted in-degree, μ converges to a finite value; if equal, it tends to infinity at least logarithmically. With initialization at 0, whether it rises or falls also depends on the comparison between F(0) and d*. The authors state that when F(0) is less than d* and d* is less than the maximum in-degree, the limit is positive. Theorem 2 states that when μ is greater than 0, the next-step support is exactly the one-step reachable set, with coefficients β_v=λ_v1{v∈ N_c}+μ∑_{u∈ Nc}λ_u1{(u→ v)∈ E}. When μ=0, nodes outside the reachable set have the same logit, which the authors say is not effective exploration.

Answer prediction

• The final thought is a superposition of nodes reachable within C steps, with weights in (0,1]. An answer token is appended at the end. In the prediction stage, W is written as μR times the outer product of the answer embedding and the reasoning embedding; the authors say the two sets of parameters can be viewed as orthogonal subspaces, or as one generation head plus one prediction head.
• Lemma 2: the vertex logit equals μA times that node's weight in the final thought (only within the C-ball), plus μR times "whether it is a candidate." Only c* has both positive residual carrying and candidate boosting.
• The loss is cross-entropy on the reachable candidate. λ* is the minimum weight of the reachable candidate over the training set, and Δtrain is the maximum positive excess of reachable-but-not-that-candidate nodes relative to it. Theorem 3 states that the gradient flow direction converges to the maximum-margin direction, with the ratio of the two components equal to λ* plus Δtrain, and the norm tending to infinity.
• For unseen instances satisfying the same weight range and with maximum excess not exceeding Δtrain, the softmax probability of c* tends to 1 over time.

Experimental Setup

• Model: GPT-2-style decoder trained from scratch. Main experiment with two layers, dmodel=768, 8 heads. AdamW with β1=0.9, β2=0.95, weight decay 10−2, constant learning rate 1×10−4, global batch size 256.
• Data: ProsQA subset used by Zhu et al. (2025), ProsQA from Hao et al. (2024). Vertex indices randomly permuted for both training and testing. Averages in Table 3: training 14785 problems, |V|=22.8, |E|=36.5, solution length 3.5; validation 257 problems, 22.7, 36.3, 3.5; test 419 problems, 22.7, 36.0, 3.5.
• Curriculum: Stage c uses c consecutive thoughts before predicting the c-th node on the predicted path; if c exceeds the chain-of-thought length l, the answer is predicted after l consecutive thoughts and the answer token. Stage 1 is 150 epochs, each subsequent stage 25 epochs, totaling 350 epochs. Data from earlier stages is mixed in with probability 0.1.
• Metrics: Test set accuracy. In the generation phase, the logit difference between frontier and non-frontier edges on the second layer proxies μv; frontier refers to edges whose source node is in Nc excluding Nc−1. In the prediction phase, the answer token's first-layer attention to the last thought proxies μA, and the answer token's second-layer attention to reasoning tokens proxies μR.
• Review and Repetition: The paper does not use a separate review model and does not report repetitions of the main accuracy. A footnote states that under different settings and random seeds, candidate elevation may not pass through reasoning tokens.

Main Results

Section 5 states that 96.2% and 99.0% are test accuracies. The extracted text of Table 4 lacks percent signs; Appendix F.1 calls the same set of experiments accuracy.

Reported Location	Setting	Reported Value
Section 5	Two-layer main experiment, total 350 epochs	96.2%
Appendix E.2	COCONUT-BFS-exp, other training settings unchanged	99.0%
Table 4	L=2, otherwise per that table's defaults	98.8
Table 4	L=12	67.4
Table 4	dmodel=384	62.0
Table 4	Learning rate 2×10−4	58.1
Appendix F.1	L=12, first stage 400 epochs, learning rate 5×10−5	99.6%

• The authors report that the main experimental pipeline can predict reachable candidates on the test set. The BFS variant is a loss control; logit curves are in the next section.
• Section 5 and Table 4 give different values for the two-layer default setting; the paper does not explain whether they are the same run. In that table, H=8, dmodel=768, learning rate 1×10−4, and tied are the same as the L=2 row and are not listed again above.
• The authors state that each ablation changes only one hyperparameter, but hyperparameters interact, so a low single row cannot alone represent the result after joint tuning.

Logits in the Thought Generation Phase

• Measurement: Track the logit difference between frontier and other edges under Lcoco on the test set. Figure 3 uses background colors to distinguish stages; the text says Stage 1 is blue and Stage 2 is purple.
• Results: The authors state that the difference for c=1 rises in Stage 1 and saturates around 60 after ∼125 epochs. Stage 2 achieves positive values for c=2 with fewer epochs; c=3 and c=4 also show this pattern, although they are not explicitly trained to generate more than two consecutive thoughts. The paper does not give the endpoints of these later differences.
• Interpretation: The authors state this is consistent with Theorem 1. They call the reuse across lengths length generalization. Appendix E.2 changes the BFS loss to encourage only the current frontier, avoiding predicting the total root. This variant rises faster for c=1 and shows no sign of saturation after 150 epochs; for c=3 and c=4, the gap between the two losses narrows. The authors attribute this to staged data mixing, gradients propagating to earlier thoughts, and the growing difference between the theoretical loss (5) and the experimental variant (19). The endpoint values of Figure 5 are not given in the paper.

Answer Prediction Phase

• Measurement: After entering the prediction phase, track the attention proxies μA and μR. Figure 4's y-axis is Average Attention Logits, x-axis is epochs after entering answer prediction.
• Results: The authors state both rise quickly and stabilize after about 5 epochs. The text does not give the exact stabilized logits. Appendix E.3 gives three pathways: reasoning tokens copy the candidate in the first layer, and the answer token attends to it in the second layer; the answer token directly attends to the candidate in the first layer; consecutive thoughts copy the candidate in the first layer, and the answer token attends to the consecutive thoughts in the second layer. The authors state the three pathways are functionally equivalent, all elevating reachable candidates.
• Interpretation: The authors state this is consistent with Theorem 3, where the two strengths grow at comparable rates, giving c* the highest logit. In practice, logits plateau, while theory predicts norms tending to infinity. The authors suggest the possible reason is that in experiments the prediction phase still interacts with thought generation, whereas theory fixes the thought distribution. The authors leave finer analysis to future work.

Other Ablations and Analyses

• L=4 is 97.3, L=8 is 96.5; H=4 is 98.0, H=12 is the same as the L=2 row in Table 4 (Table 4).
• dmodel=1536 is 97.7; learning rate 5×10−5 is 62.1; tied and untied are both the same as the L=2 row (Table 4).
• With learning rate 5×10−5 and the first stage extended to 300 epochs, Appendix F.1 reports accuracy 97.0%.
• The authors state that in the four-layer model, the first layer's edge tokens attend to the corresponding source and target nodes (Figure 7).
• The authors state that from the second layer onward, hidden states have larger inner products with reachable, frontier, and optimal nodes. In Figure 8's text labels, Layer 1, Thought 1's Not Reachable, Reachable, Frontier, Optimal are 1.23, −1.54, 1.39, 1.54; Layer 2, Thought 1 are −0.43, 8.14, 15.11, 17.42 (Appendix F.2).
• Figure 9 tracks test accuracy during the answer prediction phase. The authors state it moves from near random guessing to stable high accuracy as the two signals are learned; the text does not give the curve endpoints.

Limitations and future directions noted by the authors

The paper has no separate Limitations section. The authors explicitly leave two points for future work.

• Plateau in the prediction stage: Section 5.2 states that in practice the logit enters a plateau, unlike the unbounded growth in the theory. After giving one possible reason, the authors say a finer analysis is left for future work.
• Hyperparameter interactions: Appendix F.1 states that a comprehensive study of hyperparameter interactions is beyond the scope of this paper and is left for future investigation.

Experimental coverage

• The experimental model is a GPT-2-style decoder trained from scratch. The main setting is two layers, dmodel=768, 8 heads, with the AdamW optimizer (Section 5). Table 4 additionally reports layer counts 2, 4, 8, 12, head counts 4, 8, 12, widths 384, 768, 1536, three learning rates, and whether weights are tied. The theoretical analysis is written in terms of gradient flow (Section 3, Section 4).
• The data is a ProsQA subset. Table 3 gives training 14785, validation 257, test 419, along with average vertex counts, edge counts, and solution lengths. Section 5 states that both training and testing randomly permute vertex indices.
• The curriculum budget is 150 epochs for Stage 1, 25 epochs for each subsequent stage, 350 epochs total, with a mixing probability of 0.1 for early-stage data (Section 5).
• Thought generation baselines include Lcoco and the frontier-node variant in Appendix E.2. Answer prediction reports two attention proxies, and Appendix E.3 gives three boosting pathways. Appendix F.2 checks the first-layer copy and subsequent-layer inner products for the 4-layer model.
• The paper does not report the number of repetitions for the 96.2% accuracy in Section 5. The theory analyzes the second layer after the first-layer copy is already established, using linear attention with V fixed to the identity; Remark 1 sets off-diagonal cross terms to 0 (Section 3).

The authors use gradient flow to explain where superposition in continuous chain of thought comes from. The task is directed graph reachability: decide which of two candidate endpoints is reachable from the root. Continuous thought does not project the hidden state into discrete tokens but feeds it directly back into the input. The analysis targets a simplified two-layer transformer, with training split into generating continuous thoughts and reading out the answer.

In the generation stage, each instance has only one shortest-path demonstration. Local search strength is written as the index-matching logit. Theorem 1 states: a loss that enumerates next-layer nodes makes it grow at least logarithmically; when only the demonstration node is supervised, if that node's restricted in-degree is less than the maximum in-degree, it converges to a finite value. With initialization at 0, whether it rises also depends on the comparison between F(0) and that in-degree. A positive bounded value makes the next step simultaneously carry already-reached nodes and expand one hop, rather than committing only to the highest-in-degree branch. The prediction-stage logit is the sum of residual carrying from the final thought and boosting of the two candidates. Theorem 3, under a margin condition, pushes the probability of the reachable candidate on unseen graphs toward 1.

A two-layer GPT-2-style decoder trained from scratch achieves 96.2% accuracy on the ProsQA test set. The authors state that the Stage 1 frontier-edge logit difference saturates around 60 after ∼125 epochs, with length reuse appearing at c=3 and c=4. The BFS experimental variant achieves 99.0% test accuracy, and its c=1 logit difference remains unsaturated after 150 epochs. In the prediction stage, the two proxy quantities stabilize after about 5 epochs, but in practice enter a plateau; the authors leave this discrepancy for future work. Table 4 also reports values for deeper, narrower, or different learning-rate settings, with L=12 at 67.4 and dmodel=384 at 62.0.