Paper: Is attention required for ICL? Exploring the Relationship Between Model Architecture and In-Context Learning Ability (ICLR-2024)

Reviewer comment (Presentation):
Since the macro perspective is only studied in the language modeling domain, it would make more sense to move the write-up around it under the Language Modeling section.

Author response:
We agree with the reviewer and have actually reorganized all our experiments, including putting content exclusive to language modeling into Section 6. We did end up removing mentions of micro/macro perspective as we realized it was not strictly needed to understand our draft and seemed to add confusion rather than clarity to other readers.

Reviewer comment (Evaluation):
How does this contrast to the already present information that Transformers can do ICL considerably well? The experiments in Table 2 seem to show that attention-based models are fairly poor at ICL but Garg et. al (2023) and Muller et. al (2022) show that they do quite well.

Author response:
Perhaps this is a misunderstanding. We also find that decoder-only transformers can do ICL considerably well which is the same design used by Garg et. al (2023). In fact, our revised experiments show GPT2 and Llama2 are often the top performers when given prompt lengths seen during training. The transformers which did underperform in our previous draft were BERT and T5. The prior is an encoder-only transformer and the latter is an encoder-decoder. We are, admittedly, not familiar with Muller et. al (2022). They seem to use a modified encoder-only transformer which is interesting given the poor results we saw with BERT.  

In order to improve the focus of our new draft and for practical reasons, we removed experiments for architectures not suitable for causal language modeling (the setting where in-context learning is typically used in the real world). This includes BERT and Fnet. We also removed T5 which does not treat subsequences of its encoder input as training examples. Both types of architectures are not amenable to our experimental setups in Sections 4 and 6. We considered keeping them in Section 5 but decided showing their performance in only a single setting did not add enough value.  

In case the reviewer finds this interesting, we did perform additional experiments on encoder-decoder transformers (T5) and encoder-only transformers (BERT) that did not make it into our latest revision. First, we find that how prompts are presented to an encoder-decoder transformer make a big difference. In our original experiments, the prompt is given only to the encoder stack. The decoder stack receives only a single learned start token. This implies that the network has only bidirectional attention across the prompt via its encoder. However, we see significant performance improvements when we pass the prompt to both the encoder and decoder stacks, using only the final prediction of the decoder stack to compute loss. It is still unclear to us why this change matters. Given that we know that decoder-only transformers perform well, perhaps the encoder representations are not useful and are effectively being ignored. An alternative hypothesis is that passing the prompt to the decoder stack simply adds more computation steps which is helpful for ICL to emerge. In regards to the underperformance of encoder-only transformers, we found that this was isolated to the Huggingface implementation of BERT (which we used). We found performance significantly improved when using the x-transformers implementation of an encoder-only transformer. We are still investigating why this is the case.

Reviewer comment (Experiments):
Details about how data is processed for image classification is missing. Do the experiments use embeddings obtained from ResNet on images to get a set of vectors and labels as context, or are the models given raw pixel information in some way?

Author response:
We’ve expanded Section 5 to be more explicit about how the image data is being processed. Specifically, we use a randomly initialized Resnet to embed images (raw pixels) into vectors. The image labels are mapped to vectors with a standard embedding layer. The embedded images and labels are then interleaved into a sequence of vectors and passed to the model being evaluated. Resnet and the label embedding layer are trained alongside the model being evaluated.

Reviewer comment (Reproducibility):
A lot of the performance of transformer-based models depends on how the data is represented in the context. Providing additional implementation details about how data is represented for each experiment as well as how training is done would be helpful.

Author response:
We agree with the reviewer and expanded our discussion of experimental details. Specifically, we’ve added more information about how the various tasks are ultimately mapped to a sequence of vectors in Section 2. We included more details surrounding the training and evaluation process in Appendix A.1, A.2 and Section 5.

Reviewer comment (Experiments):
There is also a lot of interesting ablations in the appendix of the paper on different kinds of transformers (eg. Decoder only vs Encoder-Decoder). Insights about differences between such models is more fundamental than differences between small architectural changes between different named architectures, and thus a detailed comment on them in the main section would be nice.

