Context-Scaling versus Task-Scaling in In-Context Learning

## Abstract

Abstract Transformers exhibit In-Context Learning (ICL), where these models solve new tasks by using examples in the prompt without additional training. In our work, we identify and analyze two key components of ICL: (1) context-scaling, where model performance improves as the number of in-context examples increases and (2) task-scaling, where model performance improves as the number of pre-training tasks increases. While transformers are capable of both context-scaling and task-scaling, we empirically show that standard Multi-Layer Perceptrons (MLPs) with vectorized input are only capable of task-scaling. To understand how transformers are capable of context-scaling, we first propose a significantly simplified transformer architecture without key, query, value weights. We show that it performs ICL comparably to the original GPT-2 model in various statistical learning tasks including linear regression, teacher-student settings. Furthermore, a single block of our simplified transformer can be viewed as data dependent “feature map” followed by an MLP. This feature map on its own is a powerful predictor that is capable of context-scaling but is not capable of task-scaling. We show empirically that concatenating the output of this feature map with vectorized data as an input to MLPs enables both context-scaling and task-scaling. This finding provides a simple setting to study context and task-scaling for ICL.

## 1 Introduction

Pre-trained large language models exhibit In-Context Learning (ICL) capabilities, allowing them to adapt to new tasks based exclusively on input without updating the underlying model parameters [].

| Input (Prompt) | Output | Task (Pattern) |
| --- | --- | --- |
| $(1,2,3),(4,5,9),(10,-9,1),(5,6,?)$ | $11$ | In each triplet $(a,b,c)$: $c=a+b$ |
| $(4,3,1),(9,0,9),(10,8,2),(17,17,?)$ | 0 0 | In each triplet $(a,b,c)$: $c=a-b$ |
| $(1,2,5),(2,3,8),(3,4,11),(5,6,?)$ | $17$ | In each triplet $(a,b,c)$: $c=a+2b$ |
| $(2,1,7),(3,4,18),(5,2,16),(4,3,?)$ | $17$ | In each triplet $(a,b,c)$: $c=2a+3b$ |

In the table above, we provide an example of ICL. We prompt the Claude language model [] using sequences of $N$ triples of numbers ($N=4$). In each row of the table, the first $N-1$ triples follow a given pattern. The model was able to infer the pattern and to fill in the missing number denoted by the question mark correctly. What makes this in-context learning possible?

Recent research analyzed various ICL problems where, for example, the task data were generated using linear regression, student-teacher neural networks, and decision trees including [, , , , , , ]. In these problems, a transformer was first pre-trained on $T$ tasks where the data in each task was generated from a given family of functions. For example, each task may involve predicting the last element in a tuple as a linear combination of the other elements, as was shown in the table above. The pre-trained transformer was then tested on $N$ samples from a new task that is drawn from the same general family but was not seen during pretraining. Such a setup allows for understanding the effect of various factors including model architecture, the number of pre-training tasks $T$, and context length $N$ on ICL.

The ability of models to learn in context has been broadly defined as their ability to generalize to unseen tasks based on context examples without updating the model parameters. We observe that there are two distinct aspects of generalization in ICL. The first, which we call context-scaling, refers to the ability of the model to improve as the context length $N$ increases while the number of pre-training tasks $T$ is fixed. The second, task-scaling refers to the ability of a model to improve as $T$ increases while $N$ is fixed.

Figure: Figure 1: Task-scaling and context-scaling of GPT-2 architecture transformers versus MLPs for ICL with linear regression tasks. (A) Task-scaling abilities of these models with $10$ in-context examples. (B) Context-scaling abilities of these models with $10^{5}$ (left) and $10^{6}$ (right) pre-training tasks. Experimental details are provided in Appendix A.
Refer to caption: x1.png

It is a priori unclear whether a model capable of context-scaling is also capable of task-scaling and vice-versa. For example, as we show in Figure 1A, both transformers and standard Multi-Layer Perceptrons (MLPs) are capable of task-scaling on ICL of linear regression tasks. In contrast, only transformers benefit from an increasing number of context examples as shown in Figure 1B. This raises the question:

What mechanism enables models such as transformers, but not MLPs, to context-scale?

To identify such a mechanism, we begin by constructing a bare-bones transformer with all key, query, and value weight matrices fixed to be the identity matrix. We refer to our simplified model as Simplified GPT (SGPT). Despite its simplicity, we find that SGPT is competitive with GPT-2 architecture transformers [] for a variety of ICL problems considered in [] including linear regression, student-teacher networks, decision trees, and sparse linear regression.

Furthermore, we find that one block of SGPT applies a data dependent feature map that enables context-scaling. To illustrate how such a feature map can be effective for context-scaling, consider the following input data for ICL:

$$ $\displaystyle A=\begin{bmatrix}x_{1}&x_{2}&\cdots&x_{N-1}&x_{N}\\ y_{1}&y_{2}&\cdots&y_{N-1}&0\end{bmatrix}^{\top}\in\mathbb{R}^{N\times(d+1)}~{};$ (1) $$

where $x_{i}\in\mathbb{R}^{d}$ and our goal is to predict $y_{N}$. SGPT first applies a feature map $\psi:\mathbb{R}^{N\times(d+1)}\rightarrow\mathbb{R}^{N\times(d+1)}$ to $A$ and then trains an MLP on the last row of $\psi(A)$, denoted $\psi(A)_{N,:}$.

By varying $\psi$, we show that the scalar $\psi(A)_{N,d+1}$ itself is an effective estimate for $y_{N}$.
For example, when $\psi(A)=(AA^{\top})A$, $\psi(A)_{N,d+1}$ implements one-step of Gradient Descent (GD) for linear regression (using context examples $(x_{i},y_{i})_{i=1}^{N-1}$) []. When the data follow an isotropic Gaussian distribution, this estimator is consistent (as both the numbers of pre-training tasks and context examples grow) and nearly matches the optimally tuned ridge regression [, ].
Furthermore, we show that $\psi(A)_{N,d+1}$ can implement the Hilbert estimate [], which provides a statistically consistent estimate for general families of tasks beyond linear regression as the context length $N$ approaches infinity. As such, our results provably establish that one attention layer alone is capable of context-scaling for any family of tasks.

Further, we empirically show that these features can enhance the capabilities of MLPs for context-scaling. We concatenate features from $\psi(A)_{N,:}$ with the vectorized input,

$$ $A_{v}:=\begin{bmatrix}x_{1}^{T}&y_{1}&x_{2}^{T}&y_{2}&\ldots&x_{N}^{T}&0\end{ bmatrix}^{T},$ $$

and provide the combined input $[A_{v}^{T},~{}\psi(A)_{N,:}^{T}]$ to an MLP. We show that the resulting model exhibits both context-scaling and task-scaling.

We summarize our findings as follows:

- •
We identify two regimes for characterizing generalization in ICL: context-scaling and task-scaling. We observe that MLPs task-scale but do not context-scale in contrast to GPT2 architecture transformers, which exhibit both context and task-scaling.
- •
We propose a simplified transformer, SGPT, with all key, query, and value weight matrices set to the identity matrix, exhibits both context and task-scaling and is competitive with GPT-2 on a range of ICL tasks.
- •
We analyze a one-layer version of SGPT, demonstrating that this model is capable of context-scaling solely through the use of a feature map, $\psi$, applied to input data. We show that $\psi$ can be selected to perform kernel smoothing to impute the missing element in each task using the other in-context examples. Choosing $\psi$ based on the Hilbert estimate [] results in a statistically optimal (consistent) estimate as the context length approaches infinity.
- •
We show empirically that concatenating the output of $\psi$ with vectorized data as an input to MLPs enables both context-scaling and task-scaling.

