Unveiling the Statistical Foundations of Chain-of-Thought Prompting Methods

TL;DR

This work offers a rigorous statistical framework for chain-of-thought prompting, casting pretrained LLMs plus CoT prompts as Bayesian model averaging over latent task concepts. It derives a decomposition of CoT error into pretraining and prompting components, with explicit exponential-rate bounds for prompting as demonstrations increase and PAC-Bayes-based analyses for pretraining. The authors connect transformer attention to BMA in a simplified model and generalize the theory to CoT variants like SC-COT, ToT, and SI, complemented by empirical validations on synthetic tasks. The results provide both theoretical guarantees and practical insights into when and why CoT improves multi-step reasoning, guiding future prompt design and analysis.

Abstract

Chain-of-Thought (CoT) prompting and its variants have gained popularity as effective methods for solving multi-step reasoning problems using pretrained large language models (LLMs). In this work, we analyze CoT prompting from a statistical estimation perspective, providing a comprehensive characterization of its sample complexity. To this end, we introduce a multi-step latent variable model that encapsulates the reasoning process, where the latent variable encodes the task information. Under this framework, we demonstrate that when the pretraining dataset is sufficiently large, the estimator formed by CoT prompting is equivalent to a Bayesian estimator. This estimator effectively solves the multi-step reasoning problem by aggregating a posterior distribution inferred from the demonstration examples in the prompt. Moreover, we prove that the statistical error of the CoT estimator can be decomposed into two main components: (i) a prompting error, which arises from inferring the true task using CoT prompts, and (ii) the statistical error of the pretrained LLM. We establish that, under appropriate assumptions, the prompting error decays exponentially to zero as the number of demonstrations increases. Additionally, we explicitly characterize the approximation and generalization errors of the pretrained LLM. Notably, we construct a transformer model that approximates the target distribution of the multi-step reasoning problem with an error that decreases exponentially in the number of transformer blocks. Our analysis extends to other variants of CoT, including Self-Consistent CoT, Tree-of-Thought, and Selection-Inference, offering a broad perspective on the efficacy of these methods. We also provide numerical experiments to validate the theoretical findings.

Figures (18)

• Figure 1: An illustration of cot and vanilla icl. Figure (a) shows the cot prompt and the corresponding output of ChatGPT (GPT-3.5-turbo-16k). The intermediate reasoning is shown in red. The output of ChatGPT follows the pattern in the prompt, which consists of a reasoning step, followed by the desired answer. Figure (b) shows the result of using the corresponding vanilla icl prompt, which includes of two input-output pairs. In this case, ChatGPT fails to provide the correct answer. Figure (c) illustrates a general pipeline of cot prompting with $n$ demonstration examples. Each example includes an input question, $H-1$ intermediate reasoning steps, and the final answer.
• Figure 2: An illustration of the multi-step latent-variable model defined in \ref{['eq:latent_var_model']}. According to this graphical model, for any $h \geq 1$, each step $z_h^i$ of $i$-th example depends on the previous steps $\{ z_{\ell}^i \}_{ \ell < h}$ and the hidden concept $\theta^*$.
• Figure 3: An instantiation of the model in \ref{['eq:latent_var_model']} in the context of arithmetic problems. Here $\theta^*$ stands for "solving an arithmetic problem with intermediate steps", and $z_0^i$ describes a new arithmetic problem generated independently from any other demonstrations. Then each subsequent step, $z_1^i$, $z_{2}^i$, and $y^i$, depends on both the previous steps and the hidden task $\theta^*$.
• Figure 4: An illustration of how cot is represented by a two-layer transformer. Layer 1 serves as a copy head, which copies the previous steps $\{z_j^i\}_{j=1}^{h-1}$ to the current position $z_h^i$. Next, the feature mappings $v$ and $k$ map the outputs of Layer 1 to values and keys, respectively. During the generation of $z_{h+1}^\mathrm{test}$, the attention mechanism takes in key and value matrices $\mathtt{keys}$ and $\mathtt{values}$ from the demonstrations to predict the result for query $q_h^\mathrm{test}$, where $\mathtt{keys}$ and $\mathtt{values}$ are formed by stacking $\{k_h^i\}_{i=1,h=1}^{n, H}$ and $\{v_h^i\}_{i=1,h=1}^{n, H}$, respectively. Note that $\mathtt{keys}$ and $\mathtt{values}$ do not contain keys and values computed from the generated reasoning steps $\{z_j^\mathrm{test}\}_{j=1}^h$. We can achieve this by masking out the corresponding positions.
• Figure 5: An illustration of the sc-cot prompting method. This method creates the final answer $y_K^*$ based on two steps. First, we sample $K$ i.i.d. reasoning paths $\{ z_{0:H}^{\mathrm{test},i} \}_{i=1}^K$ given the cot prompt, and then report $y_K^*$ by a majority vote based on $\{ y^{\mathrm{test},i } = z_{H}^{\mathrm{test},i} \}_{i=1}^K$.

Theorems & Definitions (40)

• Lemma 4.1
• Proposition 4.2
• Lemma 5.2: cot Error Decomposition
• Definition 5.3: Equivalence Classes over $\Theta$
• Theorem 5.5
• Theorem 5.7
• Proposition 5.11: Statistical Error of ToT
• Corollary 5.13: Sample Complexity of Selection-Inference
• Proposition 5.14: CoT Outperforms Vanilla ICL
• Proposition 6.3: Pretraining Error Bound