Paper: How Do Transformers Learn In-Context Beyond Simple Functions? A Case Study on Learning with Representations (ICLR-2024)

Reviewer comment (Novelty):
Why are the constructive proofs important for understanding in-context learning in transformers? I am aware that there are prior works that design transformers that are capable of in-context learning. However, I am not convinced of the importance and significance of these results. Couldn't we also find weights for other architectures (like large MLPs or LSTMs) and argue that they are capable of in-context learning. Is the existence of these model weights informative of what is learnt in practice?

Author response:
In our opinion, the importance for the constructive proofs is **not to give a yes/no answer** on “whether transformers can do this” or whether alternative architectures can do the same. Rather, the constructive proofs can
* Demonstrate the **efficiency** of transformers for implementing these in-context learning procedures, i.e. only requiring a number of transformer layers and heads. This heavily uses the specific structure of the attention and MLP layers, instead of just invoking some universal approximation result;
* Suggest concrete mechanisms (computing representation, copying, gradient descent) that can in turn guide the probing experiments.
Alternative architectures such as LSTM or MLP (on concatenated input) may be able to implement a similar learning procedure, but **much less efficiently**, as it may be difficult to use them to efficiently replicate our mechanisms such as (batch) gradient descent or copying. They may be able to implement other mechanisms efficiently (such as online gradient descent for LSTMs), though.

Reviewer comment (Experiments):
Choice of non-linear functions. I think the authors could be more rigorous in evaluating the non-linear representations used in their setup. In particular, the non-linear representations are L-layer MLPs with the matrices being random orthogonal matrices. Do the results hold for other families of functions and does it fail to work for some other classes of functions? I think it would be helpful to clarify that the results are specific to this setup.

Author response:
We mentioned “We instantiate the representation as shallow neural networks (MLPs)” in the paragraph before the list of contributions (Page 2). To clarify, we further highlighted “L-layer MLPs” in the abstract and the list of contributions in our revision. The choice of weight matrices and non-linear functions are described around Eq(4) in Section 4.1.

Reviewer comment (Experiments):
Is the synthetic setup an accurate toy-model to understand language models? Like previous work, all the results are on synthetic data. It remains unclear if the toy setup is representative of in-context learning in language models. What kind real world tasks are captured by a composition of a fixed non-linear function and a linear function that is learnt from context?

Author response:
While there is always a gap between such synthetic tasks and real tasks, we believe our fixed representation +  linear function is at least a more realistic model than prior works in this line, which primarily focuses on linear functions only. The fixed representation can model any “prior knowledge” that can be encoded in a representation (e.g. feature extractor), whereas the changing linear function can encode the specific classification task.
As an illustrating toy example, consider the in-context learning problem

"Give me the next word or phrase: ‘apple->means anger; ocean->means sadness; grass->means happiness; banana->means’”.

The representation function for this problem would be to map each object to its color (apple -> red), and the linear function would be to map the color to an emotion (red -> anger). The feature map is fixed (relies on prior knowledge about the world) whereas the linear function needs to be learned in context (another such instance could involve a different map). Of course this is a toy example, but we believe its underlying structure (prior knowledge, then learn in context) could be broaderly present in real-world problems.

Reviewer comment (Reproducibility):
Could the authors add more details on how the non-linear functions are created? Can the authors also clarify in the introduction/abstract that the functions are L-layer MLPs?

Author response:
We mentioned “We instantiate the representation as shallow neural networks (MLPs)” in the paragraph before the list of contributions (Page 2). To clarify, we further highlighted “L-layer MLPs” in the abstract and the list of contributions in our revision. The choice of weight matrices and non-linear functions are described around Eq(4) in Section 4.1.