## 2 Prior work

#### ICL in controlled settings.

The work by [] initiated the study of ICL in statistical learning tasks, such as linear regression, decision tree learning, and teacher-student neural network learning. They showed that transformers such as GPT-2 [], when pre-trained with a large number of independent tasks (around $3.2\times 10^{7}$ independent tasks, each presented once), learn in-context. Specifically, during inference time, the performance of pre-trained transformers empirically matches that of standard, specialized algorithms for these tasks. These results were later extended to various other settings see, e.g., [, , , , ]. In particular, [] empirically showed transformers can achieve nearly optimal ICL when pre-trained with multiple passes over a smaller, fixed training set of independent tasks. [] constructed transformers that implement optimal algorithms in context even if the tasks are generated from a mixture of distributions. In a recent work, [] empirically showed that pre-trained MLPs can achieve ICL when the context length is fixed during pre-training and inference. However, it was not known whether MLPs can achieve context-scaling, and our work empirically gives a negative answer.

#### Linear attention.

Most theoretical analyses of ICL has been in the setting of linear regression for fixed context lengths. Specifically, [] showed by construction that single linear attention can implement one-step gradient descent (GD) in context. [, , , ] proved that the predictor given by pre-trained transformers with single linear layer of attention is equivalent to that given by one-step GD. [] proved the predictor from one-step of GD achieves the rate of optimally-tuned ridge regression. In addition, they show that the predictor given by one step of GD can be pre-trained to achieve this rate with finite independent tasks.
Later works such as [] connected nonlinear attention to functional one-step GD in feature space. These works together have substantially furthered our understanding of single-layer linear attention transformers for ICL of linear regression with a fixed context length.
However, these results are primarily concerned with fixed context length settings and do not address the context-scaling abilities of pre-trained transformers. Recent work [] provided an asymptotic theory of ICL for linear attention transformers. Unlike our analyses, they focused on settings in which the number of unique pre-training tasks (referred to as “diversity”) scaled proportionally with the data dimension.

#### Softmax attention and kernel smoothers.

The connection between softmax attention and kernel smoothers was first pointed out by [].
Specifically, by setting the query and key matrices to be the same, an attention head can be viewed as a kernel-smoother with a learnable positive semi-definite kernel. Empirical evidence suggests that using a shared matrix for query and key matrices does not significantly impact the performance of transformers [, ].
Later, theoretical works [, ] utilized this connection to study the ICL of softmax attention in both linear [] and nonlinear regression tasks []. In these settings, softmax attention can perform ICL by implementing a kernel smoother with a bandwidth parameter obtained by training query and key matrices. Compared with these works, we demonstrate that transformers can perform ICL using an attention component that does not contain trainable parameters (that is, setting query, key, and value matrices to identity matrices). Our results suggest that the transformer architecture is capable of performing ICL without needing to explicitly learn any hyperparameters in the attention head. We explain this by connecting to a hyper-parameter-free yet statistically consistent kernel smoother given by the Hilbert estimate [].

#### Approximation ability of transformers.

The transformer is known to be a versatile architecture that can implement efficient algorithms (by forwarding passing an input prompt) in many scenarios [, , , , ].
These results make use of query, key, and value weight matrices for constructing ICL algorithms. In this work, we show that these elements of transformers are not needed for building and training models that are competitive with GPT-2 architecture transformers for ICL tasks considered in prior works [].

## 3 Preliminaries

In this section, we outline the problem setup, training details, architectural details, and mathematical preliminaries for our work.

#### Problem formulation.

For all ICL tasks studied in our work, we consider $T$ pre-training tasks, each with input data of the form:

$$ $\displaystyle A_{t}=\begin{bmatrix}x_{1}^{t}&x_{2}^{t}&\cdots&x_{N}^{t}\\ y_{1}^{t}&y_{2}^{t}&\cdots&y_{N}^{t}\end{bmatrix}^{\top}\in\mathbb{R}^{(N) \times(d+1)},$ (2) $$

where $t=1,\ldots,T$ indexes the tasks, $N$ denotes the maximum context length, and $d$ denotes the input data dimension.
We define the loss function as:

$$ $L(\theta;A_{1},\ldots,A_{T}):=\frac{1}{T}\sum_{t=1}^{T}\left[\frac{1}{N}\sum_{ i=1}^{N}(M_{\theta}(A_{t}^{i})-y_{i+1}^{t})^{2}\right],$ $$

where

$$ $A_{t}^{i}=\begin{bmatrix}x_{1}^{t}&x_{2}^{t}&\cdots&x_{i}^{t}&x_{i+1}^{t}\\ y_{1}^{t}&y_{2}^{t}&\cdots&y_{i}^{t}&0\end{bmatrix}^{\top}\in\mathbb{R}^{(i+1) \times(d+1)},$ $$

and $M_{\theta}(\cdot)$ denotes the model with trainable parameters $\theta$. The tasks $A_{t}$ are uniformly sampled from the family of tasks $\mathcal{F}$ (e.g., linear regression with a Gaussian prior), representing the distribution of tasks relevant to the in-context learning problem.

#### Attention.

Given three matrices $A_{1},A_{2},A_{3}\in\mathbb{R}^{N\times m}$, attention layers implement functions $g:\mathbb{R}^{N\times m}\times\mathbb{R}^{N\times m}\times\mathbb{R}^{N\times m
}\to\mathbb{R}^{N\times m}$ defined as follows,

$$ $\displaystyle g(A_{1},A_{2},A_{3}):=\phi\left(\frac{1}{\sqrt{m}}A_{1}A_{2}^{ \top}\right)A_{3};$ (3) $$

where $\phi:\mathbb{R}^{N\times N}\to\mathbb{R}^{N\times N}$ is a generic function that could be a row-wise softmax function [], an entry-wise activation function such as ReLU, or just an identify map. For self-attention layers, we are typically given one input matrix $A\in\mathbb{R}^{N\times m}$ and three weight matrices $W_{Q},W_{K},W_{V}\in\mathbb{R}^{m\times m}$. In this case, the matrices $A_{1},A_{2},A_{3}$ are computed respectively as $AW_{Q}$, $AW_{K}$, and $AW_{V}$.

#### Kernel functions.

Kernel functions are positive-semidefinite functions that map pairs of inputs to real values []. Formally, given $x,y\in\mathbb{R}^{d}$, a kernel $K:\mathbb{R}^{d}\times\mathbb{R}^{d}\rightarrow\mathbb{R}$ is a function of the form $K(x,y)=\langle\psi(x),\psi(y)\rangle_{\mathcal{H}}$, where $\psi:\mathbb{R}^{d}\rightarrow\mathcal{H}$ is referred to as a feature map from $\mathbb{R}^{d}$ to a Hilbert space $\mathcal{H}$. For matrix inputs $A\in\mathbb{R}^{m\times d}$ and $B\in\mathbb{R}^{n\times d}$, we let $K(A,B)\in\mathbb{R}^{m\times n}$ such that $K(A,B)_{ij}=K(A_{i},B_{j})$ where $A_{i}$ and $B_{j}$ denote the $i$^th and $j$^th rows of $A$ and $B$ respectively.

#### Kernel smoother.

Given a kernel $K:\mathbb{R}^{d}\times\mathbb{R}^{d}\to\mathbb{R}$ and a set of points $(\mathbf{x}_{i},y_{i})_{i=1}^{n}$ where $\mathbf{x}_{i}\in\mathbb{R}^{d}$ and $y_{i}\in\mathbb{R}$, the kernel smoother estimate at point $\mathbf{x}$ is a function of the form

