Sample Complexity and Representation Ability of Test-time Scaling Paradigms

Baihe Huang Shanda LiTianhao WuYiming YangAmeet TalwalkarKannan RamchandranMichael I. JordanJiantao Jiao Thanks: baihe_huang@berkeley.edu. Affiliation: University of California, Berkeley Affiliation: University of California, Berkeley Affiliation: University of California, Berkeley Affiliation: University of California, Berkeley Affiliation: University of California, Berkeley Affiliation: Carnegie Mellon University Affiliation: Carnegie Mellon University Affiliation: Carnegie Mellon University

Abstract

Test-time scaling paradigms have significantly advanced the capabilities of large language models (LLMs) on complex tasks. Despite their empirical success, theoretical understanding of the sample efficiency of various test-time strategies—such as self-consistency, best-of- $n$ , and self-correction—remains limited. In this work, we first establish a separation result between two repeated sampling strategies: self-consistency requires $\Theta(1/\Delta^{2})$ samples to produce the correct answer, while best-of- $n$ only needs $\Theta(1/\Delta)$ , where $\Delta<1$ denotes the probability gap between the correct and second most likely answers. Next, we present an expressiveness result for the self-correction approach with verifier feedback: it enables Transformers to simulate online learning over a pool of experts at test time. Therefore, a single Transformer architecture can provably solve multiple tasks without prior knowledge of the specific task associated with a user query, extending the representation theory of Transformers from single-task to multi-task settings. Finally, we empirically validate our theoretical results, demonstrating the practical effectiveness of self-correction methods.

1 Introduction

Over the past several years, Large Language Models (LLMs) have witnessed remarkable advances, achieving unprecedented performance across a broad spectrum of application [12, 13, 20]. Driven by the paradigm of chain-of-thought (CoT) reasoning [87], the outputs of LLMs have not only grown in length but also in structural complexity. In particular, recent studies have demonstrated that scaling up computational resources during test time significantly enhances the problem-solving capabilities LLMs—a phenomenon termed as the test-time scaling law [11, 89, 36, 66]. Various methods have been proposed to effectively utilize additional test-time compute, including self-consistency [84, 11, 63, 17], best-of- $n$ sampling [42, 77, 62, 70, 73], Monte Carlo Tree Search (MCTS) [80, 101, 31, 83, 15, 56], and self-correction [59, 88, 18, 35, 100, 48]. Powered by test-time scaling paradigms, several reasoning models, such as OpenAI-o1 [65] and Deepseek-R1 [24], have achieved remarkable success in many complex tasks [34, 21, 38, 75, 22, 41, 97].

Despite these empirical advancements, the theoretical foundations of test-time scaling remain underdeveloped. While recent progress has been made in understanding the expressiveness and learnability of chain-of-thought reasoning [29, 61, 53, 44], two fundamental challenges remain unresolved:

1.

Many test-time scaling approaches rely on repeated sampling from the same LLM to select a final answer [84, 11, 42, 77, 63, 17, 91, 46, 62, 70, 73]. Two dominant paradigms are: self-consistency, which marginalizes reasoning paths and selects the most frequent answer; and best-of- $n$ , which chooses the answer with the highest reward score. However, a rigorous understanding of their sample complexities is lacking. This raises the first question:

What is the sample complexity of repeated sampling methods,
particularly self-consistency and best-of- $n$ ?
2.

Theoretical analyses of Transformers’ expressiveness have largely focused on their ability to represent individual tasks [95, 8, 9, 25, 68, 26, 27, 54, 3, 102, 94, 5, 7, 32, 81, 6, 64, 51, 32, 3, 6, 82, 57, 86, 60, 55], while the ability of Transformers to express multiple tasks at the same has been under-studied. In contrast, practical LLMs are expected to perform across diverse tasks at inference time—often using more tokens and computation than theory accounts for [19]. This gap in theory limits our understanding of test-time scaling approaches that go beyond CoT, such as self-correction [59, 88, 18, 35, 100, 48] which uses reward information. As a result, we are motivated to pose the second central question:

How can we characterize the expressiveness under test-time scaling methods,
especially in multi-task settings?

Our Contributions.

This work addresses the challenges outlined above through two key contributions. First, we analyze the sample complexity of two prominent decoding strategies: self-consistency and best-of- $n$ , in terms of the probability gap between the most likely (correct) and the second most likely model outputs. Our results reveal a fundamental separation in sample efficiency that highlights the advantage of the best-of- $n$ approach.

Proposition 1.1 (Informal statement of Theorem 3.1 and Theorem 3.2).

Let $\Delta\in(0,1)$ denote the difference between the Transformer’s probability of producing the correct answer and the probability of the second most likely answer. Then, self-consistency requires $\Theta(1/\Delta^{2})$ samples to reliably produce the correct answer, whereas best-of- $n$ achieves the same with only $\Theta(1/\Delta)$ samples.

Proof Sketch. For best-of- $n$ , correctness is achieved if the correct answer appears at least once among $n$ independent samples. Since the correct answer occurs with probability at least $\Delta$ , we have:

\displaystyle\mathbb{P}(\text{Best-of-}n\text{ outputs correct answer})\geq 1-(1-\Delta)^{n}.

To ensure high probability, it suffices to take $n\asymp 1/\Delta$ .

In contrast, self-consistency relies on the correct answer being the most frequently sampled response. Let $n_{1}$ and $n_{2}$ be the counts of the correct and second most likely answers among $n$ samples, respectively. Using the Berry-Esseen theorem, the difference

\displaystyle X=\frac{n_{1}-n_{2}-n\Delta}{\sqrt{n}}

approximately follows a normal distribution with constant mean and variance. To ensure $n_{1}>n_{2}$ with high probability, we require $\mathbb{P}(X>-\Delta\sqrt{n})\approx 1$ , or equivalently $n\asymp 1/\Delta^{2}$ . ∎

Second, we investigate Transformer’s capacity for self-correction. We demonstrate that a Transformer equipped with verifier feedback at test time can implement online learning algorithms over a pool of expert models, enabling it to adaptively identify the most suitable expert and ultimately generate a response that maximizes the reward. This process is illustrated in Figure 1: given the user query (e.g. solve the PDE $\frac{1}{c(x)^{2}}\frac{\partial^{2}u}{\partial t^{2}}-\Delta u=0$ in $\Omega\times(0,T)$ with some boundary conditions), the Transformer $f$ autoregressively generates a sequence of actions (e.g., selecting the sixth expert) and responses (e.g., constructing and applying a spectral method solver), conditioned on the history of previous action-response pairs and their corresponding rewards (e.g., solution error). Notably, this process relies solely on the Transformer $f$ —whose architecture encapsulates the capabilities of all experts—and the reward function, distinguishing it from traditional routing algorithms that explicitly query experts. As such, this mechanism allows a single Transformer architecture to solve multiple tasks without prior knowledge of the specific task associated with a user query.

Refer to caption — Figure 1: An example from [50] of test-time online learning, where the Transformer progressively learns that finite-element method solves the partial differential equation with higher accuracy.

Proposition 1.2 (Informal statement of Theorem 4.7).

There exists a generic way to construct a wider transformer $f$ from any Transformer-based expert models $f_{1},\dots,f_{E}$ such that, when provided with reward-based feedback, $f$ can generate a sequence of responses where the $t$ -th response has regret $o(1)$ .

Proof Sketch. We first construct a Transformer $f_{0}$ that implements an online learning algorithm with regret $o(1)$ . At each layer of the unified Transformer $f$ , we stack the attention blocks from the corresponding layers of experts $f_{0},f_{1},\dots,f_{E}$ . When generating the $i$ -th action, our goal is to activate only the attention blocks associated with expert $f_{0}$ ; when generating the $i$ -th response, our goal is to activate only the attention blocks associated with expert $f_{k}$ , where $k$ is the expert selected by action $i$ . To achieve the above, we add an attention block and develop a generalized position encoding scheme to induce attention sink behavior [92]: the attentions of all non-selected experts sink to the token representing action $i$ (being one at <action $i$ > and zero elsewhere) and attentions of the $k$ -th expert are identical to the attentions computed by $f_{k}$ . We illustrated this mechanism in Figure 2. As a result, the action sequence achieves $o(1)$ regret and the response sequence is generated from the corresponding expert selected by the latest action. Therefore, the response sequence also achieves regret $o(1)$ . ∎

Proposition 1.2 has two key implications. First, it demonstrates that a Transformer can express multiple tasks within a single architecture, extending beyond prior theoretical results that focus on single-task expressiveness. Importantly, the construction is task-agnostic and independent of the specific expert Transformers used, making both the result and the underlying techniques of independent theoretical interest. Second, Proposition 1.2 reveals a fundamental distinction between self-correction and repeated-sampling paradigms. While repeated-sampling methods generate identically distributed responses across attempts, self-correction provably allows the model to update its attempts based on verifier feedback, thereby increasing the probability of producing the correct answer as inference progresses. We further validate this results through controlled experiments.

2 Preliminaries

Transformers.

In this work, we consider attention-only Transformers defined as follows.

Definition 2.1 (Transformer).

We define a Transformer model over vocabulary $\mathcal{V}$ as a tuple

\displaystyle(\theta,\mathrm{pe},(\mathbf{K}^{(l)}_{h},\mathbf{Q}^{(l)}_{h},\mathbf{V}^{(l)}_{h})_{h\in[H],l\in[L]},\vartheta,\mathcal{V})

where $\theta:\mathcal{V}\to\mathbb{R}^{d}$ is the tokenizer, $\mathrm{pe}:\mathbb{R}^{d}\times\mathcal{V}^{\omega}\to\mathbb{R}^{d}$ is a position encoder, $\mathbf{K}^{(l)}_{h},\mathbf{Q}^{(l)}_{h},\mathbf{V}^{(l)}_{h}\in\mathbb{R}^{d\times d}$ are the key, query, value matrices over $L$ layers and $H$ heads each layer, and $\vartheta$ is the output feature. The computation of a Transformer rolls out as follows:

1.

For each $i=1,\dots,n$

$\displaystyle X^{(1)}_{i}=\mathrm{pe}(\theta(v_{i});v_{1},\dots,v_{i}).$

For each $l=1,\dots,L$ , compute each $X^{(l+1)}_{i}$ for $i=1,\dots,n$ by

\displaystyle X^{(l+1)}_{i}=\sum_{h=1}^{H}\sum_{j=1}^{i}\frac{\exp\left(s_{h}^{(l)}(X_{i},X_{j})\right)}{Z^{(l)}_{h}}\cdot\mathbf{V}^{(l)}_{h}X^{(l)}_{j},

(1)

where $s_{h}^{(l)}(\cdot)$ is the attention score defined by $s_{h}^{(l)}(X_{i},X_{j})=(\mathbf{Q}^{(l)}_{h}X^{(l)}_{i})^{\top}(\mathbf{K}^{(l)}_{h}X^{(l)}_{j})$ and $Z^{(l)}_{h}=\sum_{j=1}^{i}\exp\left(s_{h}^{(l)}(X_{i},X_{j})\right)$ is the normalizing constant.

3.

The output probability is given by

$\displaystyle p_{f}(y|v_{1},\dots,v_{n})=\mathrm{Softmax}(\vartheta(y)^{\top}X^{(L)}_{n}),\penalty\ y\in\mathcal{V}.$

In particular, we assume the softmax attention layer has precision $\epsilon$ : if two attention scores $s_{1},s_{2}$ satisfy $e^{s_{1}}<\epsilon\cdot e^{s_{2}}$ , then $e^{s_{1}}$ is treated as zero in the attention computation of Eq. (1).

While classical positional encoders is solely dependent on the index of the current token (i.e. we may write $\mathrm{pe}(\theta(v_{i});v_{1},\dots,v_{i})=\mathrm{pe}(\theta(v_{i});i)$ ), recent advance [37, 98, 33] has extended this notion to incorporate set membership information of preceding tokens. This generalization proves crucial for enhancing the long-context capability required for effective self-correction. Motivated by this insight, we introduce the following notion of a generalized position encoder.

Definition 2.2 (Generalized Position Encoder).

We say that $\mathrm{pe}:\mathbb{R}^{d}\times\mathcal{V}^{\omega}\to\mathbb{R}^{d}$ is a generalized position encoder w.r.t. a partition $\mathcal{V}_{1},\dots,\mathcal{V}_{K}$ of $\mathcal{V}$ if it maps an input feature in $\mathbb{R}^{d}$ and a token sequence (of arbitrary length) $v_{1},\cdots,v_{i}$ to a vector in $\mathbb{R}^{d}$ , so that it only depends on the input feature and the membership of each $v_{i}$ in the sets $\mathcal{V}_{1},\dots,\mathcal{V}_{K}$ , i.e.

\displaystyle\mathrm{pe}(\theta(v_{i});v_{1},\dots,v_{i})=\mathrm{pe}\left(\theta(v_{i});\left(\mathbbm{1}(v_{j}\in\mathcal{V}_{k})\right)_{j\in[i],k\in[K]}\right).

Test-time scaling.

In this work, we study the following three strategies for test-time scaling.

1.

Self-consistency samples $n$ i.i.d. responses from the language model and chooses the most consistent answer, while marginalizing over the reasoning paths.
2.

Best-of- ${n}$ samples $n$ i.i.d. responses from the language model and chooses the answer with the highest score given by the reward model.
3.

In the Self-Correction paradigm, the Transformer autonomously generates a sequence of $n$ responses, each conditioned on the previous responses and their respective reward scores.

3 Separation between Self-Consistency and Best-of-n

In this section, we study the sample complexity of self-consistency and best-of- $n$ . Let $q$ denote the user query (e.g. a math problem) and $\mathcal{O}$ denote the answer space; then for each answer $o\in\mathcal{O}$ we define $p(o)$ as the marginalized probability of generating $o$ over all possible reasoning paths

\displaystyle p(o)=\sum_{\mathrm{reasoning\penalty\ path}}p_{f}(\mathrm{reasoning\penalty\ path},o|q)

where $p_{f}$ denotes the probability distribution of Transformer $f$ .

To understand the sample complexity, we focus on the dependence on the following probability gap:

\displaystyle\Delta:=p(o^{*})-\max_{o\in\mathcal{O},o\neq o^{*}}p(o)

where $o^{*}$ denotes the correct answer¹¹ 1 If there are multiple correct answers, we can let $o^{*}$ to denote the set, and our results continue to hold in this setting.. If $\Delta\leq 0$ , then self-consistency fails to find the correct answer with high probability and the separation becomes trivial. Therefore, we focus on the setting where $\Delta>0$ (i.e., the most likely answer is correct), which is also considered in prior theoretical work [40]. Under this setting, we may assume without loss of generality that the reward function $r$ is maximized (only) at the correct answer, because $p$ itself is such a reward function satisfying this condition. Note that since $p(o)$ is marginalized over reasoning paths, $\Delta>0$ does not imply that the correct answer can be derived easily from greedy decoding.

Theorem 3.1 (Sample Complexity of Self-Consistency).

When $n\geq\frac{2\log(1/\delta)}{\Delta^{2}}$ , self-consistency with $n$ i.i.d. samples is able to produce the correct answer with probability at least $1-\delta$ ; When $n\leq\frac{1}{\Delta^{2}}$ , there exists a hard instance where self-consistency with $n$ i.i.d. samples fails to produce the correct answer with constant probability.

Theorem 3.2 (Sample Complexity of Best-of- $n$ ).

When $n\geq\frac{2\log(1/\delta)}{\Delta}$ , best-of- $n$ with $n$ i.i.d. samples is able to produce the correct answer with probability at least $1-\delta$ ; When $n\leq\frac{1}{\Delta}$ , there exists a hard instance where best-of- $n$ with $n$ i.i.d. samples fails to produce the correct answer with constant probability.

By providing matching (up to logarithmic factors) upper and lower bounds on the number of samples, the above results establishes the separation between self-consistency and best-of- $n$ . While self-consistency requires $\Theta(1/\Delta^{2})$ samples to produce the correct answer, best-of- $n$ shows advantage by only requiring $\Theta(1/\Delta)$ samples. Therefore, this theory corroborates the empirical findings that best-of- $n$ generally leads to better problem solving accuracy on reasoning tasks compared with self-consistency [79, 90].

4 Expressiveness under Self-Correction

A key distinction between self-correction and the repeated sampling strategies discussed in the previous section lies in the dependence structure of the generated responses: unlike repeated sampling, the outputs produced by self-correction are not i.i.d.. Consequently, to analyze the sample efficiency of self-correction, we must first address a fundamental question: can a large language model (LLM), through self-correction, increase the likelihood of generating the correct answer? At its core, this question is one of expressiveness—whether the Transformer architecture’s representation capacity is sufficient to support such improvement.

In this section, we take a first step toward analyzing the expressiveness of Transformers under the self-correction paradigm. Unlike prior work that focuses on expressiveness in the context of a single task, we study what we call general-purpose expressiveness: the ability to solve a broad range of tasks. To this end, we introduce the concept of a General-Purpose Transformer—a construction that maps any collection of task-specific Transformers (experts) into a single unified Transformer.

Definition 4.1 (General-Purpose Transformer).

We say that $\phi$ is a General-Purpose Transformer of type $(t_{1},t_{2})$ if it maps any set of Transformers with hidden size $d$ and depth $L$ into another ‘unified’ Transformer with hidden size $t_{1}\cdot d+t_{2}$ and depth $L+O(1)$ .

A general-purpose Transformer provides a principled framework for constructing more powerful Transformer architectures by composing simpler, task-specific components. This meta-architecture enables a single model to solve multiple tasks at inference time, representing a significant advancement in our theoretical understanding of the expressive power of modern machine learning systems. Our goal is to investigate the general-purpose expressiveness of self-correction paradigms through the lens of general-purpose Transformers: specifically, how a Transformer can adaptively solve different tasks during inference without prior knowledge of the task identity.

4.1 General-purpose expressiveness

In this section, we present two auxiliary results that serve as building blocks for constructing general-purpose Transformers capable of solving multiple tasks. These results may also be of independent interest beyond expressiveness of self-correction.

The first result addresses the setting in which multiple Transformers operate over distinct vocabularies, with each vocabulary corresponding to a specific task. The objective is to construct a unified Transformer that uses the final token in the input sequence to infer which task to perform, and subsequently solves the task by attending only to the task-relevant tokens. This paradigm is illustrated in Figure 3.

Proposition 4.2 (General-purpose Expressiveness over Different Token Spaces).

For any $H,L,K,N_{\max}\in\mathbb{Z}_{+}$ , $\mathcal{V}_{i}\cap\mathcal{V}_{j}=\emptyset\penalty\ (\forall i\neq j\in\{0\}\cup[K])$ , there exists a general-purpose Transformer $\phi$ of type $(O(K),O(\log N_{\max}))$ such that for any Transformers $f_{k}=(\theta,\mathrm{pe},(\mathbf{K}^{(l)}_{k;h},\mathbf{Q}^{(l)}_{k;h},\mathbf{V}^{(l)}_{k;h})_{h\in[H],l\in[L]},\vartheta,\mathcal{V}_{k})$ for $k\in[K]$ , the Transformer $\widetilde{f}=\phi(f_{1},\dots,f_{K})$ satisfies the following property: for any token sequence $v=v_{1}\cdots v_{n}$ such that $n\leq N_{\max}$ and there exists one $v_{i_{0}}\in\mathcal{V}_{0}$ , we have

\displaystyle p_{\widetilde{f}}(\cdot|v)=p_{f_{\kappa}}(\cdot|u)

where $\kappa$ is the task indicated by the last token: i.e., $v_{n}\in\mathcal{V}_{\kappa}$ , and $u=v_{i_{1}}\cdots v_{i_{m}}$ , where $\{i_{1}<\cdots<i_{m}\}=\{i:v_{i}\in\mathcal{V}_{\kappa}\}$ , is the sequence of tokens relevant to task $\kappa$ .

Remark 4.3.

The existence of $v_{i_{0}}$ which does not belong to any $\{\mathcal{V}_{i}\}_{i\in[K]}$ serves the technical purpose of inducing attention sink of all irrelevant experts to $v_{i_{0}}$ . It may be achieve by assuming the user query always ends with the special token <eos>.

The following result considers a more challenging scenario in which multiple Transformers operate across different tasks but share a common vocabulary space. A set of indicator tokens, denoted by $\Omega$ , is used to specify the intended task. The objective is to determine which task to execute based on the most recent indicator token. It then proceeds to solve the task by attending exclusively to the task-relevant tokens appearing before the first indicator token and after the last indicator token in the input sequence. This paradigm is closely related to self-correction, and is illustrated in Figure 4.

Proposition 4.4 (Multi-Task Representation over the Same Token Space).

For any $H,L,K,N_{\max}\in\mathbb{Z}_{+}$ , token spaces $\Omega\cap\mathcal{V}=\emptyset$ , there exists a general-purpose Transformer $\phi$ of type $(O(K),O(\log N_{\max}))$ such that for any Transformers $f_{k}=(\theta,\mathrm{pe},(\mathbf{K}^{(l)}_{k;h},\mathbf{Q}^{(l)}_{k;h},\mathbf{V}^{(l)}_{k;h})_{h\in[H],l\in[L]},\vartheta,\mathcal{V}),k\in[K]$ over $\mathcal{V}$ , the Transformer $\widetilde{f}=\phi(f_{1},\dots,f_{K})$ satisfies the following property: for any token sequence $v=v_{1}\cdots v_{n}$ such that

