Emergence of Superposition: Unveiling the Training Dynamics of Chain of Continuous Thought
Authors: Hanlin Zhu
Shibo Hao
Zhiting Hu
Jiantao Jiao
Stuart Russell
Yuandong Tian
Venue: ICLR 2026 Poster (2026)
Keywords: chain of continuous thought
training dynamics
reasoning
superposition
Primary area: learning theory

Abstract: Previous work shows that the chain of continuous thought (continuous CoT) improves the reasoning capability of large language models (LLMs) by enabling implicit parallel thinking, and a subsequent work provided theoretical insight by showing that a two-layer transformer equipped with continuous CoT can efficiently solve directed graph reachability by maintaining a superposition of multiple reasoning traces in the continuous thought. However, it remains unclear how the superposition mechanism is naturally learned from gradient-based training methods. To fill this gap, we theoretically analyze the training dynamics of a simplified two-layer transformer on the directed graph reachability problem to unveil how the superposition mechanism emerges during training in two training stages -- (i) a *thought-generation* stage that autoregressively expands the continuous thought, and (ii) a *prediction* stage that converts the thought into the final answer. Our analysis reveals that during training using continuous thought, the index-matching logit, an important quantity which reflects the strength of the model's local search ability, will first increase and then remain bounded under mild assumptions. The bounded index-matching logit effectively balances exploration and exploitation during the reasoning process: the model will exploit local problem structures to identify plausible search traces, and assign comparable weights to multiple such traces to explore when it is uncertain about which solution is correct, which results in superposition. Our experimental results tracking the growth of logits further validate our theory.

Published as a conference paper at ICLR 2026
EMERGENCE OF SUPERPOSITION: UNVEILING THE TRAINING DYNAMICS OF CHAIN OF CONTINUOUS THOUGHT
Hanlin Zhu' Shibo Hao" Zhiting Hu Jiantao Jiao Stuart Russell Yuandong Tian UC Berkeley
UCSD 3 Nvidia 4 Meta AI fhanlinzhu, j jiantao, russell/eberkeeey. edu
(s5hao, zhh01 Jouesd. edu, yuandong@meta.cc com
ABSTRACT
Previous work shows that the chain continuous thought (continuous CoT) im- proves the reasoning capability of large language models (LLMss) by enabling implicit parallel thinking, and subsequent work provided theoretical insight by showing that a two-layer transformer equipped with continuous CoT can effi- ciently solve directed graph reachability by maintaining a superposition of mul- tiple reasoning traces in the continuous thought. However, remains unclear how the superposition mechanism is naturally learned from gradient- based train- ing methods. To fill this gap, we theoretically analyze the training dynamics of a simplified two-layer transformer on the directed graph reachability problem to unveil how the superposition mechanism emerges during training in two training stages (i) thought- generation stage that autoregressively expands the contin- uous thought, and (ii) a prediction stage that converts the thought into the final answer. Our analysis reveals that during training using continuous thought, the index-ma logit, an important quantity which reflects the strength of the model's local search ability, will first increase and then remain bounded under mild assumptions. The bounded index-matching logit effectively balances explo- ration and exploitation during the reasoning process: the model will exploit lo- cal problem structures to identify plausible search traces, and assign comparable weights to multiple such traces explore when it uncertain about which solu- tion is correct, which results in superposition. Our experimental results tracking the growth of logits further validate our theory.
INTRODUCTIONS
Large Language Models (LLMs) show great reasoning capabilities in many complex tasks, espe- cially when equipped with chain-of-thought (CoT) (Wei et al., [022]. However, due the large inference cost of long CoT for complex tasks, many recent works seek alternative test-time scal- ing methods more efficiently improve LLMs' reasoning ability (Goyal al. 2023 Wang al.
2023b PTau et al. 2024, Su et al. (025).
One promising method is to use thain-of-continuuu-ttooog
(COCONUT, continuous CoT) Hao
et al. 2024), where the reasoning trace of an LLM is kept in a continuous latent space instead of pro jected back to the discrete token space. Continuous CoT exhibits both theoretical advantages (Zuu et al. 2025) and empirical performance gains (Hao et al. 2024) in many tasks. To more efficiently and reliably scale up continuous CoT' to solve more challenging tasks, requires a deeper under- standing of its internal mechanism.
Previous work (Zhu et al.) 2025) theoretically shows that one of the most important advantages of continuous CoT is that it enable the model to reason by superposition: when the model encounters multiple plausible search traces and uncertain about which one is correct, it can keep all plausible traces in parallel since the CoT is in continuous space instead of discrete tokens. In particular, Zhu et al. (2025) abstracted a family reasoning tasks as a directed graph reachability problem, i.e., whether there exists path from a given start node to given destination node, and

---

