Training Nonlinear Transformers for Chain-of-Thought Inference: A Theoretical Generalization Analysis

TL;DR

This work theoretically analyzes training nonlinear Transformers to acquire Chain-of-Thought (CoT) reasoning, addressing the gap between empirical CoT success and gradient-based generalization guarantees. It studies a one-layer, single-head attention Transformer trained on supervised CoT prompts and proves that CoT capability can be learned with context-size $l_{tr}=\Omega(\alpha^{-1})$ and iterations scaling as $\alpha^{-2}$, while showing that the attention concentrates on context columns sharing the same TSR pattern as the query. It further derives CoT generalization guarantees under noisy and distribution-shifted testing data, and contrasts CoT with In-Context Learning (ICL), showing that CoT does not require a dominant correct input-label ratio in prompts whereas ICL may, depending on conditions. The results are supported by numerical experiments on synthetic data and extended to a multi-layer setting, offering theoretical insight into when CoT can outperform ICL and guiding prompt design under realistic noise conditions.

Abstract

Chain-of-Thought (CoT) is an efficient prompting method that enables the reasoning ability of large language models by augmenting the query using multiple examples with multiple intermediate steps. Despite the empirical success, the theoretical understanding of how to train a Transformer to achieve the CoT ability remains less explored. This is primarily due to the technical challenges involved in analyzing the nonconvex optimization on nonlinear attention models. To the best of our knowledge, this work provides the first theoretical study of training Transformers with nonlinear attention to obtain the CoT generalization capability so that the resulting model can inference on unseen tasks when the input is augmented by examples of the new task. We first quantify the required training samples and iterations to train a Transformer model towards CoT ability. We then prove the success of its CoT generalization on unseen tasks with distribution-shifted testing data. Moreover, we theoretically characterize the conditions for an accurate reasoning output by CoT even when the provided reasoning examples contain noises and are not always accurate. In contrast, in-context learning (ICL), which can be viewed as one-step CoT without intermediate steps, may fail to provide an accurate output when CoT does. These theoretical findings are justified through experiments.

Figures (8)

• Figure 1: An example of a two-step inference
• Figure 2: Concentration of attention weights for CoT inference.
• Figure 3: CoT testing error with different (A) $\alpha'$ (B) $\tau^f$ (C) $\rho^f$.
• Figure 4: ICL testing error with different (A) $\alpha'$ (B) $\tau_o^f$ (C) $\rho_o^f$.
• Figure 5: Comparison between CoT and ICL w./w.o. Condition \ref{['cond: icl']}

Theorems & Definitions (22)

• Example 1
• Theorem 1
• Definition 1
• Theorem 2: CoT generalization
• Remark 1
• Theorem 3: ICL generalization
• Remark 2: Comparison between CoT and ICL
• Proposition 1
• Lemma 1: Multiplicative Chernoff bounds, Theorem D.4 of MRT18
• Definition 2: V10