\displaystyle\{\xi_{1}<\cdots<\xi_{m}\}=\{j:v_{j}\in\Omega\},\penalty\ \xi_{m}<n\leq N_{\max}

then we have

\displaystyle p_{\widetilde{f}}(\cdot|v)=p_{f_{\kappa}}(\cdot|u)

(2)

where $u=v_{1}\cdots v_{\xi_{1}-1}v_{\xi_{m}+1}\cdots v_{n}$ is the token sequence obtained by omitting tokens from position $\xi_{1}$ to $\xi_{m}$ , and $\kappa$ is the task indicated by token $v_{\xi_{m}}$ .

Remark 4.5.

We observe that in both results above, reducing the type parameters is generally not feasible. The dependence on $K$ arises from the need to compute features for all $K$ experts corresponding to the user query. Since the model lacks prior knowledge of the task, it must encode all task-relevant information to preserve the ability to invoke any expert at inference time. The $\log(N_{\max})$ scaling stems from the positional encoding: in order to construct $N_{\max}$ nearly orthogonal vectors, the positional embedding must have dimension at least $\log(N_{\max})$ .

4.2 General-purpose expressiveness of Transformers with self-correction

In this section we state the main result that establishes general-purpose expressiveness of Transformers with self-correction. We rely on the following notion of regret-minimization Transformer, which expresses the single task of finding the most rewardable action.

Definition 4.6 (Regret-Minimization Transformer).

We say that a Transformer $f$ achieves simple regret $\mathrm{reg}(\cdot)$ over reward function $r$ and action space $\mathcal{A}$ , if for any $T\in\mathbb{Z}_{+}$ , we have $\max_{a^{*}\in\mathcal{A}}r(a^{*})-\mathbb{E}[r(a_{T})]\leq\mathrm{reg}(T)$ where $a_{1},\dots,a_{T}$ are generated in the following way:

	$\displaystyle a_{t}\sim$	$\displaystyle\penalty\ p_{f}(\cdot\|a_{1},r_{1},\dots,a_{t-1},r_{t-1}),\penalty\ \forall t=1,\dots,T,$
	$\displaystyle r_{t}=$	$\displaystyle\penalty\ r(a_{t}),\penalty\ \forall t=1,\dots,T.$

Essentially, the goal of a regret-minimization Transformer is to learn from a reward oracle and ultimately recommend an action that is near-optimal, which is related to a concept commonly referred to as simple regret in the bandit literature [28, 14, 43]. To achieve this, the Transformer may implement strategies such as mirror descent, upper confidence bounds, or search-based algorithms, depending on the problem structure. As these procedures rely only on basic arithmetic operations, such Transformers can be constructed by applying the universal approximation capabilities of Transformers [95, 58, 29, 53]: for example, [55] provides constructions to approximate upper confidence bounds and Thompson sampling algorithms with regret $O(\sqrt{T})$ . Consequently, their construction is not the primary focus of this work.

Algorithm 1 Self-correction with verifier

1: procedure Generation(

q

)

\triangleright

q=q_{1}\dots q_{n_{0}}

denotes the user query.

\mathrm{prompt}\leftarrow q

3: for

t=1,\dots,T

a^{(t)}\sim p_{\widetilde{f}}(\cdot\mid\mathrm{prompt})

\triangleright

a^{(t)}

designates which expert to use in

t

-th iteration

\mathrm{prompt}\leftarrow\mathrm{prompt}|a^{(t)}

\triangleright

Update the prompt autoregressively,

|

represents token concatenation.

6: for

i=1,\dots

u^{(t)}_{i}\sim p_{\widetilde{f}}(\cdot\mid\mathrm{prompt})

\triangleright

Generate

t

-th response autoregressively

\mathrm{prompt}\leftarrow\mathrm{prompt}|u^{(t)}_{i}

\triangleright

Update the prompt autoregressively

9: if

u^{(t)}_{i}=\mathrm{EOS}

then

10: Break

11: end if

12: end for

13:

r^{(t)}\leftarrow r(q,u^{(t)}),\penalty\ \mathrm{prompt}\leftarrow\mathrm{prompt}|r^{(t)}

\triangleright

Query verifier to obtain reward of

t

-th response

14: end for

15: Return

16: end procedure

The following theorem establishes the existence of a general-purpose Transformer that can simulate the behavior of a set of expert Transformers (not necessarily over the same token space) through self-correction. Specifically, it shows that such a unified Transformer can, at inference time, identify and invoke the appropriate expert to solve any task that the original experts can solve. The self-correction protocol is described in Algorithm 1, wherein the unified Transformer autoregressively generates actions and responses, after which the verifier is queried to obtain reward signals. Through this process of trial and error, the model effectively “learns” at inference time, using the verifier to minimize regret and adaptively select the correct expert.

Theorem 4.7 (Regret Minimization via Self-Correction).

For any $H,L,K,N_{\max}\in\mathbb{Z}_{+}$ , token spaces $\mathcal{V}_{0},\mathcal{V}_{1},\dots,\mathcal{V}_{K},\mathcal{A}\penalty\ (|\mathcal{A}|=K)$ such that $\mathcal{V}_{0},\mathcal{V}=(\cup_{k=1}^{K}\mathcal{V}_{k}),\text{ and }\mathcal{A}$ are disjoint, and reward function $r$ , there exists a general-purpose Transformer $\phi$ of type $(O(K),O(\log N_{\max}))$ such that given any set of Transformers denoted as follows,

•

$K$ expert Transformers: $f_{k}=(\theta,\mathrm{pe},(\mathbf{K}^{(l)}_{k;h},\mathbf{Q}^{(l)}_{k;h},\mathbf{V}^{(l)}_{k;h})_{h\in[H],l\in[L]},\vartheta,\mathcal{V}_{k})$ for $k\in[K]$ , such that one of the expert $f_{k^{*}}$ achieves $\lambda$ -suboptimal reward:

$\displaystyle\mathbb{E}_{u\sim f_{k^{*}}(\cdot|q)}[r(q,u)]\geq\max_{u^{*}\in\mathcal{V}^{\omega}}r(q,u^{*})-\lambda$
•

Regret-Minimization Transformer: $f_{0}=(\theta,\mathrm{pe},\mathbf{K}^{(l)}_{0;h},\mathbf{Q}^{(l)}_{0;h},\mathbf{V}^{(l)}_{0;h})_{h\in[H],l\in[L]},\vartheta,\mathcal{V}_{0}\cup\mathcal{A})$ that implements a bandit algorithm over the reward function $r_{0}$ and action space $\mathcal{A}$ with simple regret $\mathrm{reg}(t)$ , where $r_{0}(a)=\mathbb{E}_{u\sim f_{a}(\cdot|q)}[r(q,u)]$ denotes the average reward of responses generated by the $a$ -th expert,

then the Transformer $\widetilde{f}=\phi(f_{0},f_{1},\dots,f_{K})$ satisfies the following property: for any prompt $v=v_{1}\cdots v_{n}$ , if the response sequence $u^{(1)},\dots,u^{(T)}$ generated by the protocol in Algorithm 1 has total length $\leq N_{\max}$ , then we have

\displaystyle\max_{u^{*}\in\mathcal{V}^{\omega}}r(q,u^{*})-\mathbb{E}[r(q,u^{(T)})]\leq\lambda+\mathrm{reg}(T)

Remark 4.8.

While the general-purpose Transformer $\phi$ can be applied to construct the brutal-force Transformer $\widetilde{f}$ that simply tries every expert, we note that the generality of Definition 4.6 allows us to construct more powerful Transformers beyond brutal search. Leveraging the structures in the problem and the expert pool, it is entirely possible to identify the correct expert using $\ll K$ trials [72, 30].

As a consequence of Theorem 4.7, we obtain a Transformer architecture that can provably produce a final answer that nearly maximizes the reward. This means that the unified transformer can solve $K$ distinct tasks at inference time, without requiring prior knowledge of which task the user query pertains to. Notably, the construction of such an architecture is general-purpose, in that it is independent of the specific tasks, reward functions, or expert policies. To the best of our knowledge, this constitutes the first theoretical result of its kind in the study of Transformer architectures. Furthermore, our theory aligns with the empirical finding that LLMs are able to progressively optimize outcome rewards during test-time [71].

5 Experiments

In this section, we conduct synthetic experiments to show that Transformers can self-correct with verifier feedback.

5.1 Experimental Setup

Data generation.

We aim to construct a test problem with complex prompts such that correctly solving the problem in the single-term generation is challenging. In this case, self-correction can play a critical role if Transformers have such capacities. Specifically, in our synthetic problem, the prompt is the concatenation of the following two components:

•

Instruction: A 3-SAT problem, e.g.,

(\sim x_{3}\lor\sim x_{1}\lor\sim x_{2})\land(\sim x_{1}\lor\sim x_{3}\lor x_{2})\land(\sim x_{4}\lor x_{2}\lor\sim x_{3})\land\cdots

•

Data: A string composed of characters from the set {a, b}.

Model	Depth	Heads	Width
GPT-nano	3	3	48
GPT-micro	4	4	128
GPT-mini	6	6	192
Gopher-44M	8	16	512

Table 1: Model configuration hyperparameters.

The ground truth target is defined as follows: If the 3-SAT problem in the instruction is satisfiable, the model should copy the string in the data part in the output; otherwise, the model should reverse the string in the output.

Model configuration.

We train Transformer models of various sizes. The configurations are detailed in Table 1.

Implementation details.

Our code are implemented based on PyTorch [67] and minGPT²² 2 https://github.com/karpathy/minGPT (MIT license).. All the models are trained on one NVIDIA GeForce RTX 2080 Ti GPU with 11GB memory.

Figure 5: Accuracy comparisons of different models with/without self-correction at test time.

In our experiment, we construct datasets using 3-SAT problems with 4 variables and 20 clauses. The lengths of the data strings are set to 5. We generate 10000 instances for training and 512 instances for evaluation. In the training set, we control the ratio of satisfiable and unsatisfiable 3-SAT instructions to 9:1, while in the test set, the ratio is set to 1:1.

All our models are trained with the Adam optimizer [47] for 5 epochs. Following common practice, the learning rate goes through the warm-up stage in the first 5% of training iterations, and then decays linearly to 0 until training finishes. We set the peak learning rate to $10^{-4}$ and find that all the models are stably trained under this learning rate schedule. We do not apply drop out or weight decay during training. We repeat the experiments for 3 times under different random seeds and report the average accuracy with error bars.

5.2 Results

Test set accuracy across different inference settings is shown in Figure 5. We note that model performance plateaus at $63.19\%$ when there is no self-correction at test time, with no improvement from increased model size. By contrast, when models are equipped with verifier signals to enable self-correction, test accuracy improves substantially, demonstrating the efficacy of this mechanism. Crucially, larger models – such as GPT-mini and Gopher-44M – achieve near-perfect accuracy under self-correction, suggesting that sufficiently expressive Transformers are capable of implementing effective self-correction strategies. This empirical result supports our theoretical findings.

6 Related Works

Theories of Transformers and Large Language Models.

The success of Transformers and LLMs has motivated the study on their expressiveness. Existing research has shown that Transformers can implement simple functions such as sparse linear functions, two-layer neural networks, and decision trees [32], gradient descent [3, 6, 82], automata [57, 102], Dyck languages [8, 94], Turing machines [25, 9, 96, 68, 86], variational inference [60], and bandit algorithms [55]. [95, 58, 4, 69] establish universal approximation results under various settings. [26, 27, 49, 54] study representational capabilities and properties of self-attention, the core component in Transformers. [29, 53] study the expressiveness of auto-regressive Transformers with chain-of-thought. [26, 52, 10] studies the sample complexity of Transformers. Recently, a growing body of work has begun to explore the theoretical foundations of self-improvement in large language models (LLMs). [78] introduces the generation-verification gap as a key quantity governing scaling behavior. [40] proposes a progressive sharpening framework in which the policy gradually shifts toward more confident responses. [74] draws on reinforcement learning theory to formally establish the advantages of verifier-based methods. In contrast to these works, our results provide explicit sample complexity rates and tangible representation architectures, enabling a more concrete understanding of the fundamental capabilities and limitations of test-time scaling paradigms.

Test-time scaling.

Recent research has established the test-time scaling law of LLMs, illuminating a new scaling axis beyond training-time scaling laws [45, 39]. Existing approaches of scaling up test-time compute of LLMs can be broadly classified into two categories: (1) applying test-time algorithms (aka inference-time algorithms) during LLM decoding [11, 90, 76]; and (2) explicitly training LLMs to output long chain-of-thought traces [36, 46, 66, 93]. Many recent works focus on understanding and improving the effectiveness of test-time scaling empirically: [19, 1, 23, 85] study under-thinking, over-thinking, and length control in LLM reasoning. [16] proposes to integrates self-verification and self-correction into sampling. [71] analyzes optimizing test-time compute by introducing a meta reinforcement learning formulation. [74] demonstrates that verification/RL is important for optimal test-time scaling. [99] provides an extensive review of the test-time scaling landscape. In contrast, our work focuses on theoretical analyses of test-time scaling.

7 Discussions

In this work, we present a theoretical analysis of test-time scaling paradigms, focusing on two core aspects: sample efficiency and representational capacity. Our investigation reveals a fundamental separation in sample complexity between self-consistency and best-of- $n$ , providing theoretical support for the empirically observed superiority of the latter method. Furthermore, by introducing the framework of general-purpose expressiveness, we construct generic Transformer architectures capable of emulating online learning algorithms at test time. This capability enables a single model to provably solve multiple tasks without task-specific adaptation, thus extending our understanding of expressiveness to multi-task settings. Our results highlight the theoretical advantage of self-correction paradigms, which iteratively refine predictions to increase the likelihood of correct answers—surpassing the limitations of i.i.d. responses by repeated sampling approaches. This finding is validated through experiments and we observe that it requires additional model capacities for Transformer to implement self-correction.

Despite these contributions, our work comes with limitations: our construction in Theorem 4.7 only applies to attention-only Transformers and relies on a slightly generalized position encoding method. Relaxing these constraints constitutes interesting problems for future research.

References