Published as a conference paper at ICLR 2026
showed that a two-layer transformer with ((n) (where n is the number of vertices in the graph) continuous thought decoding steps can efficiently solve the task by providing a construction of the parameters. Therefore, natural next question is:
Do gradient-bate based methods naturally lead to such a construction, and can we theoretically prove it?
This paper answers the above question affirmatively by theoretically analyzing the training dynamics of (simplified) two-layer transformer on the graph reachability problem in two training stages: (i) thought generation stage where the model autoregressively generates a chain continuous thoughts and (ii) prediction stage where the model predicts the final answer using the generated thought.
Importantly, our analysis of the thought generation training stage reveals why the superposition can emerge even if the training data only presents one demonstration for each training sample. Our theoretical analysis as well as experimental results show that when training with continuous thought (i.e., the COCONUT training method in Section [3] and Section (5). the index-matching logit, an important quantity that measures the strength of model's local search capability, will remain bounded under mild conditions, which is in contrast many previous analysis on transformer logit without continuous thought (e.g.. Tian et al.] (2023a); Nichani et al. (024a); Nguyen & Nguyen- Tang (025)) where the logit will grow logarithmically and thus unbounded. A bounded index- matching logit can balance exploration and exploitation: if the logit is too small, the model cannot even perform local search, resulting in nearly random guess in the next step; if the logit is too large, the model might over-confidently commit to one of the plausible search traces merely depending on local features (e.g., the indegree of a node in our case) even if is uncertain about the solution, and thus early discard the correct path. A bounded index-matching logit encourages the model to exploit the local structure while explore multiple plausible solutions by assigning comparable weights to them, which naturally results in superposition. This answers the question of al. (2025) why superposition can emerge during training.
1.1 RELATED WORKS
Reasoning with chain of thought. Chain of thoughttt (CoT) (Wei et al. 2022) is simple yet effective test time scaling method to enhance LLM's 's reasoning capability. can either be prompt-
based only (Khot et al. 2022; Zhou et al. 2022) or be included in the training sample to create high-quality training data Yue al. 2023; Yu et al. 2023 Wang et al.] 2023a Shao et al. 2024). Beyond empirical study, many theoretical works also explore the advantages of the CoT method. For example. Liu et al. 2022); Feng et al. (2023); MMrrill & Sabharwal (2023); Li et al. (2024b) shows that CoT can improve the expressivity of transformers, Zhu et al. (2024) studies the importance of CoT for two-hop reasoning via training dynamics. Wen et al. (2024); Kim & Suzuki| (2024) studies how CoT in the training data improve the sample efficiency of transtormers. Instead of the text- based CoT, this paper studies continuous COT where the "thinking tokens" lie in a latent continuous space and do not need to be converted to discrete tokens.
Latent space reasoning. recent line of work studies latent space reasoning, novel paradigm
beyond text-based CoT (Goyal et al. 2023 Wang et al. 2023b Pfau et al. 2024 Su et al. 2025 Hao
al. 2024). For example, Goyal et al. (2023) proposed to use pause tokens, which are learnable tokens that are inserted into the original text to increase the computation space. Later, London &
Kanade (025) theoretically shows that the pause token can strictly increase the expressivity of the
transformer. Similarly, |fau et al. (2024) studies filler tokens, which also increase the computation space of LLMs. WWang et al. (2023b) proposed to use planning tokens the beginning of the response generation to improve the reasoning capability. [Su et al. (2025) proposed to use abstract tokens in a latent space to enhance the reasoning performance while reducing the inference cost, The most related work is Hao et al. (2024), which proposes
continuous CoT for reasoning. A follow-up work Zhu et al. (2025) theoretically shows the advantage of continuous CoT via expressivity. Our work takes a further step by analyzing the training dynamics of continuous COT.
Training dynamics of transformers. There is a rich line of literature studying the optimization of transformer- -based models (Jelassi et al. 2022 Bietti et al. 2023 Mahankal et al. 20233 Fu et al,
2023 Tian et al. 2023a Zhhng al. 2024 al. 2024a Huang al. 2024 Guo
2024
2

---

Published as a conference paper at ICLR 2026
More recent works also focus on understanding reasoning abilities or patterns of transformers via training dynamics. including induction heads (Nichani et al. (024a), the reversal curse (Zhu et al.
2024
Ma et al.] 2 CoT (Wen et al. 2024: & Suzuki 2224] Huang et al. 2025b). factual Nichani et al. 2024b). two hop reasoning (Guo et al. [025 Lin et al. (025). out of context
recall
reasoning (Huang al.] 2025a), etc. Along the line, our paper aims understand the internal
mechanism of continuous CoT and why superposition emerges via the analysis of training dynamics.
2 PROBLEM FORMULATION
Basic notations. We use [N] to denote the set {1, (1. for any integer N > 0 and use [i: yl to denote {i, + ......j - 1.3} for integers S j. For any finite set X. we use | to denote its cardinality and use Unif(Xt) to denote the uniform distribution over X. We use R to denote the set of real numbers and denote x4 = max(x, 0) for x E R. For any vector x = (Ii,..., (.a) E Rd, the softmax function SoftMax(-) Rdd - Rd is defined as SoftMax(#); exp(x;) exp(æ;)). (x)p(x))) Let Ia E Rddd denote the identity matrix. Let Voc [M] denote a vocabulary of size M for a fixed integer M > O. For each token Voc, there is associated embedding E(v) E Rd. Let U [E(1), E(, E(M)) E [RdxM be the token embedding matrix.
Graph and permutation. For a directed graph g = (V,E) with the vertex set | = and the edge set & = {(s: - t:]] we fix root node E V and two candidate destination nodes C12 C2 E V such that exactly one node, denoted c., is reachable from r and denote the other с that is unreachable. For radius
N, define the e-ball as (z) = { E V v is reachable from r within steps). For a subset V C V, we define the restricted in-degree as degg.v.(u) = (u E V : (u - v) E EHI- We also fix a shortest path from to C, as = (Po. ..o.....,22 ...,..))) Pc) with Po = I, PC = C.. (pe-1 + po) E E for any c C [C]. For any permutation ' over V. we define (9) = (V. (E)), where x(E)= = {[T((s) - a()) (s - t) E), and define π (p) =
)). We also denote the set of all permutations over V as Sy.
Chain of continuous thought. Let TFo(-) : (Rd)) + Rdd be a transformer which receives an input embedding sequence h = hig A (h1, le2...., h,)) E kdxt and outputs (F) E Rd. For convenience, we assume weight tying. A traditional decoder using a discrete CoT will sample the next token U/+1 ~ SoftMax U (F))))). Then the embedding of 0141 will be appended to the end of the input, i.e., h4+1 = E(u+41). For continuous CoT, one directly appends the output of the transformer to the end of the input sequence without converting it to a token, i.e., setting h,+1 =
TF& (h)). Assume the prompt x = [xI..... Lto] E Voc* and its corresponding embedding sequence is hpp = [hi...., ] = [E(x1), E(za,) E(xta)) For notation convenience,
use [tc]] hog+o
to denote the continuous thought generated at decoding step c, where [t.] TFa TFs(h((+t--)) In particular, [to] ht. After C decoding steps, one can append a special token <A> at the end of the sequence to trigger the transformer to switch the mode and generate the final answer.
Specifically, one can set by = E(<A>) where T = to + C' + and generate the final answer
TF TFo,c.u(h)..) arg max U U TFoo(r))
Graph reachability and prompt format. In this paper, we mainly focus on the directed graph reachability problem as studied Zhu et al. (2025). where we are given a graph g = =(,, E), two candidate destination nodes C1 and C2, and root node The task is to identify which of the two nodes can reached by (denoted as c.). The prompt structure is illustrated Figure followwngg Zhu et al.) (2025). The prompt consists of (1) BOS (beginning of sentence) token 
; (2) the graph description part, which contains m edges where each edge is represented by a source node Si, target node t,, and special edge token <e>; (3) the task description part that contains special question token <Q>, two candidate destination nodes 3j and C22 a special reasoning token <R> and root node r. See [or the full list of token notations. Note that to 3m + 6 is the prompt length, and let hapo = (h1, h2,... ..,,)) be the input embedding sequence. Following Zhu et al. 2005)). we use Idx(v) to denote the position of a token in the input sequence (e.g., Idx (dx(
)=
1, (x)) 3i - 1, Idx(c1) 3m +,, Idx( <R>) 3m + 5), use Idx( dx(<<>>,) i) 3i + to denote the position of the i-th <e> token, and use Idx( [t, = to + i to denote the position of the continuous thought at step i. See Table 2 for the complete list of position indices.
3

---

Published as a conference paper at ICLR 2026
directed edge
special edge token
question token
reasoning token
<
<e>
F2
<E>
Sm
<e>
<Q>
G1
<R>
1
start token
source node
target node
candidate nodes
root
Figure 1: Prompt format of the graph reachability problem (Figure of Zhu al. 2025
Zhu et al. 2025 provided a construction of transformer parameters such that TF T C. (I.e., the transformer can predict the reachable candidate node using continuous CoT) for any graph and root-candidate node tuples, where (ito]] corresponds to the prompt of the graph and task descriptions. However, they did not theoretically study whether the constructed solution can be naturally learned via gradient-based methods. In the following sections, we theoretically show that the solution can be learned via gradient flow in both the thought generation stage (Section (3) and the prediction stage (Section 4. We also provide empirical results showing that the training dynamics in our theoretical analysis align well with the experiments (Section 5
2.1 INDEX-MATCHING LOGITS AND LOCAL SEARCH CAPABILITY
Before we delve into technical details in the following sections, we first provide an intuitive explanation of the dynamics of the main mechanism.
One step expansion via local search
[4]
d
[t.] =
Uhp
I
(+)
PEN
(a) Left: The continuous thought at step [t,] encodes embeddings of nodes that are reachable from the node within step. Middle: One-step expansion via local search where the strength is quantified by modex--atching logit p. Right: After one-step expansion, continuous thought at step 2 [t> encodes nodes reachable within two steps.
small attention v3 $ &.
embedding space
large u,
Superposition of all nodes that can be reached within steps
buffer 1
Vs
V1
j B(P))
buffer2
1/4
V2
content
[tc]
[t]]
|
One-step expansion of N.
Illustration how one-step expansion implemented (adapted from Figure 3 of Zhu et (2025)). the first layer of the transformer, each special edge token <e> copics its corresponding source and target nodes to its buffer spaces. In the second layer, as illustrated in the figure, the current thought [t.] pays large attention to edge if its source node has been explored, and adds its target node the superposition, where the strength of the added node controlled by the index- - V3 corresponds to edges Figure 2a
matching logit μr. The two edges Us + u4
Figure 2: Pictorial illustration of the superposition mechanism and the index index-matching matching logit μ.
4

---

Published as a conference paper at ICLR 2026
Global planning local search. In the context of graph reachability, the global planning refers to a model's 's capability to analyze the structure of the whole graph and then determine a path from the root node to the destination node. In contrast, local search focuses only on which nodes are reachable in one step from the current node, which is much easier to learn than global planning. When using discrete CoT, the model can choose only one path at a time. Therefore, the model needs global planning to select the correct path or to backtrack from the wrong one. When using continuous CoT, the model can keep multiple plausible paths simultaneously. Therefore, the model can rely solely on local search to perform parallel BFS, solving the task with only simple skills.
Index-matching logits. We use the mudex-matchinn logit μr to quantify the strength of the model's local search capability, which is illustrated in Figure 2 and will be formally defined in (3) in Sec- tion 3 In Theorem [ we will prove that under mild conditions, the index-matching logit l will first increase and then remain bounded. Note that positive, bounded logit μ effectively balances exploration and exploitation in node expansion: if μ is too small, each edge will receive similar attention in Figure 2b and thus the model even lacks the local search capability graph structure; if f is too large, the model will put too much weights on nodes with large in-degree (e.g., in Figure 2a 2/y weights 2 and other frontier nodes such 14 weights μ. where the difference in weights will be significant under large μ and commonly used softmax attention) and thus lacks exploration of different plausible paths.
exploit the local
3 ANALYSIS OF THE THOUGHT GENERATION STAGE
In this section, we analyze the training dynamics of the thought generation stage. We consider any graph 9 (V,E), E), a root node x E V. two candidate destination nodes C1 C2 E V, where (c1,c2) E } with c, reachable from and с unreachable. We are also given (discrete) CoT demonstration, which is a shortest path p = (Po.... .P)) from r to C, where po = r, PC = C..
We use curriculum learning following Hao et al. (2024); Zhu et al. (025 where at stage (c + 1) for any 0 < c < C, upon receiving the prompt embeddings [cco]- the model will first generate c continuous thoughts [t.]..... [te] autoregressively without supervision (i.e., no loss calculated on the first e continuous thoughts at stage e-+ 1), and then be trained to generate the next continuous thought [te+1] = (ttt)o)))))) Since the learning procedure at each stage is similar, focus
below on a fixed c.
Zhu et al. 2025) constructs solution for a two-layer transformer, where the first layer mainly performs copy (e.g., the i-th special edge token <e> will copy the information of its corresponding source node s; and target node t.). Since the copy mechanism has been widely studied (Nguyen & Nguyen- Tang 2025), as well as its formation via training dynamics (Nichani et al., 2024a), we mainly focus on the dynamics after the copy mechanism has been established. Thus, we analyze the dynamics of the second layer of the transformer.
In particular, let the hidden states of each special edge token <e> and the current thought [te] after the first transformer layer be
hidx(ce>,s) = (s:) + Et+)t) E Rd,
hidx(t=.1)
Σ
E(v) ER", (I)
(EAT(E) VWN((r)
where (v) E Rd and E. (v) Rd map token E Voc different subspaces of Rd. For example, as in the construction of Zhu et al. (2025), we can set d = 3.M, and E.(.), (-) and E(.) each corresponds .M different non-zero zero entries. This is also similar to previous work Chen et al.
Nguyen & Nguyen-Tang (2025) where E.) and E. can be viewed as previous token heads. We make the following assumptions on the embedding E, (.). E.) and E(-):
2025
Assumption (Orthonormal embeddings). Assume E, (-)
E(-). For any u,u E Voc, Eu( (v) = (u) (v) = 1fu = t} and E.()) (e(()) =
means after the first layer, each special edge token <e> will copy the embeddings of its corre- sponding source and target nodes 8; and t; to the same position in different subspaces. Also, we assume by induction that after training stages 1,
the current thought generated by the well- trained model Maxx(t.))) is a normalized superposition of token embeddings of all nodes reachable from within e steps. Below, we study the training dynamics of the current stage (i.e., stage + 1).
5

---