$$ $\displaystyle\widehat{f}_{K,n}(\mathbf{x}):=\frac{\sum_{i=1}^{n}K(\mathbf{x}, \mathbf{x}_{i})y_{i}}{\sum_{i=1}^{n}K(\mathbf{x},\mathbf{x}_{i})}~{}.$ (4) $$

We will reference kernel smoothers in Section 5.

#### Hilbert estimate.

The Hilbert estimate is a kernel smoother using the kernel

$$ $\displaystyle H(x,x^{\prime}):=\frac{1}{\left\lVert x-x^{\prime}\right\rVert_{ 2}^{d}},$ (5) $$

where $\left\lVert\cdot\right\rVert_{2}$ denotes the $\ell_{2}$-norm in $\mathbb{R}^{d}$. The key property of the Hilbert estimate that we use is that it is a asymptotically optimal (consistent) estimate. In particular, at almost all $x$, as the number of samples $n$ goes to infinity, $\widehat{f}_{H,n}\to f^{*}$ in probability where $f^{*}(x)=\mathbb{E}[y|X=x]$ denotes the optimal predictor [].

## 4 Simplified transformer model performs ICL

In this section, we introduce our simplified transformer model, SGPT, and demonstrate that it is competitive with GPT-2-type architectures on various ICL tasks. To construct SGPT, we fix all key, query, and value weights to be the identity matrix in the GPT-2 architecture. Consequently, the attention mechanism (defined in equation 3) reduces to the function:

$$ $\displaystyle g(H):=\phi(HH^{\top})H\in\mathbb{R}^{N\times(d+1)}.$ (6) $$

We define $\phi$ to be a function that performs row-wise $\ell_{1}$ normalization on its argument. To further simplify our model, we remove the final linear layer of each MLP block (i.e., our MLP blocks have one linear layer), along with batch normalization and the skip connection after the MLP layer. These details are further outlined in Appendix A.

We consider the following ICL tasks from prior work: (1) linear regression with a single noise level [], (2) linear regression with multiple noise levels used in [], (3) sparse linear functions used in [], (4) two-layer ReLU neural networks [], and (5) decision trees []. Below, we explain the problem setup and state our results for each of these five synthetic tasks.

#### Linear regression with fixed noise level.

The problem setting is as follows:

$$ $\displaystyle x\in\mathbb{R}^{d}\sim\mathcal{N}(0,I_{d}),\quad y=\beta^{\top}x +\epsilon~{}\text{with }\beta\sim\mathcal{N}\left(0,\frac{I_{d}}{d}\right),\; \epsilon\sim\mathcal{N}(0,\sigma^{2}).$ $$

Figure: Figure 2: Linear regression with a single noise level. Left panel. Performance across varying context lengths (context-scaling). Right panel. Effect of regularization on performance for a fixed number of in-context examples. Experimental details are given in Appendix A.
Refer to caption: x2.png

In this setting, prior work by [], showed that on all context lengths, GPT-2 architecture transformers can perform comparably to task-specific, optimally-tuned ridge regression. In Figure 2, we provide evidence that SGPT matches the performance of these GPT-2 models.

#### Linear regression with multiple noise levels.

The problem setting is as follows:

$$ $\displaystyle x\in\mathbb{R}^{d}\sim\mathcal{N}(0,I_{d}),\quad y_{i}=\beta^{ \top}x_{i}+\epsilon$ $\displaystyle\text{with }\beta\sim\mathcal{N}\left(0,\frac{I_{d}}{d}\right) \text{ and }\epsilon\sim\begin{cases}\mathcal{N}(0,\sigma_{1}^{2}),&\text{with probability }\frac{1}{2}\\ \mathcal{N}(0,\sigma_{2}^{2}),&\text{with probability }\frac{1}{2}\end{cases}.$ $$

Figure: Figure 3: Linear regression with multiple noise levels. Left and middle panels: Performance across varying context lengths (context-scaling). Right panel: Effect of regularization on performance for a fixed number of in-context examples. Experimental details are given in Appendix A.
Refer to caption: x3.png

In this setting, prior work by [] demonstrated that GPT-2 architecture transformers can achieve performance comparable to that of task-specific, optimally tuned ridge regression across all context lengths and for both noise levels. They refer to the model’s ability to adapt to the noise level as algorithm selection. In Figure 3, We demonstrate that SGPT performs comparably to GPT-2 architecture transformers in this setting, exhibiting similar algorithm selection capabilities.

#### Two-layer ReLU Neural Networks.

Following the work of [], we consider the following nonlinear problem setting where data for each task are generated using two-layer neural networks. In particular, data are generated according to

$$ $\displaystyle x\in\mathbb{R}^{d}\sim\mathcal{N}(0,I_{d}),\quad y=\sum_{j=1}^{r }\alpha_{j}\phi(w_{j}^{\top}x);$ $$

where $\alpha_{j},w_{j}$ are randomly initialized parameters of a fixed two-layer neural network, $\phi$ denotes the element-wise ReLU activation function, and $r=100,d=20$ (as selected in prior work). The work [] demonstrated that GPT-2 architecture transformers can match the performance of student networks (i.e., networks of the same architecture initialized differently and trained using Adam optimizer []). In Figure 4(B), we show that SGPT can match the performance of GPT-2 architecture transformers on this task.

#### Decision Tree.

Following the work of [], we consider a nonlinear problem setting where data for each task are generated using depth-four trees. For a task corresponding to a tree $f$, we have:

$$ $\displaystyle x\sim\mathcal{N}(0,I_{d}),\quad y=f(x),$ (7) $$

Previously, [] demonstrated that GPT-2 architecture transformers can perform in-context learning on this family of non-linear functions, outperforming XGBoost as a baseline. In Figure 4(A), we show that SGPT is also capable of in-context learning (ICL) in this setting, performing comparably to GPT-2 architecture and similarly outperforming XGBoost. We trained XGBoost models using the same hyperparameters as in [].

#### Sparse linear functions.

Following the work of [], we consider the class sparse linear regression problems. In this setting, data are generated according to

$$ $\displaystyle x\in\mathbb{R}^{d}\sim\mathcal{N}(0,I_{d}),\quad y=\beta^{\top}x,$ (8) $$

where $\beta\sim\mathcal{N}\left(0,I_{d}\right)$ and we zero out all but $s$ coordinates of $\beta$ uniformly at random for each task. As in prior work, we select $d=20,s=3$. In Figure 4C
, we demonstrate that SGPT is capable of ICL for this class of functions, performing comparably to GPT-2 architecture transformers and closely to the Lasso estimator [], while significantly outperforming the ordinary least square (OLS) baseline.

Figure: Figure 4: Nonlinear ICL tasks. Context-scaling capability of SGPT versus GPT-2 architecture transformers when trained on 2 million pre-training tasks. In all cases, the errors are normalized so that the trivial zero predictor achieves an error of 1.^∗ Experimental details are given in Appendix A.
Refer to caption: x4.png

Thus far, we have demonstrated SGPT is comparable to GPT-2 across various ICL tasks. In Appendix C, we revisit the experiment introduced in the introduction and show that GPT-2 and SGPT are capable of both context and task-scaling.

## 5 Kernel smoothing can perform context-scaling

We begin this section by demonstrating that even one layer of SGPT is capable of context scaling. In particular, in Figure 5, we train a one-layer model on five different context lengths and test on the same lengths for the tasks considered in the previous section. In all four problem settings, it is evident that using more context improves performance. Below, we analyze this simplified one layer model in order to pinpoint how it is capable of context scaling.

Figure: Figure 5: Context-scaling with one-layer SGPT. Experimental details are provided in Appendix A.
Refer to caption: x5.png