[1] P. Aggarwal and S. Welleck. L1: Controlling how long a reasoning model thinks with reinforcement learning. In arXiv, 2025.
[2] S. Agrawal and R. Jia. Optimistic posterior sampling for reinforcement learning: worst-case regret bounds. Advances in neural information processing systems, 30, 2017.
[3] E. Akyürek, D. Schuurmans, J. Andreas, T. Ma, and D. Zhou. What learning algorithm is in-context learning? investigations with linear models. arXiv preprint arXiv:2211.15661, 2022.
[4] S. Alberti, N. Dern, L. Thesing, and G. Kutyniok. Sumformer: Universal approximation for efficient transformers. In T. Doster, T. Emerson, H. Kvinge, N. Miolane, M. Papillon, B. Rieck, and S. Sanborn, editors, Proceedings of 2nd Annual Workshop on Topology, Algebra, and Geometry in Machine Learning (TAG-ML), volume 221 of Proceedings of Machine Learning Research, pages 72–86. PMLR, 28 Jul 2023.
[5] C. Anil, Y. Wu, A. Andreassen, A. Lewkowycz, V. Misra, V. Ramasesh, A. Slone, G. Gur-Ari, E. Dyer, and B. Neyshabur. Exploring length generalization in large language models. arXiv preprint arXiv:2207.04901, 2022.
[6] Y. Bai, F. Chen, H. Wang, C. Xiong, and S. Mei. Transformers as statisticians: Provable in-context learning with in-context algorithm selection. arXiv preprint arXiv:2306.04637, 2023.
[7] B. Barak, B. Edelman, S. Goel, S. Kakade, E. Malach, and C. Zhang. Hidden progress in deep learning: Sgd learns parities near the computational limit. Advances in Neural Information Processing Systems, 35:21750–21764, 2022.
[8] S. Bhattamishra, K. Ahuja, and N. Goyal. On the ability and limitations of transformers to recognize formal languages. In Proceedings of the 2020 Conference on Empirical Methods in Natural Language Processing (EMNLP), pages 7096–7116, 2020.
[9] S. Bhattamishra, A. Patel, and N. Goyal. On the computational power of transformers and its implications in sequence modeling. In Proceedings of the 24th Conference on Computational Natural Language Learning, pages 455–475, 2020.
[10] E. Botta, Y. Li, A. Mehta, J. T. Ash, C. Zhang, and A. Risteski. On the query complexity of verifier-assisted language generation. arXiv preprint arXiv:2502.12123, 2025.
[11] B. Brown, J. Juravsky, R. Ehrlich, R. Clark, Q. V. Le, C. Ré, and A. Mirhoseini. Large language monkeys: Scaling inference compute with repeated sampling. arXiv preprint arXiv:2407.21787, 2024.
[12] T. Brown, B. Mann, N. Ryder, M. Subbiah, J. D. Kaplan, P. Dhariwal, A. Neelakantan, P. Shyam, G. Sastry, A. Askell, S. Agarwal, A. Herbert-Voss, G. Krueger, T. Henighan, R. Child, A. Ramesh, D. Ziegler, J. Wu, C. Winter, C. Hesse, M. Chen, E. Sigler, M. Litwin, S. Gray, B. Chess, J. Clark, C. Berner, S. McCandlish, A. Radford, I. Sutskever, and D. Amodei. Language models are few-shot learners. In H. Larochelle, M. Ranzato, R. Hadsell, M. Balcan, and H. Lin, editors, Advances in Neural Information Processing Systems, volume 33, pages 1877–1901. Curran Associates, Inc., 2020.
[13] S. Bubeck, V. Chandrasekaran, R. Eldan, J. Gehrke, E. Horvitz, E. Kamar, P. Lee, Y. T. Lee, Y. Li, S. Lundberg, et al. Sparks of artificial general intelligence: Early experiments with gpt-4. arXiv preprint arXiv:2303.12712, 2023.
[14] A. Carpentier and M. Valko. Simple regret for infinitely many armed bandits. In International Conference on Machine Learning, pages 1133–1141. PMLR, 2015.
[15] G. Chen, M. Liao, C. Li, and K. Fan. Alphamath almost zero: Process supervision without process. In The Thirty-eighth Annual Conference on Neural Information Processing Systems, 2024.
[16] J. Chen, J. Ren, X. Chen, C. Yang, R. Sun, and S. Ö. Arık. Sets: Leveraging self-verification and self-correction for improved test-time scaling. arXiv preprint arXiv:2501.19306, 2025.
[17] L. Chen, J. Q. Davis, B. Hanin, P. Bailis, I. Stoica, M. Zaharia, and J. Zou. Are more LLM calls all you need? towards the scaling properties of compound AI systems. In Conference on Neural Information Processing Systems, 2024.
[18] X. Chen, M. Lin, N. Schärli, and D. Zhou. Teaching large language models to self-debug. In International Conference on Learning Representations, 2024.
[19] X. Chen, J. Xu, T. Liang, Z. He, J. Pang, D. Yu, L. Song, Q. Liu, M. Zhou, Z. Zhang, et al. Do not think that much for 2+ 3=? on the overthinking of o1-like llms. arXiv preprint arXiv:2412.21187, 2024.
[20] A. Chowdhery, S. Narang, J. Devlin, M. Bosma, G. Mishra, A. Roberts, P. Barham, H. W. Chung, C. Sutton, S. Gehrmann, et al. Palm: Scaling language modeling with pathways. arXiv preprint arXiv:2204.02311, 2022.
[21] K. Cobbe, V. Kosaraju, M. Bavarian, M. Chen, H. Jun, L. Kaiser, M. Plappert, J. Tworek, J. Hilton, R. Nakano, et al. Training verifiers to solve math word problems. arXiv preprint arXiv:2110.14168, 2021.
[22] codeforce. Codeforces, 2025.
[23] A. Cuadron, D. Li, W. Ma, X. Wang, Y. Wang, S. Zhuang, S. Liu, L. G. Schroeder, T. Xia, H. Mao, et al. The danger of overthinking: Examining the reasoning-action dilemma in agentic tasks. arXiv preprint arXiv:2502.08235, 2025.
[24] DeepSeek-AI. Deepseek-r1: Incentivizing reasoning capability in llms via reinforcement learning. In arXiv, 2025.
[25] M. Dehghani, S. Gouws, O. Vinyals, J. Uszkoreit, and Ł. Kaiser. Universal transformers. arXiv preprint arXiv:1807.03819, 2018.
[26] B. L. Edelman, S. Goel, S. Kakade, and C. Zhang. Inductive biases and variable creation in self-attention mechanisms. In International Conference on Machine Learning, pages 5793–5831. PMLR, 2022.
[27] N. Elhage, N. Nanda, C. Olsson, T. Henighan, N. Joseph, B. Mann, A. Askell, Y. Bai, A. Chen, T. Conerly, et al. A mathematical framework for transformer circuits. Transformer Circuits Thread, 1:1, 2021.
[28] E. Even-Dar, S. Mannor, Y. Mansour, and S. Mahadevan. Action elimination and stopping conditions for the multi-armed bandit and reinforcement learning problems. Journal of machine learning research, 7(6), 2006.
[29] G. Feng, B. Zhang, Y. Gu, H. Ye, D. He, and L. Wang. Towards revealing the mystery behind chain of thought: a theoretical perspective. Advances in Neural Information Processing Systems, 36:70757–70798, 2023.
[30] D. J. Foster, S. M. Kakade, J. Qian, and A. Rakhlin. The statistical complexity of interactive decision making. arXiv preprint arXiv:2112.13487, 2021.
[31] Z. Gao, B. Niu, X. He, H. Xu, H. Liu, A. Liu, X. Hu, and L. Wen. Interpretable contrastive monte carlo tree search reasoning. In arXiv, 2024.
[32] S. Garg, D. Tsipras, P. S. Liang, and G. Valiant. What can transformers learn in-context? a case study of simple function classes. Advances in Neural Information Processing Systems, 35:30583–30598, 2022.
[33] O. Golovneva, T. Wang, J. Weston, and S. Sukhbaatar. Contextual position encoding: Learning to count what’s important. arXiv preprint arXiv:2405.18719, 2024.
[34] Google. Aime problems and solutions, 2025.
[35] Z. Gou, Z. Shao, Y. Gong, yelong shen, Y. Yang, N. Duan, and W. Chen. CRITIC: Large language models can self-correct with tool-interactive critiquing. In International Conference on Learning Representations, 2024.
[36] D. Guo, D. Yang, H. Zhang, J. Song, R. Zhang, R. Xu, Q. Zhu, S. Ma, P. Wang, X. Bi, et al. Deepseek-r1: Incentivizing reasoning capability in llms via reinforcement learning. arXiv preprint arXiv:2501.12948, 2025.
[37] Z. He, G. Feng, S. Luo, K. Yang, L. Wang, J. Xu, Z. Zhang, H. Yang, and D. He. Two stones hit one bird: Bilevel positional encoding for better length extrapolation. arXiv preprint arXiv:2401.16421, 2024.
[38] D. Hendrycks, C. Burns, S. Kadavath, A. Arora, S. Basart, E. Tang, D. Song, and J. Steinhardt. Measuring mathematical problem solving with the MATH dataset. In Conference on Neural Information Processing Systems Datasets and Benchmarks Track, 2021.
[39] J. Hoffmann, S. Borgeaud, A. Mensch, E. Buchatskaya, T. Cai, E. Rutherford, D. de las Casas, L. A. Hendricks, J. Welbl, A. Clark, T. Hennigan, E. Noland, K. Millican, G. van den Driessche, B. Damoc, A. Guy, S. Osindero, K. Simonyan, E. Elsen, O. Vinyals, J. W. Rae, and L. Sifre. An empirical analysis of compute-optimal large language model training. In A. H. Oh, A. Agarwal, D. Belgrave, and K. Cho, editors, Advances in Neural Information Processing Systems, 2022.
[40] A. Huang, A. Block, D. J. Foster, D. Rohatgi, C. Zhang, M. Simchowitz, J. T. Ash, and A. Krishnamurthy. Self-improvement in language models: The sharpening mechanism. arXiv preprint arXiv:2412.01951, 2024.
[41] Z. Huang, Z. Wang, S. Xia, X. Li, H. Zou, R. Xu, R.-Z. Fan, L. Ye, E. Chern, Y. Ye, Y. Zhang, Y. Yang, T. Wu, B. Wang, S. Sun, Y. Xiao, Y. Li, F. Zhou, S. Chern, Y. Qin, Y. Ma, J. Su, Y. Liu, Y. Zheng, S. Zhang, D. Lin, Y. Qiao, and P. Liu. Olympicarena: Benchmarking multi-discipline cognitive reasoning for superintelligent AI. In Conference on Neural Information Processing Systems Datasets and Benchmarks Track, 2024.
[42] R. Irvine, D. Boubert, V. Raina, A. Liusie, Z. Zhu, V. Mudupalli, A. Korshuk, Z. Liu, F. Cremer, V. Assassi, C.-C. Beauchamp, X. Lu, T. Rialan, and W. Beauchamp. Rewarding chatbots for real-world engagement with millions of users. In arXiv, 2023.
[43] K. Jamieson, M. Malloy, R. Nowak, and S. Bubeck. lil’ucb: An optimal exploration algorithm for multi-armed bandits. In Conference on Learning Theory, pages 423–439. PMLR, 2014.
[44] N. Joshi, G. Vardi, A. Block, S. Goel, Z. Li, T. Misiakiewicz, and N. Srebro. A theory of learning with autoregressive chain of thought. arXiv preprint arXiv:2503.07932, 2025.
[45] J. Kaplan, S. McCandlish, T. Henighan, T. B. Brown, B. Chess, R. Child, S. Gray, A. Radford, J. Wu, and D. Amodei. Scaling laws for neural language models. arXiv preprint arXiv:2001.08361, 2020.
[46] Kimi. Kimi k1.5: Scaling reinforcement learning with llms. In arXiv, 2025.
[47] D. P. Kingma and J. Ba. Adam: A method for stochastic optimization. In ICLR (Poster), 2015.
[48] A. Kumar, V. Zhuang, R. Agarwal, Y. Su, J. D. Co-Reyes, A. Singh, K. Baumli, S. Iqbal, C. Bishop, R. Roelofs, et al. Training language models to self-correct via reinforcement learning. arXiv preprint arXiv:2409.12917, 2024.
[49] S. Li, X. Chen, D. He, and C.-J. Hsieh. Can vision transformers perform convolution? arXiv preprint arXiv:2111.01353, 2021.
[50] S. Li, T. Marwah, J. Shen, W. Sun, A. Risteski, Y. Yang, and A. Talwalkar. Codepde: An inference framework for llm-driven pde solver generation. arXiv preprint arXiv:2505.08783, 2025.
[51] S. Li, Z. Song, Y. Xia, T. Yu, and T. Zhou. The closeness of in-context learning and weight shifting for softmax regression. arXiv preprint arXiv:2304.13276, 2023.
[52] Y. Li, A. Kirchmeyer, A. Mehta, Y. Qin, B. Dadachev, K. Papineni, S. Kumar, and A. Risteski. Promises and pitfalls of generative masked language modeling: theoretical framework and practical guidelines. arXiv preprint arXiv:2407.21046, 2024.
[53] Z. Li, H. Liu, D. Zhou, and T. Ma. Chain of thought empowers transformers to solve inherently serial problems. In The Twelfth International Conference on Learning Representations, 2024.
[54] V. Likhosherstov, K. Choromanski, and A. Weller. On the expressive power of self-attention matrices. arXiv preprint arXiv:2106.03764, 2021.
[55] L. Lin, Y. Bai, and S. Mei. Transformers as decision makers: Provable in-context reinforcement learning via supervised pretraining. arXiv preprint arXiv:2310.08566, 2023.
[56] Q. Lin, B. Xu, Z. Li, Z. Hao, K. Zhang, and R. Cai. Leveraging constrained monte carlo tree search to generate reliable long chain-of-thought for mathematical reasoning. In arXiv, 2025.
[57] B. Liu, J. T. Ash, S. Goel, A. Krishnamurthy, and C. Zhang. Transformers learn shortcuts to automata. arXiv preprint arXiv:2210.10749, 2022.
[58] S. Luo, S. Li, S. Zheng, T.-Y. Liu, L. Wang, and D. He. Your transformer may not be as powerful as you expect. In A. H. Oh, A. Agarwal, D. Belgrave, and K. Cho, editors, Advances in Neural Information Processing Systems, 2022.
[59] A. Madaan, N. Tandon, P. Gupta, S. Hallinan, L. Gao, S. Wiegreffe, U. Alon, N. Dziri, S. Prabhumoye, Y. Yang, S. Gupta, B. P. Majumder, K. Hermann, S. Welleck, A. Yazdanbakhsh, and P. Clark. Self-refine: Iterative refinement with self-feedback. In Thirty-seventh Conference on Neural Information Processing Systems, 2023.
[60] S. Mei and Y. Wu. Deep networks as denoising algorithms: Sample-efficient learning of diffusion models in high-dimensional graphical models. arXiv preprint arXiv:2309.11420, 2023.
[61] W. Merrill and A. Sabharwal. The expressive power of transformers with chain of thought. arXiv preprint arXiv:2310.07923, 2023.
[62] T. Munkhbat, N. Ho, S. H. Kim, Y. Yang, Y. Kim, and S.-Y. Yun. Self-training elicits concise reasoning in large language models. In arXiv, 2025.
[63] A. Nguyen, D. Mekala, C. Dong, and J. Shang. When is the consistent prediction likely to be a correct prediction? In arXiv, 2024.
[64] C. Olsson, N. Elhage, N. Nanda, N. Joseph, N. DasSarma, T. Henighan, B. Mann, A. Askell, Y. Bai, A. Chen, et al. In-context learning and induction heads. arXiv preprint arXiv:2209.11895, 2022.
[65] OpenAI. Openai o1 system card. In arXiv, 2024.
[66] OpenAI. Openai o3-mini, 2024.
[67] A. Paszke, S. Gross, F. Massa, A. Lerer, J. Bradbury, G. Chanan, T. Killeen, Z. Lin, N. Gimelshein, L. Antiga, et al. Pytorch: An imperative style, high-performance deep learning library. Advances in neural information processing systems, 32:8026–8037, 2019.
[68] J. Pérez, P. Barceló, and J. Marinkovic. Attention is turing-complete. Journal of Machine Learning Research, 22(75):1–35, 2021.
[69] A. Petrov, P. H. Torr, and A. Bibi. Prompting a pretrained transformer can be a universal approximator. In Proceedings of the 41st International Conference on Machine Learning, pages 40523–40550, 2024.
[70] J. Qiu, Y. Lu, Y. Zeng, J. Guo, J. Geng, H. Wang, K. Huang, Y. Wu, and M. Wang. Treebon: Enhancing inference-time alignment with speculative tree-search and best-of-n sampling. arXiv preprint arXiv:2410.16033, 2024.
[71] Y. Qu, M. Y. Yang, A. Setlur, L. Tunstall, E. E. Beeching, R. Salakhutdinov, and A. Kumar. Optimizing test-time compute via meta reinforcement fine-tuning. arXiv preprint arXiv:2503.07572, 2025.
[72] D. Russo and B. Van Roy. Learning to optimize via information-directed sampling. Operations Research, 66(1):230–252, 2018.
[73] P. G. Sessa, R. Dadashi, L. Hussenot, J. Ferret, N. Vieillard, A. Ramé, B. Shariari, S. Perrin, A. Friesen, G. Cideron, S. Girgin, P. Stanczyk, A. Michi, D. Sinopalnikov, S. Ramos, A. Héliou, A. Severyn, M. Hoffman, N. Momchev, and O. Bachem. Bond: Aligning llms with best-of-n distillation. In arXiv, 2024.
[74] A. Setlur, N. Rajaraman, S. Levine, and A. Kumar. Scaling test-time compute without verification or rl is suboptimal. arXiv preprint arXiv:2502.12118, 2025.
[75] B. Shi, M. Tang, K. R. Narasimhan, and S. Yao. Can language models solve olympiad programming? In Conference on Language Modeling, 2024.
[76] C. V. Snell, J. Lee, K. Xu, and A. Kumar. Scaling LLM test-time compute optimally can be more effective than scaling parameters for reasoning. In The Thirteenth International Conference on Learning Representations, 2025.
[77] Y. Song, G. Wang, S. Li, and B. Y. Lin. The good, the bad, and the greedy: Evaluation of llms should not ignore non-determinism. In arXiv, 2024.
[78] Y. Song, H. Zhang, C. Eisenach, S. Kakade, D. Foster, and U. Ghai. Mind the gap: Examining the self-improvement capabilities of large language models. arXiv preprint arXiv:2412.02674, 2024.
[79] Z. Sun, L. Yu, Y. Shen, W. Liu, Y. Yang, S. Welleck, and C. Gan. Easy-to-hard generalization: Scalable alignment beyond human supervision. In The Thirty-eighth Annual Conference on Neural Information Processing Systems, 2024.
[80] Y. Tian, B. Peng, L. Song, L. Jin, D. Yu, L. Han, H. Mi, and D. Yu. Toward self-improvement of LLMs via imagination, searching, and criticizing. In Conference on Neural Information Processing Systems, 2024.
[81] J. Von Oswald, E. Niklasson, E. Randazzo, J. Sacramento, A. Mordvintsev, A. Zhmoginov, and M. Vladymyrov. Transformers learn in-context by gradient descent. arXiv preprint arXiv:2212.07677, 2022.
[82] J. Von Oswald, E. Niklasson, E. Randazzo, J. Sacramento, A. Mordvintsev, A. Zhmoginov, and M. Vladymyrov. Transformers learn in-context by gradient descent. In International Conference on Machine Learning, pages 35151–35174. PMLR, 2023.
[83] Z. Wan, X. Feng, M. Wen, S. M. McAleer, Y. Wen, W. Zhang, and J. Wang. Alphazero-like tree-search can guide large language model decoding and training. In Forty-first International Conference on Machine Learning, 2024.
[84] X. Wang, J. Wei, D. Schuurmans, Q. V. Le, E. H. Chi, S. Narang, A. Chowdhery, and D. Zhou. Self-consistency improves chain of thought reasoning in language models. In The Eleventh International Conference on Learning Representations, 2023.
[85] Y. Wang, Q. Liu, J. Xu, T. Liang, X. Chen, Z. He, L. Song, D. Yu, J. Li, Z. Zhang, et al. Thoughts are all over the place: On the underthinking of o1-like llms. arXiv preprint arXiv:2501.18585, 2025.
[86] C. Wei, Y. Chen, and T. Ma. Statistically meaningful approximation: a case study on approximating turing machines with transformers. Advances in Neural Information Processing Systems, 35:12071–12083, 2022.
[87] J. Wei, X. Wang, D. Schuurmans, M. Bosma, F. Xia, E. Chi, Q. V. Le, D. Zhou, et al. Chain-of-thought prompting elicits reasoning in large language models. Advances in neural information processing systems, 35:24824–24837, 2022.
[88] S. Welleck, X. Lu, P. West, F. Brahman, T. Shen, D. Khashabi, and Y. Choi. Generating sequences by learning to self-correct. In The Eleventh International Conference on Learning Representations, 2023.
[89] Y. Wu, Z. Sun, S. Li, S. Welleck, and Y. Yang. Scaling inference computation: Compute-optimal inference for problem-solving with language models. In Workshop on Mathematical Reasoning and AI at NeurIPS’24, 2024.
[90] Y. Wu, Z. Sun, S. Li, S. Welleck, and Y. Yang. Inference scaling laws: An empirical analysis of compute-optimal inference for LLM problem-solving. In The Thirteenth International Conference on Learning Representations, 2025.
[91] Y. Wu, Y. Wang, T. Du, S. Jegelka, and Y. Wang. When more is less: Understanding chain-of-thought length in llms. In arXiv, 2025.
[92] G. Xiao, Y. Tian, B. Chen, S. Han, and M. Lewis. Efficient streaming language models with attention sinks. arXiv preprint arXiv:2309.17453, 2023.
[93] A. Yang, A. Li, B. Yang, B. Zhang, B. Hui, B. Zheng, B. Yu, C. Gao, C. Huang, C. Lv, C. Zheng, D. Liu, F. Zhou, F. Huang, F. Hu, H. Ge, H. Wei, H. Lin, J. Tang, J. Yang, J. Tu, J. Zhang, J. Yang, J. Yang, J. Zhou, J. Zhou, J. Lin, K. Dang, K. Bao, K. Yang, L. Yu, L. Deng, M. Li, M. Xue, M. Li, P. Zhang, P. Wang, Q. Zhu, R. Men, R. Gao, S. Liu, S. Luo, T. Li, T. Tang, W. Yin, X. Ren, X. Wang, X. Zhang, X. Ren, Y. Fan, Y. Su, Y. Zhang, Y. Zhang, Y. Wan, Y. Liu, Z. Wang, Z. Cui, Z. Zhang, Z. Zhou, and Z. Qiu. Qwen3 technical report. arXiv preprint arXiv:2505.09388, 2025.
[94] S. Yao, B. Peng, C. Papadimitriou, and K. Narasimhan. Self-attention networks can process bounded hierarchical languages. arXiv preprint arXiv:2105.11115, 2021.
[95] C. Yun, S. Bhojanapalli, A. S. Rawat, S. Reddi, and S. Kumar. Are transformers universal approximators of sequence-to-sequence functions? In International Conference on Learning Representations, 2020.
[96] M. Zaheer, G. Guruganesh, K. A. Dubey, J. Ainslie, C. Alberti, S. Ontanon, P. Pham, A. Ravula, Q. Wang, L. Yang, et al. Big bird: Transformers for longer sequences. Advances in neural information processing systems, 33:17283–17297, 2020.
[97] D. Zhang, S. Zhoubian, Z. Hu, Y. Yue, Y. Dong, and J. Tang. ReST-MCTS*: LLM self-training via process reward guided tree search. In The Thirty-eighth Annual Conference on Neural Information Processing Systems, 2024.
[98] K. Zhang, G. Li, H. Zhang, and Z. Jin. Hirope: Length extrapolation for code models using hierarchical position. arXiv preprint arXiv:2403.19115, 2024.
[99] Q. Zhang, F. Lyu, Z. Sun, L. Wang, W. Zhang, Z. Guo, Y. Wang, I. King, X. Liu, and C. Ma. What, how, where, and how well? a survey on test-time scaling in large language models. arXiv preprint arXiv:2503.24235, 2025.
[100] Y. Zhang, M. Khalifa, L. Logeswaran, J. Kim, M. Lee, H. Lee, and L. Wang. Small language models need strong verifiers to self-correct reasoning. In ACL (Findings), 2024.
[101] Y. Zhang, S. Wu, Y. Yang, J. Shu, J. Xiao, C. Kong, and J. Sang. o1-coder: an o1 replication for coding. In arXiv, 2024.
[102] H. Zhao, A. Panigrahi, R. Ge, and S. Arora. Do transformers parse while predicting the masked word? In H. Bouamor, J. Pino, and K. Bali, editors, Proceedings of the 2023 Conference on Empirical Methods in Natural Language Processing, pages 16513–16542, Singapore, Dec. 2023. Association for Computational Linguistics.

Appendix A Proofs

A.1 Proof of Theorem 3.1

Proof.

Write $\mathcal{O}=\{1,\dots,O\}$ ( $O\in\mathbb{Z}_{+}$ ) where $i$ is the $i$ -th most likely answer and let $n_{i}$ denote the number of occurrences of $i$ . Then we have

\displaystyle\hat{p}=\frac{1}{n}(n_{1},\dots,n_{O})\sim\frac{1}{n}\mathrm{Multinomial}(n,p),

where $p=(p(1),\dots,p(O))$ .

Upper bound.

When $n\geq\frac{2\log(1/\delta)}{\Delta^{2}}$ we apply Claim A.5 to obtain that with probability at least $1-\delta$ ,

\displaystyle\|\hat{p}-p\|_{1}\leq\sqrt{\frac{2\ln(1/\delta)}{n}}\leq\Delta.

Under this event, we have that for any $i>1$

	$\displaystyle n_{1}-n_{i}=$	$\displaystyle\penalty\ n\cdot(\hat{p}_{1}-\hat{p}_{i})$
	$\displaystyle\geq$	$\displaystyle\penalty\ n\cdot({p}_{1}-{p}_{i}-\\|\hat{p}-p\\|_{1})$
	$\displaystyle\geq$	$\displaystyle\penalty\ 0$

and hence the correct answer $1$ is the most consistent answer. It follows that self-consistency can produce the correct answer with probability at least $1-\delta$ .

Lower bound.

When $n\leq\frac{1}{\Delta^{2}}$ , we construct the hard instance where $p_{1}=(1+\Delta)/2,p_{2}=(1-\Delta)/2$ and $\Delta<0.00001$ . If $n\leq\frac{1}{\Delta}$ then by the proof of Theorem 3.2, with constant probability the correct answer is not generated at all and hence self-consistency fails to produce the correct answer. Otherwise $n\geq\frac{1}{\Delta}\geq 10000$ . We may write $X:=\frac{n_{1}-n_{2}-n\Delta}{\sqrt{n}}$ as a sum of i.i.d. random variables divided by $\sqrt{n}$ :

\displaystyle X=\frac{\sum_{i=1}^{n}Y_{i}}{\sqrt{n}},

where $\mathbb{E}(Y_{i})=0,\sigma^{2}=\mathbb{E}(Y_{i}^{2})\geq 1/2,\rho=\mathbb{E}(|Y_{i}|^{3})\leq 1$ . By Claim A.6, we have that

	$\displaystyle\mathbb{P}(n_{1}<n_{2})=$	$\displaystyle\penalty\ \mathbb{P}(X<-1)$
	$\displaystyle\geq$	$\displaystyle\penalty\ \Phi(-1)-\frac{8\rho}{\sigma^{3}\sqrt{n}}$
	$\displaystyle\geq$	$\displaystyle\penalty\ 0.01.$

Thus in both cases, self-consistency fails to produce the correct answer with constant probability. ∎

A.2 Proof of Theorem 3.2

Proof.

Write $\mathcal{O}=\{1,\dots,O\}$ where $i$ is the $i$ -th most likely answer and let $n_{i}$ denote the number of occurrences of $i$ . Then we have

\displaystyle p(1)\geq p(2)+\Delta\geq\Delta.

Note that for best-of- $n$ , correctness is achieved if the correct answer appears at least once among $n$ independent samples.

Upper bound.

When $n\geq\frac{2\log(1/\delta)}{\Delta}$ , we have

	$\displaystyle\mathbb{P}(\text{Best-of-}n\text{ outputs correct answer})=$	$\displaystyle\penalty\ 1-(1-p(1))^{n}$
	$\displaystyle\geq$	$\displaystyle\penalty\ 1-(1-\Delta)^{\frac{2\log(1/\delta)}{\Delta}}$
	$\displaystyle\geq$	$\displaystyle\penalty\ 1-\delta.$