Published as a conference paper at ICLR 2026
The forward path and reparameterization. We consider the setting where the second layer is attention-only. The forward pass can be formulated as
(h: (h.):)) _ Wh:) his
= U UT(IIa((tLL))) + d(MMax(iLE)) (e) E RRM,
where V, W E [Rdxd are attention parameters and o()) IR + R is an activation function that determines the range of attention scores, and & = (Eo)ueVoc E RR is the output logit vector for each token in the vocabulary. Similar to the analysis in Nguyen & Nguyen- Aguyen-Tang (2025), we adopt the linear attention o (h Wh;) = h h., fix = I and use the index-matching reparameterization
W Σ Mr (()) (L(() for 0. VEV
(3)
Remark 1. Note that a more general form of the atlention weight matrix can be
W = plgs E(<A>)E(<R>)
Σ Mer, E(r)E,(u')*
(4)
UN'VV
The first term only takes effect in the prediction stage (Section 4. we can set pLez> = 0 for now. The second term involves n x n cross terms. The symmetry of the vertices, which can be enforced by permuting vertex labels during training, makes the
parameters {Me Jo,' effectively two parameters {fr, J22} where pa,s = fn and pr,,'' = P2 for u + v'. Moreover, if we focus on the relative value between μr and p2g, we can further simplify the attention weight matrix by assuming J12 0.
For notation simplicity, we h; to denote ce>.)) use httel to denote Mtx((t.,1) and use No. No+1 to denote No (z). (x). respectively when the graph & and root node is clear from the context. We also denote du = degg.N. (u) which is the indegree of u with source nodes restricted in No. Finally, we denote K [Wol and X VK*
Loss functions. An ideal model should be able directly output the shortest path from the start node r to the desired candidate destination node C.. i.e., the prediction of the (c+ 1)-th continuous thought [ter1] exactly corresponds to the (c + 1)-th step of the shortest path Po+1- However, experiments in Zhu et al. (2025) show that even for a 12-layer transformer, it is hard to predict the shortest path even if the length of the shortest path is only or 4. Therefore, we take step back and pursue more practical goal we expect the model to at least be able to generate an arbitrary path starting from the start node r, which only requires local search ability that much easier than the global planning ability. In the context of continuous thought, we expect the model to include information of all vertices that are reachable from r within (c + 1) steps in the generated thought [teii]. We consider the following two loss functions:
9. = - logg EveNo exp(Ex)
COCONUT-BES:
(5)
Evey exp(sx)
exp (xp(ee))))
COCONUT:
G.r.p log
(6)
E exp(Ex)
with permutatioonaaveraeed dataset losses
CBFS 1 and 0 Unittt()))
froce ( (r0))))))))))
Note that, intuitively, the cermutation-averagee loss will lead similar behavior across different parameters. The first loss CBFS explicitly encourages the model to predict any nodes in Nc+1- However, in practice, haustively; instead, usually present only one demonstration for each task instance during training (in our setting, only one path per instance (G, x, C1, C2)), which corresponds to the second loss feoco and aligns with the practical setting where chain of thought data can be used for supervision.
costly and even impossible to search over the entire solution space ex-
Zhu et al. (2025) observed in experiments that superposition emerges even without explicit guid- ance during training, i.e., using the loss _onen_ In this paper, we investigate the emergence of superposition by analyzing its training dynamics. The following lemma gives the gradient of the index-matching strength parameter (t. using gradient flow under the loss function fooco

---

Published as a conference paper at ICLR 2026
Lemma 1 (Gradient of pyy under L informal version of Theorem 4 in Appendix B Under permutation-avervvs training from symmetric initialization and gradient flow jy -a Me. Lenen we have pt(t) = pu(t) for all u and times t, and the gradient of pu is
EuENeH dy AA (udd))
1000(1a ((t) nyK (dee F(u(t))), F(1)= EuENoti ed(((EN))tpd)) (n (N++1l)
Moreover F is smooth, strictly increasing, with F(-00) = 0. F(+00) = max.ey d, and 0 < F(H) maxxev for all finite μ.
The proof is deferred to Appendix B Note that as long as dp... * maxTev de, put) will converge
14 < In contrast, under COCONUT-BFS with loss CBFS, ju()) will diverge infinity. We
formalize the comparison into the following theorem and defer the proof to Appendix B
Theorem (Bounded vs. divergent attention logits under COCONUT vs. COCONUT-BFS; informal
version of Theorem & Lemma 5 in Appendix
B). Let d, := dpets and dinax := max, d'
(i) Under COCONUT-BFS ( p(t) grows at least logarithmically in t, leading to unbounded attention logits.
(ii) Under COCONUT 66 if d2 < dmax then p(l) - 11* < Dc, so all attention logits remain uniformly bounded. If d, dmax then u(t) - 00 at least in a logarithmic rate.
<x,,
Emergence of Superposition via Bounded Attention Logits. By Theorem as long as F(O) <
d dmax, we have p(l) n²
Compared to many previous work (Tian al. 2023a
Nichani et al. 2024a; Nguyen & Nguyen-TTane (025) that analyze the dynamics of attentron logits
in "discrete" settings where the attention logits diverge infinity, the COCONUT training method in continuous setting usually result in bounded attention logits. The bounded attention logits lead to a more smooth probability distribution over next tokens, which is beneficial especially under uncertainty: when the model is uncertain about the next step, a more smooth probability distribution under continuous CoT mechanism results in a superposition of different plausible next steps, which implements effective exploration; distribution and thus the model will over-confidently commit to a plausible branch and likely to discard the ground-trun truth branch even when the evidence is weak.
the contrary, an unbounded logit will result in one-hot-like
Finally, we show that with a positive value of μ, the continuous thought [teri] implements a one- step expansion from No to No+1 for any graph G and root node
Theorem 2 (One-step frontier expansion; informal version of Theorem 5 in Appendix By. For any graph G and root node r, if the current thought is any positive superposition on No (r), iee. [tc] = EvEN, da E(u) with Xa > 0, then the next thought [te+i] satisfies that its token- projected output U [tei2] supported on the one-step expansion No-1 and has strictly positive mass on every node in Not1 if >> 0. In particular,
UT [to++]
Σ
By Cv
VEN
with
B. A. 1.
+ 1 Σ Au ((( 2
WEN.
carryover
one-hop expansion
The proof is deferred to Appendix B Note that at initialization where μ = O, we have 3, = 0 for w Nc+1 We. This means every node outside N₂ has the same attention logits and thus the same next token probability. However, such an exploration is not an effective exploration since it blindly puts the same weight on almost every node in the graph without exploiting the graph structure. Therefore, appropriate μ³" > 0 effectively balances the exploration and exploitation: (1) it has positive value so the model can exploit the graph structure and can distinguish nodes within the one- step expansion set; (2) it has bounded value it will not overconfidently commit to a plausible branch while discarding other branches merely relying on local structure (such as the indegree of the node) without global planning.
7

---

Published as a conference paper at ICLR 2026
4 ANALYSIS OF THE PREDICTION STAGE
In this section, we study how the transformer learns to make the correct prediction among {ci, c2] by utilizing the generated continuous thought. Note that according to Section 3 the model is able to generate [tc]] = EeENo A(E() with X E (0, 1J, superposition of all reachable nodes within C' steps, via balanced exploration and exploitation. We denote À = {>u hee. At the final stage, one appends special answer token <A> at the end of the continuous CoT,, i.e., hT = hears
and make the final prediction TFo,c.u(h(.)) == arg maXveVoc ('o
The forward path and reparameterriztiion. Similar to ( we formulate the forward pass in the prediction stage as
(h:; (h:):) = Vooh
Vo(h (h;) h;,
U UHesssaaaaa + Φ ((hidd(sa)); (hex>))) RMM,
where
(()((<B>) E(<R>) E(ci) + E(c2), h(x((a>) = h hta + E(<A>).
Note that after the first transformer layer; the hidden state of <R> contains information of two candi- date nodes C1 and 22 and the hidden state of <A> contains the representation of the last thought [tc]] both due to the copy mechanism in the first layer. Again, we adopt the linear attention oh Wh;) = h Wh;, fix V = I and use the reparameterization
w HlBo E(<A>)E(<R>)""
(8)
The scalar peRs> denotes the atlention logit strength from 10 <R>, The scalar pleaz represents the signal strength of the residual stream from the first layer. Also, note that the repa- rameterization W the prediction stage has a different form from (3 in the thought generation stage. One can either view both B and (8) as special cases of more general version 44) in or- thogonal subspaces, or view them as two different attention heads (a thought generation head and prediction head).
Remark 2.
The loss function. In the prediction stage, the goal of the model to predict the reachable candi- date node C., and thus the loss function can be written
Grr.noA = -log exp(Ec.) EVEY exp(E.))
pred
Ceeed E(g.r.en.o
l ppeed (ere
(9)
where the D = (( 200,, c c2 A {))} denote the training set. The following lemma provides closed-form logits for each vertex where the proof is deferred to Appendix C
Lemma 2 (Closed-form logits; informal version of Lemma in Appendix
The logit of each
vertex u E V has the form
Eo = Meao No 1w E No) + PeRao 1 (0 E (2).
residual carryover
candidute lift
According Lemma 2 only the candidate node c. has both positive residual carryover and can- didate lift, and an appropriate relative growth rate of Megs and (leas ensures that c, has the largest logit. We formalize the result in the following theorem with proof in Appendix
Theorem 3 (Prediction of the reachable candidate node; informal version of Theorem in Ap- pendix C Denote μA = peas and un = pAas Let (I<as)()) (M (t)) follow gradient flow on loss
defined 19. Suppose
1 E (0,1),
Auaan = max, max
VEVO (299 - x(2) E [0.1].
Then we have
(pa(t). H(((()
+ u², (((((t)) + 00,
8

---

Published as a conference paper at ICLR 2026
with / X, + Atrain and uR u'y 0. Consequently, for any unseen instance (g. r, C1. C2. 2) satisfying xu E (0, 1] on No and otherwise, and max, Xo dc S Atrain holds that:
exp (Ec, (PA(t)
Pe. Pc.(t) (t)
1.
En exp (R)
EXPERIMENTS
In this section, we present experimental results validating the theoretical analysis. We first describe the setup and overall results, then analyze training dynamics in the thought generation and answer prediction stages.
Model. We adopt a GPT-2 style decoder with two transformer layers (dmode)==66 (headd= The model is trained from scratch with AdamW =0.9..9, 82=0.95, weight decay 10 10-2), constant learning rate of 10-4, and a global batch size of 256.
Dataset. We follow the dataset from Zhu et al. (025). which is a subset of ProsQA (Hao et al, 2024]. Different from Zhuu et al. (2025). randomly permute the vertex indices in both training and testing to avoid prediction blas and validate the symmetry assumption in (4). Dataset statistics summarized in Table
multi-stage training strategy with supervision from chain-of- -f-hugghttthhttttth demonstrations. At stage c, the model learns to use c con- tinuous thoughts before predicting the e-th node on the reasoning path (thought-generatic)) stage). If e > (the CoT length), the model predicts the final answer after / continuous thoughts and the <A> token (prediction stage). We train for 150 epochs at Stage and 25 epochs for each subsequent stage, totaling 350 epochs. At each stage, data from carlier stages is mixed in with probability 0.1, which prevents the model from forgetting abilities learned from previous stages. The final accuracy of this model on the test set is 96.2%.
Training. Following Hao et al. 2024) Zhu al. (2025).
5.1 THOUGHT GENERATION
To examine the training dynamics of Me under L track the second-layer attention logits. When generating the e-th continuous thought, μ. corresponds to the logit on an edge token <e> whose source lies in No- In practice, _ooc encourages the model to predict the current search frontier rather than revisiting explored nodes, so most attention concentrates on frontier edges, i.e., edges with sources in Na N For theoretical simplicity, however, the model does assign non-zero attention logits to other edges. Therefore, we report the logit difference between frontier and non non-frontier edges on the test set, which more faithfully reflects the effective value of pt.
assume #22 0 in (4). In practice,
Average Attention ifference
60
40
20
Ce2
-20
25
50
75
100
125
150
175
200
225
250
275
300
325
350
Training Epoch
Figure 3: attention logits difference between frontier edges others under -0000 as proxy for Hay. The background colors indicate different training stages.
Figure 3 shows the results. Stage 1 (blue background), the model gradually learns to attend to frontier edges when predicting the first continuous thought (c 1). The logit difference increases

---

Published as a conference paper at ICLR 2026
steadily and saturates around 60 after ~25 epochs. This matches the theoretical prediction in under _enen, fA first grows and then stabilizes at bounded value.
Theorem I
When switching to Stage 2 (purple background), the model requires far fewer epochs to establish a positive μ for с = 2. Moreover, this pattern generalizes to c and c = 4, even though the model was never explicitly trained generate more than two continuous thoughts. This "length generalization" " indicates that once superposition emerges in earlier stages, later stages can quickly reuse it to expand the frontier further.
We also trained with a variant of LBPS Compared to _ the attention logit difference when 1 did not saturate but kept increasing much higher values, consistent with the analysis in Detailed experiments and plots are provided in Appendix E.2
Theorem I
5.2 ANSWER PREDICTION
We next analyze how the model predicts the final answer. According to Lemma 2 the prediction relies on two signals. The first is the residual carryover, which brings the explored nodes in the last thought [tc] into the answer token with strength μ... Concretely, this corresponds to the first-layer
attention from <A> to [t=], which copies the superposition of reachable nodes. The second is the candidate lift, which raises the logits of the two candidate nodes with strength μn- Since <R> copies for PR
the candidate nodes in the first layer, the second-layer attention from <A> to <R> serves proxy
Figure 4 shows the dynamics of these two proxies. Once training enters the prediction stage, both AA and uR increase rapidly and stabilize after roughly epochs. This observation is consistent with Theorem 3] which states that μld and μn grow at comparable rates, ensuring that the reachable candidate c* attains the highest logit. In contrast to the unbounded growth predicted in theory, we observe the logits plateau in practice. A possible reason is that, in practice, prediction-stage training also interacts with thought generation, whereas the theory assumes fixed thought distributions to focus the relationship between μn and yAa- We leave more detailed analysis to future work.
Awerage Atcention Logits
40
30
20
10
fresidual carryever)
in (candidate lift)
10
Training Epoch Answer Prediction
Figure 4: Training dynamics of the proxies for μa (residual carryover) and uu (candidate lift).
6 CONCLUSIONS
In this paper, we study the emergence of superposition when training with continuous COT. In par- ticular, we theoretically analyze the training dynamics of a simplified two-layer transformer on the directed graph reachability problem. Our analysis shows that under mild assumptions, the index- matching logit, an important quantity showing the strength of the model's local search ability, re- mains bounded during training. bounded index-matching logit effectively balances exploration and exploitation during the reasoning process and thus enables implicit parallel thinking, which nat- urally results in superposition. Our experimental results, which track the growth of logits, further validate our theory. We expect our theoretical analysis to bring new insights into a deeper under- standing of the mechanism of continuous CoT and ultimately scaling up this promising paradigm more efficiently and reliably.
We observe that under different experimental settings and random seeds, the candidate lift effect is not always mediated by the <R> token; alternative attention routes presented in the Appendix F.3
10

---

Published as a conference paper at ICLR 2026
ACKNOWLEDGMENTS
This work was partially supported by gift from Open Philanthropy to the Center for Human- Compatible AI (CHAI) at UC Berkeley, Berkeley Existential Risk Initiative (BERI) and by NSF Grants IIS- 1901252 and CCF-221 1209. SH grateful for the support provided by the Bloomberg Data Science Ph.D. Fellowship.
REFERENCES
Alberto Bietti, Vivien Cabannes, Diane Bouchacourt, Herve Jegou, and Leon Bottou, Birth of transformer: memory viewpoint. Advances Neural Information Processing Systems, 36: 1560-1588, 2023.
Lei Chen, Joan Bruna, and Alberto Bietti. Distributional associations in-context reasoning: A study of feed-forward and attention layers. In The Thirteenth International Conference on Learn- ing Representations, 2025.
Guhao Feng, Bohang Zhang. Yuntian Gu, Haotian Ye, Di He, and Liwei Wang. Towards revealing the mystery behind chain of thought: theoretical perspective. Advances in Neural Information Processing Systems, 36:70757-70798, 2023.
Hengyu Fu, Tianyu Guo, Yu Bai, and Song Mei. What can a single attention layer learn? a study through the random features lens. Advances in Neural Information Processing Systems, 36:1 1912-1911,, 2023.
Sachin Goyal, Ziwei Ji, Ankit Singh Rawat, Aditya Krishna Menon, Sanjiv Kumar, and Vaishnavh Nagarajan. Think before you speak: Training language models with pause tokens. arXiv preprint arXiv:2310.02226, 2023.
Tianyu Guo, Druv Pai, Yu Bai, Jiantao Jiao, Michael Jordan, and Song Mei. Active-dormant attention heads: Mechanistically demystifying extreme-token phenomena in llms. arXiv preprint arXiv:2410.1388 13835, 2024.
Tianyu Guo, Hanlin Zhu, Ruiqi Zhang, Jiantao Jiao, Song Mei, Michael I Jordan, and Stuart Russell. How do lims perform two-hop reasoning in context? arXiv preprint 133. 2025.
Shibo Hao, Sainbayar Sukhbaatar, DiJia Su, Xian Li, Zhiting Hu. Jason Weston, and Yuandong Tian. Training large language models to reason in continuous latent space. arXiv preprint arXiv:2412.06770 2024.
Yixiao Huang, Hanlin Zhu, Tianyu Guo, Jiantao Jiao, Somayeh Sojoudi, Michael I Jordan, Stuart Russell, and Song Mei. Generalization or hallucination? understanding oft-of-context reasoning in transformers. arXiv preprint 2rXiv:2506.1088 2025a.
Yu Huang, Yuan Cheng, and Yingbin Liang. In-context convergence of transformers. In Proceedings of the 41st International Conference on Machine Learning, pp. 19660-19722, 2024.
Yu Huang, Zixin Wen, Aarti Singh, Yuejie Chi, and Yuxin Chen. Transformers provably learn chain-of-thought reasoning with length generalization. arXiv preprint arXiv:2511.07377, 2025b.
Samy Jelassi, Michael Sander, and Yuanzhi Li. Vision transformers provably learn spatial structure. Advances in Neural Information Processing Systems, 35:37822-378833 2022.
Tushar Khot, Harsh Trivedi, Matthew Finlayson, Yao Fu, Kyle Richardson, Peter Clark, and Ashish Sabharwal. Decomposed prompting: A modular approach for solving complex tasks. arXiv preprint arXiv:2210.02406, 2022.
Juno Kim and Taiji Suzuki. Transformers provably solve parity efficiently with chain of thought. arXiv preprint arXiv:2410.086. 2024.
Yingcong Li, Yixiao Huang, Muhammed E Ildiz, Ankit Singh Rawat, and Samet Oymak. Mechanics of next token prediction with self-attention. In International Conference on Artificial Intelligence and Statistics, pp. 685-693. PMLR, 2024a.

---

Published as a conference paper at ICLR 2026
Zhiyuan Li, Hong Liu, Denny Zhou, and Tengyu Ma. Chain of thought empowers transformers to solve inherently serial problems. arXiv preprint 11X::22422.12875,, 1, 2024b.
Pengxiao Lin, Zheng-An An Chen, and Zhi-Qin John Xu. Identity bridge: Enabling implicit reasoning via shared latent memory. arXiv preprint arXiv:2509.2465. 2025.
Bingbin Liu, Jordan T Ash, Surbhi Goel, Akshay Krishnamurthy,, and Cyril Zhang. Transformers learn shortcuts to automata. arXiv preprini urXiv:2210.10744)) 2022.
Charles London and Varun Kanade. Pause tokens strictly increase the expressivity of constant-depth transformers. arXiv preprini arXiv:2505.21024 2025.
Xutao Ma, Yixiao Huang, Hanlin Zhu, and Somayeh Sojoudi. Breaking the reversal curse in autore- gressive language models via identity bridge. arXiv preprint urXiv:2602.024700 2026.
Arvind Mahankali, Tatsunori Hashimoto, and Tengyu Ma. One step of gradient descent is provably the optimal in-context learner with one layer of linear self-attention. arXiv preprint arXiv:2307.033: 2023.
William Merrill and Ashish Sabharwal. The expressive power of transformers with chain of thought. arXiv preprint arXiv:2310.07925: 2023.
Quan Nguyen and Thanh "guyen-Tang, One-layer transformers are provably optimal for in in-context reasoning and distributional association learning next-token prediction tasks. arXiv preprint arXiv:2505. 15009, 2025.
Eshaan Nichani, Alex Damian, and Jason Lee. How transformers learn causal structure with gradient descent. arXiv preprint 1rXiv:2402.147335, 2024a.
Eshaan Nichani, Jason D Lee, and Alberto Bietti. Understanding factual recall in transformers via associative memories. arXiv preprint arXiv:2412.06538, 2024b.
Jacob Pfau, William Merrill, and Samuel R Bowman. Let's think dot by dot: Hidden computation in transformer language models. arXiv preprint arXiv:2404.15758, 2024.
Zhihong Shao, Peiyi Wang, Qihao Zhu, Runxin Xu, Junxiao Song. Xiao Bi. Haowei Zhang, Mingchuan Zhang, YK Li, Y Wu, et al. Deepseekmath: Pushing the limits of mathematical reasoning in open language models. arXiv preprint AtV:22422.03300. 2024.
Daniel Soudry, Elad Hoffer, Mor Shpigel Nacson, Suriya Gunasekar, and Nathan Srebro. The im- plicit bias of gradient descent on separable data. Journal of Machine Learning Research, 19(70): 1-57, 2018
DiJia Su, Hanlin Zhu, Yingchen Xu, Jiantao Jiao, Yuandong Tian, and Qinging Zheng. Token assorted: Mixing latent and text tokens for improved language model reasoning. arXiv preprint arXiv:2502.03275, 2025.
Yuandong Tian, Yiping Wang, Beidi Chen, and Simon Du. Scan and snap: Understanding training dynamics and token composition in I-layer transformer. Advances in neural information process- ing systems, 36:71911-71947 2023a.
Yuandong Tian, Yiping Wang, Zhenyu Zhang, Beidi Chen, and Simon Du. Joma: Demystifying mul- tilayer transformers via joint dynamics of mlp and attention. arXiv preprint 1rX0v:2310000555, 2023b.
Peiyi Wang, Lei Li, Zhihong Shao, RX Xu, Damai Dai, Yifei Li, Deli Chen, Yu Wu, and Zhifang Sui. Math-shepherd: Verify and reinforce Ilms step-by-step without human annotations. arXiv preprint arXiv:2312.08935, 2023a.
Xinyi Wang, Lucas Caccia, Oleksiy Ostapenko, Xingdi Yuan, William Yang Wang, and Alessan- dro Sordoni. Guiding language model reasoning with planning tokens. arXiv preprint
arXiv:2310.05707, 2023b.
12

---

Published as a conference paper at ICLR 2026
Jason Wei, Xuezhi Wang, Dale Schuurmans, Maarten Bosma, Fei Xia, Ed Chi, Quoe V Le, Denny Zhou, et al. th-tout prompting elicits reasoning in large language models. Advances in neural information processing systems,
2022.
Kaiyue Wen, Huaqing Zhang, Hongzhou Lin, and Jingzhao Zhang. From sparse dependence to sparse attention: unveiling how chain-of-thought enhances transformer sample efficiency. arXiv preprint arXiv:241..0545, 2024.
Longhui Yu, Weisen Jiang. Han Shi, Jincheng Yu, Zhengying Liu, Yu Zhang, James Kwok, Zhen- guo Li, Adrian Weller, and Weiyang Liu. Metamath: Bootstrap your own mathematical questions for large language models. arXiv preprint arXiv:2309. 12284, 2023.
Xiang Yue, Xingwei Qu, Ge Zhang, Yao Fu, Wenhao Huang, Huan Sun, Yu Su, and Wenhu Chen. Mammoth: Building math generalist models through hybrid instruction tuning. arXiv preprint arXiv:2309.05653, 2023.
Ruiqi Zhang, Spencer Frei, and Peter L Bartlett. Trained transformers learn linear models in-context. Journal of Machine Learning Research, 25(49): 2024.
Denny Zhou, Nathanael Schärli, Le Hou, Jason Wei, Nathan Scales, Xuezhi Wang, Dale Schuur- mans, Claire Cui, Olivier Bousquet, Quoc Le, et al. Least-to-most prompting enables complex reasoning in large language models. arXiv preprint arXiv:2205 2022.
Hanlin Zhu, Baihe Huang, Shaolun Zhang, Michael Jordan, Jiantao Jiao, Yuandong Tian, and Stu- art J Russell. Towards theoretical understanding of the reversal curse'vi via training dynamies. Advances in Neural Information Processing Systems, 37:90473-90513, 2024.
Hanlin Zhu, Shibo Hao, Zhiting Hu, Jiantao Jiao, Stuart Russell, and Yuandong Tian. Reason- ing by superposition: A theoretical perspective on chain of continuous thought. arXiv preprint 125554 2025.
13

---

Published as a conference paper at ICLR 2026
A NOTATION DETAILS
The notation and meaning of each token and the position index are the same as Zhu et al. For completeness,
(2025)
provide detailed descriptions of different tokens in Table | (which is Table
in Zhu et al. (2025)), and the position index different tokens or continuous thoughts Table (which is Table 3 in Zhu et al. (2025)).
Tokens
Meanings


a special token denoting the beginning of the sentence the source node of edge i the target node of edge
Si
ti
<e>
special token marking the end of an edge special token followed by two candidate nodes two candidate destination nodes special token marking the start of reasoning the root node the i-th continuous thought (represented by d-dimensional vector) special token driving the model to make the final prediction
<Q>
C1 C2
<R>
[t:]
<A>
Table 1: Meaning of each token (Table 2 in Zhu et al. (2025)
Notations
Position indices
Idx(
)
Idx(s;)
3i 1
Idx(t;)
3i
Idx(<e>,1)
3i
Idx(<Q>)
3m + 2
Idx(c1)
3m +3
Idx(c2)
3m +4 4
Idx( <R>)
3m + 5
Idx(r)
3m + 6 to
Idx( [t;]]
In i
Idx( <A>)
to C + = =T
Table 2: Position indices of different tokens continuous thoughts in the input sequence (Table in Zhu et al. (2025) )
B MISSING PROOFS FOR SECTION 3
In this section, we provide the full proof of theoretical results in Section We first provide theoret- ical analysis of COCONUT-BFS and COCONUT in Appendix B.I Appendix B B.2 respectively, and provide results for continuous thought expansion in Appendix B.3
B. ANALYSIS COCONUT-BFS
In this section, we analyze the training dynamics of COCONUT-BFS, We first provide the closed- form formulation of the gradient V,, G.r. pBFs V., (BFS (H), where μ (ur]vEv the set of parameters. We omit the superscript or subscript when the context is clear.
Lemma (Per-instance gradient of ph for COCONUT-BESS
Under the loss function of COCONUT- BFS as given in (5) and the forward pass as in (2) the per-instance gradient is
1 No.. VIW.
Evtee exp (Er)
71 [N
(x+
exp (5+) + N
for any v E V, where E, = log (EveNots exp (Ex)
14

---

Published as a conference paper at ICLR 2026
Proof. First, note that for any E V, according to 2 the logit can be calculated as S, =E()) (hette)
di (h [t-1] (=.))=))
=E(u)
(= V
hh (trel Wh,) he
m
=E())
(=c) Σ
(L Whi)hi
m
=E(o)
h [tte]
+
Σ
Σ
XE
Σ My (O")) (E. ((,))) + +E((t.)) h4 N'EV
i=1
L'ENe
=E(o)
h(te] +
Σ
Σ AE((')
Ms, ps, E(si) h
i=1
N'EN
=E(n)
Mas 1 (ss E No
(A-((uENa) + Σ Aua 1ss NoJ 1(v = tij.
where x VNNol
Note that by the definition of No. we have &. 0f Net1- Therefore,
EmeNen exp (Ex)
((w) = log
Eveyes exp (Ex)
EvENen exp (A 1f0 E No) + E Ada 1(s; E 1(v t:))
log
Ever (A 1(U E No. + E Ats, 1ss E 1fc t:l)
n Wo+11
log 1
Evey exp (A 1 1 Nc) Σm Apse 1( E No 1[r t;))
For simplicity, define exp (5-) EvENei exp (5x) and thus
n n We+1
((n) - log (1- exp (E-) W
Then the per-instance gradient can be calculated as Vp. E(14
exp (E+) Wot1l
n Wo+11 (exp (5+) +n (N++1l)²
Vn, exp
exp (E+)
EveNCo exp (5x') E No) Em Al ti = ='}
n1 WWo+1
exp (5+)
exp (5-) + n We+11
En exp (Evr)
Wo
=- X- = Al-luENN) No)
exp
exp (E+) + We+1l
Now we calculate the gradient of pa over the whole dataset, where the nodes of the graphs are randomly shuffled. We also write LBPS [BFS(()) (μ) and omit the superscript when the context is clear.
Lemma 4 (Whole-dataser gradient of th for COCONUT-BFS), Under the loss function of COCONUT-BFS as given in (5) and the forward pass as in ( and assuming all p have the same value, the gradient w.r.t. the whole dataset is
exp (-E+)
n Wotil n. W.. exp (E+) We-1l
p. C(H)
Σ do exp (5v) wev
15

---

Published as a conference paper at ICLR 2026
for any v E V which is independent v. where &+ = log (EveNets exp (5o)).
Proof. Denote cl8,z) log EveNg, (e) exp (5))).))) where 5o g( is the logit of when the
graph in the prompt is 9 and the start node is
According to Lemma 3 and the condition that all jr have the same value, for any permutation
Sy, we have
(=(())=(r)) (x)) ((T(U) e NO (T(x))) + Em= (a(a.) (r(a)) e N.T 1(x(v) = x(t,))
No
(T(())
(1IEEN8(t))
m
Na 1s Σ
v(Ng(=)) -c(0.2)
This also implies
exp ((T())
E
exp (<htu) (c((())
a(o)
(x(e)) (a(t))
Σ
exp ((.x))
VENV (z)
exp ((G.t)
Therefore, by Lemma 3 we can obtain that
V Mat(o l (0)))()))(
1(T(V) (O (T(T)))
Entonooo--a exp (5((T))
((((0))(()))
VINTOO (T(2))
exp (1)()))(())
n (a(P((T(E))
(T(x))
exp ((((G))=+))) +n N (T(())
Ensor-eee exp (G())))
Ng)))
n | (x)
(=)
exp (El (((C)t)
exp
( ((.)))
+ n
(x)
-V (1).
Therefore, we can calculate the gradient with respect to the whole dataset as
V Va, L(v) =F =Fx-Uni((V)) / na en (10))
=EInUnH(Ss) M - (H)( =EorrUni(V)) n. (4))
16

---

Published as a conference paper at ICLR 2026
which is independent of and thus the gradients for fou are equal for all E V. Furthermore, can calculate that
t/0 ((s)
E Evsco- exp (Eer)
Wo+11 exp (E+) + n Wetzl
VWNll
exp (E+)
NEV
1
Eettennee exp (Eor) exp (5+)
n WWa
exp (5+) Wet1
vENo
exp (-++)
|No+1
Σ
exp (E.r)
v.v.l
exp (5+) + n - -Wa+1
MEN, o'((n=
exp (-E+)
Wetil exp (E+) + n - -Wo+1l
Σ d, (Ev)
n. V.ol
DEV
According to the gradient of pv, we finally show that by diverges to infinity at least logarithmically in I.
Lemma 5 (Dynamics of pay for COCONUT-BES) Let (o(t) be the value of poo at time t. Assume zero-initialization, i.e., (0) for all E V. Under gradient flow
V Vu, CBFS (p)
(10)
where a > 0 is the learning rate, we have
Pe 2 c1 (1 + ac2t)
for all u E V where 61
2//N.11 C2 = n 3
Proof. First, by Lemma all juy have the same value if all have the same value. Given that (0) for all E V. we can obtain that for any time t, (t) have the same value for all e V using similar argument as in Lemma 15 of Huang et al. (2025a).
Now, given any fixed time 2 0, holds that (t) has the same value for all v E We omit for notation convenience, i.e., using He to represent (t). Below, we first provide a lower bound of the gradient jly..
Since we are guaranteed that one of and cannot be reached from r, Ne+ cannot contain all the vertices in for any e, and thus n [W 2 Also, since one of ci and C2 is guaranteed to be reachable from r, there exists u E V such that d, 2 for any This is because the start node r E No S Nr. for any c 2 0, and we can take v = Pı which is on the shortest path from r to e*. Therefore, obtain that i. 0. Moreover, we have
j&v =0 exp (-E+)
n WWetl
Σ do exp (Er) VEV
VWW
exp (E+) + Wettll
2a: exp (-t+)
du exp (5v)
n VWa
exp (5+) +1
NEY
ΣVEV do exp (Ev)
n. W. (exp (E+) + 1) exp (E+)*
Note that by definition, for any vertex E NNo1 Mr). there must exists another vertex v' e No such that (u' + v) ε, which implies that d, 2 Therefore.
Σ d. exp (5x) 2
Σ vENotrifE)
d exp (5v)
NEV
Σ ((E)
exp (5.)
exxp (E+) exp (Ex)
17

---

Published as a conference paper at ICLR 2026
which further implies that
ΣnEy do exp (Ev)
Ne 2 a
/Waal n VWW
(exp(5+) + + exp (E+)
exp (5+) exp (E=) (exp (5+) + + exp (E+)
Now recall from Lemma that
1fee + Σ
Ma. Na. 1ss Hs, EN) 1fv (()=)))
VWd
(1fr N + de Mb)
Wa
(1 Nal pa). Ma)-
Wa
Therefore,
exp (E+)
Σ
exp (Ev)
VEN.11
Wo exp VWVa
(1 + Wol pe)
Sn. exp [
Wal Pa)
Also, since P1 E Not1 and degg,N. N. (p1) 2 1, we can obtain that
exp (E+) exp (E.) 2 exp (c(),e)
exx
(1 NNo + dpi He)
VWNd
2 exp
VWV
Combining the above two inequalities, we can obtain that
a
exp (5+) exp (Ex)
ju
71 [
(exp(E+)
1) exp (E+)
exp
Bu
a
VIV.i
n
N
2n exp
(2 (2+2v Wel pv
(-2V Nol (s)
exp
3n
W
Finally, by applying Lemma
13 can obtain that
In (1 + on 30- 32)))
2/a
B.2 ANALYSIS OF COCONUT
In this section, analyze the training dynamics of COCONUT. Similarly, first provide the gencn
closed-form formulation of the gradient V,.. G.e,p V u. G.r. (H), where μ === {Mp hoey is the set of parameters. We omit the superscript subscript when the context is clear.
18

---

Published as a conference paper at ICLR 2026
Lemma 6 (Per-instance gradient of de for CoCONUT), Under the loss function of COCONUT given in 6 and the forward pass as in the per-instance gradient is
you ((()) = 110 E NNo
(-1((- - po+1) E + Eviio-rre exp (5ur) exp (5+) + n |W++1]
VWV.l
for any E where E+ log
(EmeNeis exp (f.)).
Proof. First, according to the proof of Lemma we have
&. = 1 +
Σ Nu 1ss; Nc) 1f2 = tih,
where x
Note that by the definition of No. we have &o if & No+1. Therefore,
ViooT
cxp (Spo+1) log
exp (Spo.1) exp (E+) + n |Not1]
(()) = - llog "E exp (Ev)
where exp (E+) - E eeN exp (Ev).
Then the per-instance gradient can be calculated as
V pr E(1)
exp(E+) + WN++1l
exp (Spot))
Mr exp (Ere.1) (exp (E+) n (o+1l) exp (Ereta) Vno exp (E+)
(exp (E+) (N+z112
Vu exp (Eret1)
Va, exp (5_)
exp (Spo-1)
exp (5+) + n Wo+11
Nr exp (E+)
V Epets
exp (E+) + W
Since
exp (E+)
Σ
No) Σ -,
v'ENet1
1=1
=x -
Σ
exp (Eor)
vE((n=w")EE
and
Vp.Spors A 1 1f(1 + Pet1) EE),
we can finally obtain that
(()) = 1f ViNal
(-1((o- + po+1) EE Eesoorn exp (for) (E+) + Wotil
Now calculate the gradient of pa over the whole dataset. where the nodes of the graphs randomly shuffled. We also write _onen = e (p)) and omit the superscript when the context is clear.
19

---

Published as a conference paper at ICLR 2026
Lemma (Whole-dataset gradient of da for CoCONUT). Under the loss function of COCONUT given in (6) and the forward pass as in (2) and assuming all Py have the same value, the gradient w.M. the whole dataset is
(-dren
EVEN. de exp(Ee)
C(p) vWNal
exp +n WWotal
for any E V which is independent u, where &+ = log (EweNets exp (Ew))
Proof. Similar Lemma denote <( log (Ewevg, exp < (!...)) where <(9.=)) is the logit of when the graph in the prompt is 9 and the start node is r. According to the proof of Lemma for any permutation Sv, we have
(((()) ((9,2)
sx(v)
for any v E V and
exp (cf-(())(())) exp (s!6.r).
Therefore, by Lemma 6 we can obtain that
Vanco (1()) (FL) 1(T(V) E No (x(E)
(-1((T(v)) (
No (T(z))
(=)))))<()))
Ento)1oo)-= exp
(x(0')
exp (c(-10)((7) + - |vot (T(T))
Enso-nneE exp ((9
1f (N)
-1((v PP41) EE}
(=)
exp (c(01c) + NN9 (x)1(x))
-Vp.lg,sp(1))
Therefore, calculate the gradient with respect the whole dataset
V po L(u)) =F u--uri(ss))) V (1)
=Etr V -,- (p)
=Em'>Unif((y)> (V [ (g.
which is independent of and thus the gradients for μ., equal for all v E V. Furthermore, similar to Lemma 4 we can calculate that
V ((4))
1l0 EN ViVal
Edsoo-nee exp (Eor)
-1((0 + P
n NEV
exp
(1x) + n Nottl
EveNion do (xp
-deei +
n W.T
exp (e((it) + (ot1
Finally, we derive the dynamics of My. Recall that we denote K WWl, VR We also make the following notation for Theorem Let d. dp and d max Moreover, Let
20

---

Published as a conference paper at ICLR 2026
cp |e+11 and denote
Eu(m) (( NcT + pdd),
E+(p) Σ clc(1))
mEN.11
S(1))
S(1) =
du esull))
F(u) = (+(1) ( +
MENo+1
Theorem (Dynamics of j4. for COCONUT). Let pe (1) be the value of Pho at time t. Assume zero- initialization, ie., (0) 0 for all v V. Under gradient flow
Nn = V Croe (fL)
(11)
where or is the learning rate, suppose the initialization satisfies (o) 0 for all Then:
1. Scalar reduction. For all t 2 O. polt) = ju(t) is shared across v, and p(t) satisfies
a (d. -F(((t))).
i(t)
(12)
nvK
2. Regularity of F. The function F: R + IR is C strictly increasing, and satisfies
lim F(n) O, 11-00
lim 110000
F(R)= = dmax: 0 F< < dmax for all finite jr.
3. Finite fixed point when d. < Amaz. df < dmax. there exists a unique μ E R such that F(P²) d,. The solution jut) of 12) with μu0) O converges monotonically to p*:
p(t) 7 4tt if F()) < dd.,
p(t) > n¹ if d..
and the equilibrium μ² is locally exponentially stable, ie., there exists ? such that for all large enought t, it holds that
(m(() - 10 S $5770 - =
4. Logarithmic divergence when d, = dmax- If d. dmax. then ju(l) > 0 for all I and
n(1) + +00. Moreover, for all 2 0,
a d2 max co
t).
(tt) 2
In 1
(13)
Admax
2n² KK
Proof. (1) Scalar reduction. By Lemma and similar argument as in the proof of Lemma we have po put) for all 2 0. Therefore, the gradient V u. C(u) is independent of and equals
ΣveNee, du cbulma)
d, (d,, EmeNen eSu(Hu) +n [Ntiil +n-WWotll
1 (-d, F(u.))
V
Mo C(p)
nvK
nVK
Thus, we have
jut) -V -Vp.C(u((t) nvK (d. FF((())
(2) Regularity and limits of By the proof of Lemma and the condition that (u.(1) = p(() for all E V and / 2 0, we have
E4 (14)
Σ
exp()(1(uEMM) + pdd)) S(4)
Σ du exp(A(I {u NNa + pdm)))
RENe11
ueA
Both functions are finite sums t functions of μ, hence F(μ) = S(p)/ (p) - co) is also
21

---

Published as a conference paper at ICLR 2026
Now show the strict monotonicity of F(:) μ by differentiation. We further write Eu Ea(u) for brevity. Then
(4) -
Σ
duefx
S'(n) =
Σ
MENe11
NENet1
Therefore, obtain that
F'(n) (F+ c0) ((1)
(F+ (1)) coo)²
27
A (E+ co)2
Σ NEN
of (E+(1) + co)
Σ daesa
-
uENo11
Note that by the Cauchy-Schwarz inequality, have
2
Σ uENo+1
Le F..(F=+()))))
Σ
dretu
2 0.
"ENot1
and hence
F()) 2 (E+(1) + co) E+(1))
Σ
Lees
0
uENoti
since co and there exists at least one node u e No+1 (e.g., Pc+1) such that du 2 1 by definition. Thus, F(-) is strictly increasing.
Now we consider the limits of F(-). First, note that
S(n)
Σ
du el.
uENot1
and for each u E No either d, 0 or du > 0 and thus
lim due (u()) lim du exp(A(L E No) + udw.) =0. 0. 10--0
Therefore, have - S(n) 0. Moreover, since E+(n) + co 0, we have
lim F(A) =0
2
Now we consider the case when + -00 Since
(5u(m)
(Xp(X( + pdw))
(+) + EvEN.,, exp(X(1fe E + ude)) + Co
EvEN., exp()(1 No) 1 (uENN + p(d. du))) co
As μ +00, we can obtain that if d2 < dmaxx then
Ea(a)
lim +00
lim # exp(X(-1 u(dmax d.))) + co
0.
E. (M) + co
d, = dmax, then
(L((m)
lim
400 E+(p) c
exp 1fu
lim 1+0 EveD exp() 1lv E N.)) + EvENeH Duna exp(A(1 E N₂) p(ds dmax))) cAnumax
exp (X 1lu E NoI) EvED exp() 1I [N))
22

---

Published as a conference paper at ICLR 2026
where Dmax= {u No+1 dy dmax). Therefore,
exp (A Liu NoH) Σ du Ever exp(1 1lv E (No))
lim 10000
F(N)
dmax.
of Daax
Finally, the inequality F(p) dmaax for finite μ follows from S(n) S d. (1) and co > 0:
F()) E-(n) S dmax dmaxE+(p) < dimax.
S(p)
(+(u) +
(3) Finite fixed point and monotone convergence for dh dmmax. By (2), F(.) is continuous, strictly increasing, with range (0, dmax). Therefore, there exists unique j E IR such that F(u*) d, Now we first argue v(() - 1r"
S μr' and (thus) F(u(O)) & F(2*) d.. then ju(t) 2 0 as long as put) S μ², hence is non-decreasing and bounded above by j=; monotone convergence implics μu() + f S μ for some p. By the continuity of F(-) and the fact that j (u()= > ee is analogous with a non-increasing trajectory.
Consider the ODE jt (d. - F(u)) with c c= a/(nn/ K) 0. If (())
obtain that F(a) d., which implies μ². The
For local exponential stability,
set μ μs and write
A(t) = - (F + (u(t) - F(u"))
By the mean value theorem, +(u*+ū)= F(
(g) for some & between μ² and +f.. Since F' ((*)) > and Ft continuous, there exists 1 > and m > such that FE(E) 2 m whenever I p= S vy. Hence, as long as [a(()) S 7,. we have
di (+(0) = I((t) (()) = -c FF(()( [,(1) S -cm (7(1)--
Applying Gronwall's inequality, we have (ia(t) S e-mtt (00) in this neighborhood, which estab- lishes local exponential convergence.
(4) Divergence and logarithmic lower bound for d, dmaxx When d, dmaxx since F(u) < dmax for all finite μs, we have ii()= c (dmax (2(())) > 0 where e nvK and thus ((() is
strictly increasing. We now lower bound the growth rate similar Lemma 5
Since S(n) S dmax E- (u), we have
S(p)
2 dmax (1-
E- (M) E+(u) (4))) E_ (1) co
dmax Co
dmax F(u)= dmax
E+ (1) + c0
Moreover, for each u E Net1, it holds that
(La(m) S exp(A(1 + pdmax)).
Therefore, (M) < [No+1 A(L+admax) < e2++ and thus we obtain that (x) + c S ne Al t S ( + +1) < 2n e AlL+udmax)
where used e for 2 O. Combining the above derivation, can obtain that
dmax F(P) 2 dmax Co - e
2n
We can then plug this into ja = C. (dmax - F(u)) with == an/nV// K) to get
dmax Co - Admaxu(()) (2H(1)
iutt > RvK
2n
a dmax e and 02 Xdmax 0.
where C₁
2n² VK
Applying Lemma
13 we can obtain exactly
13 This shows ((t) + +oo at least logarithmically
fast.
23

---

Published as a conference paper at ICLR 2026
B.3 THOUGHT EXPANSION
Finally, we provide results for continuous thought expansion. Note that the following results hold for any directed graph that differs from the graphs in the training set.
Theorem 5 (One-hop expansion of continuous thoughts). Let G (V.E) be any directed graph (which can differ from the graphs in the training set) and r E V be a root node. Assume the current thought is any positive superposition No (r):
h (te)
Σ E(u,,, NEN.
A > O.
Then the next continuous thought [to+]] 1 h [cons generated by the forward pass
satisfies
- Tiihitl
Σ
B. Evs
DEN++1
with coefficients
B. 1. + 1 Σ 1() 2
(14)
wEN.
carryaver
one-hop expansion
where we assume the model has been trained until time t and the trained model satisfies (t > 0 in (3) In particular, 3. > 0 for every E Not1 if >> O, so the support of & is exactly Net1. and the output is superposition of Net1.
Proof. Note that h; We can calculate that
and W EVVEV (() =En ME(D)E
Wh:
Σ E(D)), (Es(()s) + EE.(t.) = - E((s)), DEV
where we used E, Therefore, with h [t.]] = EvEN. X,E(u), we calculate
= 1fc = si} and E, (E) = according to Assumption
0₂ = h( Wh; μ Σ Au E(u) E(s:) = Mb, 1ss, Nc). ueNo
The value aggregation becomes
m
( (hh)) Σ h p Σ
As (Eo(s:) + Et(tt))
i: i,,,N,
Furthermore, we have
UT (h (ta) (hi)) =p
Σ
Xa, Ct. = Σ
Au 1[(u u) E E Σ
8: BIENo
PEV
MEN,
Similarly,
U Uhh = Σ M Σ No 1lue Ao
VEN.
vEV
Adding the two parts yields & = Σ Bo en with 8, as in hite) + (11 (h:3) and & U (tai)).
11 since by definition we have [
C MISSING PROOFS FOR SECTION 4
In this section, analyze the training dynamics of the prediction stage, i.e., after thought gener- ation, how the model extracts the information from the generated continuous thought to make the final prediction.
24

---

Published as a conference paper at ICLR 2026
Recall that the i-th training sample consists of (g(() r(0, (1) (6) ci ,0 ))))), and denote C* as
(6)
the reachable candidate and ( the unreachable candidate. We also use No .V denote the C'-all for the i-th training sample. We assume is large enough so that c E -
(r(1) to
for any i. Note that cl & Ne for any C by definition. For notation convenience, also Meas, MR = MeRs, and denote g($) = (5% JeEVoc as the logits calculated by forward pass
(6)
for the i-th training sample. We denote c(0) 2,09 (! A convenience.
= for t E {1,2, -3 for notation
To start with, we first provide a closed-form logit expression. Lemma (Closed-form logits in prediction stage). Under repurameterization ( and forward pass for the prediction stage f every V we have
(p MARR)) = pe Aa frontier residiaal
Meao Mo E (c1, o2)1.
(15)
candidate lift
In particular,
Ee, - plea Ae.
(16)
Proof. For the reasoning token Mdxx(
) = E(<R>) + E(CI) + E((22), we have
Widdx(sws) MeR= E(<A>) E(<R>)
hldx(<w>) MeRa E(<A>).
Therefore,
U (T(ide((<as))) Whids(<R>) H(xx(<<>)) )
MeRRo (eago + eei + ecz)
Also,
U UTiuass hixx(=a))) = Hegx
Σ
Au Cu + Moa Cexs
wENo
Combining the above two expressions yields 15) For TT $ No we have dc, 0, and (6 follows.
For each training sample, we can construct a two-dimensional feature for every node v E
(2(), 1fv E (o( c E
20:
so that cl0) (MA, MR) (w. zt with (MA, MR) E R2 For instance i, correct classification ( ( za for {1,2, x, L}. We
means (w, (1) x > 0 for all + (t) where we denote x (1) =
further define the difference of features with respect to c, for later use:
(A 0), (19 (A
==
Di,o
(5) Iec.
1
1),
VE (c.
(17)
1),
u U Ufo
10
C.1 LINEARLY SEPARABLE STRUCTURE AND MAX-MARGIN PROBLEM
Lemma 9 (Separation by a nonnegative direction). For every instance i and each competitor v + (1)
V (()
((11)) 1), Ai,))
+1,
>0
+1,
Hence, the training data are linearly separable by direction in R2
Proof. The result holds because 1! 00 A..0 (
$1.
25

---

Published as a conference paper at ICLR 2026
Define the hard-man hard-margin value of unit direction u E S! n R2 (where S' {u E RR? |ull2 1}) as
(u)) min min (u, Ai.)). uxtt
The corresponding maximum-murgin direction is
arg max 7(u).
(18)
We characterize u* using the following two quantities of the training sets:
X, min X
(), Atrain max; max, "ENE (A
[0.
where (z)+ := max 0}. Intuitively, A, the smallest mass ever placed on reachable candidate across the training set, and Atrain is the largest overshoot of a non-candidate but reachable node's weight relative the reachable candidate.
Lemma 10 (Closed-form lower envelope of the margin). For any unil u = (MA, uR) E ES n R
(u) min u X., uR UA trains UA UA + uRR min u Xa ug U A Arain
Proof. According (7), can directly obtain the the lower bounds u,a X! un uA (2 x(e))
and u _ +up- Minimizing over and according to the definition X., Atrain yields the desired result.
Li
Proposition C.1 (Properties of the maximum -arginn direction). Let u² = (u unp) be solution of (18). Then the unique maximizer satisfies
"R X2 + Atrain.
1
A + Atrain
VI+ (A. + Atrain)² "n
"
VI+ + Atrain)
Proof. By Lemma 10 we can maximize (uu) minfuAd. uR UAAuiii over the unit vector by equalizing the two arguments (otherwise one can rotate u to increase the minimum). Therefore, we can equalize the two arguments, which yields un (A + Atrain). ), and obtain the desired result.
C.2 IMPLICIT BIAS OF GRADIENT FLOW AND DIRECTIONAL CONVERGENCE
(E.
gpred G.r.
log
Cred E(G.n.er.ss)) l (.
Evey (E)
Recall the loss function on the prediction stage over the training set ( We can rewrite it as follows
N
exp(s!)
C( =
1
e0 (HA.HR))
(HA.PR) = - -log
N
Evey exp(s!")
and run the gradient-flow dynamics vi(() = -o ((w(()) with wo()) (Ma(t)) (t); (t)) and o > 0. By Lemma the data are linearly separable, so the implicit bias gradient flow directly yields the following lemma.
Lemma 11 (Implicit bias of gradient flow). Along gradient flow from any bounded initialization w()), we have
w(t)
[w(()]| - 00,
u²
(tw(t]
where un is the unique solution to the maximun-margin problem (88.. there exists scalar radius r(t) such that
Combining Proposition C.1
(HA(t),uR(t) r(t) u + o(r(t)),
and for any E 0.
(t) (t)
2 X. + Atrain for all sufficiently large t.
The proof can be straightforwardly adapted from its gradient descent counterpart (Soudry et al. 2018).
26

---

Published as a conference paper at ICLR 2026
C.3 PREDICTION ON UNSEEN GRAPHS
Finally, show that after sufficient training, the model can correctly predict the reachable candidate node even for unseen graphs, showcasing its generalization capability.
Fix any unseen test graph along with the exploration set No and weights Atest such that Atrst (0, 1] on No and otherwise. The test graph also satisfies
max
An
lest
test
S A
with A S Arain-
Therefore, for every non-candidate Nor
\ [otstt
J. it holds that Atest S A + A.
The following lemma shows that as long as the test graph satisfies the above condition, has positive margin using the maximum margin direction for the training set u².
Lemma 12 (Positive time margins from the trained direction). Let u* (u', u²)) be the unique max-margin direction with ug/ / = X, + Auaan > A. Then for every competitor v # chest,
(u', test b 2 min das uAA X
test
0.
Proof. For = c the difference is (Atrrt 0); since A test O.
0 we have {u", x 200t)) 2
WAA X
For o & Noestt the difference is (A 1) and the bound is even larger.
For E Nostt (ctest), we have Arsst & Abeat + A, hence
(u", 1eat teat) = u'a (A Xteat) +uR 2 uR uA (s -2) 2 uny d, > 0.
Finally, show that after sufficient training, the model can correctly predict the reachable candidate node..
Theorem 6 (Generalization for unseen graphs). Let (1)) (*(()) follow gradient flow on the loss from any bounded initialization. Suppose the training set is linearly separable and X. Atrain are defined as above. Then, for any unseen instance satisfying Aeest E (0, on Nesst and 0 otherwise, and
max
test
A test
with
Arrain
WEN
Ox
we have for all sufficiently larget. t:
exp (Eresh (MA (HA(t),
Peteal (t) =
E, exp (Clest (MA(s,
Proof. By Lemma 11
we have
(Ha(t),pR(t) r(t)u* + r(t) + co.
Then, by Lemma 12 for every competitor + cleat,
c! (Ha(())) - gt (n (U)())
testt + o(r(t)) =r(t) (u", test
2r(1) min (wn Xo, uA Xtest) + (r((1) t-300 400.
Hence the arg max is test and its softmax probability tends to 1.
27

---

Published as a conference paper at ICLR 2026
D AUXILIARY LEMMAS
Lemma 13 (ODE lower bound), Let c1, > O be two constants. Assume the function f - satisfies f(0) and
()) dt 2 exp(-022 s()), v2 0.
Then holds
1 In(1 cc C2
S(1) 2
for all
Proof. We define g(() e Note that dg(t) = d (c22(()) = (f) dt Therefore, dg(1)) 2 G22dl for 2 0, and thus
enl(() 22 exp(-c2 f(t)) exp(ee2(()) = c102.
dt
dt
[ dg(t) 2 L ac2dt = g(t) 9(0) 2 ce2t.
Therefore,
s(t) ex1( 2 g0) + c1e2t = c00(00) + = + a022,
which implies
1 In(1 + cc
S(1) 2
E EXPERIMENT DETAILS
E. DATASET
Table 3: ProsQA statistics. Numbers averaged over problem instances.
#Problems
V
[E|
Sol. Len.
Train
14785
22.8
36.5
3.5
Val 257
22.7
36.3
3.5
Test
419
22.7
36.0
3.5
The statistics of the ProsQA dataset is shown in Table
F.2 EXPERIMENT WITH COCONUT-BFS
As comparison, also train model with modified version of LBFS, Recall that the original LBFS (5) encourages the model to predict any nodes within Net1- To avoid the trivial solution of al- ways predicting the root node, introduce experimental variant that only encourages predicting nodes on the current frontier:
EveNe exp(E.)
COCONUT-BES
BFS.exp &.
-lo
(19)
Evey exp(E.)
All other training settings remain unchanged. The answer accuracy of this model on the test set is 99.0%. We then track the logit difference between frontier and non-frontier edges as a proxy for pas with results shown in Figure 5)
28

---

Published as a conference paper at ICLR 2026
Number of layers
Number of heads
Width
98.8
H 98.0
dmnnee 384
62.0
L 97.3
H = 8
98.8
dmadell = 768
98.8
= 8
96.5
H = 12 98.8
dmodel 1536
97.7
12
67.4
Learning rate
Weight tying
1) = 2 x
58.1
Tied Untied
98.8
1 = 1 10
98.8
98.8
7 5 X 10
62.1
Table 4: Ablation on depth, heads, width, learning rate, and weight tying. By default, other hyper- parameters follow the main experiments.
In Stage 1, the logit difference for 1 grows much faster than under ceoco and shows no sign of saturation even after 150 epochs. This agrees with the theoretical prediction in Theorem 1 under COCONUT-BFS, j4y diverges rather than stabilizing. At later steps (c 3, 4), the gap between COCoNUT-BFS and CoCONUT becomes smaller. We attribute this to practical factors such as stage- wise data mixing, gradient propagation across earlier thoughts, and a larger discrepancy between losses and (99 in the later stage.
Average Attention Difference
150
125
100
75
50
125
175
225
350
Training Epoch
Figure 5: The attention logits difference between frontier edges and others. The model is trained with modified version of CBPS
E.3 ALTERNATIVE ATTENTION ROUTES CANDIDATE LIFT
Our theoretical analysis in Lemma assumes that <R> copies the candidate nodes in the first layer, and <A> then attends to <R> in the second layer. practice, however, we observe three distinct yet functionally equivalent attention routes that realize the same candidate lift. Example attention maps for each route are shown in Figure
F ADDITIONAL EXPERIMENTAL RESULTS
In this section, provide additional experiments to complement our main analysis of training dynamics.
F.I ABLATION STUDY: ARCHITECTURAL OPTIMIZATION SENSITIVITY
We evaluate the sensitivity of COCONUT training model depth, number of attention heads, hidden width, and learning rate. The results are summarized Table 4
We observe that models with L {4, 8) layers maintain high accuracy, while L 12 is harder to optimize. The performance remains comparable when dmmate1 E {768, 1536). but degrades when the width is too small (e.g., d 384). Varying the number of heads does not have major effects
29

---

Published as a conference paper at ICLR 2026
42
qq
80
20
1.5
20
daterts
2
Irpat Ney!
Youn Tokans (8ay
(a) Layer-1 attention map for Pattern A
(h) Layer-2 attention map for Pattern A
1.75
125
20
2.0
hos
0.20
deaant
Arats
daterts
daterb
0.10
Inpui Thwwee Dey!
Iour Takes rely
(c) Layer-1 attention map for Pattern B
Layer-2 attention map for Pattern
10-
12-
Yosk Tebrr (bary)
doadd Hegge:
<aterb
[A]
AA
(e) Layer-l attention map for Pattern C
(f) Layer-2 attention map for Pattern C
Figure 6: Example attention maps illustrating three alternative routes for candidate lift. For clarity, we omit earlier tokens in the sequence and only visualize the final segment containing some of edges, the candidate nodes, latent thoughts, and answer tokens. Pattern (consistent with the theoretical assumption); <R> copies candidate nodes in Layer 1, and <A> attends to <R> in Layer 2. Pattern B: <A> directly attends to candidate nodes Layer 1. Pattern C: continuous thoughts copy candidate nodes in Layer 1, and <A> attends to the continuous thoughts in Layer 2. All three patterns achieve the same functional effect of lifting the reachable candidate.
30

---

Published as a conference paper at ICLR 2026
8
<5>
<
11
Input Tokens (Query)
16
Attention Weight
3
<6>
17
-2
<e>
0-


15
<e>
11
16
14
<e>
Input Tokens (Key)
Figure 7: The first-layer attention patterns in 4-layer transformers. <e> tokens attend the corre- sponding source and target nodes to aggregate the information. This is consistent with the analysis of the two-layer transformer in Zhu et al. 2025
final accuracy, whereas too large or too small learning rates tend to degrade performance. Weight tying setting does not affect model performance.
We emphasize that each ablation in Table f varies only a single hyperparameter at a time, keeping all other settings identical to our main experiment. In practice, these hyperparameters interact in coupled manner. For instance, with a smaller learning rate of x 10-5, we can extend the first-stage training to 300 epochs and get 97.0% accuracy. For deeper models with L 12. prolonging first- stage training to 400 epochs and reducing the learning rate to 10 improves accuracy to 99.6%. A comprehensive hyperparameter interaction study is beyond the scope of this work and is left for future investigation.
F.2 MULTI-LAYER TRANSFORMERS AND MECHANISTIC PATTERNS
We use the COCONUT model with L to analyze the reasoning pattern beyond two-layer trans- formers. The results shown in Figure 7 and Figure and summarize the reasoning patterns below.
First layer (induction head): The first layer performs token-level copying, propagating node information into edge tokens <e>, consistent with the copy mechanism derived in previous theoretical analysis (Zhu al. 2025).
Second layer and beyond : From the second layer onward, the model aggregates over reachable nodes in a superpositional representation that enables parallel breadth-first exploration.
F.3 ACCURACY DYNAMICS IN THE ANSWER-PREDICTION STAGE
We track the test accuracy during the final answer wwwerrppeeictionn stage following the setting in Figure The result is shown in Figure 9 which shows a rapid transition from near-random guessing to stable high accuracy once the model integrates residual carryover and candidate lift signals.
31

---

Published as a conference paper at ICLR 2026
Cert muous thought?
Contrnuu though
Cortnunsstoough 4
2530
2003
330
23:03
643
030
300
270
330
1930
2030
1253
1590
503
370
192
3753
1205
3300
1003
10:0
330
3750
13.0
(17)
3300
3000
630
400
230
1111l Froduct POLR
150 Noduct Value
Figure 8: Inner product between layer-wise hidden states and different types of nodes 4-layer transformer. The experimental setting follows Zhu et al. (2025). From the second layer onward, hidden states exhibit larger inner products with reachable, frontier, and optimal nodes, indicating that superpositional representations emerge early as layer 2 in the 4-layer transformer.
Average Attention Logits
40
80
Accuracy (9%
60
20
40
(residual carryover)
0-
Aia (candidate lift)
220
-10
Accuracy
Training Epoch after Answer Prediction
Figure 9: Accuracy curve during the answer-prediction stage. The accuracy shows rapid improve- ment corresponding to the learning of residual carryover and candidate lift signals.
G THE USE OF LARGE LANGUAGE MODELS (LLMs)
We used LLMs mainly for grammar checking and polishing in paper writing.
32


Reviews:
[{"review_id": "96YsAif8kn", "reviewer": "Reviewer_Aa1k", "summary": "This paper investigates the superposition mechanism underlying continuous chain-of-thought (continuous CoT, or COCONUT) in large language models (LLMs), a paradigm where reasoning traces are maintained in a continuous latent space rather than discrete tokens. Building on prior work that theoretically showed a two-layer transformer with continuous CoT can solve directed graph reachability via superposition of reasoning traces, the authors address a critical open question: Do gradient-based training methods naturally induce this superposition?\n\nTo answer this, the authors analyze the training dynamics of a simplified two-layer transformer on the directed graph reachability task, splitting training into two stages: (1) a thought-generation stage, where the model autoregressively expands continuous thoughts; and (2) a prediction stage, where it converts these thoughts into a final answer. Their core theoretical insight is that the index-matching logit—a metric of the model’s local search ability—remains bounded under mild assumptions during COCONUT training. This bounded logit balances exploration (assigning comparable weights to multiple plausible reasoning traces) and exploitation (leveraging local graph structure), directly enabling superposition. The authors validate these claims with experiments on a GPT-2-style model, showing logits stabilize at bounded values (consistent with theory) and the model achieves 96.2% test accuracy on graph reachability.", "questions": "1. The paper notes that logits plateau in practice but grow unbounded in theory, attributed to interactions between thought-generation and prediction stages. Could you elaborate on why these interactions cause plateauing? For example, does the prediction stage’s loss implicitly regularize the thought-generation logits, or is it a numerical effect (e.g., gradient clipping in practice)?\n\n2. The model exhibits \"length generalization\" (e.g., generating c=4 continuous thoughts despite training on c≤2). Do you have evidence that this generalization relies on superposition? For example, if you perturb the superposition (e.g., zero out weights for non-optimal traces), does length generalization break?\n\n3. The theoretical analysis assumes linear attention. How might non-linear attention (e.g., softmax attention, the standard in LLMs) affect the index-matching logit’s dynamics? Would softmax attention’s normalization further bound logits, or introduce new complexities (e.g., vanishing gradients)?", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1. Mechanistic insight: The paper demystifies a core advantage of continuous CoT—superposition—by tying it to a measurable quantity (index-matching logit) and its training dynamics. This is a rare example of \"reverse-engineering\" why a latent reasoning method works, rather than just showing it works.\n\n2. Theory-experiment alignment: Experiments are not just \"proof of concept\" but directly validate theoretical predictions (e.g., logit saturation, length generalization to c=3,4 despite training only on c=1,2). This strengthens confidence in the mechanistic claims.\n\n3. Practical relevance: The focus on single demonstrations per sample (COCONUT) reflects real-world training constraints (where large CoT datasets are costly to annotate). This makes the results more applicable than methods relying on exhaustive BFS-style supervision.", "weaknesses": "1. Task generality: The analysis is limited to directed graph reachability, a synthetic task with well-defined \"reasoning traces\" (paths). It is unclear if the superposition mechanism extends to tasks with more ambiguous traces (e.g., open-domain QA, where reasoning steps are less structured). The paper should acknowledge this limitation and suggest future work on diverse tasks.\n\n2. Model scalability: The experiments use a small two-layer model. While this is standard for mechanistic work, the paper could include a brief discussion of how deeper models might interact with superposition (e.g., whether layers stack superposition or refine it) to address scalability concerns.\n\n3. Alternative mechanisms: The paper does not rule out competing explanations for superposition. For example, could the model’s weight tying (assumed for convenience) or curriculum learning (multi-stage training) also contribute to superposition? A short ablation on weight tying or curriculum learning would strengthen the causal link between COCONUT training and superposition.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "j5savjhsbr", "reviewer": "Reviewer_ohkE", "summary": "Building on recent work in continuous Chain-of-Thought, the authors seek to explain why the training procedure yields strong reasoning behavior. They decompose the model’s prediction into two stages, develop a mathematical analysis focusing on the index-matching logit, an important quantity in the model, and present experiments that support their claims.", "questions": "See weaknesses.", "limitations": "", "rating": 6, "confidence": 2, "soundness": 3, "presentation": 2, "contribution": 2, "strengths": "1. The paper formulates a clear, important problem in continuous Chain-of-Thought, explaining why such capability is enabled during training.\n2. The claims are supported by precise mathematical derivations and well-designed experiments that substantiate the analysis.", "weaknesses": "1. The writing is careful but not intuitive. As a non-expert reader in this domain, it is important to provide a more intuitive background introduction. In this paper, it is especially important to introduce index-matching logic and related concepts gradually and intuitively. \n2. The paper is focused on one single task of graph reachability. Could it be tested on more tasks and how would it be generalized to other tasks?", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "xLZ02oKv7s", "reviewer": "Reviewer_iW6r", "summary": "This work theoretically studies how the superposition mechanism in continuous CoT is learned from gradient-based methods. It derives theoretical results by formulating the training dynamics of a two-layer transformer on the graph reachability problem in two training stages of thought generation and prediction. It also reveals the importance of the bounded index-matching logit and its connection to the tradeoff between exploration and exploitation.", "questions": "Please refer to the Weaknesses.", "limitations": "", "rating": 6, "confidence": 2, "soundness": 3, "presentation": 2, "contribution": 3, "strengths": "1. This work addresses an important problem: as the previous work (Zhu et al., 2025) theoretically shows that the core superiority of continuous CoT lies in reason by superposition, this work solves a next question: how the superposition can be learned by gradient-based methods. \n\n2. The contributions of this work are sufficient as it derives multiple novel theoretic results on the gradient-based learned superposition mechanism.", "weaknesses": "1. The presentation should be improved. \n(1) The draft contains many typos and inconsistency, which should be thoroughly revised. \n(2) As the content is mostly hard to follow, more intuitive visuals may need to be added for better understanding. \n\n2. This work discusses only two-layer transformer. It would be better to discuss about the generalization to models with more layers. \n\n3. Practical implication may be supplemented. Based on the theorems proposed in this work, does one design a better training recipe for continuous CoT?", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "T6durq1vyh", "reviewer": "Reviewer_eBsL", "summary": "This paper offers a theoretical and empirical account of how the superposition principle arises in transformers under gradient-based training. Focusing on a two-layer transformer with continuous-thought decoding, the authors delineate two curriculum stages: (i) autoregressive generation of continuous thoughts and (ii) final prediction from the aggregated latent sequence. They show that, under an appropriate loss, attention logits remain bounded in Stage I, which expands the next thought. Leveraging max-margin analysis, they establish direction convergence and obtain accuracy guarantees on unseen examples.", "questions": "Please refer to the Weaknesses.", "limitations": "", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 2, "contribution": 3, "strengths": "- Theoretical rigor: Under some assumptions, the paper provides rigorous theory about how bounded attention logits and superposition naturally arise from gradient-based training. From a methodological point of view, this paper simplifies the difficulties of theoretical analysis by closing the dynamics through some assumptions.\n- Insight into reasoning processes: Through the focus on the graph reachability problem, the analysis reveals concrete relationships between loss design and the evolution of attention weights.\n- Experimental verification: The theoretical results were partially verified through some experiments. Section 5 provide the evidence that the theoretically predicted logit behaviors manifest during actual learning.", "weaknesses": "- The clarity of the article's presentation needs significant improvement. For example:\n\n  - In the third subsection of section 3, the article introduces prior knowledge about the first decoder layer (Eq. (1)) based on some related work. This should be stated as an assumption and verified through experiments.\n  - Similarly, the reparameterization introduced in Eq. (3) and Eq. (7) should also be considered as an assumption and partially verified through experiments.\n  -  In section 5.1, the explanation of $\\mu$ is confusing. More intuition or discussion is needed to explain why the chosen logit can be a proxy for $\\mu$. In contrast, the explanation in 5.2 does not have this problem.\n  - On page 5, the fourth line from the bottom considers two losses under permutation. Intuitively, the permutation invariance of the data will lead to similar properties of the parameters (as shown in Lem. 1), but the relevant insight needs to be pointed out more explicitly.\n  - The statement of Theorem 2 needs to be improved. According to the expression of $\\beta$, all parameters should be regarded as vectors, and the theorem should show the boundedness of the operator between vectors.\n\n- The plausibility of the assumptions or motivations is unclear.\n  - In Eq. (3), the initialization of $\\mu$ is setted as 0. However, it is in an inductive process. Specifically, this is step c of the course learning. Assuming $\\mu$ is 0 does not seem reasonable considering that the attention parameters have been changed during the previous training.  Additional discussion or proof is needed to complete the inductive process.\n  - As mentioned before, priors about the behavior of reparameterization and the first layer of attention need to be empirically supported.\n\n- The significance of some experiments is unclear.\n\nIn Section 5.2, the article analyzes the attention of the answer token and the reasoning token. Although it shows similar results to the first half of Theorem 3, it does not show how the behavior of accuracy changes as the direction converges.\n\n- Scalability of experiments.\nThe paper experiments on a two-layer transformer. The relevant experiments, definitions, and detailed analysis are not yet clear in their guidance for training with more layers or for attention behavior. A supplementary example with more layers would be helpful.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "8afVBy4wuR", "reviewer": "Reviewer_dwvn", "summary": "This paper focuses on the gap of whether superpositional reasoning (superposition) in continuous Chain-of-Thought (continuous CoT) naturally emerges under standard gradient training, chooses directed graph reachability, and provides a theoretical analysis of the training dynamics of a simplified two-layer Transformer, explicitly distinguishing the two phases of “thought generation” and “answer prediction” to explain the source of the mechanism; the core conclusion is that under continuous-thought training, the “index-matching logit” (which characterizes local search strength) first increases and then, under mild assumptions, remains bounded, thereby achieving an exploration–exploitation balance under uncertainty: it both leverages local structure to advance and assigns comparable weight to multiple candidate paths, so superposition arises naturally; the authors also validate this by experiments that track the growth of logits.", "questions": "Based on your theoretical understanding, can you propose any targeted insights or approaches for enhancing a model’s reasoning ability during the training process?", "limitations": "", "rating": 8, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1. Focused problem that fills a gap: the paper centers on the core mechanism of why superposition would naturally arise during training in continuous Chain-of-Thought (continuous CoT) and provides an explanation from the perspective of training dynamics, going beyond prior work that only proves the existence of feasible solutions/expressivity without discussing whether they can be learned by gradient methods. The abstract and introduction clearly define the goal and differentiated positioning.\n\n2. Solid key theoretical results: it formally compares two objectives (COCONUT vs. COCONUT-BFS), proving that under the former the index-matching logit is bounded whereas under the latter it diverges; it thus concludes that a bounded logit leads to a smoother next-step distribution, thereby retaining multiple candidate paths and promoting superposition. This conclusion offers a convincing explanation for why training does not over-commit to a single path.\n\n3. Mutual support between empirics and theory: a two-layer GPT-2 decoder trained from scratch achieves 96.2% accuracy on a subset of ProsQA; during training the “frontier vs. non-frontier” logit gap first rises and then stabilizes, consistent with the “bounded” theory; length extrapolation appears at longer steps. These observations support the chain “bounded logit → superposition/exploration–exploitation balance.”", "weaknesses": "1. Ablation and robustness analyses are somewhat limited: there is a lack of systematic ablations over factors such as number of layers/heads/width and learning rate; nor are results under different random seed settings presented (there is only a verbal note in a footnote that different random seeds were used). The dataset is limited to a subset of ProsQA. It is recommended that the authors conduct more detailed experimental validation on more related datasets and under more settings.\n\n2. (Small suggestion) To facilitate readers’ quick understanding of the paper’s contributions, it is recommended to describe an informal version of the conclusions in the first half of the paper.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (Poster)