In particular, the model we analyze is identical to one layer of SGPT up to the omission of remaining skip connections (for more details, see Appendix A). The model implements a function $f:\mathbb{R}^{N\times(d+1)}\to\mathbb{R}$ and takes the form below:

$$ $\displaystyle f(A):=\left[\sigma\left(\psi\left(A\right)W^{(1)}\right)W^{(2)} \right]_{N}=\sigma\left(\psi\left(A\right)_{N,:}W^{(1)}\right)W^{(2)};$ (9) $$

where $\psi:\mathbb{R}^{N\times(d+1)}\rightarrow\mathbb{R}^{N\times(d+1)}$ is a feature map (generalizing the attention function defined in equation 6), $A\in\mathbb{R}^{(N)\times(d+1)}$ denotes the input data, and $W^{(1)}\in\mathbb{R}^{d+1\times k},W^{(2)}\in\mathbb{R}^{k\times 1}$.

Figure: Figure 6: Comparison between using a row of features given by $\psi_{K}(\cdot)_{N,:}$ and using the scalar given by $\psi_{K}(\cdot)_{N,d+1}$ for $K\in\{L,H\}$. Here, the feature maps $\psi_{L},\psi_{H}$ are defined in equation 10 and equation 5 respectively.
Refer to caption: x6.png

### 5.1 Feature-map that enables context-scaling

The key aspect distinguishing the model in equation 9 from a standard MLP operating on vectorized inputs $A\in\mathbb{R}^{N(d+1)}$ is the feature map $\psi$. As such, we analyze how the feature map $\psi$ transforms an input

$$ $A=\begin{bmatrix}x_{1}&x_{2}&\cdots&x_{N-1}&x_{N}\\ y_{1}&y_{2}&\cdots&y_{N-1}&0\end{bmatrix}^{\top}\in\mathbb{R}^{N\times(d+1)}.$ $$

First, we note that upon varying the function $\psi$, the bottom-right element of $\psi(A)$, denoted $\psi(A)_{N,d+1}$, is capable of implementing several well-known estimators, which we describe below. To ease notation, we let $\mathbf{X}:=[x_{1},x_{2},\cdots,x_{N}]^{\top}\in\mathbb{R}^{N\times d}$ and $\mathbf{y}:=[y_{1},\cdots,y_{N}]^{\top}$. Detailed derivations of the explicit forms for $\psi(A)_{N,d+1}$ below are presented in Appendix B.

(1) 1-step GD estimate. Let $\psi_{L}(A):=(AA^{\top})A$. Then,

$$ $\psi_{L}(A)_{N,d+1}=x_{N}^{\top}\mathbf{X}^{\top}\mathbf{y}.$ (10) $$

Thus, $\psi_{L}$ computes the estimate arising from a linear predictor trained for one step of gradient descent on the data $(\mathbf{X},\mathbf{y})$. This estimate has been previously considered as a mechanism through which transformers performed ICL, but there have been no theoretical guarantees for this approach for general ICL tasks beyond linear regression [, , , , ].

(2) Kernel smoother. Given a kernel $K$, let $\psi_{K}(A)=\widehat{K}(\mathbf{X},\mathbf{X})A$, where

$$ $\widehat{K}(\mathbf{X},\mathbf{X})_{i,j}=\begin{cases}\frac{K(x_{i},x_{j})}{ \sum_{j\neq i}K(x_{i},x_{j})}&\text{ if $i\neq j$}\\ 0&\text{ if $i=j$}\end{cases}.$ $$

In this case, $\psi_{K}(A)_{N,:}$ has the following form

$$ $\displaystyle\psi_{K}(A)_{N,:}=\underbrace{\bigg{[}\frac{\sum_{i=1}^{N-1}K(x_{ N},x_{i})x_{i}^{\top}}{\sum_{i=1}^{N-1}K(x_{N},x_{i})}}_{\text{smoothed $d$- dimensional features}},\underbrace{\frac{\sum_{i=1}^{N-1}K(x_{N},x_{i})y_{i}}{ \sum_{i=1}^{N-1}K(x_{N},x_{i})}\bigg{]}}_{\text{smoothed estimate}}\in\mathbb{ R}^{d+1},$ (11) $$

and the last element $\psi_{K}(A)_{N,d+1}$ is the kernel smoother estimate,

$$ $\psi_{K}(A)_{N,d+1}=\frac{\sum_{i=1}^{N-1}K(x_{N},x_{i})y_{i}}{\sum_{i=1}^{N-1 }K(x_{N},x_{i})}.$ (12) $$

Below, we provide key examples of kernel smoothers that can be implemented by equation 12 upon changing the kernel $K$.

- 1.
When $K$ is the exponential kernel, i.e., $K(z,z^{\prime})=e^{-z^{\top}z}$, then $\psi_{K}$ implements softmax attention, and $\psi_{K}(A)_{N,d+1}$ is the kernel smoother corresponding to the exponential kernel.
- 2.
When using the kernel $H$ defined in equation 5, then, $\psi_{H}(A)_{N,d+1}$ implements the Hilbert estimate, which is consistent as the number of in-context examples goes to infinity [].

In our experiments in Figure 5, we trained an MLP on features computed using $\psi_{K}(A)_{N,:}$. Yet, the results above suggest that the scalar $\psi_{K}(A)_{N,d+1}$ alone should be sufficient for context scaling. Indeed, in the case of the Hilbert estimate, this entry alone provides a consistent estimate as the context length goes to infinity. To this end, in Figure 6, we compare the performance of two MLPs when the number of tasks is fixed and the context length increases. The first MLP is trained using $\psi_{K}(A)_{N,:}\in\mathbb{R}^{d+1}$, and the second is trained on only $\psi_{K}(A)_{N,d+1}$. The results in this figure confirm that using $\psi_{K}(A)_{N,d+1}$ is as good as using $\psi_{K}(A)_{N,:}$ for context-scaling.

### 5.2 Training MLPs that simultaneously context-scale and task-scale

Figure: Figure 7: Comparison of MLPs trained using (1) vectorized inputs; (2) features from $\psi_{K}$ for $K\in L,H$ defined in equation 10 and equation 5; (3) both features from $\psi_{K}$ and vectorized inputs. We compare performance across two ICL tasks: linear regression and two-layer teacher-student neural networks. (A) Task-scaling ability of MLPs using different inputs. (B) Context-scaling ability of MLPs using different inputs. MLPs trained on both vectorized inputs and features from $\psi_{K}$ are able to simultaneously context-scale and task-scale. Experimental details are provided in Appendix A.
Refer to caption: x7.png

As the Hilbert estimate provides a consistent estimate, our results above show that transformers provably generalize to unseen tasks, when the context length approaches infinity. Nevertheless, the issue with using the Hilbert estimate alone is that the Hilbert estimate is only computed using examples provided in a context. As such, it cannot task-scale unlike MLPs trained on vectorized inputs. We now show that training MLPs on vectorized inputs concatenated with features estimated using $\psi_{K}(A)$ result in MLPs that can both context-scale and task-scale.

Namely, we revisit the experiment presented in Figure 1 and extend our analysis by training MLPs on three distinct input configurations: (1) vectorized input data; (2) features from $\psi_{K}(A)_{N,:}$; and (3) the concatenation of vectorized input data and features from $\psi_{K}(A)_{N,:}$. In our experiments, we consider the feature maps $\psi_{L}$ and $\psi_{H}$ discussed in the previous section. The results of training these MLPs is presented in Figure 7 and we summarize the results below.

