How Do Transformers Learn In-Context Beyond Simple Functions? A Case Study on Learning with Representations

We thank the reviewer for their valuable feedback. We respond to the specific questions as follows.

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?

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.

Choice of non-linear functions… I think it would be helpful to clarify that the results are specific to this setup.

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?

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.

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?

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.

We thank the reviewer for their valuable feedback. We respond to the specific questions as follows.

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.

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.

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.

Labels in figures and the figure captions can be more clear. For example, items like "TF_upper+1_layer_TF_embed" in Figure 4b are not very readable.

Thank you for pointing this out. We have added a short explanation about this in the caption of Figure 4 in our revision.

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$.

We thank the reviewer for their valuable feedback. We respond to the specific questions as follows.

The theory only encompasses representational results by providing some settings of the parameters in a way that a transformer performs an ICL task. Given the transformer's highly expressive capability, these types of constructions are generally relatively straightforward.

While transformers are known to be highly expressive approximators of sequence-to-sequence functions, the efficiency and the internal mechanisms of the transformer still matters a lot. Our transformer constructions in Theorem 1&2 are efficient (transformer has low number of layers and heads). Further, both its high-level structure (such as computing representation + using gradient descent to do linear regression) + lower-level operations such as copying show up experimentally in trained transformers via probing. We believe these make the construction itself much more interesting and reflective of reality than a simple universal approximation type argument.

There's no assurance that these theoretical constructs are truly internalized by the model during the training process. Although probing experiments gave us some confidence that, in specific instances, the theory can predict the model's behavior, these types of experiments generally don't offer robust guarantees. As a result, while the theory is logical and sometimes mirrors empirical events, it could be counterproductive to lean too heavily on these theoretical constructs.

We appreciate this thoughtful question. First of all, we agree that theoretical construct is just one way to guide the probing experiments, there may be other ways to formalize hypotheses about the actual mechanism, and we don’t necessarily need to rely solely on the theory. Nevertheless, we still believe that the probing results on their own are actually a strength of our work, in that we are the first to give evidence of these mechanisms (especially lower-level ones such as copying) for in-context learning on trained transformers. In addition to probing, we also conducted a “pasting” experiment (Figure 4) as a more controlled way to test the linear ICL capability of the upper module of the trained transformer. We believe further probing studies, potentially going beyond the theoretical constructs, is an important direction which we would like to leave as future work.

It might be necessary to carry out an analysis of training dynamics in order to theoretically determine under which conditions the model actually aligns with the theoretical constructs.

We agree that the training dynamics is an interesting future direction as well. In our ablations we already had some preliminary experiments on the training trajectory (Appendix F.1, Figure 9) where we found that the representation module and the linear ICL module gets simultaneously learned in the early training stage. Further studies in this direction would be interesting.

We thank the reviewer for their valuable feedback. We respond to the specific questions as follows.

Novelty of the technical contributions is limited... 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… could authors comment on the significance of the key contributions of the paper?

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.

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.

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.

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?

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.

Could the authors comment on generalizing this to broader nonlinearities in the true representation map? How would it increase the required number of layers?

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.

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

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.

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

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