This confirms that best-of- $n$ achieves the correct answer with $1-\delta$ probability.

Lower bound.

When $n\leq\frac{1}{\Delta}$ , we construct the hard instance where $p(1)=\Delta+(1-\Delta)/O,p(2)=\cdots=p(O)=(1-\Delta)/O$ and $\Delta<0.0000001$ . Since the correct answer occurs with probability at least $\Delta$ , we have:

	$\displaystyle\mathbb{P}(\text{Best-of-}n\text{ outputs correct answer})=$	$\displaystyle\penalty\ 1-(1-p(1))^{n}$
	$\displaystyle\leq$	$\displaystyle\penalty\ 1-(1-2\Delta)^{\frac{1}{\Delta}}$
	$\displaystyle\leq$	$\displaystyle\penalty\ 0.99.$

This confirms that best-of- $n$ fails to produce the correct answer with constant probability. ∎

A.3 Proof of Proposition 4.2

We first introduce the following result that extends any Transformer to a larger vocabulary, so that it only attends to tokens in its original vocabulary.

Proposition A.1 (Extended Representation to Multiple Token Spaces).

For any $H,L,N_{\max}\in\mathbb{Z}_{+}$ , $\mathcal{V}_{1}\cap\mathcal{V}_{0}=\emptyset$ , there exists a general-purpose Transformer $\phi$ of type $(O(1),O(\log N_{\max}))$ such that for any Transformers $f=(\theta,\mathrm{pe},(\mathbf{K}^{(l)}_{h},\mathbf{Q}^{(l)}_{h},\mathbf{V}^{(l)}_{h})_{h\in[H],l\in[L]},\vartheta,\mathcal{V}_{1})$ over vocabulary $\mathcal{V}_{1}$ , the Transformer $\widetilde{f}=\phi(f_{1})$ satisfies the following property: for any token sequence $v=v_{1}\cdots v_{n}$ such that $n\leq N_{\max}$ , denote $\{i_{1}<\cdots<i_{m}\}=\{i:v_{i}\in\mathcal{V}_{1}\}$ , then we have

\displaystyle p_{\widetilde{f}}(\cdot|v)=p_{f}(\cdot|u),

where $u=v_{i_{1}}\cdots v_{i_{m}}$ .

Proof.

Set constants $B_{v},B_{qk},B_{\theta}$ such that for any layer $l$ and head $h$ , it holds that $\left\|(\mathbf{Q}^{(l)}_{h})^{\top}\mathbf{K}^{(l)}_{h}\right\|_{2}\leq B_{qk}$ , $\left\|\mathbf{V}^{(l)}_{h}\right\|_{2}\leq B_{v}$ , and $\|\theta(v)\|_{2}\leq B_{\theta}$ holds for all $v\in\mathcal{V}$ . Let $B=(HB_{v})^{L}B_{qk}B_{\theta},C=4B^{2}+\log(1/\epsilon),C_{0}=4C$ . By Lemma A.3, there exists $\alpha_{1},\dots,\alpha_{N_{\max}},\beta_{0},\beta_{1}\in\mathbb{R}^{d_{0}}$ and $A_{0},A_{1},A\in\mathbb{R}^{d_{0}\times d_{0}}$ for $d_{0}\leq O(\log N_{\max})$ such that

For any $i\geq j_{1},j_{2},j_{3}$ :

	$\displaystyle(\alpha_{i}+\beta_{1})^{\top}A_{0}(\alpha_{j_{1}}+\beta_{1})=$	$\displaystyle\penalty\ (\alpha_{i}+\beta_{1})^{\top}A_{0}(\alpha_{j_{2}}+\beta_{1})\geq(\alpha_{i}+\beta_{1})^{\top}A_{0}(\alpha_{j_{1}}+\beta_{0})+C_{0}$
	$\displaystyle(\alpha_{i}+\beta_{0})^{\top}A_{0}(\alpha_{i}+\beta_{0})\geq$	$\displaystyle\penalty\ (\alpha_{i}+\beta_{0})^{\top}A_{0}(\alpha_{j_{1}}+\beta_{1})+C_{0},$		(3)

For any $i>j$

	$\displaystyle(\alpha_{i}+\beta_{1})^{\top}A(\alpha_{i}+\beta_{1})\geq$	$\displaystyle\penalty\ (\alpha_{i}+\beta_{1})^{\top}A(\alpha_{j}+\beta_{1})+C_{0}$
	$\displaystyle\geq$	$\displaystyle\penalty\ (\alpha_{i}+\beta_{1})^{\top}A(\alpha_{j}+\beta_{0})+2C_{0},$		(4)

For any $i\geq j,j_{1}$

	$\displaystyle(\alpha_{i}+\beta_{1})^{\top}A_{1}(\alpha_{j}+\beta_{0})=$	$\displaystyle\penalty\ (\alpha_{i}+\beta_{1})^{\top}A_{1}(\alpha_{j_{1}}+\beta_{1})+C_{0}$
	$\displaystyle(\alpha_{i}+\beta_{1})^{\top}A_{1}(\alpha_{i}+\beta_{1})\geq$	$\displaystyle\penalty\ \max\{(\alpha_{i}+\beta_{1})^{\top}A_{1}(\alpha_{j_{1}}+\beta_{1}),(\alpha_{i}+\beta_{1})^{\top}A_{1}(\alpha_{j_{1}}+\beta_{0})\}+C_{0}.$		(5)

We define $\phi$ as follows: for any Transformers $f=(\theta,\mathrm{pe},(\mathbf{K}^{(l)}_{h},\mathbf{Q}^{(l)}_{h},\mathbf{V}^{(l)}_{h})_{h\in[H],l\in[L]},\vartheta,\mathcal{V}_{1})$ , the Transformer $\widetilde{f}=\phi(f)$ is given by

\displaystyle(\widetilde{\theta},\widetilde{\mathrm{pe}},(\widetilde{\mathbf{K}}^{(l)}_{h},\widetilde{\mathbf{Q}}^{(l)}_{h},\widetilde{\mathbf{V}}^{(l)}_{h})_{h\in[H+1],l\in[L]},\widetilde{\vartheta},\mathcal{V}_{1}\cup\mathcal{V}_{0}),

where the tokenizer is given by

\displaystyle\widetilde{\theta}(v)=\mathbbm{1}(v\in\mathcal{V}_{1})\cdot\begin{pmatrix}\theta(v)\\ \beta_{1}\end{pmatrix}+\mathbbm{1}(v\in\mathcal{V}_{0})\cdot\begin{pmatrix}0\\ \beta_{0}\end{pmatrix},

the positional encoder is given by

\displaystyle\widetilde{\mathrm{pe}}\left(\begin{pmatrix}x\\ y\end{pmatrix};v_{1},\dots,v_{i}\right)=\begin{pmatrix}\mathrm{pe}\left(x;u\right)\\ \alpha_{i}+y\end{pmatrix},

where $u=v_{i_{1}}\cdots v_{i_{m}}$ and $x\in\mathbb{R}^{d}$ ; for $l=1,\dots,L$ the key, query, value matrices are given by

		$\displaystyle\penalty\ \widetilde{\mathbf{K}}^{(l)}_{h}=\begin{pmatrix}\mathbf{K}^{(l)}_{h}&\\ &A_{0}\end{pmatrix},\penalty\ \widetilde{\mathbf{Q}}^{(l)}_{h}=\begin{pmatrix}\mathbf{Q}^{(l)}_{h}&\\ &I\end{pmatrix},$
		$\displaystyle\penalty\ \widetilde{\mathbf{V}}^{(l)}_{h}=\begin{pmatrix}\mathbf{V}^{(l)}_{h}&\\ &0\end{pmatrix},$
		$\displaystyle\penalty\ \widetilde{\mathbf{K}}^{(l)}_{H+1}=\begin{pmatrix}0&\\ &A\end{pmatrix},\penalty\ \widetilde{\mathbf{Q}}^{(l)}_{H+1}=\begin{pmatrix}0&\\ &I\end{pmatrix},\penalty\ \widetilde{\mathbf{V}}^{(l)}_{H+1}=\begin{pmatrix}0&\\ &I\end{pmatrix}.$

The output feature is given by $\widetilde{\vartheta}(y)=\begin{pmatrix}\vartheta(y)\\ 0\end{pmatrix}$ . Since $i_{1},\dots,i_{m}$ only depends on whether $v_{i}$ ’s belong to the set $\mathcal{V}_{1}$ , the generalized position encoding $\mathrm{pe}$ is well-defined. It can be verified that $\phi$ is indeed a general-purpose Transformer of type $(O(1),O(\log N_{\max}))$ .

We show that for any $l=1,\dots,L$ ,

\displaystyle\widetilde{X}^{(l)}_{i}=\begin{pmatrix}X^{(l)}_{i}\\ \widetilde{\alpha}_{i}\end{pmatrix},\penalty\ \forall i=i_{1},\dots,i_{m}

(6)

where $X^{(l)}_{i}$ is the $l$ -th layer of Transformer $f$ at position $i$ (attending only to positions $i_{1},\dots,i_{m}$ ) such that

\displaystyle\|X^{(l)}_{i}\|_{2}\leq B_{\theta}(HB_{v})^{l},

(7)

and

\displaystyle\widetilde{X}^{(l)}_{j}=\begin{pmatrix}0\\ \widetilde{\alpha}_{j}\end{pmatrix},\penalty\ \forall j\notin\{i_{1},\dots,i_{m}\}

(8)

where $\widetilde{\alpha}_{i}=\alpha_{i}+\mathbbm{1}(v\in\mathcal{V}_{0})\cdot\beta_{0}+\mathbbm{1}(v\in\mathcal{V}_{1})\cdot\beta_{1}$ .

We prove these results by induction. The case $l=1$ folows directly from the definitions of the tokenizer.

Prove Eq. (6).

Suppose Eq. (6) and Eq. (8) hold for $1,\dots,l-1$ -th layer, and consider $l$ -the layer. We have

	$\displaystyle\widetilde{X}^{(l+1)}_{i}=$	$\displaystyle\penalty\ \underbrace{\sum_{h=1}^{H}\sum_{j=1}^{i}\frac{\exp\left((\widetilde{\mathbf{Q}}^{(l)}_{h}\widetilde{X}^{(l)}_{i})^{\top}(\widetilde{\mathbf{K}}^{(l)}_{h}\widetilde{X}^{(l)}_{j})\right)}{\widetilde{Z}^{(l)}_{h}}\cdot\widetilde{\mathbf{V}}^{(l)}_{h}\widetilde{X}^{(l)}_{j}}_{\text{term 1}}$
		$\displaystyle\penalty\ +\underbrace{\sum_{j=1}^{i}\frac{\exp\left((\widetilde{\mathbf{Q}}^{(l)}_{H+1}\widetilde{X}^{(l)}_{i})^{\top}(\widetilde{\mathbf{K}}^{(l)}_{H+1}\widetilde{X}^{(l)}_{j})\right)}{\widetilde{Z}^{(l)}_{H+1}}\cdot\widetilde{\mathbf{V}}^{(l)}_{H+1}\widetilde{X}^{(l)}_{j}}_{\text{term 2}}.$

Eq. (1) ensures that for any $i,i^{\prime}\in\{i_{1},\dots,i_{m}\},j\notin\{i_{1},\dots,i_{m}\}$ :

	$\displaystyle(\widetilde{\mathbf{Q}}^{(l)}_{h}\widetilde{X}^{(l)}_{i})^{\top}(\widetilde{\mathbf{K}}^{(l)}_{h}\widetilde{X}^{(l)}_{i^{\prime}})=$	$\displaystyle\penalty\ (\mathbf{Q}^{(l)}_{h}\widetilde{X}^{(l)}_{i})^{\top}(\mathbf{K}^{(l)}_{h}\widetilde{X}^{(l)}_{i^{\prime}})+(\alpha_{i}+\beta_{1})^{\top}A_{0}(\alpha_{i^{\prime}}+\beta_{1})$
	$\displaystyle\geq$	$\displaystyle\penalty\ (\mathbf{Q}^{(l)}_{h}X^{(l)}_{i})^{\top}(\mathbf{K}^{(l)}_{h}X^{(l)}_{j})+(\alpha_{i}+\beta_{1})^{\top}A_{0}(\alpha_{j}+\beta_{0})+C$
	$\displaystyle=$	$\displaystyle\penalty\ (\widetilde{\mathbf{Q}}^{(l)}_{h}\widetilde{X}^{(l)}_{i})^{\top}(\widetilde{\mathbf{K}}^{(l)}_{h}\widetilde{X}^{(l)}_{j})+C,$

and if $i,j_{1},j_{2}\in\{i_{1},\dots,i_{m}\}$

		$\displaystyle\penalty\ (\widetilde{\mathbf{Q}}^{(l)}_{h}\widetilde{X}^{(l)}_{i})^{\top}(\widetilde{\mathbf{K}}^{(l)}_{h}\widetilde{X}^{(l)}_{j_{1}})-(\widetilde{\mathbf{Q}}^{(l)}_{h}\widetilde{X}^{(l)}_{i})^{\top}(\widetilde{\mathbf{K}}^{(l)}_{h}\widetilde{X}^{(l)}_{j_{2}})$
	$\displaystyle=$	$\displaystyle\penalty\ (\mathbf{Q}^{(l)}_{h}X^{(l)}_{i})^{\top}(\mathbf{K}^{(l)}_{h}X^{(l)}_{j_{1}})+(\alpha_{i}+\beta_{1})^{\top}A_{0}(\alpha_{j_{1}}+\beta_{1)})-(\mathbf{Q}^{(l)}_{h}X^{(l)}_{i})^{\top}(\mathbf{K}^{(l)}_{h}X^{(l)}_{j_{2}})-(\alpha_{i}+\beta_{1})^{\top}A_{0}(\alpha_{j_{2}}+\beta_{1})$
	$\displaystyle=$	$\displaystyle\penalty\ (\mathbf{Q}^{(l)}_{h}X^{(l)}_{i})^{\top}(\mathbf{K}^{(l)}_{h}\widetilde{X}^{(l)}_{j_{1}})-(\mathbf{Q}^{(l)}_{h}\widetilde{X}^{(l)}_{i})^{\top}(\mathbf{K}^{(l)}_{h}X^{(l)}_{j_{2}}),$

where we use the fact that $C_{0}\geq C+2\max_{h,l,i,j}\left|(\mathbf{Q}^{(l)}_{h}X^{(l)}_{i})^{\top}(\mathbf{K}^{(l)}_{h}X^{(l)}_{j})\right|$ . Since the transformers have precision $\epsilon$ and $C\geq 2\max_{h,l,i,j}\left|(\mathbf{Q}^{(l)}_{h}X^{(l)}_{i})^{\top}(\mathbf{K}^{(l)}_{h}X^{(l)}_{j})\right|+\log(1/\epsilon)$ , it follows that the attention weights of head $(k-1)H+h$ is identical to the attention weights of expert $k$ , i.e.

\displaystyle\frac{\exp\left((\widetilde{\mathbf{Q}}^{(l)}_{h}\widetilde{X}^{(l)}_{i})^{\top}(\widetilde{\mathbf{K}}^{(l)}_{h}\widetilde{X}^{(l)}_{j})\right)}{\widetilde{Z}^{(l)}_{h}}=\mathbbm{1}(j\in\{i_{1},\dots,i_{m}\})\cdot\frac{\exp\left((\mathbf{Q}^{(l)}_{h}X^{(l)}_{i})^{\top}(\mathbf{K}^{(l)}_{h}X^{(l)}_{j})\right)}{Z^{(l)}_{h}}.

Therefore

\displaystyle\text{term 1}=\sum_{h=1}^{H}\sum_{j=i_{1},\dots,i_{m}}\frac{\exp\left((\mathbf{Q}^{(l)}_{h}X^{(l)}_{i})^{\top}(\mathbf{K}^{(l)}_{h}X^{(l)}_{j})\right)}{Z^{(l)}_{h}}\cdot\begin{pmatrix}\mathbf{V}^{(l)}_{h}X^{(l)}_{j}\\ 0\end{pmatrix}=\begin{pmatrix}X^{(l+1)}_{j}\\ 0\end{pmatrix}.

Furthermore, by Eq. (2) we have for any $j<i$

	$\displaystyle(\widetilde{\mathbf{Q}}^{(l)}_{H+1}\widetilde{X}^{(l)}_{i})^{\top}(\widetilde{\mathbf{K}}^{(l)}_{H+1}\widetilde{X}^{(l)}_{i})=$	$\displaystyle\penalty\ \widetilde{\alpha}_{i}^{\top}A\widetilde{\alpha}_{i}$
	$\displaystyle\geq$	$\displaystyle\penalty\ \widetilde{\alpha}_{i}^{\top}A\widetilde{\alpha}_{j}+C$
	$\displaystyle=$	$\displaystyle\penalty\ (\widetilde{\mathbf{Q}}^{(l)}_{H+1}\widetilde{X}^{(l)}_{i})^{\top}(\widetilde{\mathbf{K}}^{(l)}_{H+1}\widetilde{X}^{(l)}_{j})+C,$

and hence the attention weighs concentrates on $i$ itself. Thus

\displaystyle\text{term 2}=\begin{pmatrix}0&\\ &I\end{pmatrix}\cdot\begin{pmatrix}X^{(l)}_{i}\\ \widetilde{\alpha}_{i}\end{pmatrix}=\begin{pmatrix}0\\ \widetilde{\alpha}_{i}\end{pmatrix}.

Combining, we derive Eq.(6) for $(l+1)$ -th layer.

Prove Eq. (7).

From above,

	$\displaystyle\\|X^{(l+1)}_{i}\\|_{2}=$	$\displaystyle\penalty\ \left\\|\sum_{h=1}^{H}\sum_{j=1}^{i}\frac{\exp\left((\widetilde{\mathbf{Q}}^{(l)}_{h}\widetilde{X}^{(l)}_{i})^{\top}(\widetilde{\mathbf{K}}^{(l)}_{h}\widetilde{X}^{(l)}_{j})\right)}{\widetilde{Z}^{(l)}_{h}}\cdot\mathbf{V}^{(l)}_{h}X^{(l)}_{j}\right\\|_{2}$
	$\displaystyle\leq$	$\displaystyle\penalty\ HB_{v}\cdot\max_{j\leq i}\\|X^{(l)}_{j}\\|_{2}$
	$\displaystyle\leq$	$\displaystyle\penalty\ B_{\theta}(HB_{v})^{l+1}.$

This confirms Eq. (24) for $l+1$ .

Prove Eq. (8).

Notice that Eq. (1) ensures that for any $j,j^{\prime}\notin\{i:v_{i}\in\mathcal{V}_{1}\}$ and $i\in\{i:v_{i}\in\mathcal{V}_{1}\}$ :

	$\displaystyle(\widetilde{\mathbf{Q}}^{(l)}_{h}\widetilde{X}^{(l)}_{j})^{\top}(\widetilde{\mathbf{K}}^{(l)}_{h}\widetilde{X}^{(l)}_{j^{\prime}})=$	$\displaystyle\penalty\ (\mathbf{Q}^{(l)}_{h}X^{(l)}_{j})^{\top}(\mathbf{K}^{(l)}_{h}X^{(l)}_{j^{\prime}})+(\alpha_{j}+\beta_{0})^{\top}A_{0}(\alpha_{j^{\prime}}+\beta_{0})$
	$\displaystyle\geq$	$\displaystyle\penalty\ (\mathbf{Q}^{(l)}_{h}X^{(l)}_{j})^{\top}(\mathbf{K}^{(l)}_{h}X^{(l)}_{i})+(\alpha_{j}+\beta_{0})^{\top}A_{0}(\alpha_{i}+\beta_{1})+C$
	$\displaystyle=$	$\displaystyle\penalty\ (\widetilde{\mathbf{Q}}^{(l)}_{h}\widetilde{X}^{(l)}_{j})^{\top}(\widetilde{\mathbf{K}}^{(l)}_{h}\widetilde{X}^{(l)}_{i})+C.$

It follows that the attention weights is concentrated on the compliment of $\{i:v_{i}\in\mathcal{V}_{1}\}$ itself, and therefore Eq. (8) follows by a simple induction argument.

Finally, at the output layer

	$\displaystyle p_{\widetilde{f}}(y\|v_{1},\dots,v_{n})=$	$\displaystyle\penalty\ \mathrm{Softmax}(\widetilde{\vartheta}(y)^{\top}\widetilde{X}^{(L)}_{n})$
	$\displaystyle=$	$\displaystyle\penalty\ \mathrm{Softmax}(\vartheta(y)^{\top}X^{(L)}_{m})$
	$\displaystyle=$	$\displaystyle\penalty\ p_{f_{\kappa}}(y\|u).$

This establishes the desired statement. ∎

Now we return to the proof of Proposition 4.2.

Proof.

By Proposition A.1, it suffices to construct general-purpose Transformer $\phi$ such that

\displaystyle p_{\widetilde{f}}(\cdot|v)=p_{f_{\kappa}}(\cdot|u),

where $u=v_{1}\cdots v_{i_{0}-1}v_{i_{0}+1}\cdots v_{n}$ , because then the $\widetilde{\phi}$ given by

\displaystyle\widetilde{\phi}(f_{1},\dots,f_{K})=\phi(\phi_{e}(f_{1}),\dots,\phi_{e}(f_{K}))