- 1.
MLPs with vectorized input data: Figure 7A demonstrates that these MLPs exhibit task-scaling. Yet, Figure 7B reveals that these MLPs fail to context-scale and performance can even deteriorate with increased context length.
- 2.
MLPs with features from $\psi_{K}(\cdot)_{N,:}$: Figure 7A illustrates that these MLPs do not task-scale, as performance does not improve with an increasing number of tasks. This behavior matches intuition as the Hilbert smoother and 1-step gradient descent features are task-specific and do not leverage inter-task relationships. Yet, Figure 7B shows that these MLPs successfully context-scale, which happens provably for the particular case of $\psi_{H}$.
- 3.
MLPs with both vectorized inputs and features from $\psi_{K}(\cdot)_{N,:}$: Figure 7A demonstrates that these MLPs are capable of task-scaling, consistent with the performance of MLPs on vectorized data alone. Moreover, in Figure 7B, we now observe that these MLPs are now capable of context-scaling, consistent with the performance of MLPs using the features from $\psi_{K}(\cdot)_{N,:}$ alone.

These results underscore the importance of the feature map $\psi_{K}$ for context-scaling and highlight the effectiveness of using both vectorized inputs and features from $\psi_{K}$ in improving the ability of models to learn in-context.

## 6 Conclusion

#### Summary.

In this work, we observed that transformers, unlike MLPs, are able to simultaneously context-scale (improve performance as the context length increases) and task-scale (improve performance as the number of pre-training tasks increases). To better understand this property of transformers, we first identified a simplified transformer (SGPT) that could solve ICL tasks competitively with GPT-2 architecture transformers despite having no trainable key, query, and value weights in attention layers. By studying a one-layer version of SGPT, we identified that the attention operator of SGPT applied a feature map, $\psi$, on input data that enabled context-scaling. In particular, we showed that this feature map could implement kernel smoothers such as the Hilbert estimate, which is a statistically consistent estimator as the context length goes to infinity. As such, our work provably demonstrates that transformers can context-scale, generalizing to new, unseen tasks when provided a large context. We demonstrated the effectiveness of the feature map, $\psi$, for context-scaling by showing that MLPs trained on both features from $\psi$ and vectorized inputs could simultaneously context-scale and task-scale.

#### Future work and limitations.

While we have provably established that one-layer transformers can context-scale, we empirically observe that one-layer transformers are not as sample-efficient as deep transformers for both context-scaling and task-scaling. Thus, an important future direction is understanding how depth improves the sample complexity of transformers in both context-scaling and task-scaling settings. Exploring this aspect remains a promising avenue for future research and could provide a comprehensive understanding of ICL, and more broadly, a better understanding of how transformers are able to generalize to new tasks when provided large contexts.

## Appendix A Experiments details

We provide all experimental details below.

#### Problem formulation.

For all ICL tasks studied in our work, we consider T pretraining tasks, each with input data of the form:

$$ $A_{t}=\begin{bmatrix}x_{1}^{t}&x_{2}^{t}&\cdots&x_{N}^{t}\\ y_{1}^{t}&y_{2}^{t}&\cdots&y_{N}^{t}\end{bmatrix}^{\top}\in\mathbb{R}^{(N) \times(d+1)}$ $$

where $t\in{1,\ldots,T}$ indexes the tasks, $N$ denotes the maximum context length, and $d$ denotes the input data dimension.
We define the loss function $L(\theta)$ as:

$$ $L(\theta;X_{1},\ldots,X_{T}):=\frac{1}{T}\sum_{t=1}^{T}\left[\frac{1}{N}\sum_{ i=1}^{N}(M_{\theta}(X_{t}^{i})-y_{i+1}^{t})^{2}\right],$ $$

where

$$ $A_{t}^{i}=\begin{bmatrix}x_{1}^{t}&x_{2}^{t}&\cdots&x_{i}^{t}&x_{i+1}^{t}\\ y_{1}^{t}&y_{2}^{t}&\cdots&y_{i}^{t}&0\end{bmatrix}^{\top}\in\mathbb{R}^{(i+1) \times(d+1)},$ $$

and $M_{\theta}$ denotes the model with trainable parameters $\theta$.

#### Vectorized input.

By vectorizing input, we mean flattening the input into a vector. After vectorization, $A_{t}^{i}$ defined above becomes,

$$ $A_{t}^{i}=\begin{bmatrix}{x_{1}^{t}}^{\top}&{x_{2}^{t}}^{\top}&\cdots&{x_{i}^{ t}}^{\top}&{x_{i+1}^{t}}^{\top}&y_{1}^{t}&y_{2}^{t}&\cdots&y_{i}^{t}&0\end{ bmatrix}^{\top}\in\mathbb{R}^{(i+1)(d+1)}.$ $$

#### GPT-2.

We used the GPT-2 implementation from prior work [, ], which is based on the Hugging Face implementation []. Following the approach in these prior works, we modified the embedding layer with a learnable linear layer that maps from the ambient dimension to an embedding dimension.

#### SGPT.

To construct SGPT, we make the following modifications to the GPT-2 architecture: (1) we fix all key, query, value weights to the identity; (2) we eliminate all batch-normalization layers; and (3) we remove the second linear layer from each MLP. Following prior work [, ], we modified the embedding layer with a linear layer that maps from the ambient dimension to an embedding dimension. In SGPT, this linear layer is not trainable and serves as a fixed random map. We outline the architecture below.

Let $A\in\mathbb{R}^{N\times(d+1)}$ be the input of the model. We initialize a random matrix $W_{0}\in\mathbb{R}^{(d+1)\times k}$, where $k$ is the embedding dimension. Defining the input of the $i$-th layer as $H^{(i)}$, we have $H^{(0)}:=AW_{0}$,

$$ $\displaystyle H^{(i)}=\sigma\Big{(}\big{(}g(H^{(i-1)})W^{(i)}_{proj}+H^{(i-1)} \big{)}W^{(i)}_{MLP}\Big{)}+g(H^{(i-1)})W^{(i)}_{proj}+H^{(i-1)}$ (13) $$

where:

- •
$\sigma$ is the activation function, chosen to be GeLU,
- •
$g$ is as defined in Equation 6, $\phi$ is row wise $l1$ normalziation.
- •
$W^{(i)}_{proj}\in\mathbb{R}^{k\times k}$ is the projection matrix,
- •
$W^{(i)}_{MLP}\in\mathbb{R}^{k\times k}$ is the MLP weight matrix for the $i$-th layer.

The last layer of the network is a linear layer with weights $W_{O}\in\mathbb{R}^{k\times 1}$.

#### MLP architectures.

In all experiments, we use a standard 2-layer ReLU MLP with a width of 1024 units. Given an input vector $x\in\mathbb{R}^{d_{in}}$, the MLP implements a function $f$ of the form

$$ $\displaystyle f(x):=\sigma(xW_{0})W_{1}$ (14) $$

where $W_{0}\in\mathbb{R}^{d_{in}\times 1024}$, $W_{1}\in\mathbb{R}^{1024\times 1}$, and $\sigma$ is the ReLU activation function.

#### Zero-padding input for MLP.

In all experiments, we always trained a single MLP for all context lengths by zero-padding to the largest context length. For example, if the input data is in $\mathbb{R}^{d}$ and the largest context length is $N_{max}$, then the input dimension of the MLP is $dN_{max}$.

#### Expeirmental details for Figure 1 .

In this experiment, we trained an 8-layer GPT-2 model with 8 attention heads and a width of 256. The MLP configuration is the same as that in equation 14. We trained and tested the models on the same set of context lengths: 5, 10, 20, 30, and 40.

#### Experimental details for Section 4 .

We trained both standard GPT-2 architecture transformers and our proposed SGPT with the following configurations:

- •
Widths: {256, 512, 1024},
- •
Number of layers: {2, 4, 6, 8},
- •
Number of attention heads for GPT-2: {2, 4, 8}.

