Transformers are Minimax Optimal Nonparametric In-Context Learners

Thank you for your detailed review and positive assessment of our theoretical contributions! Here are our responses to the comments.

Weakness 1 & Limitations.

While the reviewer mentioned in the Limitations section that our model does not include the MLP layer, we find it illuminating to consider our setup as a deep Transformer with all attention layers (except the last one) and layernorm/output layers removed, so that the MLP layers have combined into a feedforward DNN. Since this simplified model already achieves optimal ICL, we can expect a full Transformer to achieve the same. Moreover, our techniques can be partially applied to settings with multiple attention layers or nonlinear (e.g. softmax) attention; please see Item 1 of the global response where we address this in depth.

Weakness 2.

Motivated by the reviewers' comments, we have conducted new numerical experiments justifying our simplified model setup (Weakness 1) and the assumption of empirical risk minimization (Weakness 2). Please see Item 3 of the global response for details.

Question 1.

Indeed, if we completely remove the DNN and set both the features and basis to be simply the coordinate mappings $x_1,\cdots,x_d$, this reduces to the class of linear maps in prior works, e.g. Zhang et al. (2023). In a nutshell, our contribution is extending this to infinite-dimensional nonparametric task classes via learnable representations, enabling fine-grained sample complexity analysis and guaranteeing optimality of ICL with deep architectures.

Question 2.

If the features are fixed, the dynamics are arguably less interesting. Since the attention layer output is linear in the attention matrix $\Gamma$, the pretraining loss function is always convex (ignoring clipping). Hence gradient descent is easily shown to converge exponentially fast to a minimizer regardless of the number of samples, random initialization, etc. which can also be written down explicitly by computing the matrix derivative of line 145 to zero and solving for $\Gamma$.

The linear regression setup of the previous work by Zhang is not quite as simple, in that they also include a scalar multiplier $\sigma$ representing the value matrix and study the joint dynamics. In that paper, convergence is proved by showing a PL inequality under restrictive assumptions on initialization. Rough calculations show that a similar result can be obtained in the fixed feature setting as well, justifying the fixed value matrix assumption in our model. Some other works also indicate that the attention mechanism may possess structures favorable for optimization; see our discussion in Section 1.1. Nevertheless, the main novelty of our work comes from incorporating the representational capability of the MLP layers, for which a dynamical analysis is unfortunately still out of reach.

Thank you again for your time and effort in reviewing our paper! As the discussion period is ending soon, we are following up to see if our response was satisfactory in addressing the reviewer's questions. If not, we would be happy to discuss any remaining concerns.

Thank you for your through review and helpful suggestions, which have helped us greatly to improve our paper!

Response to Weakness 1.

Unfortunately, Section 5 was completed last minute which resulted in the flow of the paper being rather disjoint. Together with improved lower bounds, we restate our take-home message (besides in-context optimality in $n$) as follows. We believe this novel understanding will be helpful to both theoreticians and practitioners.

• (In the Besov space setting) The obtained upper bound $n^{-\frac{2\alpha}{2\alpha+d}}$ when $T\gtrsim nN$ which is minimax optimal in $n$, has also been shown to be jointly optimal in $n,T$. Hence ICL is provably optimal in the large $T$ regime.
• (In the coarser basis setting) We obtained an improved lower bound $n^{-\frac{2\alpha}{2\alpha+d}} + (nT)^{-\frac{2\tau}{2\tau+d}}$ using the method in Appendix E.3. If $T=O(1)$, this gives the slower $\tau$ rate, while if $T$ is larger this gives the faster $\alpha$ rate. This aligns with Corollary 4.9 (which only gave an upper bound), hence ICL is provably suboptimal in the small $T$ regime.

When combined, these results are stronger than optimality in only $n$. (Since the Transformer requires $n$ in-context samples plus $n\times T$ pretraining samples, one could argue that standard minimax rates do not apply. However, the above lower bounds do apply rigorously to any meta-learning algorithm in the $n\times T$ samples.) They also align with experimental observations of task diversity threshold (Raventos 2023), rigorously supporting the importance of varied tasks over increasing prompt length.

Response to Weakness 2 & Question 1.