Author response:
Since we made substantial revisions to our draft to improve the focus of our manuscript, we decided to remove most of the appendix experiments found in the original draft. Instead, we spend much more time deeply investigating our new results in Sections 4, 5, and 6.

Reviewer comment (Novelty):
The authors should consider citing the following works which are directly related to ICL.  

**Citations missing**  

Garnelo, M., Schwarz, J., Rosenbaum, D., Viola, F., Rezende, D. J., Eslami, S. M., & Teh, Y. W. (2018). Neural processes. arXiv preprint arXiv:1807.01622.  

Garnelo, M., Rosenbaum, D., Maddison, C., Ramalho, T., Saxton, D., Shanahan, M., ... & Eslami, S. A. (2018, July). Conditional neural processes. In International conference on machine learning (pp. 1704-1713). PMLR.  

Müller, S., Hollmann, N., Arango, S. P., Grabocka, J., & Hutter, F. (2021). Transformers can do bayesian inference. arXiv preprint arXiv:2112.10510.  

Mittal, S., Bracher, N. L., Lajoie, G., Jaini, P., & Brubaker, M. A. (2023, July). Exploring Exchangeable Dataset Amortization for Bayesian Posterior Inference. In ICML 2023 Workshop on Structured Probabilistic Inference {\&} Generative Modeling.

Author response:
Thank you for bringing these works to our attention. We will review these works and figure out how best to incorporate them into our draft. The works appear to represent an interesting perspective on ICL that we had not considered.

Reviewer comment (Writing):
*Thus, each train and test example is a unique learning problem, but of a consistent type (e.g. linear regression)* : In this particular setup, “consistent type” does not imply linear regression but just regression in general. Would that be correct?

Author response:
By “consistent type” we just mean that each prompt (train, test examples) are members of the same class of functions. For example, in our linear regression task, let's say we generate a prompt which is associated with some implicit w_1 (defining a specific instance of a linear regression problem). Now we generate another prompt which is associated with w_2. Both prompts are unique problems (different w) but of the same type: linear regression. That said, you are correct in that the space of predictions by the model is not strictly linear.

Reviewer comment (Writing):
The difference between "bursty" and "non-bursty" prompts was not clear, at least to me, from the write-up. Could the authors clarify this distinction?

Author response:
We’ve revised Section 5 to explain the differences between “bursty” and “non-bursty” more clearly. We’ve also included visual examples in Table 1 and Figure 6.

Reviewer comment (Experiments):
For Transformer based experiments (eg. BERT, GPT2, etc.), are positional encodings used? This is primarily because observations seen in linear regression, for eg., are exchangeable and thus observation i and j should not be treated differently (i.e. there should be permutation invariance in the predictive distribution).

Author response:
Yes, all transformers in our experiments use positional encodings. To be precise, GPT2 uses absolute learned positional embeddings (just another embedding layer) and Llama2 uses rotary positional embeddings. You are correct in that there should be permutation invariance among the input-output pairs. However, the pairs themselves are not permutation invariant as the model needs to be able to distinguish between the input from the output. Still it would be interesting to experiment with positional embeddings removed as prior work has shown that they may not be necessary in causal transformers [1, 2].

Reviewer comment (Writing):
For the transformer models, can the authors clarify what an encoder-only transformer model is?

Author response:
An encoder-only transformer is a “BERT-style” transformer, i.e., a single stack of transformer blocks that do not employ causal masking.

Reviewer comment (Evaluation):
It is not clear why the ``ICL Score” is a good metric to care about. Could the authors explain what is the relevance of this score, and why is it important other than validation loss?

Author response:
ICL score can be thought of as a simple heuristic for measuring statistical efficiency. It differs from validation loss in that it’s effectively the slope of the validation loss curve plotted against token index (context length). Prior work [3] has used ICL score as a signal for when “induction heads” develop, a pair of attention heads that are theorized to be responsible for ICL ability in transformers. We include the metric in our study to see if a similar signal emerges in non-transformers and to gauge whether it correlates well to other measures of ICL ability.