The best performance was achieved with an 8-layer model with a width of 256. For the original GPT-2, the optimal configuration used 8 attention heads.

We outline per-task observations and configurations below:

- 1.
Linear regression with single noise level: Following the prior work [], the input dimension is $d=20$ and noise level is $\sigma=0.5$. We trained on context lengths from 10 to 40 with a step size of 5.
- 2.
Linear regression with two noise levels: Following the prior work of [], the input dimension is $d=20$ and noise levels are $\sigma_{1}=0.1,\sigma_{2}=0.5$. We trained on context lengths from 1 to 40 with a step size of 1.
- 3.
Decision tree: Following the prior work of [], the input dimension is $d=20$ with a tree depth of 4. We note that with the input structure equation 2, GPT-2 performs poorly, so in our figure we used the pretrained model from [].
- 4.
Two-layer ReLU Neural Networks. As mentioned before, we chose the width of this family of neural networks to be $r=100$ and the input dimension to be $d=20$. We trained on context lengths from 1 to 100 with a step size of 10. We observed that unlike our model, GPT-2 does not generalize well for context lengths that it has not been trained on.
- 5.
Sparse Linear Regression As mentioned previously, the ambient dimension of the input is $d=20$, consistent with prior work[], and the effective dimension is $s=3$. We used scikit-learn [] for the Lasso and Ordinary least sauare performances.

#### Experimental details for Figure 5 .

In all tasks, we used input dimension $d=8$ following the setting in []. We trained and tested both models on context lengths of 10, 20, 30, 40, and 50.

#### Experimental details for Figure 6 .

We train an MLP (equation 14) on features extracted using $\psi_{K}(\cdot)_{N,:}$ and linear regression on the scalar $\psi_{K}(\cdot)_{N,d+1}$.

#### Experimental details for Figure 7 :

- •
Linear regression: We used the same settings as used for the model in equation 3 with $d=8$ and $\sigma=0.22$. We trained and tested models on the context lengths 5, 10, 20, 30, and 40.
- •
2-layer NN task: We used the same settings as used for the model in 4 with $d=8,r=100$. Trained and tested models on the context lengths 10, 20, 30, 40, 50, 60, 70, and 80.

#### Hardware.

We used machines equipped with NVIDIA A100 and A40 GPUs, featuring V-RAM capacities of 40GB. These machines also included 8 cores of Intel(R) Xeon(R) Gold 6248 CPU @ 2.50GHz with up to 150 GB of RAM. For all our experiments, we never used more than one GPU, and no model was trained for more than two days.

## Appendix B Feature map derivation

Let $A=\begin{bmatrix}x_{1}^{\top}&y_{1}\\
x_{2}^{\top}&y_{2}\\
\vdots&\vdots\\
x_{N-1}^{\top}&y_{N-1}\\
x_{N}^{\top}&0\\
\end{bmatrix}\in\mathbb{R}^{N\times(d+1)}$.

#### 1-step of GD.

In this case, we have

$$ $\displaystyle\psi_{L}(A)$ $\displaystyle=(AA^{\top})A$ $\displaystyle=\begin{bmatrix}x_{1}^{\top}&y_{1}\\ x_{2}^{\top}&y_{2}\\ \vdots&\vdots\\ x_{N-1}^{\top}&y_{N-1}\\ x_{N}^{\top}&0\end{bmatrix}\begin{bmatrix}x_{1}&x_{2}&\cdots&x_{N-1}&x_{N}\\ y_{1}&y_{2}&\cdots&y_{N-1}&0\end{bmatrix}A$ $\displaystyle=\begin{bmatrix}x_{1}^{\top}x_{1}+y_{1}y_{1}&\cdots&x_{1}^{\top}x _{N-1}+y_{1}y_{N-1}&x_{1}^{\top}x_{N}\\ x_{2}^{\top}x_{1}+y_{2}y_{1}&\cdots&x_{2}^{\top}x_{N-1}+y_{2}y_{N-1}&x_{2}^{ \top}x_{N}\\ \vdots&\cdots&\vdots&\vdots\\ x_{N-1}^{\top}x_{1}+y_{N-1}y_{1}&\cdots&x_{N-1}^{\top}x_{N-1}+y_{N-1}y_{N-1}&x _{N-1}^{\top}x_{N}\\ x_{N}^{\top}x_{1}&\cdots&x_{N}^{\top}x_{N}&x_{N}^{\top}x_{N}\end{bmatrix} \begin{bmatrix}x_{1}&y_{1}\\ x_{2}&y_{2}\\ \vdots&\vdots\\ x_{N-1}&y_{N-1}\\ x_{N}&0\\ \end{bmatrix}.$ $$

Thus, $\psi_{L}(A)_{N,d+1}=x_{N}^{\top}\mathbf{X}^{\top}\mathbf{y}$, where $\mathbf{X}:=\begin{bmatrix}x_{1}^{\top}\\
\vdots\\
x_{N}^{\top}\end{bmatrix}\in\mathbb{R}^{N\times(d+1)}$ and $\mathbf{y}:=\begin{bmatrix}y_{1}\\
\vdots\\
y_{N}\end{bmatrix}$. This value is equivalent to the prediction given by using one-step of gradient descent to solve linear regression.

#### Kernel smoothers.

In this case, we have

$$ $\displaystyle\psi_{K}(A)$ $\displaystyle=\widehat{K}(\mathbf{X},\mathbf{X})A$ $\displaystyle=\begin{bmatrix}0&\frac{K(x_{1},x_{2})}{\sum_{\begin{subarray}{c} i=1\ i\neq 1\end{subarray}}^{N}K(x_{1},x_{i})}&\cdots&\frac{K(x_{1},x_{N})}{ \sum_{\begin{subarray}{c}i=1\ i\neq 1\end{subarray}}^{N}K(x_{1},x_{i})}\\ \frac{K(x_{2},x_{1})}{\sum_{\begin{subarray}{c}i=1\ i\neq 2\end{subarray}}^{N} K(x_{2},x_{i})}&0&\cdots&\frac{K(x_{2},x_{N})}{\sum_{\begin{subarray}{c}i=1\ i \neq 2\end{subarray}}^{N}K(x_{2},x_{i})}\\ \vdots&\vdots&\vdots&\vdots\\ \frac{K(x_{N},x_{1})}{\sum_{\begin{subarray}{c}i=1\ i\neq N\end{subarray}}^{N} K(x_{N},x_{i})}&\frac{K(x_{N},x_{2})}{\sum_{\begin{subarray}{c}i=1\ i\neq N \end{subarray}}^{N}K(x_{N},x_{i})}&\cdots&0\\ \end{bmatrix}\begin{bmatrix}x_{1}&y_{1}\\ x_{2}&y_{2}\\ \vdots&\vdots\\ x_{N-1}&y_{N-1}\\ x_{N}&0\\ \end{bmatrix}.$ $$

Now last row of the $\psi_{K}(A)$ is given by

$$ $\displaystyle\psi_{K}(A)_{N,:}=\begin{bmatrix}\frac{\sum_{i=1}^{N-1}K(X_{N},x_ {i})x_{i}}{\sum_{i=1}^{N-1}K(X_{N},x_{i})}&\frac{\sum_{i=1}^{N-1}K(X_{N},x_{i} )y_{i}}{\sum_{i=1}^{N-1}K(X_{N},x_{i})}\end{bmatrix}\in\mathbb{R}^{d+1}.$ $$

Thus, $\psi_{K}(A)_{N,d+1}$ is equvialent to the prediction for $x_{N}$ given by using a kernel smoother with kernel $K$.

## Appendix C Additional experiments