Motivated by your comments on practical relevance, we have conducted new numerical experiments justifying our simplified model setup and the assumption of empirical risk minimization. Figures can be found in the attached PDF of the global response. We implement and compare the following toy models:

• Linear Transformer: the model studied in our paper, where a DNN (2 hidden layers) feeds into a simplified linear attention layer
• Softmax Transformer: the same model as (a) but with softmax attention
• Full Transformer: an ordinary transformer with 2 stacked encoder layers

Architectures are compared in Figure 1. The number of feedforward layers, width of hidden layers, learning rate etc. are set to equal for a fair comparison.

• Figure 2 shows training (solid lines) and test loss curves (dashed lines) during pretraining. All 3 architectures exhibit similar behavior and converge to near zero training loss, justifying the use of our simplified model and also supporting the assumption that the empirical risk can be minimized.
• Moreover, Figure 3 shows the converged losses over a wide range of $N,n,T$ values (median over 5 runs). We verify that increasing $N,n$ leads to decreasing train and test error, corresponding to the approximation error of Theorem 3.1. We also observe that increasing $T$ tends to improve the pretraining generalization gap up to a threshold, confirming our theoretical analyses. Again, this behavior is consistent across the 3 architectures.
• Although the dimensions and architectures are relatively small due to limited time and compute, we plan to scale up our experiments and add them to the final version of the paper.

Please also see Item 1 of the global response where we discuss how our theoretical results may be extended to more complex transformer architectures. We also humbly ask the reviewer to consider raising their score if their concerns were addressed.

(continued from rebuttal)

• Isn't is unrealistic that $N$ should scale in $n$ for a real transformer?

In practice we agree that the architecture should not depend on $n$, especially if prompts of any length are allowed. However since we assume all prompts (during both pretraining and ICL) are of fixed length $n$, it is reasonable to choose a more powerful architecture for larger $n$.

Strictly speaking, $N$ does not need to scale in $n$ since our bounds hold for all $N,n,T$. But the bound itself is natural: as we mentioned, approximation error must decrease in $N$ due to the infinite dimensionality of the target class, and generalization error should of course decrease in $n,T$. For example if $N$ is considered fixed, the bound is interpreted as an excess risk of $O(1/n+1/T)$. However we want to obtain the overall sample complexity rate in $n$, which necessitates $N$ to scale in $n$. The rate itself is less important than the fact that it is the best rate attainable by any estimator, and the optimality result should be viewed more as guaranteeing tightness of the derived bounds (further justified by Section 5) than a prescription of how $N,T$ should scale.

We also mention that this approach is standard in minimax analyses. For example, rate-optimal scaling for ordinary supervised regression with DNNs (Suzuki, 2019) also requires the width to scale as $n^{\frac{1}{2s+1}}$ even though the smoothness of the target is unknown.

• Rebuttal End

Finally, we have also obtained improved lower bounds which reinforce the message of our paper, and newly conducted numerical experiments justifying our assumptions and results. Please see Item 2 and Item 3 of the global response for details. We also humbly ask the reviewer to consider raising their score if some of their concerns were addressed.

Thank you for the detailed review and insightful questions! Our responses are as follows.

• It would be helpful to instantiate a target function class in the main paper, specify what $\alpha$ is, etc.

(We answer this question first to help in understanding our paper.) Throughout Section 4, we consider concrete task classes including the Besov space, anisotropic Besov space, and piecewise $\gamma$-smooth function class. Our main result for DNNs (Theorem 4.5) is stated when $\mathcal{F}^\circ$ is the unit ball in Besov space $B_{p,q}^\alpha(\mathcal{X})$, which is a very wide function class generalizing Holder and Sobolev spaces; please see Section 4.1 for details. Here $\alpha$ is the smoothness and $p,q$ are additional parameters controlling spatial inhomogeneity. These relate to the decay rate $s$ of Assumption 2 by $s=\alpha/d$; the high-level Assumptions 1-3 all follow from Assumption 4 in this setting.

We also show the curse of dimensionality can be avoided when $\mathcal{F}^\circ$ is extended to the anisotropic Besov space in Section 4.3. Furthermore, for the multilayer Transformer setting (Section 4.5) the inputs $x$ are sequences and $\mathcal{F}^\circ$ is the piecewise $\gamma$-smooth class, which is a more realistic model for text by allowing the position of important tokens to depend on input.