satisfies the requirement, where $\phi_{e}$ is the general-purpose Transformer that extends the $K$ Transformers to the larger vocabulary $\mathcal{V}:=\cup_{k=1}^{K}\mathcal{V}_{k}$ as given by Proposition A.1.

Set constants $B_{v},B_{qk},B_{\theta}$ such that for any layer $l$ and head $h$ , it holds that $\left\|(\mathbf{Q}^{(l)}_{h})^{\top}\mathbf{K}^{(l)}_{h}\right\|_{2}\leq B_{qk}$ , $\left\|\mathbf{V}^{(l)}_{h}\right\|_{2}\leq B_{v}$ , and $\|\theta(v)\|_{2}\leq B_{\theta}$ holds for all $v\in\mathcal{V}$ . Let $B=(KHB_{v})^{L}B_{qk}B_{\theta},C=4B^{2}+\log(1/\epsilon),C_{0}=4C$ . By Lemma A.3, there exists $\alpha_{1},\dots,\alpha_{N},\beta_{0},\beta_{1},\dots,\beta_{K}\in\mathbb{R}^{d_{0}}$ and $A_{1},\dots,A_{K}\in\mathbb{R}^{{d_{0}}\times{d_{0}}}$ for ${d_{0}}\leq O(K+\log N_{\max})$ such that

For any $i\geq j_{1},j_{2},j_{3}$ and $k,k^{\prime},k^{\prime\prime}\neq 0$ :

	$\displaystyle(\alpha_{i}+\beta_{k})^{\top}A_{0}(\alpha_{j_{1}}+\beta_{k^{\prime}})=$	$\displaystyle\penalty\ (\alpha_{i}+\beta_{k})^{\top}A_{0}(\alpha_{j_{2}}+\beta_{k^{\prime\prime}})\geq(\alpha_{i}+\beta_{k})^{\top}A_{0}(\alpha_{j_{1}}+\beta_{0})+C_{0}$
	$\displaystyle(\alpha_{i}+\beta_{0})^{\top}A_{0}(\alpha_{i}+\beta_{0})\geq$	$\displaystyle\penalty\ (\alpha_{i}+\beta_{0})^{\top}A_{0}(\alpha_{j_{1}}+\beta_{k})+C_{0},$		(9)

For any $i>j$ and $k\neq k^{\prime}\neq 0$

	$\displaystyle(\alpha_{i}+\beta_{k})^{\top}A(\alpha_{i}+\beta_{k})\geq$	$\displaystyle\penalty\ (\alpha_{i}+\beta_{k})^{\top}A(\alpha_{j}+\beta_{k^{\prime}})+C_{0}$
	$\displaystyle\geq$	$\displaystyle\penalty\ (\alpha_{i}+\beta_{k})^{\top}A(\alpha_{j}+\beta_{0})+2C_{0},$		(10)

For any $i\geq j,j_{1}$ and $k\neq k^{\prime},k^{\prime\prime}$

	$\displaystyle(\alpha_{i}+\beta_{k})^{\top}A_{k^{\prime}}(\alpha_{j}+\beta_{0})=$	$\displaystyle\penalty\ (\alpha_{i}+\beta_{k})^{\top}A_{k^{\prime}}(\alpha_{j_{1}}+\beta_{k^{\prime\prime}})+C_{0}$
	$\displaystyle(\alpha_{i}+\beta_{k})^{\top}A_{k}(\alpha_{i}+\beta_{k})\geq$	$\displaystyle\penalty\ \max\{(\alpha_{i}+\beta_{k})^{\top}A_{k}(\alpha_{j_{1}}+\beta_{k^{\prime\prime}}),(\alpha_{i}+\beta_{k})^{\top}A_{k^{\prime}}(\alpha_{j_{1}}+\beta_{0})\}+C_{0},$		(11)

We define $\phi$ as follows: for any Transformers

\displaystyle f_{k}=(\theta_{k},\mathrm{pe}_{k},(\mathbf{K}^{(l)}_{k;h},\mathbf{Q}^{(l)}_{k;h},\mathbf{V}^{(l)}_{k;h})_{h\in[H],l\in[L]},\vartheta_{k},\mathcal{V}_{k}),

over $\mathcal{V}_{k},\penalty\ k\in[K]$ , the Transformer $\widetilde{f}=\phi(f_{1},\dots,f_{K})$ is given by

\displaystyle(\widetilde{\theta},\widetilde{\mathrm{pe}},(\widetilde{\mathbf{K}}^{(l)}_{h},\widetilde{\mathbf{Q}}^{(l)}_{h},\widetilde{\mathbf{V}}^{(l)}_{h})_{h\in[KH+1],l\in[L+1]},\widetilde{\vartheta},\mathcal{V}),

where the tokenizer is given by

\displaystyle\widetilde{\theta}(v)=\mathbbm{1}(v\notin\mathcal{V}_{0})\cdot\begin{pmatrix}\theta_{1}(v)\\ \vdots\\ \theta_{K}(v)\\ 0\end{pmatrix}+\begin{pmatrix}0\\ \vdots\\ 0\\ \beta_{\mathcal{E}(v)}\end{pmatrix},

where $\mathcal{E}(v)=k$ iff $v\in\mathcal{V}_{k}$ . Let the positional encoder be given by

\displaystyle\widetilde{\mathrm{pe}}\left(\begin{pmatrix}x\\ y\end{pmatrix};v_{1},\dots,v_{i}\right)=\begin{pmatrix}\mathrm{pe}_{1}\left(x;u\right)\\ \vdots\\ \mathrm{pe}_{K}\left(x;u\right)\\ \alpha_{i}+y\end{pmatrix},

where $x\in\mathbb{R}^{d}$ and $u$ is the sub-sequence of $v$ that omits $v_{i_{0}}$ (if any); for $l=1,\dots,L$ the key, query, value matrices are given by

		$\displaystyle\penalty\ \widetilde{\mathbf{K}}^{(l)}_{(k-1)H+h}=\begin{pmatrix}0&&&&\\ &\ddots&&&\\ &&\mathbf{K}^{(l)}_{k;h}&&\\ &&&\ddots&\\ &&&&A_{0}\end{pmatrix},\penalty\ \widetilde{\mathbf{Q}}^{(l)}_{(k-1)H+h}=\begin{pmatrix}0&&&&\\ &\ddots&&&\\ &&\mathbf{Q}^{(l)}_{k;h}&&\\ &&&\ddots&\\ &&&&I\end{pmatrix},$
		$\displaystyle\penalty\ \widetilde{\mathbf{V}}^{(l)}_{(k-1)H+h}=\begin{pmatrix}0&&&&\\ &\ddots&&&\\ &&\mathbf{V}^{(l)}_{k;h}&&\\ &&&\ddots&\\ &&&&0\end{pmatrix},$

\displaystyle\widetilde{\mathbf{K}}^{(l)}_{KH+1}=\begin{pmatrix}0&&&\\ &\ddots&&\\ &&0&\\ &&&A\end{pmatrix},\penalty\ \widetilde{\mathbf{Q}}^{(l)}_{KH+1}=\begin{pmatrix}0&&&\\ &\ddots&&\\ &&0&\\ &&&I\end{pmatrix},\penalty\ \widetilde{\mathbf{V}}^{(l)}_{KH+1}=\begin{pmatrix}0&&&\\ &\ddots&&\\ &&0&\\ &&&I\end{pmatrix},

where the submatrices $\mathbf{K}^{(l)}_{k;h},\mathbf{Q}^{(l)}_{k;h},\mathbf{V}^{(l)}_{k;h}$ are located in the $k$ -th diagonal block, and for the final layer

\displaystyle\penalty\ \widetilde{\mathbf{K}}^{(L+1)}_{k}=\begin{pmatrix}0&&&\\ &\ddots&&\\ &&0&\\ &&&A_{k}\end{pmatrix},\penalty\ \widetilde{\mathbf{Q}}^{(L+1)}_{k}=\begin{pmatrix}0&&&\\ &\ddots&&\\ &&0&\\ &&&I\end{pmatrix},\penalty\ \widetilde{\mathbf{V}}^{(L+1)}_{k}=\begin{pmatrix}0&&&&\\ &\ddots&&&\\ &&I&&\\ &&&\ddots&\\ &&&&0\end{pmatrix},

where the identity sub-matrix in $\widetilde{\mathbf{V}}^{(L+1)}_{k}$ is located in the $k$ -th block. The output feature is given by $\widetilde{\vartheta}(y)=\begin{pmatrix}\vartheta_{1}(y)\\ \vdots\\ \vartheta_{K}(y)\\ 0\end{pmatrix}$ . Since $u^{(k)}$ ’s only depend on set membership information of $v_{i}$ ’s, the generalized position encoding $\mathrm{pe}$ is well-defined. We can easily verify that $\phi$ is indeed a general-purpose Transformer of type $(O(K),O(\log N_{\max}))$ .

We show that for any $l=1,\dots,L$ ,

\displaystyle\widetilde{X}^{(l)}_{i}=\begin{pmatrix}X^{(l)}_{1;i}\\ \vdots\\ X^{(l)}_{K;i}\\ \widetilde{\alpha}_{i}\end{pmatrix},\penalty\ \forall i\neq i_{0}

(12)

where $X^{(l)}_{k;i}$ is the $l$ -th layer of Transformer $k$ at position $i$ (attending to all positions but $i_{0}$ ) such that

\displaystyle\|X^{(l)}_{k;i}\|_{2}\leq B_{\theta}(KHB_{v})^{l}.

(13)

and

\displaystyle\widetilde{X}^{(l)}_{i_{0}}=\begin{pmatrix}0\\ \vdots\\ 0\\ \widetilde{\alpha}_{i_{0}}\end{pmatrix}

(14)

where $\widetilde{\alpha}_{i}=\alpha_{i}+\beta_{\mathcal{E}(v_{i})}$ .

We prove these results by induction. The case $l=1$ folows directly from the definitions of the tokenizer.

Prove Eq. (12).

Suppose Eq. (12) and Eq. (14) hold for $1,\dots,l-1$ =th layer, and consider $l$ -the layer. We have

	$\displaystyle\widetilde{X}^{(l+1)}_{i}=$	$\displaystyle\penalty\ \underbrace{\sum_{k=1}^{K}\sum_{h=1}^{H}\sum_{j=1}^{i}\frac{\exp\left((\widetilde{\mathbf{Q}}^{(l)}_{(k-1)H+h}\widetilde{X}^{(l)}_{i})^{\top}(\widetilde{\mathbf{K}}^{(l)}_{(k-1)H+h}\widetilde{X}^{(l)}_{j})\right)}{\widetilde{Z}^{(l)}_{(k-1)H+h}}\cdot\widetilde{\mathbf{V}}^{(l)}_{(k-1)H+h}\widetilde{X}^{(l)}_{j}}_{\text{term 1}}$
		$\displaystyle\penalty\ +\underbrace{\sum_{j=1}^{i}\frac{\exp\left((\widetilde{\mathbf{Q}}^{(l)}_{KH+1}\widetilde{X}^{(l)}_{i})^{\top}(\widetilde{\mathbf{K}}^{(l)}_{KH+1}\widetilde{X}^{(l)}_{j})\right)}{\widetilde{Z}^{(l)}_{KH+1}}\cdot\widetilde{\mathbf{V}}^{(l)}_{KH+1}\widetilde{X}^{(l)}_{j}}_{\text{term 2}}.$

Eq. (1) ensures that for any $j_{1}<j_{2}\leq i$ such that $i_{0}\notin\{i,j_{1},j_{2}\}$ :

	$\displaystyle(\widetilde{\mathbf{Q}}^{(l)}_{(k-1)H+h}\widetilde{X}^{(l)}_{i})^{\top}(\widetilde{\mathbf{K}}^{(l)}_{(k-1)H+h}\widetilde{X}^{(l)}_{j_{1}})=$	$\displaystyle\penalty\ (\mathbf{Q}^{(l)}_{k;h}X^{(l)}_{k;i})^{\top}(\mathbf{K}^{(l)}_{k;h}X^{(l)}_{k;j_{1}})+(\alpha_{i}+\beta_{\mathcal{E}(i)})^{\top}A_{0}(\alpha_{j_{1}}+\beta_{\mathcal{E}({j_{1}})})$
	$\displaystyle\geq$	$\displaystyle\penalty\ (\mathbf{Q}^{(l)}_{k;h}X^{(l)}_{k;i})^{\top}(\mathbf{K}^{(l)}_{k;h}X^{(l)}_{k;j_{1}})+(\alpha_{i}+\beta_{\mathcal{E}(i)})^{\top}A_{0}(\alpha_{i_{0}}+\beta_{\mathcal{E}({i_{0}})})+C$
	$\displaystyle=$	$\displaystyle\penalty\ (\widetilde{\mathbf{Q}}^{(l)}_{(k-1)H+h}\widetilde{X}^{(l)}_{i})^{\top}(\widetilde{\mathbf{K}}^{(l)}_{(k-1)H+h}\widetilde{X}^{(l)}_{i_{0}})+C.$

and

		$\displaystyle\penalty\ (\widetilde{\mathbf{Q}}^{(l)}_{(k-1)H+h}\widetilde{X}^{(l)}_{i})^{\top}(\widetilde{\mathbf{K}}^{(l)}_{(k-1)H+h}\widetilde{X}^{(l)}_{j_{1}})-(\widetilde{\mathbf{Q}}^{(l)}_{(k-1)H+h}\widetilde{X}^{(l)}_{i})^{\top}(\widetilde{\mathbf{K}}^{(l)}_{(k-1)H+h}\widetilde{X}^{(l)}_{j_{2}})$
	$\displaystyle=$	$\displaystyle\penalty\ (\mathbf{Q}^{(l)}_{k;h}X^{(l)}_{k;i})^{\top}(\mathbf{K}^{(l)}_{k;h}X^{(l)}_{k;j_{1}})+(\alpha_{i}+\beta_{\mathcal{E}(i)})^{\top}A_{0}(\alpha_{j_{1}}+\beta_{\mathcal{E}({j_{1}})})$
		$\displaystyle\penalty\ -(\mathbf{Q}^{(l)}_{k;h}X^{(l)}_{k;i})^{\top}(\mathbf{K}^{(l)}_{k;h}X^{(l)}_{k;j_{2}})-(\alpha_{i}+\beta_{\mathcal{E}(i)})^{\top}A_{0}(\alpha_{j_{2}}+\beta_{\mathcal{E}({j_{2}})})$
	$\displaystyle=$	$\displaystyle\penalty\ (\mathbf{Q}^{(l)}_{k;h}X^{(l)}_{k;i})^{\top}(\mathbf{K}^{(l)}_{k;h}X^{(l)}_{k;j_{1}})-(\mathbf{Q}^{(l)}_{k;h}X^{(l)}_{k;i})^{\top}(\mathbf{K}^{(l)}_{k;h}X^{(l)}_{k;j_{2}}).$

It follows from the precision $\epsilon$ of the transformers that the attention weights of head $(k-1)H+h$ is identical to the attention weights of expert $k$ , i.e.

\displaystyle\frac{\exp\left((\widetilde{\mathbf{Q}}^{(l)}_{(k-1)H+h}\widetilde{X}^{(l)}_{i})^{\top}(\widetilde{\mathbf{K}}^{(l)}_{(k-1)H+h}\widetilde{X}^{(l)}_{j})\right)}{\widetilde{Z}^{(l)}_{(k-1)H+h}}=\frac{\exp\left((\mathbf{Q}^{(l)}_{k;h}X^{(l)}_{k;i})^{\top}(\mathbf{K}^{(l)}_{k;h}X^{(l)}_{k;j})\right)}{Z^{(l)}_{k;h}}.

Therefore

\displaystyle\text{term 1}=\sum_{k=1}^{K}\sum_{h=1}^{H}\sum_{j=1}^{i}\frac{\exp\left((\mathbf{Q}^{(l)}_{k;h}X^{(l)}_{k;i})^{\top}(\mathbf{K}^{(l)}_{k;h}X^{(l)}_{k;j})\right)}{Z^{(l)}_{k;h}}\cdot\begin{pmatrix}0\\ \vdots\\ \mathbf{V}^{(l)}_{k;h}X^{(l)}_{k;j}\\ \vdots\\ 0\end{pmatrix}=\begin{pmatrix}X^{(l)}_{1;i}\\ \vdots\\ X^{(l)}_{K;i}\\ 0\end{pmatrix}.

Furthermore, by Eq. (2) we have for any $j<i$

	$\displaystyle(\widetilde{\mathbf{Q}}^{(l)}_{KH+1}\widetilde{X}^{(l)}_{i})^{\top}(\widetilde{\mathbf{K}}^{(l)}_{KH+1}\widetilde{X}^{(l)}_{i})=$	$\displaystyle\penalty\ \widetilde{\alpha}_{i}^{\top}A\widetilde{\alpha}_{i}$
	$\displaystyle\geq$	$\displaystyle\penalty\ \widetilde{\alpha}_{i}^{\top}A\widetilde{\alpha}_{j}+C$
	$\displaystyle=$	$\displaystyle\penalty\ (\widetilde{\mathbf{Q}}^{(l)}_{KH+1}\widetilde{X}^{(l)}_{i})^{\top}(\widetilde{\mathbf{K}}^{(l)}_{KH+1}\widetilde{X}^{(l)}_{j})+C$

and hence the attention weighs concentrates on $i$ itself. Thus

\displaystyle\text{term 2}=\begin{pmatrix}0&&&\\ &\ddots&&\\ &&0&\\ &&&I\end{pmatrix}\cdot\begin{pmatrix}X^{(l)}_{1;i}\\ \vdots\\ X^{(l)}_{K;i}\\ \widetilde{\alpha}_{i}\end{pmatrix}=\begin{pmatrix}0\\ \vdots\\ 0\\ \widetilde{\alpha}_{i}\end{pmatrix}

Combining these two terms, we confirm that Eq.(12) holds for $(l+1)$ -th layer.

Prove Eq. (13).

From above,

	$\displaystyle\\|X^{(l+1)}_{k;i}\\|_{2}=$	$\displaystyle\penalty\ \left\\|\sum_{k=1}^{K}\sum_{h=1}^{H}\sum_{j=1}^{i}\frac{\exp\left((\widetilde{\mathbf{Q}}^{(l)}_{(k-1)H+h}\widetilde{X}^{(l)}_{i})^{\top}(\widetilde{\mathbf{K}}^{(l)}_{(k-1)H+h}\widetilde{X}^{(l)}_{j})\right)}{\widetilde{Z}^{(l)}_{(k-1)H+h}}\cdot\mathbf{V}^{(l)}_{k;h}X^{(l)}_{k;j}\right\\|_{2}$
	$\displaystyle\leq$	$\displaystyle\penalty\ KHB_{v}\cdot\max_{j\leq i}\\|X^{(l)}_{k;j}\\|_{2}$
	$\displaystyle\leq$	$\displaystyle\penalty\ B_{\theta}(KHB_{v})^{l+1}.$

This confirms Eq. (13) for $l+1$ .

Prove Eq. (14).

Notice that Eq. (1) ensures that for any $j\leq i_{0}$ :

	$\displaystyle(\widetilde{\mathbf{Q}}^{(l)}_{(k-1)H+h}\widetilde{X}^{(l)}_{i_{0}})^{\top}(\widetilde{\mathbf{K}}^{(l)}_{(k-1)H+h}\widetilde{X}^{(l)}_{i_{0}})=$	$\displaystyle\penalty\ (\mathbf{Q}^{(l)}_{k;h}X^{(l)}_{k;i_{0}})^{\top}(\mathbf{K}^{(l)}_{k;h}X^{(l)}_{k;i_{0}})+(\alpha_{i_{0}}+\beta_{\mathcal{E}({i_{0}})})^{\top}A_{0}(\alpha_{i_{0}}+\beta_{\mathcal{E}({i_{0}})})$
	$\displaystyle\geq$	$\displaystyle\penalty\ (\mathbf{Q}^{(l)}_{k;h}X^{(l)}_{k;i_{0}})^{\top}(\mathbf{K}^{(l)}_{k;h}X^{(l)}_{k;j})+(\alpha_{i_{0}}+\beta_{\mathcal{E}({i_{0}})})^{\top}A_{0}(\alpha_{j}+\beta_{\mathcal{E}({j})})+C$
	$\displaystyle=$	$\displaystyle\penalty\ (\widetilde{\mathbf{Q}}^{(l)}_{(k-1)H+h}\widetilde{X}^{(l)}_{i_{0}})^{\top}(\widetilde{\mathbf{K}}^{(l)}_{(k-1)H+h}\widetilde{X}^{(l)}_{j})+C.$

It follows that the attention weights of head $(k-1)H+h$ is concentrated on $i_{0}$ itself, therefore

\displaystyle\text{term 1}=\sum_{k=1}^{K}\sum_{h=1}^{H}\begin{pmatrix}0\\ \vdots\\ \mathbf{V}^{(l)}_{k;h}\cdot 0\\ \vdots\\ 0\end{pmatrix}=0.

By the same argument, for $i=i_{0}$ we have

\displaystyle\text{term 2}=\begin{pmatrix}0&&&\\ &\ddots&&\\ &&0&\\ &&&I\end{pmatrix}\cdot\begin{pmatrix}0\\ \vdots\\ 0\\ \widetilde{\alpha}_{i_{0}}\end{pmatrix}=\begin{pmatrix}0\\ \vdots\\ 0\\ \widetilde{\alpha}_{i_{0}}\end{pmatrix}.

Combining these confirms Eq. (14).

