PaperlayerView exact revision on arXiv

Source-linked AI summary

Towards Revealing the Mystery behind Chain of Thought: A Theoretical Perspective

Guhao Feng, Bohang Zhang, Yuntian Gu, Haotian Ye, Di He, Liwei Wang

arXiv:2305.15408v5cs.LGcs.CCcs.CLstat.ML

TL;DR

The paper asks why Chain-of-Thought prompting improves LLMs on mathematical and decision-making tasks despite unclear underlying mechanisms. Using circuit-complexity arguments, constructive Transformer proofs, and experiments, it shows that CoT supports solutions that direct prediction cannot generally produce, while identifying unresolved questions about prompt triggering, scaling, and generalization.

  • Problem

    The paper examines why CoT improves LLM performance on mathematical and reasoning tasks and whether direct answer generation has inherent limitations.

  • Method

    The paper uses circuit complexity theory, constructive autoregressive-Transformer proofs, and experiments on mathematical and Dynamic Programming tasks.

  • Results

    Experiments show direct prediction always fails, whereas CoT models learn complete solutions and generalize to longer inputs; theory proves constant-size autoregressive Transformers suffice for arithmetic and equations.

  • Takeaways & Limitations

    CoT’s step-by-step generation provides a theoretically sufficient route for these tasks where bounded-depth direct prediction requires prohibitively large models or is impossible in general.

  • Takeaways & Limitations

    The paper does not theoretically explain how prompts trigger CoT, how model size affects CoT, or how CoT solutions generalize.

Abstract

from arXiv · show

Recent studies have discovered that Chain-of-Thought prompting (CoT) can dramatically improve the performance of Large Language Models (LLMs), particularly when dealing with complex tasks involving mathematics or reasoning. Despite the enormous empirical success, the underlying mechanisms behind CoT and how it unlocks the potential of LLMs remain elusive. In this paper, we take a first step towards theoretically answering these questions. Specifically, we examine the expressivity of LLMs with CoT in solving fundamental mathematical and decision-making problems. By using circuit complexity theory, we first give impossibility results showing that bounded-depth Transformers are unable to directly produce correct answers for basic arithmetic/equation tasks unless the model size grows super-polynomially with respect to the input length. In contrast, we then prove by construction that autoregressive Transformers of constant size suffice to solve both tasks by generating CoT derivations using a commonly used math language format. Moreover, we show LLMs with CoT can handle a general class of decision-making problems known as Dynamic Programming, thus justifying its power in tackling complex real-world tasks. Finally, an extensive set of experiments show that, while Transformers always fail to directly predict the answers, they can consistently learn to generate correct solutions step-by-step given sufficient CoT demonstrations.

1 Introduction

The paper investigates why Chain-of-Thought prompting improves mathematical and reasoning performance, focusing on its mechanisms and the limits of direct answer generation. It combines theoretical expressivity results with experiments on mathematical and dynamic-programming tasks.

  • 1 Introduction: CoT improves generated-answer correctness on arithmetic and reasoning tasks by eliciting intermediate derivations.The paper notes that this can be triggered by special phrases or few-shot CoT demonstrations.
  • 1 Introduction: The paper asks whether LLMs have inherent limitations in directly answering math and reasoning questions and why CoT succeeds.
  • 1 Introduction: Bounded-depth Transformers cannot directly solve arithmetic and linear-equation tasks without super-polynomial model growth in input length.The paper attributes this limitation to the tasks’ parallel complexity rather than their serialized computational cost.
  • 1 Introduction: Constant-size autoregressive Transformers can solve both mathematical tasks by generating intermediate derivations step by step.The constructions use a commonly used mathematical language format and rely on CoT’s recursive nature.
  • 1 Introduction: CoT enables LLMs to generate complete solution chains for general Dynamic Programming problems, although direct generation is impossible in general.The paper uses Context-Free Grammar Membership Testing as a counterexample for bounded-depth polynomial-size Transformers.
  • 1 Introduction: Direct answer prediction always fails experimentally, while CoT models learn entire solutions and generalize to longer input sequences.The experiments cover arithmetic, linear equations, longest increasing subsequence, and edit distance; direct accuracy is mostly below 60%.