• Is there any hope to extend this analysis to more attention layers/non-linear attention/constant $N$?

Please see Item 1 of the global response where we address this question in depth.

• When $N$ is constant, the error is constant, which seems vacuous.

While our bounds (e.g. Theorem 4.5) hold for all $N$, the risk dependency is not an artifact but a natural result of how our task class is set up. After pretraining our transformer has learned $N$ feature maps $\phi_1,\cdots,\phi_N$ and its output $f_\Theta(\mathbf{X},\mathbf{y},\cdot)$ is always contained in the span of these $N$ functions. However, the true test task $F_\beta^\circ$ is randomly chosen from an infinite-dimensional space, and so the $L^2$ regression error is fundamentally lower bounded by its Kolmogorov width -- the minimum approximation error by an $N$-dimensional linear subspace -- which scales as $N^{-\alpha/d}$ for the Besov space. Hence this error is unavoidable for any fixed-basis approximation scheme; an implication of our analysis is that Transformers attain this optimal rate (the first term in Theorem 4.5).

Besides the approximation error, the pretraining and in-context generalization errors decrease as $O(1/T)$ and $O(1/n)$, which is still a useful result if one is concerned with improving generalization.

• Is it possible to have a setting where we removed the DNN?

Indeed, if we completely remove the DNN and set both the features and basis to be simply the coordinate mappings $x_1,\cdots,x_d$, this reduces to the class of linear maps in prior works. In a nutshell, our contribution is extending this to infinite-dimensional nonparametric task classes via learnable representations, enabling fine-grained sample complexity analysis and guaranteeing optimality of ICL with deep architectures.

• It seems very important to justify how we can find a $\phi$ satisfying Assumption 3 for DNNs.

That is exactly what Appendix C.1.1 (Verification of Assumptions) and Lemma 4.4 are for! We show that the abstract Assumptions 1,2,3 are all replaced by the single concrete Assumption 4 when the target class is set to the Besov space. For the multilayer Transformer setting (Section 4.5), Assumption 3 instead follows from Theorem D.2. We will make this point clearer in the paper.

• Why aren’t the basis functions the first N functions?