Next, we show that the last layer satisfies

\displaystyle\widetilde{X}^{(L+1)}_{n}=\begin{pmatrix}0\\ \vdots\\ X^{(L+1)}_{\kappa;n}\\ \vdots\\ 0\end{pmatrix}

(15)

where $X^{(L+1)}_{\kappa;n}$ is the $\kappa$ -th block. To see this, we notice that Eq. (3) implies the followings (the proofs are identical to the above):

Attention sink to dummny token $v_{i_{0}}$ for mismatch expert: for any $k^{\prime}\neq\kappa$ and $j\leq n$ we have

$\displaystyle(\widetilde{\mathbf{Q}}^{(L)}_{(k^{\prime}-1)H+h}\widetilde{X}^{(L)}_{n})^{\top}(\widetilde{\mathbf{K}}^{(L)}_{(k^{\prime}-1)H+h}\widetilde{X}^{(L)}_{j})=$	$\displaystyle\penalty\ (\alpha_{n}+\beta_{\mathcal{E}({n})})^{\top}A_{k^{\prime}}(\alpha_{j}+\beta_{\mathcal{E}({j})})$
$\displaystyle\leq$	$\displaystyle\penalty\ (\alpha_{n}+\beta_{\mathcal{E}({n})})^{\top}A_{k^{\prime}}(\alpha_{i_{0}}+\beta_{\mathcal{E}({i_{0}})})-C$
$\displaystyle=$	$\displaystyle\penalty\ (\widetilde{\mathbf{Q}}^{(L)}_{(k^{\prime}-1)H+h}\widetilde{X}^{(L)}_{n})^{\top}(\widetilde{\mathbf{K}}^{(L)}_{(k^{\prime}-1)H+h}\widetilde{X}^{(L)}_{i_{0}})-C.$	(16)

Attention to oneself for matching expert: for any $j\neq i_{0}$ we have

$\displaystyle(\widetilde{\mathbf{Q}}^{(L)}_{(\kappa-1)H+h}\widetilde{X}^{(L)}_{n})^{\top}(\widetilde{\mathbf{K}}^{(L)}_{(\kappa-1)H+h}\widetilde{X}^{(L)}_{j})=$	$\displaystyle\penalty\ (\alpha_{n}+\beta_{\mathcal{E}({n})})^{\top}A_{\kappa}(\alpha_{j}+\beta_{\mathcal{E}({j})})$
$\displaystyle\geq$	$\displaystyle\penalty\ (\alpha_{n}+\beta_{\mathcal{E}({n})})^{\top}A_{\kappa}(\alpha_{i_{0}}+\beta_{\mathcal{E}({i_{0}})})+C$
$\displaystyle=$	$\displaystyle\penalty\ (\widetilde{\mathbf{Q}}^{(L)}_{(\kappa-1)H+h}\widetilde{X}^{(L)}_{n})^{\top}(\widetilde{\mathbf{K}}^{(L)}_{(\kappa-1)H+h}\widetilde{X}^{(L)}_{i_{0}})+C,$	(17)

and

$\displaystyle(\widetilde{\mathbf{Q}}^{(L)}_{(\kappa-1)H+h}\widetilde{X}^{(L)}_{n})^{\top}(\widetilde{\mathbf{K}}^{(L)}_{(\kappa-1)H+h}\widetilde{X}^{(L)}_{n})=$	$\displaystyle\penalty\ (\alpha_{n}+\beta_{\mathcal{E}({n})})^{\top}A_{\kappa}(\alpha_{n}+\beta_{\mathcal{E}({n})})$
$\displaystyle\geq$	$\displaystyle\penalty\ (\alpha_{n}+\beta_{\mathcal{E}({n})})^{\top}A_{\kappa}(\alpha_{j}+\beta_{\mathcal{E}({j})})+C$
$\displaystyle=$	$\displaystyle\penalty\ (\widetilde{\mathbf{Q}}^{(L)}_{(\kappa-1)H+h}\widetilde{X}^{(L)}_{n})^{\top}(\widetilde{\mathbf{K}}^{(L)}_{(\kappa-1)H+h}\widetilde{X}^{(L)}_{j})+C.$	(18)

Combining Eq. (1), Eq. (2), and Eq. (2), we have

\displaystyle\frac{\exp\left((\widetilde{\mathbf{Q}}^{(L)}_{(k-1)H+h}\widetilde{X}^{(L)}_{n})^{\top}(\widetilde{\mathbf{K}}^{(L)}_{(k-1)H+h}\widetilde{X}^{(L)}_{j})\right)}{Z^{(l)}_{k}}=\begin{cases}\delta^{i_{0}}_{j},&\penalty\ k\neq\kappa\\ \delta^{n}_{j},&\penalty\ k=\kappa\end{cases}

It follows that

	$\displaystyle\widetilde{X}^{(L+1)}_{n}=$	$\displaystyle\penalty\ \widetilde{\mathbf{V}}^{(L)}_{(\kappa-1)H+h}\cdot\widetilde{X}^{(L)}_{n}+\sum_{k\neq\kappa}\mathbf{V}^{(L)}_{(\kappa-1)H+h}\cdot\widetilde{X}^{(L)}_{i_{0}}$
	$\displaystyle=$	$\displaystyle\penalty\ \begin{pmatrix}0&&&&\\ &\ddots&&&\\ &&I&&\\ &&&\ddots&\\ &&&&0\end{pmatrix}\cdot\begin{pmatrix}X^{(L)}_{1;i}\\ \vdots\\ X^{(L)}_{K;i}\\ \widetilde{\alpha}_{i}\end{pmatrix}=\begin{pmatrix}0\\ \vdots\\ X^{(L)}_{\kappa;n}\\ \vdots\\ 0\end{pmatrix}.$

Therefore we establish Eq. (15).

Finally, at the output layer

	$\displaystyle p_{\widetilde{f}}(y\|v_{1},\dots,v_{n})=$	$\displaystyle\penalty\ \mathrm{Softmax}(\widetilde{\vartheta}(y)^{\top}\widetilde{X}^{(L+1)}_{n})$
	$\displaystyle=$	$\displaystyle\penalty\ \mathrm{Softmax}(\vartheta(y)^{\top}Y^{(L)}_{n-1})$
	$\displaystyle=$	$\displaystyle\penalty\ p_{f_{\kappa}}(y\|u).$

This establishes the desired statement. ∎

A.4 Proof of Proposition 4.4

Proof.

Set constants $B_{v},B_{qk},B_{\theta}$ such that for any layer $l$ and head $h$ , it holds that $\left\|(\mathbf{Q}^{(l)}_{h})^{\top}\mathbf{K}^{(l)}_{h}\right\|_{2}\leq B_{qk}$ , $\left\|\mathbf{V}^{(l)}_{h}\right\|_{2}\leq B_{v}$ , and $\|\theta(v)\|_{2}\leq B_{\theta}$ holds for all $v\in\mathcal{V}$ . Let $B=(KHB_{v})^{L}B_{qk}B_{\theta},C=2B^{2}+\log(1/\epsilon),C_{0}=4C$ . Define $\iota(i)=u$ iff $\xi_{u}\leq i<\xi_{u+1}$ ( $\xi_{0}=-1,\xi_{m+1}=\infty$ by default). Let $\mathcal{E}(\cdot)$ denote the task id indicated by the special token. By Lemma A.2, there exists $\alpha_{1},\dots,\alpha_{N},\beta_{1},\dots,\beta_{K}\in\mathbb{R}^{d_{0}}$ and $A_{1},\dots,A_{K}\in\mathbb{R}^{{d_{0}}\times{d_{0}}}$ for ${d_{0}}\leq O(K+\log N_{\max})$ such that for any $n\leq N$ we have

For any $k\neq k^{\prime}$ :

\displaystyle\alpha_{n}^{\top}A_{k}(\alpha_{n}+\beta_{k^{\prime}})\geq C_{0}+\begin{cases}\alpha_{n}^{\top}A_{k}\alpha_{n}\\ \alpha_{n}^{\top}A_{k}\alpha_{j}\\ \alpha_{n}^{\top}A_{k}(\alpha_{j}+\beta_{k^{\prime\prime}})\end{cases},\penalty\ \forall 0\leq j\leq n,1\leq k^{\prime\prime}\leq K.

(19)

For any $k\in[K]$ :

\displaystyle\alpha_{n}^{\top}A_{k}\alpha_{n}=\alpha_{n}^{\top}A_{k}\alpha_{0}\geq C_{0}+\begin{cases}\alpha_{n}^{\top}A_{k}(\alpha_{n}+\beta_{k})\\ \alpha_{n}^{\top}A_{k}\alpha_{j}\\ \alpha_{n}^{\top}A_{k}(\alpha_{j}+\beta_{k^{\prime}})\end{cases},\penalty\ \forall 0<j<n,k^{\prime}\neq k.

(20)

For any $k,k^{\prime},k^{\prime\prime}\in[K]$ :

\displaystyle(\alpha_{n}+\beta_{k^{\prime}})^{\top}A_{k}(\alpha_{n}+\beta_{k^{\prime}})\geq C_{0}+(\alpha_{n}+\beta_{k^{\prime}})^{\top}A_{k}\alpha_{j},\penalty\ \forall 0\leq j\leq n.

(21)

For any $0<j<n$ :

	$\displaystyle\alpha_{n}^{\top}A\alpha_{n}\geq$	$\displaystyle\penalty\ \alpha_{n}^{\top}A(\alpha_{n}+\beta_{k})+C_{0}$
	$\displaystyle\geq$	$\displaystyle\penalty\ C_{0}+\max\{\alpha_{n}^{\top}A\alpha_{j},\alpha_{n}^{\top}A(\alpha_{j}+\beta_{k^{\prime}})\},\penalty\ \forall k,k^{\prime\prime}\in[K].$		(22)

We define $\phi$ as follows: for any Transformers

\displaystyle f_{k}=(\theta_{k},\mathrm{pe}_{k},(\mathbf{K}^{(l)}_{k;h},\mathbf{Q}^{(l)}_{k;h},\mathbf{V}^{(l)}_{k;h})_{h\in[H],l\in[L]},\vartheta_{k},\mathcal{V}),k\in[K]

over $\mathcal{V}$ , the Transformer $\widetilde{f}=\phi(f_{1},\dots,f_{K})$ is given by

\displaystyle(\widetilde{\theta},\widetilde{\mathrm{pe}},(\widetilde{\mathbf{K}}^{(l)}_{h},\widetilde{\mathbf{Q}}^{(l)}_{h},\widetilde{\mathbf{V}}^{(l)}_{h})_{h\in[KH+1],l\in[L]},\widetilde{\vartheta},\mathcal{V}\cup\Omega),

where the tokenizer is given by

\displaystyle\widetilde{\theta}(v)=\begin{pmatrix}\theta_{1}(v)\\ \vdots\\ \theta_{K}(v)\\ 0\end{pmatrix},\penalty\ v\in\mathcal{V},\penalty\ \widetilde{\theta}(\omega)=\begin{pmatrix}0\\ \vdots\\ 0\\ \beta_{\mathcal{E}(\omega)}\end{pmatrix},\penalty\ \omega\in\Omega,

the positional encoder is given by

\displaystyle\widetilde{\mathrm{pe}}\left(\begin{pmatrix}x\\ y\end{pmatrix};v_{1},\dots,v_{i}\right)=\begin{pmatrix}\mathrm{pe}_{1}\left(x;v_{1},\cdots,v_{\xi_{1}-1},v_{\xi_{m}+1},\cdots,v_{n}\right)\\ \vdots\\ \mathrm{pe}_{K}\left(x;v_{1},\cdots,v_{\xi_{1}-1},v_{\xi_{m}+1},\cdots,v_{n}\right)\\ \alpha_{\iota(i)}+y\end{pmatrix},

where $x\in\mathbb{R}^{d}$ ; for $l=1,\dots,L$ the key, query, value matrices are given by

		$\displaystyle\penalty\ \widetilde{\mathbf{K}}^{(l)}_{(k-1)H+h}=\begin{pmatrix}0&&&&\\ &\ddots&&&\\ &&\mathbf{K}^{(l)}_{k;h}&&\\ &&&\ddots&\\ &&&&A_{k}\end{pmatrix},\penalty\ \widetilde{\mathbf{Q}}^{(l)}_{(k-1)H+h}=\begin{pmatrix}0&&&&\\ &\ddots&&&\\ &&\mathbf{Q}^{(l)}_{k;h}&&\\ &&&\ddots&\\ &&&&I\end{pmatrix},$
		$\displaystyle\penalty\ \widetilde{\mathbf{V}}^{(l)}_{(k-1)H+h}=\begin{pmatrix}0&&&&\\ &\ddots&&&\\ &&\mathbf{V}^{(l)}_{k;h}&&\\ &&&\ddots&\\ &&&&0\end{pmatrix},$
		$\displaystyle\penalty\ \widetilde{\mathbf{K}}^{(l)}_{KH+1}=\begin{pmatrix}0&&&\\ &\ddots&&\\ &&0&\\ &&&A\end{pmatrix},\penalty\ \widetilde{\mathbf{Q}}^{(l)}_{KH+1}=\begin{pmatrix}0&&&\\ &\ddots&&\\ &&0&\\ &&&I\end{pmatrix},\penalty\ \widetilde{\mathbf{V}}^{(l)}_{KH+1}=\begin{pmatrix}0&&&\\ &\ddots&&\\ &&0&\\ &&&I\end{pmatrix},$

where the submatrices $\mathbf{K}^{(l)}_{k;h},\mathbf{Q}^{(l)}_{k;h},\mathbf{V}^{(l)}_{k;h}$ are located in the $k$ -th diagonal block.

The output feature is given by $\widetilde{\vartheta}(y)=\begin{pmatrix}\vartheta_{1}(y)\\ \vdots\\ \vartheta_{K}(y)\\ 0\end{pmatrix}$ . Since $\xi_{1},\xi_{m}$ only depends on whether $v_{i}$ ’s belong to the set $\Omega$ , the generalized position encoding $\mathrm{pe}$ is well-defined. We can easily verify that $\phi$ is indeed a general-purpose Transformer of type $(O(K),O(\log N_{\max}))$ .

Let $\widetilde{X}^{(l)}_{1},\dots,\widetilde{X}^{(l)}_{n}$ represent the $l$ -th hidden layer. Our goal is to show that for any $l=1,\dots,L$ , $\widetilde{X}^{(l)}_{i}$ can be written as:

\displaystyle\widetilde{X}^{(l)}_{i}=\begin{pmatrix}X^{(l)}_{1;i}\\ \vdots\\ X^{(l)}_{K;i}\\ \widetilde{\alpha}_{i}\end{pmatrix},\penalty\ i=1,\dots,n,

(23)

where $\widetilde{\alpha}_{i}=\alpha_{\iota(i)}+\mathbbm{1}(\iota(i)=i)\cdot\beta_{\mathcal{E}(v_{i})}$ and $X^{(l)}_{k;i}\in\mathbb{R}^{d}$ such that

\displaystyle\|X^{(l)}_{k;i}\|_{2}\leq B_{\theta}(KHB_{v})^{l}.

(24)

In particular, for $i=1,\dots,m$ we have

\displaystyle X^{(l)}_{k;\xi_{i}}=0,\penalty\ \forall k=1,\dots,K,

(25)

and for $j=1,\dots,\xi_{1}$ we have

\displaystyle X^{(l)}_{k;j}=Y^{(l)}_{k;j},\penalty\ \forall k=1,\dots,K,

(26)

and for $j=1,\dots,\xi_{1}-1,\xi_{m}+1,\dots,n$ we have

\displaystyle X^{(l)}_{\kappa;j}=Y^{(l)}_{\kappa,j-\xi_{m}-1+\xi_{1}},\penalty\ X^{(l)}_{k^{\prime};j}=0,\penalty\ \forall k^{\prime}\neq\kappa,

(27)

where $Y^{(l)}_{k;j}$ is the $l$ -th hidden layer of ${f_{k}}$ (attending only to positions $1,\dots,\xi_{1}-1,\xi_{m}+1,\dots,n$ ) .

Thus we apply induction on $l$ . The case $l=1$ holds trivially from the definition of $\widetilde{\theta}$ and $\widetilde{\mathrm{pe}}$ . Suppose the above relationship holds for all layers $1,\dots,l$ , consider layer $l+1$ . We have

	$\displaystyle\widetilde{X}^{(l+1)}_{i}=$	$\displaystyle\penalty\ \underbrace{\sum_{k=1}^{K}\sum_{h=1}^{H}\sum_{j=1}^{i}\frac{\exp\left((\widetilde{\mathbf{Q}}^{(l)}_{(k-1)H+h}\widetilde{X}^{(l)}_{i})^{\top}(\widetilde{\mathbf{K}}^{(l)}_{(k-1)H+h}\widetilde{X}^{(l)}_{j})\right)}{\widetilde{Z}^{(l)}_{(k-1)H+h}}\cdot\widetilde{\mathbf{V}}^{(l)}_{(k-1)H+h}\widetilde{X}^{(l)}_{j}}_{\text{term 1}}$
		$\displaystyle\penalty\ +\underbrace{\sum_{j=1}^{i}\frac{\exp\left((\widetilde{\mathbf{Q}}^{(l)}_{KH+1}\widetilde{X}^{(l)}_{i})^{\top}(\widetilde{\mathbf{K}}^{(l)}_{KH+1}\widetilde{X}^{(l)}_{j})\right)}{\widetilde{Z}^{(l)}_{KH+1}}\cdot\widetilde{\mathbf{V}}^{(l)}_{KH+1}\widetilde{X}^{(l)}_{j}}_{\text{term 2}},$

where

\displaystyle\widetilde{Z}^{(l)}_{(k-1)H+h}=\sum_{j=1}^{i}\exp\left((\widetilde{\mathbf{Q}}^{(l)}_{(k-1)H+h}\widetilde{X}^{(l)}_{i})^{\top}(\widetilde{\mathbf{K}}^{(l)}_{(k-1)H+h}\widetilde{X}^{(l)}_{j})\right).

By induction hypothesis,

\displaystyle\widetilde{X}^{(l)}_{i}=\begin{pmatrix}X^{(l)}_{1;i}\\ \vdots\\ X^{(l)}_{K;i}\\ \widetilde{\alpha}_{i}\end{pmatrix},

and $X^{(l)}_{k;i}=Y^{(l)}_{\zeta(i)}$ for $i=1,\dots,\xi_{1}-1,\xi_{m}+1,\dots,n$ , where $\zeta(i):=\begin{cases}i,&\penalty\ i<\xi_{1}\\ i-\xi_{m}-1+\xi_{1},&\penalty\ i>\xi_{m}\end{cases}$ .

Notice that for $j\leq i$ :

	$\displaystyle(\widetilde{\mathbf{Q}}^{(l)}_{(k-1)H+h}\widetilde{X}^{(l)}_{i})^{\top}(\widetilde{\mathbf{K}}^{(l)}_{(k-1)H+h}\widetilde{X}^{(l)}_{j})=$	$\displaystyle\penalty\ (X^{(l)}_{k;i})^{\top}(\mathbf{Q}^{(l)}_{k;h})^{\top}\mathbf{K}^{(l)}_{k;h}X^{(l)}_{k;j}+\widetilde{\alpha}_{i}^{\top}A_{k}\widetilde{\alpha}_{j},$
	$\displaystyle(\widetilde{\mathbf{Q}}^{(l)}_{KH+1}\widetilde{X}^{(l)}_{i})^{\top}(\widetilde{\mathbf{K}}^{(l)}_{KH+1}\widetilde{X}^{(l)}_{j})=$	$\displaystyle\penalty\ \widetilde{\alpha}_{i}^{\top}A\widetilde{\alpha}_{j}.$

Prove Eq (23).

By properties of $\alpha,\beta,A$ , for any $j_{2}<\xi_{u}<j_{1}<i<\xi_{u+1}$ notice that:

	$\displaystyle(\widetilde{\mathbf{Q}}^{(l)}_{KH+1}\widetilde{X}^{(l)}_{i})^{\top}(\widetilde{\mathbf{K}}^{(l)}_{KH+1}\widetilde{X}^{(l)}_{j_{1}})\geq$	$\displaystyle\penalty\ (\widetilde{\mathbf{Q}}^{(l)}_{KH+1}\widetilde{X}^{(l)}_{i})^{\top}(\widetilde{\mathbf{K}}^{(l)}_{KH+1}\widetilde{X}^{(l)}_{\xi_{u}})+C$
	$\displaystyle\geq$	$\displaystyle\penalty\ (\widetilde{\mathbf{Q}}^{(l)}_{KH+1}\widetilde{X}^{(l)}_{i})^{\top}(\widetilde{\mathbf{K}}^{(l)}_{KH+1}\widetilde{X}^{(l)}_{j_{2}})+2C.$

Due to $\epsilon$ -precision of transformers, this implies that

\displaystyle\frac{\exp\left((\widetilde{\mathbf{Q}}^{(l)}_{KH+1}\widetilde{X}^{(l)}_{i})^{\top}(\widetilde{\mathbf{K}}^{(l)}_{KH+1}\widetilde{X}^{(l)}_{j})\right)}{Z^{(l)}_{KH+1}}=\begin{cases}\frac{\mathbbm{1}(j>\xi_{u})}{i-\xi_{u}},&\penalty\ \xi_{u}<i<\xi_{u+1}\\ \delta^{j}_{\xi_{l}},&\penalty\ i=\xi_{u}\end{cases},

