Source-linked AI summary
Reasoning by Superposition: A Theoretical Perspective on Chain of Continuous Thought
Hanlin Zhu, Shibo Hao, Zhiting Hu, Jiantao Jiao, Stuart Russell, Yuandong Tian
TL;DR
The paper asks why continuous CoT can outperform discrete CoT on graph reachability, a problem whose continuous-thought mechanism lacks theoretical understanding. It constructs a two-layer transformer using D continuous steps and validates the resulting superposition-based search empirically. The construction solves D-diameter graphs in D steps, while discrete CoT requires O(n^2) steps in the cited constant-depth result.
Problem
Continuous CoT shows empirical gains, but its expressive power and mechanism remain theoretically underexplained for reasoning tasks such as directed graph reachability.
Method
The paper constructs a two-layer transformer whose continuous thought vectors maintain superposition states encoding multiple reachable search frontiers.
Results
D continuous-thought steps solve reachability on n-vertex graphs of diameter D, compared with O(n^2) steps for the best known constant-depth discrete-CoT result.
Takeaways & Limitations
Continuous thought supports parallel breadth-first-style search, and experiments show that this superpositional representation emerges in trained models and matches the theoretical construction.
Takeaways & Limitations
The construction uses uniform node weights, whereas training signals and internal heuristics can bias practical superposition states toward particular frontiers or optimal nodes.
Abstract
from arXiv · showhide
Large Language Models (LLMs) have demonstrated remarkable performance in many applications, including challenging reasoning problems via chain-of-thoughts (CoTs) techniques that generate ``thinking tokens'' before answering the questions. While existing theoretical works demonstrate that CoTs with discrete tokens boost the capability of LLMs, recent work on continuous CoTs lacks a theoretical understanding of why it outperforms discrete counterparts in various reasoning tasks such as directed graph reachability, a fundamental graph reasoning problem that includes many practical domain applications as special cases. In this paper, we prove that a two-layer transformer with $D$ steps of continuous CoTs can solve the directed graph reachability problem, where $D$ is the diameter of the graph, while the best known result of constant-depth transformers with discrete CoTs requires $O(n^2)$ decoding steps where $n$ is the number of vertices ($D<n$). In our construction, each continuous thought vector is a superposition state that encodes multiple search frontiers simultaneously (i.e., parallel breadth-first search (BFS)), while discrete CoTs must choose a single path sampled from the superposition state, which leads to sequential search that requires many more steps and may be trapped into local solutions. We also performed extensive experiments to verify that our theoretical construction aligns well with the empirical solution obtained via training dynamics. Notably, encoding of multiple search frontiers as a superposition state automatically emerges in training continuous CoTs, without explicit supervision to guide the model to explore multiple paths simultaneously.
1 Introductions
The section motivates continuous chain-of-thought as a response to limits in complex reasoning and develops a theoretical and empirical account of its graph-reachability advantages.
- Continuous CoT remains theoretically underexplained despite empirical gains on synthetic reasoning tasks.
- Graph reachability provides a general test problem spanning theoretical applications such as halting and practical uses such as knowledge graphs.
- D steps of continuous thought solve reachability on n-vertex graphs, where D is the diameter and D < n, versus O(n^2) steps for the best known constant-depth discrete-CoT result.
- Continuous latent vectors represent multiple search traces and frontiers simultaneously, enabling implicit parallel search, whereas discrete tokens realize one branch at a time.
- Experiments show that trained continuous-CoT models encode multiple plausible frontiers in parallel, aligning learned representations with the theoretical construction.
- Related work: The work extends prior transformer-expressivity studies, which largely focus on discrete CoT, by analyzing continuous CoT theoretically.
2 Preliminaries
The preliminaries define the transformer, positional-encoding, token-embedding, and continuous-thought machinery used in the paper’s construction.
- Tokens and embeddings: Token embeddings are partitioned into content, two buffers, and effective positional-encoding dimensions.
- Tokens and embeddings: The notation assumes orthonormal content embeddings, with U^⊤U = I_V.
- Transformer architectures: The transformer computation is organized as repeated attention and MLP layers, producing the embedding at the final sequence position.
- Transformer architectures: The transformer processes token and thought embeddings through causal self-attention and position-wise multilayer perceptrons.
- Transformer architectures: Continuous CoT appends the transformer’s output embedding directly as the next input instead of sampling a discrete token.
- Positional encodings: Positional encodings are added at each sequence position, with the construction covering sinusoidal encodings and also discussing RoPE.
3 Problem Formulations
The paper formulates directed graph reachability as a tokenized transformer task and specifies how continuous thoughts are generated after the graph prompt.
- Graph reachability: The task gives a directed graph, two candidate destinations, and a root, with exactly one candidate reachable from the root.
- Input structures: Figure 1 presents the prompt format for the graph reachability problem.
- Input structures: Each graph edge is represented by source, target, and a special edge token, followed by the question, candidates, reasoning token, and root.
- Chain of continuous thought: Continuous thoughts are generated autoregressively by feeding each transformer output back as the next sequence embedding.
- Position index: The notation indexes each continuous thought by its position after the prompt, enabling the sequence of reasoning steps to be tracked.
- The following sections evaluate continuous thought on graph reachability both theoretically and empirically.
4 Theoretical Results
The theoretical construction uses attention choosers, layered attention, and filtering to maintain continuous-thought superpositions of reachable vertices. Each step expands the represented search frontier, while the final readout measures candidate signals; practical weights may vary from the uniform idealization.
- Attention chooser: The construction uses an attention chooser to select positions based on the current token and relative offset, enabling shared parameters across input lengths.For matching tokens it attends to the position i−ℓ; otherwise it attends to the first position.
- Superposition states: Each continuous thought is the normalized superposition of all vertices reachable from the root within c steps.The result is established inductively, beginning with the root at c = 0.
- First layer attention: The first attention layer uses five attention-chooser heads to copy edge, reasoning, and answer-related values into designated buffer spaces.The heads target edge and special-token positions using different relative offsets.
- Second layer attention: The second attention layer expands the represented vertices by attending to edges whose source lies in the current superposition and adding their target nodes.This produces the next reachable set Vc+1 from Vc.
- Filtering and readout: An MLP filters noise tokens and equalizes retained node weights, while the final answer step measures candidate signals in the superposition state.The construction separates content and buffer spaces, though these can be projected into a more compact shared space.
- Discussion: In practice, superposition weights need not be uniform and may favor exact-step frontier nodes or nodes heuristically aligned with the destination.The theoretical construction maintains uniform weights, whereas training signals and internal heuristics can bias them.
5 Experiments
Experiments show that a two-layer COCONUT model nearly solves ProsQA and develops the predicted superpositional search behavior. Attention and representation analyses indicate parallel frontier expansion, while BFS-style supervision produces a similar strategy without optimal-node guidance.
- 5.2 Overall Results: COCONUT achieves near-perfect accuracy on ProsQA, whereas CoT and No CoT reach about 75%, and 12-layer CoT reaches 83%.Random guessing is 50%.
- 5.3 Visualising Latent Reasoning: Layer 1 attention copies edge endpoints onto each edge token, matching the theoretical construction.The representative attention map places nearly all mass on the source and target nodes.
- 5.3 Visualising Latent Reasoning: Layer 2 concentrates attention on reachable edges, with additional emphasis on frontier and optimal edges.The analysis groups edges as Reachable, Not Reachable, Frontier, and Optimal.
- 5.3 Visualising Latent Reasoning: Continuous thoughts are more similar to nodes within i hops, especially frontier and optimal nodes, indicating a soft parallel representation of reachable states.Frontier nodes are a subset of reachable nodes, while optimal nodes are a subset of frontier nodes.
- 5.3 Visualising Latent Reasoning: Together, the learned attention and representations instantiate the intended search: Layer 1 establishes context, Layer 2 expands the frontier, and latent vectors encode reachable states.The search pattern is consistent across multiple random seeds.
- 5.4 Exploration Priority: COCONUT-BFS matches original COCONUT at near-perfect accuracy, and both supervision methods converge to similar exploration strategies.Original COCONUT still emphasizes non-optimal frontier nodes despite training on optimal nodes.
6 Conclusions
The paper concludes that continuous CoT solves graph reachability efficiently by maintaining superposition states encoding multiple search traces, and validates this construction experimentally.
- A two-layer transformer solves reachability on an n-vertex, D-diameter graph in D continuous-thought steps, compared with O(n^2) discrete-CoT steps.
- Superposition states encode multiple search traces simultaneously, supporting the paper’s explanation of COCONUT’s reasoning capability.
- Experiments show that the theoretical construction matches solutions obtained through training dynamics.
- The paper identifies lower bounds for discrete CoT and broader continuous-space reasoning as future directions.
B Missing Proofs
The attention chooser construction uses positional encoding and token-content conditions to concentrate attention on a designated prior position or the BOS token.
- The attention chooser constructs key and query matrices so attention selects position i − ℓ when the current token matches <x>.
- Under the stated embedding and prefix assumptions, the selected attention score exceeds 1 − ε.
- When the current token does not match <x>, the construction directs attention to the BOS token as an attention sink.
- The desired position receives the maximum query-key score with gap ηε_T, yielding a softmax concentration bound.
- Choosing sufficiently large η makes the concentration bounds exceed 1 − ε and completes the proof.
B.2 Proof of Lemma 2
The proof avoids induction by exploiting autoregressive causality: later appended vectors cannot alter computations at earlier positions, so one forward pass suffices.
- Autoregressive computation ensures that appending vectors does not affect the transformer outputs at earlier positions.
- This causal property replaces separate inductive proofs for each step with a single forward-pass construction.
- The construction sets designated final positions to the answer token and proves the required hidden embedding directly.
B.3 Construction of transformer parameters
The transformer construction uses attention heads, buffer spaces, and MLPs to copy graph information, update continuous search frontiers, and produce the final prediction.
- First-layer attention: First-layer attention heads are constructed so selected attention terms exceed 1 − ε for every position.
- First-layer MLP: Each edge token attends to its corresponding source and target nodes, while subsequent MLP processing produces the first-layer representation.
- First-layer attention: The first layer copies attended token contents into separate buffer spaces using head-specific output matrices.
- Second-layer attention: The second-layer attention makes each current thought attend to edges whose source nodes occur in the current superposition.
- Final prediction: The answer token attends to the reasoning token, which stores the candidate destination nodes in its buffer space.
- Second-layer MLP: After second-layer attention, each thought contains the current frontier and its one-step reachable vertices, forming the next frontier after noise removal.
B.4 Proof of the main theorem
The proof extends the attention-chooser construction to rotary positional embeddings under explicit sequence-length and embedding assumptions. It establishes that attention can focus on a designated earlier position or the BOS token with arbitrarily high concentration.
- RoPE construction: RoPE is used because positional encoding functions only in the first attention layer, where the construction builds attention choosers.The construction omits positional, buffer, and extra embedding dimensions when they are not needed for the RoPE-based chooser.
- RoPE construction: The RoPE construction assumes M ≥ T, where T is the maximum input-sequence length, and uses ω = M^-2/d.The positional frequencies are chosen so distinct positions remain distinguishable over the permitted sequence length.
- Attention chooser: The construction applies to input embeddings that are either equal to <x>'s embedding or orthogonal to it, with the first ℓ tokens orthogonal to <x>.Under these conditions, the resulting query and key scores satisfy the attention-chooser guarantee for every sequence position.
- Attention chooser: For a token <x>, the attention chooser targets position i−ℓ when the current token is <x> and otherwise targets the BOS token.The proof uses a positive score gap and sufficiently large η to make both attention lower bounds exceed 1−ε.
C.1 Dataset
This section reports ProsQA dataset statistics and evaluates whether COCONUT learns the intended superpositional search behavior. The reported inner-product results are consistent across three random-seed runs.
- Dataset: ProsQA dataset statistics are reported in Table 4, with numbers averaged over problem instances.The supplied passage identifies the table and averaging procedure but does not provide the individual statistics.
- Empirical evaluation: COCONUT is evaluated across three random seeds using mean inner products between continuous thoughts and each node group.The measurements follow the setting in Figure 6 and are reported in Table 5.
- Empirical evaluation: The inner-product results are consistent across multiple runs.Each run takes about 24 hours on two Nvidia A100 80GB GPUs.