In fact, this is to address a unique technical difficulty that we newly uncover: the necessity of the inverse covariance matrix $\Psi_N^{-1}$ to approximate the attention matrix. Wavelet systems of Besov-type spaces are fundamentally co-dependent and form a multi-resolution scheme: there are $O(2^k)$ independent basis elements at each resolution $k\in\mathbb{N}$, which can always be refined as a linear combination of higher resolution wavelets. This makes $\Psi_N$ singular and affects various decay rates, so we cannot simply take the first $N$ elements. This was not a problem in prior attention-only works ($\text{Var}(x)$ was simply assumed to be positive definite) nor any existing minimax optimality works (which don't have a product structure). Nonetheless, we were still able to prove optimality by decomposing all wavelets to the finest resolution -- numbered from $\underline{N}$ to $\overline{N}$ -- via the refinement analysis in Appendix C.3, hence the additional steps to account for the aggregated coefficients $\bar{\beta}$. While this additional notation may cause confusion upon first reading, we have left it in since specializing to the wavelet system to obtain optimality is an important part of our work.

• What role does the sparsity $S$ play, why is it set to be $O(1)$?

The sparsity bound $S=O(1)$ is merely to clarify the minimum number of essential parameters and can be completely removed. Recall that a fully connected network has $S\sim LW^2$. The dominating term from the DNN class entropy (line 675) is $SLN\log W$, so that letting $S=O(\log N)$ (fully connected depth $L$ network) would only incur an additional log factor to yield $N(\log N)^2$. Since the $N^2\log N$ term from the attention matrix entropy dominates this, the overall bound remains unchanged.

(continued in comment)

Thank you again for your time and effort in reviewing our paper! As the discussion period is ending soon, we are following up to see if our response addressed the reviewer's questions. If so, we hope that they would be willing to increase their score. If not, we would be happy to discuss any remaining concerns.

Thank you for your response. I will raise my score to a 5.

Thank you for your through review of our work! We hope our responses can help clarify some important points.

Weakness 1.

Theorem 3.1 is indeed a standard result, as we have indicated it is a straightforward adaptation of Lemma 4 of Schmidt-Hieber (2020). While it is only the first step of the proof of the main results, we included the statement to explicitly show how the in-context and pretraining generalization gaps arise from different elements of the risk: the former manifests when evaluating the approximation error w.r.t. $n$, while the latter is the usual entropy term.

Weakness 2.

The attention matrix construction can be seen as an extension of Zhang (2023), however we urge the reviewer to take the following points into consideration.

• The optimal construction is completely determined for a linear attention layer since the L2 linear regression loss is convex; even the construction in Zhang et al. (2023) is stated to be a generalization of Oswald et al. (2023). So while we agree the construction is not surprising, it is an important step allowing us to reduce to the analysis of the DNN module which facilitates the main risk analysis. Conversely, Zhang and Oswald do not consider the MLP at all and only study the attention matrix.
• The given form (line 174) is only an example construction to upper bound the empirical loss. Unlike Zhang, we do not claim the ERM or optimization dynamics must result in such a simplified form, which may not be true in a deep Transformer setting.
• The idea of approximating the target function with a well-chosen basis is of course fundamental to all functional analysis and is not unique to our approach. Indeed, our goal is to establish a strong connection between ICL and ordinary supervised regression by isolating the pretraining generalization gap. We dive deeper into this idea by pointing out a new limitation of ICL stemming from non-adaptive function approximation (Remark 4.6).
• Moreover, a unique technical difficulty that we newly uncover and address is the necessity of the inverse covariance matrix $\Psi_N^{-1}$ to approximate the attention matrix. Wavelet systems of Besov-type spaces are fundamentally co-dependent and form a multi-resolution scheme which makes $\Psi_N$ singular and affects various decay rates. This was not a problem in Zhang's work ($\Lambda$ was simply assumed to be positive definite) nor any existing minimax optimality works (which don't have a product structure). Nonetheless, we were still able to prove optimality by decomposing all splines to the finest resolution and performing the analysis strictly at that scale, hence the additional steps to account for the aggregated coefficients $\bar{\beta}$.
• We have added experiments justifying the use of our simplified transformer model; please see Item 3 of the global response.

We also implore the reviewer to judge our result not based on the perceived difficulty level of the proof, but its implications towards understanding the effectiveness of in-context learning.

Question 1.

The assumption can be substantially relaxed; the stated form is mostly to clarify the minimum number of essential parameters. Recall that a fully connected network has $S\sim LW^2$. The dominating term from the DNN entropy (line 675) is $SLN\log W$, so that letting $S=O(\log N)$ (fully connected depth $L$ network) would only incur an additional log factor to yield $N(\log N)^2$. Since the $N^2\log N$ term from the attention matrix entropy dominates this, the overall bound remains unchanged. The same holds even if width is relaxed to $W=O(N)$, and furthermore even if $S=O(N)$ in this case.

However if $W=O(N)$ and the network is fully connected, that is $S=O(N^2\log N)$, the pretraining gap will worsen to around $N^3/T$. This reflects the fact that wavelets are 'easier' to approximate with DNNs than Besov functions are with wavelets.

Question 2.

That is indeed the setup we were aiming for, however please consider our response to Weakness 2 as well as the following points.

• Linear attention has been studied in many other papers including Zhang, Oswald, Mahankali, and Ahn's works (cited in the paper), all yielding similar 'optimal matrix' constructions. As mentioned before, we do not claim originality in this regard. However, our focus is on reduction to the representational capability of the MLP layers. Except constructive works such as Bai et al. (2023), The NN+attention setup had only been studied before very recently by Kim and Suzuki (2024), and they only considered a shallow MLP.
• We also find it illuminating to consider our setup as a deep Transformer with all but the last attention layer removed. Since this simplified model is already in-context sample optimal, we can expect a full Transformer to be the same. Moreover, we indicate how to extend our results to nonlinear attention or multiple attention layers in Item 1 of the global response.

We have also obtained improved lower bounds which reinforce the message of our paper, and newly conducted numerical experiments justifying our assumptions and results. Please see Item 2 and Item 3 of the global response for details. We also humbly ask the reviewer to consider raising their score if some of their concerns were addressed.