and hence for $\xi_{u}<i<\xi_{u+1}$

	$\displaystyle\widetilde{X}^{(l+1)}_{i}=$	$\displaystyle\penalty\ \sum_{k=1}^{K}\sum_{h=1}^{H}\sum_{j=1}^{i}\frac{\exp\left((\widetilde{\mathbf{Q}}^{(l)}_{(k-1)H+h}\widetilde{X}^{(l)}_{i})^{\top}(\widetilde{\mathbf{K}}^{(l)}_{(k-1)H+h}\widetilde{X}^{(l)}_{j})\right)}{\widetilde{Z}^{(l)}_{(k-1)H+h}}\cdot\widetilde{\mathbf{V}}^{(l)}_{(k-1)H+h}\begin{pmatrix}\vdots\\ X^{(l)}_{k;j}\\ \vdots\\ 0\end{pmatrix}$
		$\displaystyle\penalty\ +\sum_{j=\xi_{u}+1}^{i}\cdot\frac{1}{i-\xi_{u}}\cdot\begin{pmatrix}0\\ \vdots\\ 0\\ \alpha_{\iota(i)}\end{pmatrix}$
	$\displaystyle=$	$\displaystyle\penalty\ \begin{pmatrix}X^{(l+1)}_{1;i}\\ \vdots\\ X^{(l+1)}_{K;i}\\ \widetilde{\alpha}_{i}\end{pmatrix},$

and for $i=\xi_{u}$

	$\displaystyle\widetilde{X}^{(l+1)}_{i}=$	$\displaystyle\penalty\ \sum_{k=1}^{K}\sum_{h=1}^{H}\sum_{j=1}^{i}\frac{\exp\left((\widetilde{\mathbf{Q}}^{(l)}_{(k-1)H+h}\widetilde{X}^{(l)}_{i})^{\top}(\widetilde{\mathbf{K}}^{(l)}_{(k-1)H+h}\widetilde{X}^{(l)}_{j})\right)}{\widetilde{Z}^{(l)}_{(k-1)H+h}}\cdot\widetilde{\mathbf{V}}^{(l)}_{(k-1)H+h}\begin{pmatrix}\vdots\\ X^{(l)}_{k;j}\\ \vdots\\ 0\end{pmatrix}+\begin{pmatrix}0\\ \vdots\\ 0\\ \alpha_{\iota(i)}+\beta_{\mathcal{E}(v_{i})}\end{pmatrix}$
	$\displaystyle=$	$\displaystyle\penalty\ \begin{pmatrix}X^{(l+1)}_{1;i}\\ \vdots\\ X^{(l+1)}_{K;i}\\ \widetilde{\alpha}_{i}\end{pmatrix},$

where

\displaystyle X^{(l+1)}_{k;i}=

\displaystyle\penalty\ \sum_{k=1}^{K}\sum_{h=1}^{H}\sum_{j=1}^{i}\frac{\exp\left((\widetilde{\mathbf{Q}}^{(l)}_{(k-1)H+h}\widetilde{X}^{(l)}_{i})^{\top}(\widetilde{\mathbf{K}}^{(l)}_{(k-1)H+h}\widetilde{X}^{(l)}_{j})\right)}{\widetilde{Z}^{(l)}_{(k-1)H+h}}\cdot\mathbf{V}^{(l)}_{k;h}X^{(l)}_{k;j}.

(28)

This confirms Eq. (23) for $l+1$ .

Prove Eq. (24).

From above,

	$\displaystyle\\|X^{(l+1)}_{k;i}\\|_{2}=$	$\displaystyle\penalty\ \left\\|\sum_{k=1}^{K}\sum_{h=1}^{H}\sum_{j=1}^{i}\frac{\exp\left((\widetilde{\mathbf{Q}}^{(l)}_{(k-1)H+h}\widetilde{X}^{(l)}_{i})^{\top}(\widetilde{\mathbf{K}}^{(l)}_{(k-1)H+h}\widetilde{X}^{(l)}_{j})\right)}{\widetilde{Z}^{(l)}_{(k-1)H+h}}\cdot\mathbf{V}^{(l)}_{k;h}X^{(l)}_{k;j}\right\\|_{2}$
	$\displaystyle\leq$	$\displaystyle\penalty\ KHB_{v}\cdot\max_{j\leq i}\\|X^{(l)}_{k;j}\\|_{2}$
	$\displaystyle\leq$	$\displaystyle\penalty\ B_{\theta}(KHB_{v})^{l+1}.$

This confirms Eq. (24) for $l+1$ .

Prove Eq. (25).

We first show $X^{(l)}_{k;\xi_{1}}=0$ . Indeed, by the properties of $\alpha_{t},\beta_{k}$ , for any $j\leq\xi_{1}$

		$\displaystyle\penalty\ (\widetilde{\mathbf{Q}}^{(l)}_{(k-1)H+h}\widetilde{X}^{(l)}_{\xi_{1}})^{\top}(\widetilde{\mathbf{K}}^{(l)}_{(k-1)H+h}\widetilde{X}^{(l)}_{\xi_{1}})$
	$\displaystyle=$	$\displaystyle\penalty\ (X^{(l)}_{k;\xi_{1}})^{\top}(\mathbf{Q}^{(l)}_{k;h})^{\top}\mathbf{K}^{(l)}_{k;h}X^{(l)}_{k;\xi_{1}}+(\alpha_{0}+\beta_{\mathcal{E}(v_{\xi_{1}})})^{\top}A_{k}(\alpha_{0}+\beta_{\mathcal{E}(v_{\xi_{1}})})$
	$\displaystyle\geq$	$\displaystyle\penalty\ (X^{(l)}_{k;\xi_{1}})^{\top}(\mathbf{Q}^{(l)}_{k;h})^{\top}\mathbf{K}^{(l)}_{k;h}X^{(l)}_{k;\xi_{1}}+(\alpha_{0}+\beta_{\mathcal{E}(v_{\xi_{1}})})^{\top}A_{k}\alpha_{0}+C$
	$\displaystyle=$	$\displaystyle\penalty\ (\widetilde{\mathbf{Q}}^{(l)}_{(k-1)H+h}\widetilde{X}^{(l)}_{\xi_{1}})^{\top}(\widetilde{\mathbf{K}}^{(l)}_{(k-1)H+h}\widetilde{X}^{(l)}_{j})+C$

It follows from Eq. (28) that

\displaystyle X^{(l+1)}_{k;\xi_{1}}=\sum_{k=1}^{K}\sum_{h=1}^{H}\mathbf{V}^{(l)}_{k;h}X^{(l)}_{k;\xi_{1}}=0.

For $\xi_{i}\penalty\ (i>1)$ , we apply the same argument again to obtain that for any $j\leq\xi_{i}$ such that $j\notin\{\xi_{1}<\cdots<\xi_{\iota(n)}\}$ and any $i^{\prime}<i$ ,

		$\displaystyle\penalty\ (\widetilde{\mathbf{Q}}^{(l)}_{(k-1)H+h}\widetilde{X}^{(l)}_{\xi_{i}})^{\top}(\widetilde{\mathbf{K}}^{(l)}_{(k-1)H+h}\widetilde{X}^{(l)}_{\xi_{k^{\prime}}})$
	$\displaystyle\geq$	$\displaystyle\penalty\ (\widetilde{\mathbf{Q}}^{(l)}_{(k-1)H+h}\widetilde{X}^{(l)}_{\xi_{1}})^{\top}(\widetilde{\mathbf{K}}^{(l)}_{(k-1)H+h}\widetilde{X}^{(l)}_{j})+C$

This implies that the attention weights are supported on $\{\xi_{1}<\cdots<\xi_{i}\}$ , and therefore

\displaystyle X^{(l+1)}_{k;\xi_{i}}=\sum_{k=1}^{K}\sum_{h=1}^{H}\sum_{j=1}^{i}\frac{\exp\left((\widetilde{\mathbf{Q}}^{(l)}_{(k-1)H+h}\widetilde{X}^{(l)}_{\xi_{i}})^{\top}(\widetilde{\mathbf{K}}^{(l)}_{(k-1)H+h}\widetilde{X}^{(l)}_{\xi_{j}})\right)}{\widetilde{Z}^{(l)}_{(k-1)H+h}}\cdot\mathbf{V}^{(l)}_{k;h}X^{(l)}_{k;\xi_{j}}=0

where we apply the induction hypothesis $k;X^{(l)}_{\xi_{j}}=0$ for all $j=1,\dots,i-1$ . This thus completes the proof of Eq. (25).

Prove Eq. (26).

When $j_{1}<j_{2}\leq i<\xi_{1}$ , we have

		$\displaystyle\penalty\ (\widetilde{\mathbf{Q}}^{(l)}_{(k-1)H+h}\widetilde{X}^{(l)}_{i})^{\top}(\widetilde{\mathbf{K}}^{(l)}_{(k-1)H+h}\widetilde{X}^{(l)}_{j_{1}})-(\widetilde{\mathbf{Q}}^{(l)}_{(k-1)H+h}\widetilde{X}^{(l)}_{i})^{\top}(\widetilde{\mathbf{K}}^{(l)}_{(k-1)H+h}X^{(l)}_{j_{2}})$
	$\displaystyle=$	$\displaystyle\penalty\ (X^{(l)}_{k;i})^{\top}(\mathbf{Q}^{(l)}_{k;h})^{\top}\mathbf{K}^{(l)}_{k;h}X^{(l)}_{k;j_{1}}+\alpha_{0}^{\top}A_{k}\alpha_{0}^{\top}$
		$\displaystyle\penalty\ -(X^{(l)}_{k;i})^{\top}(\mathbf{Q}^{(l)}_{k;h})^{\top}\mathbf{K}^{(l)}_{k;h}X^{(l)}_{k;j_{2}}-\alpha_{0}^{\top}A_{k}\alpha_{0}^{\top}$
	$\displaystyle=$	$\displaystyle\penalty\ (\mathbf{Q}^{(l)}_{k;h}Y^{(l)}_{k;i})^{\top}(\mathbf{K}^{(l)}_{k;h}Y^{(l)}_{k;j_{i}})-(\mathbf{Q}^{(l)}_{k;h}Y^{(l)}_{k;i})^{\top}(\mathbf{K}^{(l)}_{k;h}Y^{(l)}_{k;j_{2}}).$

It follows that

\displaystyle\widetilde{Z}^{(l)}_{(k-1)H+h}=\sum_{j=1}^{i}\exp\left((\mathbf{Q}^{(l)}_{k;h}Y^{(l)}_{k;i})^{\top}(\mathbf{K}^{(l)}_{k;h}Y^{(l)}_{k;j})\right),

and

	$\displaystyle X^{(l+1)}_{k;i}=$	$\displaystyle\penalty\ \sum_{k=1}^{K}\sum_{h=1}^{H}\sum_{j=1}^{i}\frac{\exp\left((\mathbf{Q}^{(l)}_{k;h}Y^{(l)}_{k;i})^{\top}(\mathbf{K}^{(l)}_{k;h}Y^{(l)}_{k;j})\right)}{\widetilde{Z}^{(l)}_{(k-1)H+h}}\cdot\mathbf{V}^{(l)}_{k;h}Y^{(l)}_{k;j}$
	$\displaystyle=$	$\displaystyle\penalty\ Y^{(l+1)}_{k;i}.$

This confirms Eq. (26).

Prove Eq. (27).

When $i>\xi_{m}$ , we rely on the following properties:

Attention sink to $v_{\xi_{m}}$ for mismatch expert: for any $k^{\prime}\neq\kappa$ and $j\leq i$ we have

\displaystyle(\widetilde{\mathbf{Q}}^{(l)}_{(k^{\prime}-1)H+h}\widetilde{X}^{(l)}_{i})^{\top}(\widetilde{\mathbf{K}}^{(l)}_{(k^{\prime}-1)H+h}\widetilde{X}^{(l)}_{j})\leq(\widetilde{\mathbf{Q}}^{(l)}_{(k^{\prime}-1)H+h}\widetilde{X}^{(l)}_{i})^{\top}(\widetilde{\mathbf{K}}^{(l)}_{(k^{\prime}-1)H+h}\widetilde{X}^{(l)}_{\xi_{m}})-C.

(29)

Attention to task-relevant tokens for matching expert: for $j\in\{1,\dots,\xi_{1}-1,\xi_{m}+1,\dots,n\}$ , and $\xi_{1}\leq j^{\prime}\leq{\xi_{m}}$ we have

\displaystyle(\widetilde{\mathbf{Q}}^{(l)}_{(\kappa-1)H+h}\widetilde{X}^{(l)}_{i})^{\top}(\widetilde{\mathbf{K}}^{(l)}_{(\kappa-1)H+h}\widetilde{X}^{(l)}_{j})\geq(\widetilde{\mathbf{Q}}^{(l)}_{(\kappa-1)H+h}\widetilde{X}^{(l)}_{i})^{\top}(\widetilde{\mathbf{K}}^{(l)}_{(\kappa-1)H+h}\widetilde{X}^{(l)}_{j^{\prime}})+C.

(30)

and for $j_{1}<j_{2}\in\{1,\dots,\xi-1-1,{\xi_{m}}+1,\dots,n\}$

		$\displaystyle\penalty\ (\widetilde{\mathbf{Q}}^{(l)}_{(\kappa-1)H+h}\widetilde{X}^{(l)}_{i})^{\top}(\widetilde{\mathbf{K}}^{(l)}_{(\kappa-1)H+h}\widetilde{X}^{(l)}_{j_{1}})-(\widetilde{\mathbf{Q}}^{(l)}_{(\kappa-1)H+h}\widetilde{X}^{(l)}_{i})^{\top}(\widetilde{\mathbf{K}}^{(l)}_{(\kappa-1)H+h}\widetilde{X}^{(l)}_{j_{2}})$
	$\displaystyle=$	$\displaystyle\penalty\ (\mathbf{Q}^{(l)}_{\kappa;h}Y^{(l)}_{\kappa;i-\xi_{m}-1+\xi_{1}})^{\top}(\mathbf{K}^{(l)}_{\kappa;h}Y^{(l)}_{\zeta(j_{1})})-(\mathbf{Q}^{(l)}_{\kappa;h}Y^{(l)}_{i-\xi_{m}-1+\xi_{1}})^{\top}\mathbf{K}^{(l)}_{\kappa;h}Y^{(l)}_{\kappa;\zeta(j_{2})}),$		(31)

To see Eq. (29), we notice that

		$\displaystyle\penalty\ (\widetilde{\mathbf{Q}}^{(l)}_{(k^{\prime}-1)H+h}\widetilde{X}^{(l)}_{i})^{\top}(\widetilde{\mathbf{K}}^{(l)}_{(k^{\prime}-1)H+h}\widetilde{X}^{(l)}_{j})$
	$\displaystyle=$	$\displaystyle\penalty\ (X^{(l)}_{k^{\prime};i})^{\top}(\mathbf{Q}^{(l)}_{k^{\prime};h})^{\top}\mathbf{K}^{(l)}_{k^{\prime};h}X^{(l)}_{k^{\prime},j}+\alpha_{m}^{\top}A_{k^{\prime}}(\alpha_{\iota(j)}+\beta_{\mathcal{E}(v_{j})}\cdot\mathbbm{1}(\iota(j)=j))$
	$\displaystyle\leq$	$\displaystyle\penalty\ (X^{(l)}_{k^{\prime};i})^{\top}(\mathbf{Q}^{(l)}_{k^{\prime};h})^{\top}\mathbf{K}^{(l)}_{k^{\prime};h}X^{(l)}_{k^{\prime};\xi_{m}}+\alpha_{m}^{\top}A_{k^{\prime}}(\alpha_{m}+\beta_{\mathcal{E}(v_{\xi_{m}})})-C$
	$\displaystyle=$	$\displaystyle\penalty\ (\widetilde{\mathbf{Q}}^{(l)}_{(k^{\prime}-1)H+h}\widetilde{X}^{(l)}_{i})^{\top}(\widetilde{\mathbf{K}}^{(l)}_{(k^{\prime}-1)H+h}\widetilde{X}^{(l)}_{\xi_{m}})-C,$

where we use Eq. (19) with $k^{\prime}\neq\kappa$ .

To see Eq. (30), we notice that

	$\displaystyle(\widetilde{\mathbf{Q}}^{(l)}_{(\kappa-1)H+h}\widetilde{X}^{(l)}_{i})^{\top}(\widetilde{\mathbf{K}}^{(l)}_{(\kappa-1)H+h}\widetilde{X}^{(l)}_{j})=$	$\displaystyle\penalty\ (\mathbf{Q}^{(l)}_{k;h}X^{(l)}_{k;i})^{\top}(\mathbf{K}^{(l)}_{k;h}X^{(l)}_{k;j})+\alpha_{m}^{\top}A_{\kappa}\alpha_{0}$
	$\displaystyle\geq$	$\displaystyle\penalty\ (\mathbf{Q}^{(l)}_{k;h}X^{(l)}_{k;i})^{\top}(\mathbf{K}^{(l)}_{k;h}X^{(l)}_{k;j^{\prime}})+\alpha_{m}^{\top}A_{\kappa}(\alpha_{\iota(j^{\prime})}+\beta_{\mathcal{E}(v_{j^{\prime}})})+C$
	$\displaystyle=$	$\displaystyle\penalty\ (\widetilde{\mathbf{Q}}^{(l)}_{(k-1)H+h}\widetilde{X}^{(l)}_{i})^{\top}(\widetilde{\mathbf{K}}^{(l)}_{(k-1)H+h}\widetilde{X}^{(l)}_{j^{\prime}})+C,$

and

	$\displaystyle(\widetilde{\mathbf{Q}}^{(l)}_{(\kappa-1)H+h}\widetilde{X}^{(l)}_{i})^{\top}(\widetilde{\mathbf{K}}^{(l)}_{(\kappa-1)H+h}\widetilde{X}^{(l)}_{j})=$	$\displaystyle\penalty\ (\mathbf{Q}^{(l)}_{k;h}X^{(l)}_{k;i})^{\top}(\mathbf{K}^{(l)}_{k;h}X^{(l)}_{k;j})+\alpha_{m}^{\top}A_{\kappa}\alpha_{0}$
	$\displaystyle\geq$	$\displaystyle\penalty\ (\mathbf{Q}^{(l)}_{k;h}X^{(l)}_{k;i})^{\top}(\mathbf{K}^{(l)}_{k;h}X^{(l)}_{k;j^{\prime}})+\alpha_{m}^{\top}A_{k}\alpha_{\iota(j^{\prime})}+C$
	$\displaystyle=$	$\displaystyle\penalty\ (\widetilde{\mathbf{Q}}^{(l)}_{(k-1)H+h}\widetilde{X}^{(l)}_{i})^{\top}(\widetilde{\mathbf{K}}^{(l)}_{(k-1)H+h}\widetilde{X}^{(l)}_{j^{\prime}})+C,$

where we use Eq. (20) and Eq. (4).

When $\xi_{m}<j_{1}<j_{2}$ , Eq. (2) follows directly from

		$\displaystyle\penalty\ (\widetilde{\mathbf{Q}}^{(l)}_{(\kappa-1)H+h}\widetilde{X}^{(l)}_{i})^{\top}(\widetilde{\mathbf{K}}^{(l)}_{(\kappa-1)H+h}\widetilde{X}^{(l)}_{j_{1}})-(\widetilde{\mathbf{Q}}^{(l)}_{(\kappa-1)H+h}\widetilde{X}^{(l)}_{i})^{\top}(\widetilde{\mathbf{K}}^{(l)}_{(\kappa-1)H+h}\widetilde{X}^{(l)}_{j_{2}})$
	$\displaystyle=$	$\displaystyle\penalty\ (\mathbf{Q}^{(l)}_{k;h}X^{(l)}_{k;i})^{\top}(\mathbf{K}^{(l)}_{k;h}X^{(l)}_{k;j_{1}})+\alpha_{m}^{\top}A_{k}\alpha_{m}^{\top}$
		$\displaystyle\penalty\ -(\mathbf{Q}^{(l)}_{k;h}X^{(l)}_{k;i})^{\top}(\mathbf{K}^{(l)}_{k;h}X^{(l)}_{k;j_{2}})+\alpha_{m}^{\top}A_{k}\alpha_{m}^{\top}$
	$\displaystyle=$	$\displaystyle\penalty\ (\mathbf{Q}^{(l)}_{\kappa;h}Y^{(l)}_{\kappa;i-\xi_{m}-1+\xi_{1}})^{\top}(\mathbf{K}^{(l)}_{\kappa;h}Y^{(l)}_{j_{1}-\xi_{m}-1+\xi_{1}})-(\mathbf{Q}^{(l)}_{\kappa;h}Y^{(l)}_{i-\xi_{m}-1+\xi_{1}})^{\top}\mathbf{K}^{(l)}_{\kappa;h}Y^{(l)}_{\kappa;j_{2}-\xi_{m}-1+\xi_{1}}).$

The other cases follow similarly due to Eq. (4).

We have hence confirmed Eq. (29), Eq. (30), Eq. (2), and therefore