2 Preliminary

The preliminary section describes autoregressive Transformers as sequence-processing architectures that generate later tokens from earlier ones, and presents CoT as stepwise intermediate reasoning. It also specifies causal masking and the model components used for this generation process.

  • 2 Preliminary: An autoregressive Transformer processes input tokens and generates tokens for subsequent positions.The architecture first embeds tokens with positional information, then applies Transformer blocks.
  • 2 Preliminary: Each Transformer block combines multi-head self-attention and a feed-forward network with residual connections.The standard setting uses an H-head softmax attention followed by a two-layer pointwise FFN.
  • 2 Preliminary: Causal masking ensures position i attends only to preceding positions j ≤ i, supporting autoregressive generation.
  • 2 Preliminary: Next-token prediction is repeated by appending each generated token until an End-of-Sentence token appears.
  • 2 Preliminary: CoT prompting induces intermediate reasoning steps before the final answer when direct generation struggles on challenging math or reasoning tasks.

3 CoT is the Key to Solving Mathematical Problems

The paper shows that bounded-depth Transformers face fundamental limitations on directly solving arithmetic and linear-equation tasks, while constant-size autoregressive Transformers can generate correct solutions through CoT derivations.

  • Problem formulation: Arithmetic CoT repeatedly performs an intermediate operation while copying unrelated expression items, reducing the expression step by step.The task evaluates expressions containing numbers, operators, brackets, and an equal sign.
  • Problem formulation: Equation CoT follows Gaussian elimination, eliminating one variable from all but one equation until each equation contains a single variable.After m − 1 elimination steps, the system is solved under the assumption that a unique solution exists.
  • Impossibility results: Without CoT, bounded-depth Transformers cannot solve arithmetic evaluation or linear equations unless model size grows super-polynomially with input length.The analysis uses finite-field formulations with integers modulo a prime p.
  • Constructive results: A constant-size autoregressive Transformer with depth 5 and 5 heads per layer can generate arithmetic CoT solutions for all inputs in Arithmetic(n, p).Its hidden size is independent of n, and parameter values are bounded by O(poly(n)).
  • Constructive results: The constructions implement CoT generation using Transformer operations such as conditional COPY and MEAN, treating autoregressive generation as a loop.The proofs also identify softmax attention, multi-head attention, feedforward networks, and residual connections as key components.

4 CoT is the Key to Solving General Decision-Making Problems

The paper extends its CoT analysis to Dynamic Programming, showing that constant-size autoregressive Transformers can solve suitably bounded DP problems by generating intermediate subproblem results, whereas direct prediction can be impossible for CFG membership testing.

  • DP framework: Dynamic Programming decomposes complex decision-making problems into sequentially solved subproblems whose results are reused by later subproblems.The framework is characterized by a state space, transition function, and aggregation function.
  • DP framework: The restricted DP setting assumes each state depends on finitely many input tokens and finitely many previous states.The transition can therefore be expressed through functions determining input access, predecessor states, and the resulting DP value.
  • CoT construction: The CoT generation format lists the input sequences, then state–value pairs in a feasible topological order, followed by the final answer.Each pair (i, dp(i)) is generated jointly as a vector.
  • Theoretical results: Under Assumptions 4.2–4.5, an autoregressive Transformer with constant depth, hidden dimension, and head count can generate correct answers for DP problems of input length n.The parameters are bounded by O(poly(n)).
  • Impossibility results: Without CoT, log-precision bounded-depth Transformers cannot solve CFG Membership Testing for all input strings at some lengths, assuming TC0 ≠ NC1.The lower bound applies when hidden dimension is at most polynomial in sequence length.
  • Implications: The paper concludes that CoT substantially improves Transformer expressivity for tasks intrinsically hard for well-parallelized computation models.Universal CFG Membership Testing is identified as a P-complete problem.

5 Experiments