In the following experiments, we trained both an 8-layer, 8-head GPT-2 model and an 8-layer SGPT model, using identical context lengths during training and testing. As shown in Figure 8, both models are capable of scale and context scaling. Here for both tasks we set $d=8$ and for the linear regression task $\sigma=0.22$ and for 2-layer neural network we set the width of the networks to $r=100$.

Figure: Figure 8: Task-scaling and context-scaling of GPT-2 architecture transformers versus SGPT for linear regression and 2-layer neural networks tasks. (A) Task-scaling abilities of these models with $10$ in-context examples. (B) Context-scaling abilities of these models with $10^{5}$ (left) and $10^{6}$ (right) pre-training tasks.
Refer to caption: x8.png

## Peer Reviews

confidence: 3
contribution: 2
initial_rating: 3
presentation: 2
questions: - Why do previous theoretical results only apply to fixed context length, as stated in l.157-159? (See weaknesses.)
- The SGPT as defined in Eq. 13 appears to be more specific than the generic feature map $\psi$ you are then introducing in Equation (9). Is that true and if so, can you explain how this generalized version connects to the SGPT as well as the generic Transformer architecture?
- Tong & Pehlevan (2024) also show that MLPs cannot context scale (Fig. 1d). This is currently not reflected in your related work section (l. 143-146). Could you please clarify and explain how your findings relate to these prior findings?

Minor comments:
L. 51: “an non-exhaustive” -> “a non-exhaustive”
L. 157: typo?
L. 213: “identify” -> “identity”
soundness: 2
strengths: - I really liked the introduction and thought that the question of context-scaling was nicely set up.
- I appreciated the authors considering a wide range of different tasks that they collected from the relevant literature.
- Building part of the SGPT into the MLP was a neat method for illustrating the origin of the distinct mechanism between either.
summary: The authors study an important question: how does in-context learning in Transformers depend on the number of in-context examples as well as the number of overall tasks? They draw a connection to kernel smoothing and demonstrate that a simplified version of the Transformer architecture can implement this algorithm equally well. They show that in contrast to Transformers, MLPs do not exhibit scaling with in-context examples, but that a feature map inspired the Transformer's mechanism can yield successful improvement with more in-context examples.
weaknesses: Unfortunately, I do not think this work is ready for publication in its current state. Primarily, I believe that the paper does not provide sufficiently novel insight from the prior literature.

Notably, improvement of Transformer performance with the number of in-context examples was already noted (as the authors lay out in the related work section), e.g. in Bai et al. (2023) and the prior theoretical literature also explains why this would be the case, as it draws the connection between ICL in Transformers and gradient descent and kernel smoothing --- both of which improve with the number of samples. I do not understand why previous theoretical results (e.g. Proposition 1 in von Oswald et al. (2023)) would only apply to fixed context length.

It is also unclear to me how SGPT provides a novel insight into the mechanism of ICL compared to, say, the construction in von Oswald et al. (2023), who also explicitly draw the connection to kernel smoothing and use a similar set of simple key, query, and value matrices (especially when $W_0=0$ in their construction in Appendix A.1). The authors argue that the simplicity of SGPT is a substantial strength of this paper, as it demonstrates that there are many problems such a simple architecture can solve. But it is unclear to me whether this a theoretical argument (which seems to make the connection to kernel smoothing, in which case I'm unsure how this is different from the insight by von Oswald et al.) or an empirical argument (in which case I think the authors would have to demonstrates concretely that SGPT outperforms kernel smoothing algorithms).

As I noted, I think the contrast to MLPs and providing the modified features to the MLPs was interesting. I think it would be important, however, to provide insight into *how* the vectorized component enables them to scale with the number of examples.

Taken together, I think the paper in its current form is not sufficiently distinct from existing work --- or at least does not explain sufficiently clearly how it is different. As I noted above, I do think that the authors focus on a really interesting question (context scaling) that provides a different angle from prior work. However, I think for the paper to be ready for publication, I think this investigation would have to further explore how this angle can change our theoretical understanding of context scaling.

---

confidence: 3
contribution: 3
initial_rating: 6
presentation: 3
questions: What is the task-scaling performance for single-layer SGPT (a task-scaling counterpart of Fig 5)?

Could you explain more on the right panel of Fig 3? 

Could you explain more on the Fig 4(B), 2-layer NN, SGD?
soundness: 3
strengths: The SGPT considered and its experiments are novel. And it is quite surprising and interesting to see that its performance is comparable to GPT2. I suppose it is due to the simplicity of the ICL tasks conducted in the paper. 

The idea of connecting ICL and kernel smoothly is clearly presented and is of insight. 

The separation of context-scaling and task-scaling via feature from kernel estimate and vectorized input is novel and can potentially help us understand their impacts better individually.
summary: This paper studies context-scaling and task-scaling of ICL, under multiple ICL regression tasks, including linear regression with fixed and multi noise level, two-layer ReLU neural networks, decision trees, and sparse linear regression. Experiments are conducted on GPT2 and Simplified GPT (SGPT) by taking key, query and value matrices to be identify matrices and removing batch normalization. Under those tasks, GPT2 and SGPT both demonstrate context-scaling ability and the performances of GPT2 and SGPT are very close to each other. SGPT allows a kernel smoothing perspective interpretation of transformer's ICL capability. Specifically, the authors demonstrate the capability of ICL capability of transformer by showing a single layer SGPT can perform kernel smoothing (including the consistent Hilbert estimate as special case) with appropriate feature map corresponding to the attention. To see what statistics are the essence of task-scaling and context-scaling, experiments on MLP with different inputs , such as vectorized input data with or without kernelized features are conducted. Specifically, task-scaling attributes to the vectorized data and context scaling attributes to the kernelized features, and combining both inputs provide both task-scaling and context scaling.
weaknesses: Though with the consistency of Hilbert estimate, how exactly the transformer performs Hilbert estimate i.e., via the construction of activation function in attention, is not straightforward. 

What is the major intuition of taking key, query, and value matrices to be identity? Such intuition is vital since it is shown that the context-scaling capability is attributed to the attention, and task-scaling is to the MLP with vectorized data. I wonder if key, query, and value matrices are learnable, will they also provide sufficient task-scaling ability? 

What is the theoretical or explanatory justification for the capability of the task-scaling ability of transformer and MLP, generalization?

---

confidence: 3
contribution: 3
initial_rating: 6
presentation: 2
questions: See weaknesses above.
soundness: 2
strengths: I found the paper overall to be very well-written and a pleasure to read. The topic is extremely important, particularly in our post-ChatGPT era, and deals with a critical ability in Transformers. The contrast to MLPs sparks a fascinating discussion about the relative merits of different architectures.
summary: This study examines task versus context scaling in various in-context regression tasks, finding that Transformers (but not MLPs) can scale with increasing context sizes, as well as tasks. The authors suggest that a Transformer's ability to context-scale stems from its ability to implement kernel smoothers in the attention matrix. Equipping an MLP with features derived from these kernel smoothers enable it to also context-scale.
weaknesses: I would love to see this manuscript published at ICLR, but there are a few oversights that prevent me from assigning a higher score. If these are able to be addressed, I will be delighted to raise my score.

The discussion on context-scaling in MLPs appears to be drawing from prior work by Tong and Pehlevan (https://arxiv.org/abs/2405.15618). The authors claim that MLPs do not context-scale, but Tong and Pehlevan seem to be showing otherwise. I may be misunderstanding both sides here, but Fig 1d and 1i of Tong and Pehlevan appear to demonstrate that at least MLP-Mixers continue to do well for arbitrary contexts. While MLP performance decays as context length increases in ICL regression, it doesn't appear to be the case for ICL classification. Were you able to look at classification tasks as well? Further, it looks like Tong and Pehlevan made the choice of plotting *excess* MSE above Bayes optimal rather than raw MSE. Because Bayes optimal MSE falls as context length increases, zero excess MSE would imply context-scaling. Looking at Figure 6b, for fewer dimensions than the $d= 8$ you tested, it does appear that MLPs adhere to the Bayes optimal level before failing at longer contexts. 

Overall, it would appear that context-scaling in MLPs does happen, but is bottlenecked by some aspect of insufficient data and long inputs, rather than some inability to implement kernel smoothers. MLP-Mixers, which do not have any product interactions that could implement a kernel smoother in an obvious way, continue to do well also. Indeed, it looks like in your Figure 7B, you do demonstrate some context-scaling in unmodified MLPs for 5M-pretraining tasks (top row), quite similar to using $\psi_L$ features.

Additionally, I thought that kernel smoothers are weak in high dimensions, and require a dataset size that is exponential in the input dimension in order to interpolate well -- a classical curse of dimensionality. However, modern Transformers routinely handle token embeddings with dimensions that number in the tens of thousands, which would presumably defeat a kernel smoother even if it were exposed to an Internet-scale corpus -- and in-context, no less! It's quite possible I misunderstand this aspect of your analysis, but it seems implausible that a kernel smoother interpretation of attention is applicable to real-world Transformers?

Finally, you mention there is a theoretical connection between having identity KQV matrices and the Hilbert estimate kernel, but I couldn't find the derivation in your manuscript. This could very well be a blatant oversight on my part, but could you point me to its location?

---

confidence: 4
contribution: 1
initial_rating: 3
presentation: 2
questions: In the weakness section, I provide both the concerns and the questions for the authors.
soundness: 2
strengths: 1. I like the lens that the authors use in the paper, i.e., separating in-context learning into task scaling and context scaling. Even though this lens is not completely new, its a useful one. 

2. The authors propose a simplification of GPT which is referred to as SGPT. This simplification is a useful one from two perspectives -- it gives way to a model that can be understood theoretically and it possibly indicates that we could simplify the architectures for in-context learning. 

3. Overall, the paper is easy to read and reasonably well presented.
summary: In recent years, there has been a growing interest in the community to understand and build better in-context learning capabilities. In this work, the authors study in-context learning through the lens of two separate abilities: context scaling and task scaling. Context scaling refers to the capability to improve predictions as the number of in-context examples grow while keeping the number of tasks fixed. Task scaling refers to the capability to improve predictions as the number of task grows while keeping the number of in-context examples fixed. The authors show that a lot of in-context learning capabilities arise from the behavior of attention as a kernel smoother. To show this, the authors consider a simplified version of the GPT architecture, which they call SGPT, where the attention blocks are not trained and key, query, and value matrices are set to identity. The authors compare GPT-2 like architecture with SGPT and show that on synthetic in-context   learning tasks from the literature, the performance of the two models match. Further, the authors study context scaling and task scaling in MLPs. First, they show that MLPs with vectorized inputs are capable of task scaling and not capable of context scaling. They show that MLPs with featurized inputs are capable of context scaling but not of task scaling. Finally, they combine the two types of inputs to show MLPs exhibit both context scaling and task scaling.
weaknesses: I have several concerns with the paper that I highlight below. 

1. **On feature maps for context-scaling:**  
a) If you construct a Kernel based mapping of the form in equation 11, then of course due to sheer consistency argument one can say that the map in equation 12 converges to the true function. This would require infinite examples in-context. The whole point of in-context learning is to learn in-context as quickly as possible. From an asymptotics point of view, all the consistent estimators in equation 4 can learn the function. Thus the experiments you conduct in Figure 5 are not really interesting. If you took the model in equation 4 itself, then so long as the Kernel you use gives a consistent estimate, one could argue that as sequence length n increases in equation 4 the performance will improve and this model will context scale. Would you elaborate, what is really the important and new insight one should take from Figure 5?

    b) In the first part of Section 5.1, the authors construct a mapping which is based on linear attention and show the equivalence to one step of gradient descent. This result is already known from previous works such as Von Oswald's work.  Can the authors explain what is truly new in Section 5.1? 

 