Reviewer comment (Experiments):
My main concern is that the paper lacks the discussion of under which conditions a certain architecture tends to succeed in in-context learning and under which conditions it fails. In other words, I think the current main claim of this paper that “all considered architectures can perform in-context learning under certain conditions” is not very interesting and surprising. Instead, presenting both the success and failure cases and carefully comparing the settings/conditions for failure cases to occur may provide more insights to the community. I do not expect the paper to find the essential reasons for a certain architecture to fail in in-context learning, but more ablation studies on various training settings/hyperparameters are encouraged and at least some reasonable hypotheses can be made to explain the failure cases. The authors have some discussions on the unexpected low performance of BERT and T5, which I greatly appreciate. However, this is not enough, and I hope to see a more comprehensive picture of how various factors affect the in-context learning ability.

Author response:
We acknowledge that the claims in our original draft were overly conservative. In our revised draft, we demonstrate that all architectures are capable of performing in-context learning across a wider range of conditions than previously documented. Specifically, we highlight how ICL emerges in RNN and LSTM for image classification, which contrasts with the findings of [1].

Furthermore, we have significantly revised our experimental setup and discussion. In Section 4, we report how ICL ability responds to varying prompt lengths, including those not seen during training. In Section 5, we show how changes in the distributional properties of training data affect the emergence of ICL. In Section 6, we focus on tasks involving natural language which more closely resemble ICL in the real world. We also converted our dense data tables into line plots which make relative performance takeaways much easier to see. These revisions allow us to more clearly delineate the circumstances under which each architecture succeeds or fails, as well as to make stronger relative comparisons. 

In our revised draft, we have expanded the discussion to include additional hypotheses that may explain the varying performances of different architecture. While we recognize the importance of understanding the 'why' behind these performance differences, we believe that first clearly documenting the observed phenomena is an equally important contribution to the research community.

Reviewer comment (Writing):
When introducing the image classification task on Page 4, the definition of “bursty” and “non-bursty” samples are unclear and not self-contained. I feel confused without referring to the original paper. I encourage the authors to give a toy example in a figure along with the description.

Author response:
We’ve revised Section 5 to explain the differences between “bursty” and “non-bursty” more clearly. We’ve also included visual examples in Table 1 and Figure 6.

Reviewer comment (Writing):
What is the difference between the micro perspective of in-context learning and the macro perspective of in-context learning? On Page 4, the paper says “While the micro perspective focuses on specific tasks, the macro perspective offers a broader view of in-context learning by focusing on loss”, but I’m not able to capture the essential difference of the macro perspective from the micro perspective from this sentence.

Author response:
We ended up removing mentions of the micro and macro perspective in our revised draft. These perspectives were introduced by [2] but we now believe that, in the context of this study, end up adding confusion rather than clarity.

Reviewer comment (Experiments):
In the task of Associative Recall, I’m concerned with the following case. If the query token never appears in the prompt sequence, how can the network be able to predict the corresponding target of this query token? For example, if the prompt is “a,1,b,3,c,2,d”, how can the network predict the target for the token “d” when the network has never seen “d” in the prompt?

Author response:
You are correct. If the query token never appears in the prompt sequence, the network cannot accurately predict the corresponding target. However, all architectures are trained and evaluated with the same datasets so their performance ceiling is the same, i.e., we can still make relative comparisons among architectures. We also find that given a modest prompt length, the probability of the query token not appearing in the prompt decays quickly. We can see in Table 9 that prompts of length 2^5 to 2^6 have scores near 100% accuracy, indicating that the query must have almost always appeared in those prompts.

There are also practical concerns with forcing the query and its completion to appear in the prompt as it makes performing the experiments where we show how accuracy changes with varying prompt length difficult (Figure 1). This introduces a dependence on what may appear in the in-context examples and the current query index we are computing loss for.