Experiments compare CoT-trained and direct-answer Transformers across four tasks, difficulty levels, data corruption, and longer inputs. CoT models learn near-perfect solutions, remain robust to corrupted demonstrations, and extrapolate beyond training lengths.

  • Tasks and datasets: The experiments evaluate Arithmetic, Equation, longest increasing subsequence, and edit distance using datasets with increasing difficulty.Difficulty varies by operator count, variable count, sequence length, or string length.
  • Main results: 3-layer Transformers with CoT achieve near-perfect accuracy across all tasks and difficulty levels.Figure 2 reports test accuracy by difficulty level and repeats each experiment five times with error bars.
  • Main results: Direct-answer Transformers typically fail even at depths L = 3, 4, 5, whereas CoT models consistently outperform them.The direct datasets remove intermediate derivations from the CoT datasets.
  • Robustness to data quality: 3-layer Transformers exceed 95% accuracy when 30% of arithmetic demonstrations omit an intermediate CoT step and include a single-token corruption.This experiment tests robustness to missing and corrupted training data.
  • Length extrapolation: A 3-layer arithmetic model trained on 1–15 operators performs well on test expressions containing 16, 17, or 18 operators.The experiment tests length extrapolation to sequences longer than those used in training.

6 Related Work

Related work studies Transformer expressivity, computational power, and in-context learning. This paper builds on those directions while examining CoT through model-capacity and computation perspectives.

  • Function approximation: Earlier work established that sufficiently large Transformers can universally approximate arbitrary continuous sequence-to-sequence functions.This line of research focuses on function approximation.
  • Relation to this paper: The paper’s related-work position connects Transformer expressivity and computation studies with its analysis of CoT-enabled mathematical and decision-making tasks.The supplied passages establish these as neighboring research directions rather than reporting a new comparison.
  • Computational power: Some studies show standard and looped Transformers are Turing-complete, but these results rely on infinite precision.The paper characterizes that assumption as unrealistic for practical scenarios.
  • In-context learning: Recent research demonstrates that autoregressive Transformers can in-context learn function classes and algorithms including linear regression, gradient descent, and Bayesian inference.These results examine Transformer power from an in-context learning perspective.

7 Limitations and Future Directions

The paper’s scope is bounded by unresolved questions about prompt-triggered CoT generation, model scaling, theoretical generalization, and learning from limited demonstrations. Its technical framework uses finite-precision Transformers and formal CoT task definitions.

  • Limitations and future directions: The paper does not explain how specific prompts trigger CoT generation or how prompt-output relations can be used to harness LLMs.These questions are identified as future work.
  • Limitations and future directions: The theoretical role of model size in CoT remains unresolved despite empirical observations that scaling improves CoT ability.The paper leaves this relationship for future theoretical study.
  • Limitations and future directions: The paper mainly studies expressivity rather than theoretically analyzing generalization from CoT demonstrations, including out-of-distribution length extrapolation.The authors identify theoretical generalization as an important future direction.
  • Limitations and future directions: The practical question of learning CoT solutions from limited demonstrations or purely direct datasets remains open.The authors explicitly leave these settings as future work.
  • Technical scope and assumptions: The analysis adopts finite O(log n)-bit precision and uses finite fields such as Zp to make arithmetic and equation tasks well-defined.The log-precision setting limits each neuron to O(log(n))-bit information and avoids infinite-precision assumptions.
  • Technical scope and assumptions: The paper formally defines CoT derivation formats for arithmetic expressions and linear equations before proving their Transformer constructions.These definitions support the task-specific theoretical analysis.

C.1 Technical lemmas for MLP

