Is attention required for ICL? Exploring the Relationship Between Model Architecture and In-Context Learning Ability

W1:

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.

W2:

We’ve expanded Section 2 to clarify the differences among the various architectures. That said, some architectures, in particular those inspired by state-space models, are quite complex and require substantial background knowledge. In our experimental results, we also spend more time discussing our hypotheses on why we see differences among architectures.

W3:

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.

W4:

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.

W5:

We made substantial changes in how we present and discuss the results of our experiments. We provide a brief summary of the changes in the general response but hope the reviewer finds value in our new presentation in Sections 4, 5, and 6.

W6:

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.

W7:

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.

W8:

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.

We agree that the proposed positional encoding experiment is interesting. We conducted the experiment and included the results in Figure 9 of our latest revision. The setup and our observations are included in the general response. We now address the reviewer’s remaining concerns.

Our findings do not contradict Muller et. al (2022) [1].

1. Prior to conducting the aforementioned experiment, our revised draft made no claims about encoder-only transformers (the base architecture used by [1]).

2. If the reviewer is referring to the results in our original draft (Table 2), there are a few variables at play.

   • First, the experimental set up of [1] is different from ours. [1] sums the embeddings of each in-context example pair to represent it as a single token. They also use a special Riemann head for regression tasks and remove positional embeddings.
   • Second, the poor performance was isolated to the Huggingface implementation of BERT, but we did find that BERT performed better with more hyperparameter tuning (Table 14, original draft). Our current position is that this implementation of BERT might not be a good representative of a basic encoder-only transformer as it was designed to specifically support masked language modeling and next sentence prediction [2]. In our original draft, we found that a generic implementation of an encoder-only transformer performed well at ICL (Table 9, original draft) which does not contradict [1].
   • Observing that a specific (and popular) implementation of an encoder-only transformer failed at our tasks is valuable in itself. At the time of writing, BERT is the 2nd most downloaded model on the Huggingface hub. It is reasonable to assume that others are using this implementation as a generic encoder-only transformer. We also do not believe that extensively comparing two specific implementations of an encoder-only architecture is within the scope of this study which seeks to compare fundamentally different architecture classes.

3. Most importantly, in light of the new experiments, our findings are indeed consistent with [1]. Our experimental set up is now comparable and we confirm that encoder-only transformers perform well on all tasks considered.

Questions regarding T5.

1. Clarity on “We also removed T5 which does not treat subsequences of its encoder input as training examples”. Apologies for wording this in a confusing way. We will lay out a concrete example. Consider trying to use T5 for the experiments in Section 4. For each architecture, we report performance with respect to context length. We can do this efficiently because the considered architectures are all auto-regressive: each in-context example and its left context is treated as a single input sequence. This is what we meant by “subsequence”. This implies that all in-context examples produce a valid loss and prediction conditioned on its left context. However, this is not straightforward with T5 because its encoder stack has bidirectional attention. Even though we compute loss only using its decoder-stack, cross-attention to the encoder allows the model to “cheat”. This is not to say that Figure 1 is impossible to produce with T5, it is just impractical.
2. The initial poor performance of T5 was not due to any specific implementation. We explained earlier that T5 only performed poorly if the prompt is only given to its encoder stack and the decoder stack receives a single learned start token. Performance improves dramatically when the prompt is given to both the encoder and decoder stack. See Table 11 of the original draft where we perform an ablation to confirm this.
3. Note that we make no claims about T5 or encoder-decoder transformers in our revised draft. We focus solely on the autoregressive architectures, the dominant paradigm where ICL is used.

We thank the reviewer for engaging in thoughtful discussion and providing constructive feedback. We have made every effort to address their concerns within the rebuttal period and hope that we succeeded.

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

[2] https://arxiv.org/abs/1810.04805

W1:

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.

W2:

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.

Q1:

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.

Q2:

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.

Q3:

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.

Q4:

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.

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

[2] https://arxiv.org/abs/2209.11895

[3] https://arxiv.org/abs/2212.14052

[4] https://arxiv.org/abs/2302.10866

[5] https://arxiv.org/abs/1610.02357

[6] https://arxiv.org/abs/1612.08083

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.

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.

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.

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

We agree with the reviewer that our data presentation can be improved. To that end, we converted all the dense tables into line plots which allow us to quickly visualize how ICL performance responds to prompt length and how it changes over the course of training. The dense tables are still accessible in the appendix and we’ve added bidirectional links for easy navigation. We found that patterns in the new visualizations are now much easier to identify. We also reorganized our results into three sections that focus on different aspects of ICL performance: extrapolation, data distributional properties, and natural language. We hope the reviewer finds more value in our revised data visualizations and discussion compared to our original draft.

We are happy to add additional experiments for linear attention and sparse attention in our final draft. Given that we’ve increased the number of training runs substantially for each architecture, the experiments may not finish before the rebuttal period. If helpful, one of the architectures we do have experiments for, RWKV, does employ a linear attention mechanism and can be viewed as a successor to An Attention Free Transformer [1].

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

As the rebuttal period comes to a close, we would like to thank the reviewers again for their time and constructive feedback. Reviewer 9CgZ proposed an interesting experiment to measure the effects of positional embeddings given that in-context examples in our tasks should be permutation invariant. This experiment may also be of interest to reviewer Rxok who asked about the effects of sequence order. We conducted this experiment and show the results in Figure 9 of the latest revision (page 29).

Specifically, we consider the following variables:

• Token representation scheme: We represent in-context example pairs as a single token (instead of two in our original experiments) which allows us to remove positional embeddings. Specifically, we either sum their respective embeddings as represented in [1] or concatenate their embeddings, as suggested by reviewer 9CgZ. The query label is masked out by setting its embedding to zero.
• Positional embeddings: whether to use learned absolute positional embeddings or no positional embeddings at all.
• Attention mask: encoder-only vs decoder-only transformer. Note that in both scenarios, the query can attend to all in-context examples. In the encoder-only transformer, each example can attend to all other examples since it does not employ a causal mask. Examples in the decoder-only transformer can only attend to examples to its left.
• 4 tasks: linear regression, associative recall, multiclass classification, image classification.
• The remaining settings are identical to Section 4 with the following changes: Our hyperparameter sweep covers 2 learning rates, 2 seeds, and 2 layer depths. We train for 50K steps and only take the loss (and evaluate) at the token index 32 (i.e., models are trained to make a single prediction given 31 example pairs and the query). We conducted 768 training runs in total.

We make the following observations:

• Token representation scheme sensitivity: Associative recall and multiclass classification are not sensitive to tokenization schemes. However, we observe that concatenating embeddings in linear regression and image classification resulted in noticeably improved performance. We suspect that it is easier for attention heads to discern in-context inputs from outputs if they initially reside in their own subspace.
• Removing positional embeddings did not impact performance. This makes intuitive sense as in-context examples in this setting are permutation invariant.
• For most tasks, encoder-only and decoder-only transformers perform on par. The exception was linear regression where the encoder-only outperformed the decoder-only in the more difficult settings (d=20, 30). For image classification, we observed that ICL emerged in both transformers in very similar windows and followed a similar decay scheduled (as discussed in Section 6).

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