\displaystyle\frac{\exp\left((\widetilde{\mathbf{Q}}^{(l)}_{(k-1)H+h}\widetilde{X}^{(l)}_{i})^{\top}(\widetilde{\mathbf{K}}^{(l)}_{(k-1)H+h}\widetilde{X}^{(l)}_{j})\right)}{\widetilde{Z}^{(l)}_{(k-1)H+h}}=\begin{cases}\delta^{\xi_{m}}_{j},&\penalty\ k\neq\kappa\\ \frac{\exp\left((\mathbf{Q}^{(l)}_{\kappa;h}Y^{(l)}_{\kappa;i-\xi_{m}-1+\xi_{1}})^{\top}(\mathbf{K}^{(l)}_{\kappa;h}Y^{(l)}_{j})\right)}{\widetilde{Z}^{(l)}_{(k-1)H+h}},&\penalty\ k=\kappa,\penalty\ j<\xi_{1}\\ 0,&\penalty\ k=\kappa,\penalty\ \xi_{1}\leq j\leq\xi_{m}\\ \frac{\exp\left((\mathbf{Q}^{(l)}_{\kappa;h}Y^{(l)}_{\kappa;i-\xi_{m}-1+\xi_{1}})^{\top}(\mathbf{K}^{(l)}_{\kappa;h}Y^{(l)}_{j-\xi_{m}-1+\xi_{1}})\right)}{\widetilde{Z}^{(l)}_{(k-1)H+h}},&\penalty\ k=\kappa,\penalty\ j>\xi_{m}\\ \end{cases}

and

\displaystyle\widetilde{Z}^{(l)}_{(k-1)H+h}=\sum_{j=1,\dots,\xi_{1}-1,\xi_{m}+1,\dots,n}\exp\left((\mathbf{Q}^{(l)}_{\kappa;h}Y^{(l)}_{\kappa;i-\xi_{m}-1+\xi_{1}})^{\top}(\mathbf{K}^{(l)}_{\kappa;h}Y^{(l)}_{j})\right).

It follows that

	$\displaystyle X^{(l+1)}_{\kappa;i}=$	$\displaystyle\penalty\ \sum_{j=1}^{\xi_{1}-1}\frac{\exp\left((\mathbf{Q}^{(l)}_{\kappa;h}Y^{(l)}_{\kappa;i-\xi_{m}-1+\xi_{1}})^{\top}(\mathbf{K}^{(l)}_{\kappa;h}Y^{(l)}_{j})\right)}{\widetilde{Z}^{(l)}_{(k-1)H+h}}\mathbf{V}^{(l)}_{k;h}Y^{(l)}_{j}$
		$\displaystyle\penalty\ +\sum_{j=\xi_{m}+1}^{i}\frac{\exp\left((\mathbf{Q}^{(l)}_{\kappa;h}Y^{(l)}_{\kappa;i-\xi_{m}-1+\xi_{1}})^{\top}(\mathbf{K}^{(l)}_{\kappa;h}Y^{(l)}_{j-\xi_{m}-1+\xi_{1}})\right)}{\widetilde{Z}^{(l)}_{(k-1)H+h}}\mathbf{V}^{(l)}_{k;h}Y^{(l)}_{j-\xi_{m}-1+\xi_{1}},$
	$\displaystyle=$	$\displaystyle\penalty\ Y^{(l+1)}_{\kappa;i-\xi_{m}-1+\xi_{1}}$
	$\displaystyle X^{(l+1)}_{k^{\prime};i}=$	$\displaystyle\penalty\ X^{(l)}_{k^{\prime};\xi_{m}}=0,\penalty\ \forall k^{\prime}\neq\kappa.$

Therefore we establish Eq. (27). This completes the induction.

At the output layer, we have

	$\displaystyle p_{\widetilde{f}}(y\|v_{1},\dots,v_{n})=$	$\displaystyle\penalty\ \mathrm{Softmax}(\widetilde{\vartheta}(y)^{\top}\widetilde{X}^{(L)}_{n})$
	$\displaystyle=$	$\displaystyle\penalty\ \mathrm{Softmax}(\vartheta(y)^{\top}Y^{(L)}_{n-\xi_{m}-1+\xi_{1}})$
	$\displaystyle=$	$\displaystyle\penalty\ p_{f_{\kappa}}(y\|u_{1},\dots,u_{n-\xi_{m}-1+\xi_{1}}).$

This establishes the desired Eq. (2). ∎

A.5 Proof of Theorem 4.7

Proof.

Let $\phi_{s},\phi_{m},\phi_{e}$ denote the general-purpose Transformers in Proposition 4.4 (with $K$ experts), 4.2 (with $K=3$ token spaces), and A.1 (extending to $\mathcal{V}$ ) respectively. We construct a dummy Transformer $f_{d}$ that outputs $\mathrm{BOS}$ immediately after a token in $\mathcal{A}$ . Then we claim that the general-purpose Transformer $\widetilde{\phi}$ defined by

\displaystyle\widetilde{\phi}(f_{0},f_{1},\dots,f_{K})=\phi_{m}(\phi_{s}(\phi_{e}(f_{1}),\dots,\phi_{e}(f_{K})),f_{d},f_{0})

achieves the desired property.

Indeed, let $g_{1}=\phi_{s}(\phi_{e}(f_{1}),\dots,\phi_{e}(f_{K}))$ , by Proposition 4.4, we have

1.

Expert following: At $t$ -th iteration,

$\displaystyle p_{g_{1}}\left(\cdot\Big|\mathrm{prompt}\right)\sim p_{f_{a^{(t)}}}\left(\cdot\Big|q|u^{(t)}_{1:i-1}\right),$

where $q|u^{(t)}_{1:i-1}$ is the token sequence obtained by concatenating the user query $q$ and prior generated part in response $t$ : $u^{(t)}_{1:i-1}$ .
2.

Regret minimization:

$\displaystyle\max_{a^{*}\in\mathcal{A}}r_{0}(a^{*})-\mathbb{E}[r_{0}(a^{(T)})]\leq\mathrm{reg}(T).$

Therefore by Proposition 4.2, we have

\displaystyle u^{(t)}_{i}\sim p_{f_{a^{(t)}}}\left(\cdot\Big|q|u^{(t)}_{1:i-1}\right).

It follows that

	$\displaystyle\max_{u^{}\in\mathcal{V}^{\omega}}r(q,u^{})-\mathbb{E}[r(q,u^{(T)})]\leq$	$\displaystyle\penalty\ \lambda+\mathbb{E}_{u\sim f_{k^{*}}(\cdot\|p)}[r(q,u)]-\mathbb{E}_{a^{(T)}}\left[\mathbb{E}_{u^{(T)}\sim f_{a^{(t)}}(\cdot\|q)}[r(q,u^{(T)})]\right]$
	$\displaystyle\leq$	$\displaystyle\penalty\ \lambda+\max_{a^{}\in\mathcal{A}}r_{0}(a^{})-\mathbb{E}[r_{0}(a^{(T)})]$
	$\displaystyle\leq$	$\displaystyle\penalty\ \lambda+\mathrm{reg}(T).$

Finally, $\widetilde{\phi}$ has type $\phi$ of type $(O(K),O(\log(N_{\max})))$ because $\phi_{s}$ has type $(O(K),O(\log(N_{\max})))$ and $\phi_{m},\phi_{e}$ has type $(O(1),O(\log(N_{\max})))$ . This completes the proof. ∎

A.6 Attention Sink Positional Encoding

In this section, we introduce positional encoding mechanisms that induce attention sink behaviors used by Theorem 4.7.

Lemma A.2 (Attention Sink Positional Encoding, Type 1).

For any $C\in\mathbb{R}_{+}$ , $K,N\in\mathbb{Z}_{+}$ , there exist vectors $\alpha_{1},\dots,\alpha_{N},\beta_{1},\dots,\beta_{K}\in\mathbb{R}^{d}$ and matrices $A,A_{1},\dots,A_{K}\in\mathbb{R}^{d\times d}$ for $d\leq O(K+\log N)$ such that for any $n\in[N]$ the followings hold

For any $k\neq k^{\prime}$ :

\displaystyle\alpha_{n}^{\top}A_{k}(\alpha_{n}+\beta_{k^{\prime}})\geq C+\begin{cases}\alpha_{n}^{\top}A_{k}\alpha_{n}\\ \alpha_{n}^{\top}A_{k}\alpha_{j}\\ \alpha_{n}^{\top}A_{k}(\alpha_{j}+\beta_{k^{\prime\prime}})\end{cases},\penalty\ \forall 0\leq j\leq n,1\leq k^{\prime\prime}\leq K.

For any $k\in[K]$ :

\displaystyle\alpha_{n}^{\top}A_{k}\alpha_{n}=\alpha_{n}^{\top}A_{k}\alpha_{0}\geq C+\begin{cases}\alpha_{n}^{\top}A_{k}(\alpha_{n}+\beta_{k})\\ \alpha_{n}^{\top}A_{k}\alpha_{j}\\ \alpha_{n}^{\top}A_{k}(\alpha_{j}+\beta_{k^{\prime}})\end{cases},\penalty\ \forall 0<j<n,k^{\prime}\neq k.

For any $k,k^{\prime},k^{\prime\prime}\in[K]$ :

\displaystyle(\alpha_{n}+\beta_{k^{\prime}})^{\top}A_{k}(\alpha_{n}+\beta_{k^{\prime}})\geq C+(\alpha_{n}+\beta_{k^{\prime}})^{\top}A_{k}\alpha_{j},\penalty\ \forall 0\leq j\leq n.

For any $0<j<n$ :

	$\displaystyle\alpha_{n}^{\top}A\alpha_{n}\geq$	$\displaystyle\penalty\ \alpha_{n}^{\top}A(\alpha_{n}+\beta_{k})+C$
	$\displaystyle\geq$	$\displaystyle\penalty\ C+\max\{\alpha_{n}^{\top}A\alpha_{j},\alpha_{n}^{\top}A(\alpha_{j}+\beta_{k^{\prime}})\},\penalty\ \forall k,k^{\prime\prime}\in[K].$

Proof.

Notice that the following relations are sufficient to guarantee the desired properties

	$\displaystyle\alpha_{n}^{\top}A_{k}\alpha_{n}=$	$\displaystyle\penalty\ \alpha_{n}^{\top}A_{k}\alpha_{0},$
	$\displaystyle\alpha_{n}^{\top}A_{k}\beta_{k^{\prime}}=$	$\displaystyle\penalty\ C,$
	$\displaystyle\alpha_{n}^{\top}A_{k}\alpha_{n}\geq$	$\displaystyle\penalty\ \alpha_{n}^{\top}A_{k}\alpha_{j}+\alpha_{n}^{\top}A_{k}\beta_{k^{\prime}}+C,$
	$\displaystyle\alpha_{n}^{\top}A_{k}\beta_{k}=$	$\displaystyle\penalty\ -C,$
	$\displaystyle\alpha_{n}^{\top}A\beta_{k}=$	$\displaystyle\penalty\ -C,$
	$\displaystyle\beta_{k^{\prime}}^{\top}A_{k}\beta_{k^{\prime}}=$	$\displaystyle\penalty\ 9C.$

By Lemma A.4, we can find $\gamma_{1},\dots,\gamma_{N}\in\mathbb{R}^{\bar{d}}$ such that $\bar{d}=O(\log N)$ , $\gamma_{i}^{\top}\gamma_{j}\leq 1/2$ for any $i\neq j\in[N]$ , and $\gamma_{i}^{\top}\gamma_{i}\geq 1$ for any $i\in[N]$ . Define

\displaystyle B_{k}=e_{k}e_{k}^{\top},\penalty\ \eta_{k}=-e_{k}.

where $e_{1},\dots,e_{K}$ form the standard basis of $\mathbb{R}^{K}$ .

We thus let

\displaystyle\alpha_{i}=\begin{pmatrix}a\gamma_{i}\\ b\mathbf{1}_{E}\\ c1\\ c1\\ 0\end{pmatrix},\penalty\ \beta_{k}=\begin{pmatrix}0\\ f\eta_{k}\\ e\\ -e\\ h\end{pmatrix},\penalty\ \alpha_{0}=\begin{pmatrix}0\\ 0\\ g1\\ -g1\\ 0\end{pmatrix}

\displaystyle A_{k}=\begin{pmatrix}I&&&&\\ &B_{k}&&&\\ &&1&&\\ &&&-1&\\ &&&&1\end{pmatrix},\penalty\ A=\begin{pmatrix}I&&&&\\ &I/K&&&\\ &&0&&\\ &&&0&\\ &&&&0\end{pmatrix},

where $b=c=f=\sqrt{C},e=\sqrt{C}/2,a=\sqrt{3C},g=2\sqrt{C},h=3\sqrt{C}$ . The dimension can be bounded by $d=\bar{d}+K+3=O(K+\log N)$ . ∎

Lemma A.3 (Attention Sink Positional Encoding, Type 2).

For any $C\in\mathbb{R}_{+}$ , $K,N\in\mathbb{Z}_{+}$ , there exist vectors $\alpha_{1},\dots,\alpha_{N},\beta_{0},\dots,\beta_{K}\in\mathbb{R}^{d}$ and matrices $A,A_{1},\dots,A_{K}\in\mathbb{R}^{d\times d}$ for $d\leq O(K+\log N)$ such that for any $n\in[N]$ the followings hold

For any $i\geq j_{1},j_{2},j_{3}$ and $k,k^{\prime},k^{\prime\prime}\neq 0$ :

	$\displaystyle(\alpha_{i}+\beta_{k})^{\top}A_{0}(\alpha_{j_{1}}+\beta_{k^{\prime}})=$	$\displaystyle\penalty\ (\alpha_{i}+\beta_{k})^{\top}A_{0}(\alpha_{j_{2}}+\beta_{k^{\prime\prime}})\geq(\alpha_{i}+\beta_{k})^{\top}A_{0}(\alpha_{j_{1}}+\beta_{0})+C$
	$\displaystyle(\alpha_{i}+\beta_{0})^{\top}A_{0}(\alpha_{i}+\beta_{0})\geq$	$\displaystyle\penalty\ (\alpha_{i}+\beta_{0})^{\top}A_{0}(\alpha_{j_{1}}+\beta_{k})+C.$

For any $i>j$ and $k\neq k^{\prime}\neq 0$

	$\displaystyle(\alpha_{i}+\beta_{k})^{\top}A(\alpha_{i}+\beta_{k})\geq$	$\displaystyle\penalty\ (\alpha_{i}+\beta_{k})^{\top}A(\alpha_{j}+\beta_{k^{\prime}})+C$
	$\displaystyle\geq$	$\displaystyle\penalty\ (\alpha_{i}+\beta_{k})^{\top}A(\alpha_{j}+\beta_{0})+2C.$

For any $i\geq j,j_{1}$ and $k\neq k^{\prime},k^{\prime\prime}$

	$\displaystyle(\alpha_{i}+\beta_{k})^{\top}A_{k^{\prime}}(\alpha_{j}+\beta_{0})\geq$	$\displaystyle\penalty\ (\alpha_{i}+\beta_{k})^{\top}A_{k^{\prime}}(\alpha_{j_{1}}+\beta_{k^{\prime\prime}})+C$
	$\displaystyle(\alpha_{i}+\beta_{k})^{\top}A_{k}(\alpha_{i}+\beta_{k})\geq$	$\displaystyle\penalty\ \max\{(\alpha_{i}+\beta_{k})^{\top}A_{k}(\alpha_{j_{1}}+\beta_{k^{\prime\prime}}),(\alpha_{i}+\beta_{k})^{\top}A_{k^{\prime}}(\alpha_{j_{1}}+\beta_{0})\}+C.$

Proof.

Following the notations in Lemma A.2, let

\displaystyle\alpha_{i}=\begin{pmatrix}\gamma_{i}\\ 0\\ 0\\ 0\end{pmatrix},\beta_{k}=\begin{pmatrix}0\\ \gamma\\ e_{k}\\ 1\end{pmatrix},\beta_{0}=\begin{pmatrix}0\\ \gamma\\ 1\\ f\end{pmatrix},

and

\displaystyle A=\begin{pmatrix}0&&&\\ &a\cdot I&&\\ &&0&\\ &&&0\end{pmatrix},A_{k}=\begin{pmatrix}b\cdot I&&&\\ &0&&\\ &&c\cdot e_{k}e_{k}^{\top}&\\ &&&1\end{pmatrix},A=\begin{pmatrix}e\cdot I&&&\\ &0&&\\ &&0&\\ &&&0\end{pmatrix},

where $a=c=e=C,f=3.5C,d=4C$ . The dimension can be bounded by $d=\bar{d}+K+3=O(K+\log N)$ . ∎

A.7 Technical Claims

Claim A.4 (Johnson-Lindenstrauss Lemma).

Given $0<\varepsilon<1$ , a set $X$ of $N$ points in $\mathbb{R}^{n}$ , and an integer $k>\frac{8(\ln N)}{\varepsilon^{2}}$ , there is a linear map $f:\mathbb{R}^{n}\to\mathbb{R}^{k}$ such that

\displaystyle(1-\varepsilon)\|u-v\|^{2}\leq\|f(u)-f(v)\|^{2}\leq(1+\varepsilon)\|u-v\|^{2}

holds for all $u,v\in X$ .

Claim A.5 (Concentration of Multinomial Distributions, adapted from [2]).

Let $p\in\Delta^{S}$ and $\hat{p}\sim\frac{1}{n}\text{Multinomial}(n,p)$ . Then, for any $\delta\in[0,1]$ :

\displaystyle\mathbb{P}\left(\|\hat{p}-p\|_{1}\geq\sqrt{\frac{2\ln(1/\delta)}{n}}\right)\leq\delta.

Claim A.6 (Berry-Esseen theorem).

If $X_{1},X_{2},\dots$ are i.i.d. random variables with $\mathbb{E}(X_{1})=0$ , $\mathbb{E}(X_{1}^{2})=\sigma^{2}>0$ , and $\mathbb{E}(|X_{1}|^{3})=\rho<\infty$ , we define

\displaystyle Y_{n}=\frac{X_{1}+X_{2}+\cdots+X_{n}}{n}

as the sample mean, with $F_{n}$ the cumulative distribution function of $\frac{Y_{n}\sqrt{n}}{\sigma}$ and $\Phi$ the cumulative distribution function of the standard normal distribution, then for all $x$ and $n$ ,

\displaystyle|F_{n}(x)-\Phi(x)|\leq\frac{8\rho}{\sigma^{3}\sqrt{n}}.

Appendix B Detailed Experiment Results

In Table 2, we report detailed test accuracy comparisons among different models with/without self-correction at test time. We note that:

•

Self-correction significantly boosts models’ test performances.
•

Larger models benefit more from self-correction, indicating that model expressiveness plays an important role in implementing self-correction.

Those empirical findings corroborate our theoretical results.

Model	Accuracy with self-correction (%)	Accuracy without self-correction (%)
GPT-nano	$1.23\pm 1.07$	$2.56\pm 0.43$
GPT-micro	$63.19\pm 0.16$	$93.09\pm 9.70$
GPT-mini	$63.19\pm 0.16$	$98.57\pm 1.85$
Gopher-44M	$63.19\pm 0.16$	$99.15\pm 0.23$

Table 2: Detailed test accuracy comparisons.

Sample Complexity and Representation Ability of Test-time Scaling Paradigms

Abstract

1 Introduction

Our Contributions.

Proposition 1.1 (Informal statement of Theorem 3.1 and Theorem 3.2).

Proposition 1.2 (Informal statement of Theorem 4.7).

2 Preliminaries

Transformers.

Definition 2.1 (Transformer).

Definition 2.2 (Generalized Position Encoder).

Test-time scaling.

3 Separation between Self-Consistency and Best-of-n

Theorem 3.1 (Sample Complexity of Self-Consistency).

Theorem 3.2 (Sample Complexity of Best-of-nn).

4 Expressiveness under Self-Correction

Definition 4.1 (General-Purpose Transformer).

4.1 General-purpose expressiveness

Proposition 4.2 (General-purpose Expressiveness over Different Token Spaces).

Remark 4.3.

Proposition 4.4 (Multi-Task Representation over the Same Token Space).

Remark 4.5.

4.2 General-purpose expressiveness of Transformers with self-correction

Definition 4.6 (Regret-Minimization Transformer).

Theorem 4.7 (Regret Minimization via Self-Correction).

Remark 4.8.

5 Experiments

5.1 Experimental Setup

Data generation.

Model configuration.

Implementation details.

5.2 Results

6 Related Works

Theories of Transformers and Large Language Models.

Test-time scaling.

7 Discussions

References

Appendix A Proofs

A.1 Proof of Theorem 3.1

Proof.

Upper bound.

Lower bound.

A.2 Proof of Theorem 3.2

Proof.

Upper bound.

Lower bound.

A.3 Proof of Proposition 4.2

Proposition A.1 (Extended Representation to Multiple Token Spaces).

Proof.

Prove Eq. (6).

Prove Eq. (7).

Prove Eq. (8).

Proof.

Prove Eq. (12).

Prove Eq. (13).

Prove Eq. (14).

A.4 Proof of Proposition 4.4

Proof.

Prove Eq (23).

Prove Eq. (24).

Prove Eq. (25).

Prove Eq. (26).

Prove Eq. (27).

A.5 Proof of Theorem 4.7

Proof.

A.6 Attention Sink Positional Encoding

Lemma A.2 (Attention Sink Positional Encoding, Type 1).

Proof.

Lemma A.3 (Attention Sink Positional Encoding, Type 2).

Proof.

A.7 Technical Claims

Claim A.4 (Johnson-Lindenstrauss Lemma).

Claim A.5 (Concentration of Multinomial Distributions, adapted from [2]).

Claim A.6 (Berry-Esseen theorem).

Appendix B Detailed Experiment Results

Theorem 3.2 (Sample Complexity of Best-of- $n$ ).