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

2025-09-27

AI Summary

This work investigates how the superposition mechanism naturally emerges in chain-of-thought (CoT) reasoning under gradient-based training. Method: We propose a theoretical framework for two-stage training—thought generation and final prediction—and analyze the logit dynamics of index-matching tokens, revealing their “increase-then-saturate” behavior as the root cause of superposition. Using a simplified two-layer Transformer on a directed graph reachability task, we empirically trace logit evolution and quantify probability mass assigned to multiple valid reasoning paths. Contribution/Results: We formally characterize how exploration–exploitation trade-offs drive weights across diverse reasoning paths toward near-equality. Under mild assumptions, we prove that superposition necessarily emerges during training. Our analysis provides novel insights into the interpretability and robustness of continuous reasoning, establishing a principled foundation for understanding emergent multi-path inference in CoT models.

Technology Category

• Knowledge Representation and Reasoning: Computational Complexity of Reasoning
• Cognitive Modeling & Cognitive Systems: Conceptual Inference and Reasoning
• Reasoning under Uncertainty: Other Foundations of Reasoning under Uncertainty

Application Category

• Graph Algorithms and Modeling for the Web: Foundation models and LLMs for Web-related graphs
• Search and Retrieval-Augmented AI: Web learning to rank, online learning, and counterfactual learning for ranking
• Semantics and Knowledge: Methods to enhance, augment, integrate or synergize semantic models such as knowledge graphs and LLMs

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.

Problem

• Analyzing how superposition emerges during transformer training dynamics
• Investigating training stages for continuous chain of thought reasoning
• Understanding how gradient-based methods learn superposition mechanisms

Innovation

• Analyzes training dynamics of two-layer transformer on graph reachability
• Reveals superposition emerges through thought-generation and prediction stages
• Shows bounded index-matching logit balances exploration and exploitation