The technical lemmas show that small GeLU MLPs and causal attention layers can approximate the basic arithmetic, selection, lookup, copying, and averaging operations used in the constructions.

  • MLP operations: Two-layer GeLU MLPs approximate scalar multiplication with hidden dimension 4 and polynomially bounded parameters.For inputs a,b ∈ [−M,M], the approximation error is at most ϵ and parameter norms are O(poly(M, 1/ϵ)).
  • MLP operations: Two-layer GeLU MLPs can approximate arbitrary linear transformations with hidden dimension 2d2 and polynomially bounded parameter norms.The approximation holds uniformly with ℓ∞ error at most ϵ.
  • MLP operations: A two-layer GeLU MLP implements conditional selection between vectors when the scalar selector is separated from zero by α.The hidden dimension is 2d + 2, and the approximation applies for bounded vector inputs and selectors outside (−α, α).
  • MLP operations: Two-layer GeLU MLPs represent discrete lookup tables over one-hot inputs while controlling perturbation error.The construction supports tables of size d^k with parameter norms O(poly(k, 1/ϵ)).
  • Attention operations: Under Assumption C.6, one causal attention head approximates COPY and MEAN operations with embedding size O(d).The parameter norms are O(poly(M, 1/δ, log(n), log(1/ϵ))).
  • Role in the construction: The lemmas provide the basic operations used to construct constant-size autoregressive Transformer solutions for arithmetic tasks.The paper states that such Transformers can evaluate arithmetic expressions with CoT but cannot do so without CoT.

D.1 Proof of Theorem 3.3

The proof constructs a constant-size autoregressive Transformer that generates arithmetic Chain-of-Thought solutions. Its layers collect positional and syntactic information, copy relevant tokens, and perform intermediate calculations.

  • Construction: Theorem D.1 gives a depth-5 autoregressive Transformer with hidden size O(poly(p)), independent of input length n, that generates arithmetic CoT solutions.The construction uses five heads per layer and parameters bounded by O(poly(n)).
  • Proof intuition: The proof relies on fixed local context: the Transformer decides at each position whether to copy a token or perform a calculation.The final construction tolerates O(1) output error while keeping intermediate errors below a chosen epsilon with parameters O(poly(n,1/epsilon)).
  • Layer 1: The first layer counts equal signs and copies the position of the latest equal sign using separate attention heads.Additional MLPs derive positional quantities used by later layers.
  • Layer 2: The second layer computes distances to equal signs and counts equal signs in strictly preceding tokens.These quantities prepare the model to identify the next reducible arithmetic expression.
  • Layer 3: The third layer copies five context tokens, checks whether the current expression is evaluable, and computes its result when appropriate.Multi-head attention performs the five copies in parallel, while an MLP applies the local evaluation rules.

D.2 Proof of Theorem 3.1

The proof shows that direct arithmetic evaluation is hard for bounded-depth, polynomial-size log-precision Transformers under TC0 ≠ uniform NC1. It establishes hardness by reducing Boolean Formula Evaluation to arithmetic-expression evaluation.

  • Hardness framework: The proof assumes TC0 ≠ uniform NC1 and uses this separation to establish an impossibility result for direct arithmetic evaluation.The target models have fixed depth, polynomial hidden dimension, and log precision.
  • Conclusion: Theorem D.4 concludes that for any fixed depth L and polynomial Q, some input size defeats every log-precision Transformer with hidden dimension at most Q(n).Thus direct arithmetic solving cannot be achieved by this bounded-depth polynomial-size model family under the stated assumption.
  • Reduction: Boolean Formula Evaluation is reduced to arithmetic expressions while preserving the formula's Boolean result.Negation maps to 1−φ, conjunction to multiplication, and disjunction to 1−(1−φ1)×(1−φ2).
  • Reduction: The translation has length O(|s|) and can be implemented by uniform TC0 circuits.The construction rewrites Boolean operators and inserts brackets using TC0-computable operations.

E.1 Proof of Theorem 3.4

The proof constructs a constant-size autoregressive Transformer for systems of linear equations with CoT. Its layers track equations, CoT steps, variables, and elimination choices to implement Gaussian elimination token by token.

  • Construction: Theorem E.1 gives a depth-4 autoregressive Transformer with hidden size O(poly(p)), independent of the variable bound m, that generates equation CoT solutions.The construction uses five heads per layer and parameters bounded by O(poly(m)).
  • Implication: The proof therefore shows that CoT enables autoregressive Transformers to solve linear equations even though direct answer prediction is impossible for the bounded-depth polynomial-size family.The direct impossibility result is stated under TC0 ≠ NC1.
  • Layer 1: The first layer counts equations and CoT-step markers and determines the number of variables in the system.Attention heads use MEAN and COPY operations to recover these quantities from preceding tokens.
  • Layer 2: The second layer identifies an equation with a nonzero coefficient for the variable eliminated at the current CoT step.It records whether that equation will be used in the next elimination step and computes auxiliary context.
  • Layers 3–4: Later layers copy the relevant equation and determine the next output tokens, including numbers, variables, separators, and elimination results.The construction uses attention for retrieval and MLPs for conditional selection and auxiliary arithmetic.