2. **On the results with simplified transformer:** As i stated in the strengths section, I found the result with simplified transformer intriguing. However, there are few important ablations that would have helped understand this section better. First of all, the authors change the attention to a linear attention l1 normalization operation. What would have happened if the authors only used linear attention and no normalization? While the model is close to standard GPT in terms of performance, how much of this is due to learning that happens due to the MLPs at different depths? Put differently, can the authors show the impact of depth on the performance, i.e., as depth increases model learns more complex hierarchy of features that allow it to match GPT? If that is the case, then this should be clearly stated in the paper too, that depth was crucial to match the performance of GPT at short context lengths. In some sense, if depth is crucial to match the performance, then the role of kernel smoothing alone is not a crucial one. 

3. **On Section 5.2**: In this section, the authors study various variants of MLPs to study context scaling and task scaling capabilities. I find somethings unclear here. 


   a) Firstly, I find the fact that vectorized MLPs with sufficient capacity not able to context scale not clear. The training data takes the form P, x, where P is the prompt that contains (x,y) pairs from the task of interest, x is the current query. In the case, where the model class is unrestricted, then the learner should ideally learn E[y| P,x], where y is the true label and expectation is over the distribution of prompts and query, label pairs. If the task of interest is linear regression (studied in Garg et al.), the coefficients of the regression are drawn from an isotropic Gaussian, and the features are drawn from an isotropic Gaussian as well, then   
   E[y| P,x] = ((X.t X + sigma I)^{-1} X.t Y).t x, where sigma is noise. This is solution to standard ridge regression. Provided the model class be it transformer or be it MLP is sufficently expressive and contains ((X.t X + sigma I)^{-1} X.t Y).t x, then in principle MLP should also be capable of context scaling. So my question to the authors is "If MLP has sufficient capacity, what stops it from implementing a ridge regression solution at different context lengths up to the maximum context length determined by size of vectorized input?" Basically, if the MLP is not able to learn, then it has to be an argument that is not explained by expressivity but by learnability. Perhaps the optimization cannot find the global minimum in the above case easily? 

     b) Secondly, the authors show that MLPs with features from kernel smoothes can context scale. I don't quite get what's the surprise here. Isn't the input feature to these MLPs itself guaranteed to be consistent in infinite context limit? 


4. **On the role of context-length:** The authors state that existing works fail to explain the context scaling capabilities of transformers.  In the example, I gave in 3a), the Bayes optimal predictor is implemented by the transformer. This Bayes optimal predictor of course improves with context length. In this sense, I don't quite get why do the authors say that existing works fail to explain context scaling capabilities. 

5. **On MLPs and context scaling**: The authors also mention that it is unclear whether MLPs scale with context and they are the first to dive into this. Perhaps the authors have missed https://arxiv.org/pdf/2311.18194.

Overall, I don't feel I learned something insightful from the paper, the only thing I liked was the experiment with SGPT.

Meta-review: This paper explores task and context scaling in in-context learning. The authors propose simplified models like SGPT and make theoretical connections to kernel smoothers. While the ICL topic is very timely and relevant, the reviewers have concerns regarding the novelty and depth of insights. The claims regarding SGPT’s advancements over prior works are not convincingly justified, and some of the findings seem incremental. Reviewers found certain theoretical contributions to be unclear and there was also a concern on insufficient experimentation to substantiate claims. For example, SGPT matching GPT-2's performance lacks thorough ablations. Furthermore, the analysis/discussion of MLPs raises valid questions that remain not fully unresolved. Given the limited novelty and incomplete discussions, I recommend rejecting this submission in its current form.

Decision: Reject