Moreover, we’ve found that this task is present in prior work on ICL [3] [4] so we felt it important to include in our study for completeness.

Reviewer comment (Experiments):
For CNNs, why do the authors use lightweight convolutions and dynamics convolutions for experiments? Why not consider traditional convolutions, which are more widely used and studied? The authors are encouraged to clarify the reason for this specific setting.

Author response:
We originally considered adding traditional convolutions to our experiments but decided to only focus on architectures capable of causal language modeling (the setting where in-context learning is typically used in the real world). There are also practical concerns that arise when conducting our experiments. Specifically, the experiments in Section 4 where we examine performance versus varying prompt lengths as well as the language modeling experiments in Section 6 are not straightforward to perform with a traditional CNN.

Given our criteria, we are not aware of a CNN based architecture that is simpler than lightweight convolutions. It is simply a (normalized) depthwise convolution which has been widely studied [5]. There exists Gated Convolutional Model [6] which predates lightweight convolutions but is arguably more “complex”. That said, if the reviewer has a specific architecture mind, we are happy to perform and add the experiments to our final draft. However, given that we’ve increased the number of training runs substantially for each architecture, the experiments may not finish before the rebuttal period.

Reviewer comment (Presentation):
What is the criterion to assign the cells red/green colors in Table 2?

Author response:
We have since removed the red/green coloring in the data tables and replaced them with line plots for clearer presentation of the data. A new view of the dense data tables are still available in the appendix. The color coding in these new data tables are simply normalized with respect to the min and max values within that table.

Reviewer comment (Novelty):
The conclusion of the paper may be trivial, as the few-shot learning task can still be seen as having a mapping between input and output, albeit with more complicated inputs. Therefore, any universal approximator should at least have the capability to perform such a task.

Author response:
While we agree that any theoretical universal approximator should have the capability to learn any input-output mapping, this may not be true in practice. Indeed, prior work has shown that it is not clear if non-transformer architectures can perform ICL [1]. That said, we agree with the reviewer that we did not place enough emphasis on the relative performance of the architectures. We have since addressed this in our revised draft by conducting both a larger hyperparameter sweep and revising our experimental set up. Specifically, in Section 4, we report how ICL ability responds  to varying prompt lengths, including those not seen during training. In Section 5, we show how changes in the distributional properties of training data affect the emergence of ICL. In Section 6, we focus on tasks involving natural language which more closely resemble ICL in the real world.

[1] https://arxiv.org/abs/2205.05055

Reviewer comment (Experiments):
The performance in few-shot learning would vary because of characteristics of the architecture and tuning, but the authors do not seem to put emphasis on the relative performance of the different architectures in their study.

Author response:
We hope that the reviewer has found our latest revision much improved. We specifically set out to address their concern that we did not place enough emphasis on the relative performance of different architectures. The reviewer will be happy to find that we thoroughly address this in Sections 4-6.

Reviewer comment (Experiments):
It sounds like in-context capability comes from the fact that data are presented in a sequence and models are able to make use of the temporal correlation between tokens when making predictions. Has the authors investigated what order of the sequence encourages the learning of the correlations?

Author response:
Q1: We agree that investigating the effects of sequence order is interesting as most ICL tasks are permutation invariant with respect to in-context examples. While we do not experiment with sequence order specifically, we do perform a new study on sequence distributional properties in Section 5. That is, we investigate how different sampling procedures used to construct training sequences affect ICL.

Reviewer comment (Experiments):
Why do you choose to only reduce depth when normalizing the parameter counts? The normal practice is typically to reduce both the width and depth in some proportion.

Author response:
Q2. Thank you for pointing this out. Indeed we found that increasing depth for smaller parameter architectures (on a per layer basis) sometimes led to unstable training. We have since redone our experiments and normalized parameter accounts by adjusting both depth and embedding size. We’ve included the specific widths and depths we selected in Tables 3 and 5.