F.2 Recurrent Neural Network(RNN)

This section contrasts recurrent models with the paper's Transformer results and describes dynamic-programming formulations for several problems. RNNs require growing hidden dimension to generate the same CoT format, while the DP formulations fit the Transformer framework.

  • RNN comparison: The RNN lower bound follows because the input is compressed into a finite hidden state, limiting the number of distinct output sequences the model can generate.The proof applies the pigeonhole principle to the number of possible arithmetic and equation solution sequences.
  • Scope boundary: The authors caution that most Transformer theorems do not generalize directly to RNNs because the architectures differ substantially.They also note that relatively little prior work studies CoT in RNNs theoretically or empirically.
  • Dynamic programming: The DP section formulates CFG Membership Testing through a CYK-style state space whose final answer is dp(|V|,0,n,n,S,m).The state tracks iteration, substring boundaries, split position, nonterminal, and rule index.
  • Dynamic programming: The paper concludes that LIS, edit distance, and CFG Membership Testing satisfy the framework's assumptions and can be solved by autoregressive Transformers with CoT.The CFG formulation addresses the input-dependent terms by extracting them before retrieving the relevant DP states.

G.2 Proof of Theorem 4.7

The proof constructs a constant-size autoregressive Transformer that solves general Dynamic Programming problems by generating intermediate states and aggregating their values.

  • A constant-depth, constant-width autoregressive Transformer solves every DP input sequence of length at most n with parameters bounded by O(poly(n)).
  • The construction embeds tokens, separators, positions, input values, DP states, DP values, and answer flags.
  • The first blocks copy separator positions and compute the next state from the current state and problem size.
  • A later block copies required input tokens and prior DP values, then computes the next DP value.
  • The final block aggregates DP values using max, min, or sum-like operations and generates the answer conditionally.

G.3 Proof of the Theorem 4.8

The section proves a conditional impossibility result for CFG Membership Testing using the problem’s NC1-hardness and the assumption TC0 ≠ NC1.

  • Assuming TC0 ≠ NC1, some context-free language cannot be recognized by fixed-depth, polynomial-width, log-precision autoregressive Transformers at all lengths.
  • The proof reduces the claim to an NC1-complete automaton-acceptance problem, which is a special case of CFG Membership Testing.

H.1 Datasets

The experiments use synthetic datasets of increasing difficulty for arithmetic, equations, LIS, and edit distance, with CoT outputs representing intermediate solution procedures.

  • Arithmetic Expression: Arithmetic expressions are generated by decomposing sampled answers with sequential operators, and CoT reverses that generation process.
  • Linear Equation: Linear-equation examples use uniquely solvable systems, with CoT showing Gaussian elimination steps before the variable assignments.
  • Longest Increasing Subsequence: LIS inputs concatenate increasing subsequences, randomly insert numbers, and use the DP array plus final answer as the CoT solution.
  • Edit Distance: Edit-distance datasets include both randomly generated string pairs and corrupted copies, covering large and small edit distances.
  • The experiments use minGPT with shared architectural and optimization settings across direct, CoT, and length-extrapolation evaluations.

I Robustness in Training Data Quality

The robustness experiments test arithmetic CoT training under missing and corrupted intermediate steps, finding high accuracy despite substantial data imperfections.

  • γ controls arithmetic-data corruption: γ = 0.1 omits 10% of intermediate steps and single-token-corrupts 10% of the remainder.
  • The robustness evaluation uses arithmetic problems with 10 operators and reports results in Table 1.
  • As corruption and omission increase, accuracy decreases slightly, while remaining high despite substantial distortions.
Loading 2305.15408v5…