Reviewer comment (Experiments):
Ruiqi et al. (https://arxiv.org/abs/2306.09927) show that in-context learning fails if the linear functions are selected to be out-of-distribution. Is this also the case here?

Author response:
Ruiqi et al. consider two out-of-distribution tasks: OOD linear weight $w$, and OOD covariates $x$. In these OOD scenarios, we expect the same kind of failure here in our task: The theoretically constructed transformer can still implement the correct algorithm in these OOD scenarios, but the trained transformers (trained in in-distribution instances) may not perform well due to the failure of the empirical risk minimization on OOD.

Reviewer comment (Theory):
Section 3.1 states that the representation function can be chosen arbitrarily, but Lemma B.3 requires a specific structure and non-linearity to work.

Author response:
> Section 3.1 states that the representation function can be chosen arbitrarily, but Lemma B.3 requires a specific structure and non-linearity to work.

Did the reviewer mean Section 4.1? Right after “The representation function $\Phi^\star$ can in principle be chosen arbitrarily” in Section 4.1, we stated “As a canonical and flexible choice for both our theory and experiments, we choose $\Phi^\star$ to be a standard L-layer MLP”, followed by Eq(4) which describes the MLP architecture in details (including the choice of non-linearity). Lemma B.3 follows the specifications made here.

Reviewer comment (Experiments):
Although the paper does a good job of illustrating the claimed mechanism, it does not analyze settings where the mechanism breaks.

Author response:
> Although the paper does a good job of illustrating the claimed mechanism, it does not analyze settings where the mechanism breaks.

> What happens when the representation function is of a different form? If either the transformer does not have enough layers or the width, is there an approximate representation function learned on which regression is performed, or does the entire mechanism fall apart?

In theory, if the transformer does not have enough layers or width, then it can still implement an approximate algorithm (such as computing an approximate version of the representation function $\Phi^\star$, as the reviewer mentioned, combined with fewer gradient descent steps to implement approximate ridge regression).

Empirically, we conducted some preliminary experiments with the same transformer architecture but more sophisticated representation functions, such as MLPs with higher depth or width. We’ve included these experiments in Appendix F.3 in our revision. Qualitatively, the trained transformer still exhibits similar mechanisms, as shown in the risk and probing curves (Figure 13-15), which gives evidence that they may indeed be implementing approximate versions of our claimed mechanism. Further questions such as what approximate representation function these are implementing could be an interesting question for future work.

Reviewer comment (Theory):
What happens when the representation function is of a different form? If either the transformer does not have enough layers or the width, is there an approximate representation function learned on which regression is performed, or does the entire mechanism fall apart?

Author response:
> Although the paper does a good job of illustrating the claimed mechanism, it does not analyze settings where the mechanism breaks.

> What happens when the representation function is of a different form? If either the transformer does not have enough layers or the width, is there an approximate representation function learned on which regression is performed, or does the entire mechanism fall apart?

In theory, if the transformer does not have enough layers or width, then it can still implement an approximate algorithm (such as computing an approximate version of the representation function $\Phi^\star$, as the reviewer mentioned, combined with fewer gradient descent steps to implement approximate ridge regression).

Empirically, we conducted some preliminary experiments with the same transformer architecture but more sophisticated representation functions, such as MLPs with higher depth or width. We’ve included these experiments in Appendix F.3 in our revision. Qualitatively, the trained transformer still exhibits similar mechanisms, as shown in the risk and probing curves (Figure 13-15), which gives evidence that they may indeed be implementing approximate versions of our claimed mechanism. Further questions such as what approximate representation function these are implementing could be an interesting question for future work.

Reviewer comment (Theory):
How robust is learning of the representation function in settings where the pretraining data contains spurious correlations? Can we say anything about the transformer's ability to compositionally generalize with either the representation function, regression, or both?

Author response:
> How robust is learning of the representation function in settings where the pretraining data contains spurious correlations? Can we say anything about the transformer's ability to compositionally generalize with either the representation function, regression, or both?

By “spurious correlations”, did the reviewer mean OOD-like scenarios where the test time representation function is different from the training time one in function space, but they happen to be identical on the training data distribution? In that case, we can still guarantee the existence of a transformer that can do well on both the test time representation function and the training data, but we won’t have a statistical guarantee on how training based on empirical risk minimization can actually find this transformer, as the statistical guarantee (e.g. the one sketched in the last paragraph of Section 4.1) only holds in-distribution.

Re the OOD question: Yes, if the transformer is trained with (non-diagnoal) Gaussian prior for the x, then we expect it to utilize this structure and it may perform suboptimally when test-time x is say spherical Gaussian. It would also be unclear whether it "reveals the true representation function".

Reviewer comment (Presentation):
What is OLS in Fig 1b?

Author response:
> What is OLS in Fig 1b?

The OLS means ordinary least squares (i.e. standard linear regression) on top of the representation function $\Phi^\star$.

Reviewer comment (Novelty):
Novelty of the technical contributions is limited given prior works of similar flavor that provide the feasibility of empirical risk minimization during forward pass. One of the aspects which authors claim to be novel is that they allow for representation-based learning. However, the underlying assumption is that the representation map is a multi-layer MLP model, which Transformers should be easily able to simulate through its MLP layers. In that sense, the results in not very surprising.

Author response:
For the MLP part of the theory, we agree the high-level idea itself is not very surprising, as mentioned. However, instantiating this idea into a full rigorous construction involves several other new ingredients different from existing work. For example, we required an efficient implementation of N parallel gradient descent simultaneously for *predicting at every token* using a decoder transformer (Proposition B.4), whereas similar existing work (Bai et al. 2023, Zhang et al. 2023a) only implements a single gradient descent for predicting at the last token only using an encoder transformer. See the “Proof techniques” paragraph after Theorem 1 for more details.
Further, another core contribution of this paper lies in using theory to guide empirical mechanistic study in a fine-grained (low-level) way, and indeed validating that the mechanisms identified from the theory show up in trained transformers (Figure 3-5). Along with the findings, we also proposed new techniques such as pasting (Figure 4) which may be of further interest.

Reviewer comment (Experiments):
The probing analysis is done only on a synthetic setup. It would be nice to get some supporting evidence for the proposed in-context learning mechanism on a real dataset.

Author response:
While our setting is synthetic, the setting (linear function on top of common representation) is arguably more realistic than most existing works in this direction that focus on linear functions. The synthetic setting also allows probing analyses to be done in conjunction with rigorous theoretical studies. Nevertheless, We agree that some probing analyses on e.g. real data or even more realistic synthetic data is an important question, which we believe is out of the scope of this paper but an interesting direction for future work.

Reviewer comment (Writing):
In the paragraph on **In-context learning** in Section 2, $\mathcal{D}^{(j)}$ and $\mathbf{w}_{\star}^{(j)}$ are not defined before their usage.

Author response:
Thanks for spotting these issues. We’ve updated them accordingly in our revision.

Reviewer comment (Novelty):
In general, the transformers are known to be universal approximators. In light of these, could authors comment on the significance of the key contributions of the paper, i.e., learning a representation map before applying gradient-descent in the representation space?

Author response:
While the high-level plan of “learning representation + applying gradient descent” can be done by universal approximation, and this paradigm is not so surprising given prior work,**how** we instantiate this construction is important and makes a difference. The way we instantiate it in our Theorem 1&2 are efficient (transformer has low number of layers and heads), and lower-level operations such as copying we use in our construction do show up experimentally via probing. We believe these make the construction itself much more interesting and reflective of reality than a simple universal approximation type argument.

Reviewer comment (Experiments):
Could the authors elaborate on using a **linear model** for their investigation on the upper module via pasting (Figure 4)?

Author response:
Is the question about why we **use a linear model as the data distribution** to investigate the upper module via pasting? Our goal is to test if the upper layers are simply implementing linear regression, or if they are implementing a more complex algorithm like linear regression with feature learning. The linear model is a sensible choice for testing this. In our experiment, we generated data using a linear model and passed it through the upper layers. The results suggest the upper layers are just doing vanilla linear regression.

Reviewer comment (Theory):
Currently the non-linearity in the true representation map is closely tied to the nonlinearity used in the Transformer network. Could the authors comment on generalizing this to broader nonlinearities in the true representation map? How would it increase the required number of layers?

Author response:
If the true representation map involves other nonlinearities beyond leaky ReLU, we could approximate such a non-linearity by a linear combination of ReLU functions. Therefore we can construct a similar transformer with potentially more heads within each layer (for approximating scalar nonlinearity using a linear combination of ReLUs), but the same number of layers.