Paper: Understanding In-Context Learning in Transformers and LLMs by Learning to Learn Discrete Functions (ICLR-2024)

Reviewer comment (Theory):
The author claimed that TFs can learn two distinct algorithms on tasks such as conjunctions, which is not that convincing to me. Let's take figure 3 as an example and let's call the four sub-figures a to d from the left to the right. The way that authors compared these four figures and draw conclusions is that: they compare a and b, and then compare c with d, and then they drew the conclusion that the trained TFs can learn two distinct algorithms. This is not a good way for comparison. This is because one can easily give a counter-example to show that a single algorithm can simultaneously achieve a and b. For example, suppose the algorithm that TFs learn is: randomly pick one decision rule (or Boolean function) which is in the conjunction function class and agrees with all examples in the test context (the x_i, y_i pairs at test time). Then, when tested on teach conjunction (fig a), since there is only one conjunction function matching all test context, it will achieve a perfect accuracy. When tested with random sequence of examples from a conjunction task, there can be many functions in the conjunction function class that satisfy all test examples, so it can achieve an imperfect performance. Similar logic applies on the comparison between c and d.

Author response:
“Reviewer: *The author claimed that TFs can learn two distinct algorithms on tasks such as conjunctions, which is not that convincing to me.
Let's take figure 3 as an example and let's call the four sub-figure a to d from the left to the right. The way that authors compared these four figures and draw conclusions is that: ….. , so it can achieve a imperfect performance. Similar logic applies on the comparison between c and d.*”

We think there could be a misunderstanding here. We are not completely sure we understand your explanation here. However, based on our interpretation, we think our argument in the paper is quite similar to what you suggest in the next paragraph.

“Reviewer: *At a high level, if you want to show the trained TFs can learn two distinct algorithms, it is not a good idea to show that they achieve different performance on different test examples but trained on same data distribution, It is better to show that they achieve different performance when trained on different data distribution, but tested on the same examples. …*”

We are not comparing Transformers trained on the same data distribution. We are comparing two Transformers trained on two different pretraining distributions (one with random examples and one with teaching sequences). Leaving aside Transformers, it is straightforward to see that there can be two types of algorithms: algo (a) is a more general purpose like FFN trained with gradient descent and algo (b) is a more specific algorithm which uses teaching sequences to find the correct conjunction. Now, the general purpose algo (a) will perform well with and without teaching sequences but it need not be optimal on examples with teaching sequences. On the other hand, the algo (b) will be optimal on examples with teaching sequences but can completely fail when such teaching sequences are not provided. 

We claim that Transformers learn two distinct algorithms in this sense. This part is unrelated to Figures 3 c and d. If you look at Figure 3 a and Table 1 (or Figure 6) then these correspond to results with Transformers trained for Conjunctions in the vanilla setting (without teaching sequences) which work well with and without teaching sequences (analogous to aglo (a) described above). Figure 2 left and Figure 3 c (center-left) correspond to results with Transformers trained with teaching sequences (different pretraining distribution). Figure 2 left shows that they can perform optimally when teaching sequences are provided but fail otherwise (Figure 3 c) (analogous to algo (b)).  Figure 3 a and b are results with two different Transformers trained on two different input distributions (mentioned in the title above the figure). Figures 3 c and d serve a different purpose which depicts the result of a single Transformer trained on a mixture of distributions and tested on individual distributions.

Reviewer comment (Experiments):
At a high level, if you want to show the trained TFs can learn two distinct algorithms, it is not a good idea to show that they achieve different performance on different test examples but trained on same data distribution. It is better to show that they achieve different performance when trained on different data distributions, but tested on the same examples. In this case, you should compare figure b and c to draw your conclusion, instead of comparing a and b. The reason behind this is that the in-context learning is specific to the pre-training distribution, and it is very natural that they learn different algorithms on different pre-training distributions.

Author response:
“Reviewer: *At a high level, if you want to show the trained TFs can learn two distinct algorithms, it is not a good idea to show that they achieve different performance on different test examples but trained on same data distribution, It is better to show that they achieve different performance when trained on different data distribution, but tested on the same examples. …*”

We are not comparing Transformers trained on the same data distribution. We are comparing two Transformers trained on two different pretraining distributions (one with random examples and one with teaching sequences). Leaving aside Transformers, it is straightforward to see that there can be two types of algorithms: algo (a) is a more general purpose like FFN trained with gradient descent and algo (b) is a more specific algorithm which uses teaching sequences to find the correct conjunction. Now, the general purpose algo (a) will perform well with and without teaching sequences but it need not be optimal on examples with teaching sequences. On the other hand, the algo (b) will be optimal on examples with teaching sequences but can completely fail when such teaching sequences are not provided. 

We claim that Transformers learn two distinct algorithms in this sense. This part is unrelated to Figures 3 c and d. If you look at Figure 3 a and Table 1 (or Figure 6) then these correspond to results with Transformers trained for Conjunctions in the vanilla setting (without teaching sequences) which work well with and without teaching sequences (analogous to aglo (a) described above). Figure 2 left and Figure 3 c (center-left) correspond to results with Transformers trained with teaching sequences (different pretraining distribution). Figure 2 left shows that they can perform optimally when teaching sequences are provided but fail otherwise (Figure 3 c) (analogous to algo (b)).  Figure 3 a and b are results with two different Transformers trained on two different input distributions (mentioned in the title above the figure). Figures 3 c and d serve a different purpose which depicts the result of a single Transformer trained on a mixture of distributions and tested on individual distributions.

Reviewer comment (Novelty):
Your definition and narratives of in-context learning and the algorithm learning perspective (the last two paragraphs in 'in-context learning' paragraph in section 2) seem to follow closely with the formulation in [1] and the formal definition in [2]. Does your definition of in-context learning resemble their definitions? Or do you have any differences? Maybe you can discuss more about your definition of ICL and how it relates to the algorithm learning process.

Author response:
**(Q1)** *Your definition and narratives …  discuss more about your definition of ICL and how it relates to algorithm learning process.*

Yes, indeed our definition of in-context learning is similar to those of [1, 2] which are also based on Garg et al [3]. There are no significant differences in our framework compared to [1] and our perspective of looking at the meta-learning-like setup as algorithmic learning is almost identical. Their [1] definition is more formal which is necessary for their theoretical results and not required to understand the main findings of our paper. Some of our experiments exploring the sample complexity of Transformers (Section D.3 and Figure 13) are also along the same lines as their [1] work. 

We are glad that you find the algorithmic learning perspective interesting. There are some theoretical insights about the representational capabilities of Transformers one can glean from such a perspective. Maybe you will find the following interesting: Recent works [5] have shown that hard-attention (hardmax instead of softmax) Transformers can represent functions in the class of AC0 circuits which are constant depth polynomial size Boolean circuits (with ‘and’ and ‘or’ gates). We found that we can represent the known PAC-learning algorithm for learning Conjunctions with a Boolean circuit in AC0 which implies that hard-attention Transformers are capable of representing the learning algorithm for Conjunctions (but does not imply that they can learn). If you find this interesting, we can discuss this in detail in our next version. Given space constraints we cannot include it in the main paper.

Reviewer comment (Experiments):
For some tasks in your paper, there should be existing algorithms to learn them, such as conjunction or disjunction. I am also wondering how the trained TFs compare to them. Also, some existing papers show that the TFs can approximate or possibly learn (in-context) the Bayesian optimal estimator over some function class [1,2,3]. I suppose under the task distribution in your paper, it is very likely that the Bayesian optimal estimator is analytically intractable, but I think it should be numerically computable under some simple tasks. So I am wondering how the trained TFs compare to the Bayesian optimal estimators (or some strong, computationally efficient baseline on specific task, if any).

Author response:
**(Q4)** *For some tasks in your paper, there should be existing algorithms to learn it, such as conjunction or disjunction. I am also wondering how does trained TFs compare to them? … Bayesian optimal estimators (or some strong, computationally efficient baseline on specific task, if any).*

We conducted such experiments and the results are present in the paper.  Figure 9 in the paper compares the trained Transformers with known PAC-learning algorithm for tasks such as Conjunction and Disjunctions. There are a few things to note here: (a) These classical algorithms are near-optimal and are guaranteed to work for any distributions over the input. They are optimal in the sense that the lower bound on the sample complexity of learning Conjunctions/Disjunctions is $O(n)$ and these PAC-learn these classes with $O(n)$ samples. (b) Transformers can perform better on the first few examples because they have knowledge about the distributions of inputs and functions since they are exposed to it in the meta-learning stage. So in that sense, they are closer to the Bayes-optimal-estimator than the classical algo in their behaviour. However, unlike the provable PAC-learning algorithms, the algorithms learned by Transformers are not guaranteed to work for arbitrary distributions. (c)  Both Transformers and the classical algo reach (near) perfect accuracy after observing almost the same number of $O(n)$ examples. 


Regarding the second part of the questions, we think the idea is interesting but we are not sure how feasible it would be to implement and compare with the Bayes optimal estimator within the discussion period. We will look into it as soon as we are done with the experiments requested by other reviewers which are relatively direct in terms of implementation.

Reviewer comment (Writing):
The takeaway messages are a bit unclear to me. [_Update_: The authors have clarified in the rebuttal. I hope the changes could be reflected in the revised paper.]

Author response:
Thank you for your thoughtful comments and time.

Response to Weakness: “Reviewer: *The takeaway messages are a bit unclear to me.*”

List of takeaways.

**(Takeaway 1)** Transformers (as well as other architectures) have concrete limitations in in-context learning certain classes such as Parities which are known to be learnable in polynomial time. They even fail when provided with a number of examples which is sufficient for a feedforward network trained with gradient descent. 

While there are various classes of functions for which Transformers have been shown to perform well, we are not aware of works which indicate a precise class of functions (such as Parities) which are known to be efficiently learnable and Transformers fail to learn such classes of functions in-context. 

**(Takeaway 2)** Attention-free architectures can perform in-context learning as well (in the sense of implementing learning algorithms). However, there still exists a gap between their performance and Transformers’ performance. We are unaware of previous works that report such a comparison on in-context learning with such a benchmark of learning problems. Reviewer bf3N found both takeaways 1 and 2 interesting and noted that they can be of interest to the community.

**(Takeaway 3)** Transformers can leverage more informative examples such as teaching sequences to learn more sample-efficiently when such prompts are provided and can switch to the vanilla version when teaching sequences are not provided. 
Your specific questions about takeaway 3 are discussed separately.

**(Takeaway 4)** Pretrained models such as GPT-2 which are trained on natural language data encode mechanisms to achieve non-trivial performance on learning problems such as Conjunctions and can implement the nearest neighbour algorithm.

**(Takeaway 5)** LLMs primarily pretrained on text data can take a sequence of inputs and labels in a manner similar to the meta-learning-like setup and perform as well as certain baselines such as the nearest neighbour algorithm. Evidence for 4 and 5 indicate that LLMs can learn from in-context examples alone apart from just indexing to tasks already seen during training. In other words, they can act as learning algorithms to some degree in a manner similar to the way we test models trained from scratch in the meta-learning-like setup.

To our knowledge, our experiments with LLMs (takeaways 4 and 5) are the first to explore the ability of LLMs to act as learning algorithms and solve learning problems.  Reviewer 9RLW found the experiments with LLMs and teaching sequences quite interesting. 

As mentioned in the introduction of our paper, our work focuses on the following research questions and the insights above are a step towards answering those. 

From Section 1:   
(Takeaway 1) -> (a) What are the limits of the in-context learning ability of Transformers?  
(Takeaway 2) -> (b) Is the attention mechanism essential for in-context learning?   
(Takeaway 3) -> (c) Can Transformers exploit high-quality and informative examples to learn more efficiently?   
(Takeaways 4 and 5) -> (d) To what extent do LLMs that are not specifically trained for these tasks carry the capacity to implement non-trivial learning algorithms on in-context examples?  


We hope this helps in clarifying the key takeaways related to the primary research questions explored in the paper. Let us know if anything regarding them is unclear.

Reviewer comment (Evaluation):
Parity-20: what if reporting the accuracy on every position, rather than the final position? i.e. whether the full-sequence accuracy is too harsh, and some more continuous progress measure might be more informative.

Author response:
“**(Q1)** Reviewer: *Parity-20: what if reporting the accuracy on every position, rather than the final position?*”

Figure 6 depicts the accuracy at every position and not just the final position for almost all of the tasks. For Parity-(10, 2), you can see that the accuracy at every position is near chance level. The performance on Parity-(20, 3) and Parity-20 is almost identical, i.e. near chance level at every position. It was omitted since the plots were almost identical to Parity-(10, 2).

Reviewer comment (Experiments):
For OOD test results (Fig 10), how about the comparison to LSTM, especially under aggressive distribution shift where Transformers fail to preserve the in-distribution accuracy?

Author response:
“**(Q2)** Reviewer: For OOD test results (Fig 10), how about the comparison to LSTM, especially under aggressive distribution shift where Transformers fail to preserve the in-distribution accuracy?”

We will try to do experiments along these lines and report back in a few days.

Reviewer comment (Evaluation):
About teaching sequences: [_Update_: the questions have been addressed during the rebuttal. I hope the updated paper could include the clarifications in the discussions; for example, which specific aspects of the empirical results suggest "two algorithms", and that it's important for the teaching sequences to be at the beginning of the list of examples, rather than mixed in as an arbitrary order.]

Author response:
"*re teaching sequences: I'd like to get some further clarification please:
The authors replied ... in the paper please?*"

The results are not present in the paper currently but we will include it in the final version. We checked for their OOD performance when we first observed that they could successfully learn from teaching sequences. We did not have any specific reason to include it in the paper earlier. While we have tried to include an extensive set of experiments in our paper, it is natural that some empirical results useful to answer specific questions are not preemptively added to the paper.

Reviewer comment (Presentation):
Fig 2: what's the value of $k$? To better show the progression of performance, could you please start the x-axis from 1?

Author response:
“**(Q3a)** Reviewer: *Fig 2: what's the value of k? To better show the progression of performance, could you please start the x-axis from 1?*”

The value of k is not fixed. It depends on the number of literals in the target Conjunctions or DNFs. For Figure 2 Left, the expected length of the teaching sequence in the Conjunction is ~27 based on the distribution of the class of functions. For function classes such as Conjunctions or DNFs, most of the examples in the teaching sequence have negative labels and hence any kind of baseline even Null accuracy can achieve near-perfect accuracy on the inputs in the teaching sequence or equivalently on the first 30 examples on average. However, the performances of the baselines drop right after the teaching sequence whereas Transformers stay at perfect accuracy on the rest of the examples. Hence, the performances on the first ~27 examples primarily containing teaching sequences are not informative of the models’ true effectiveness and were omitted to avoid confusion. Based on your comment, we will include the plot where the x-axis begins with 1 in the appendix along with the explanation of their performance in the next couple of days.

Reviewer comment (Experiments):
I'm curious about the robustness of the teaching sequence results. For example, for different values of $k$ (with dimension fixed to 20), does Transformer always able to ...

Author response:
“**(Q3b)** Reviewer: I'm curious about the robustness of the teaching sequence results. For example, for different values of  (with dimension fixed to 20), does Transformer always able to”

For classes such as Conjunctions and DNFs, the length of the teaching sequence varies so the performance of Transformers is not specific to teaching sequences of a particular length. For Parity-(20, 3), the length of the teaching sequence is fixed since the output always depends on a fixed number of variables in sparse Parities. If we change the length of the teaching sequence, then the class of functions will change as well and hence it becomes more of a problem about generalizing to unseen functions apart from generalizing to teaching sequences of different lengths.

Reviewer comment (Theory):
Appendix F: why would this not work as a teaching sequence for the entire set of parity? Note that $k$ samples suffice only if the value of $k$ is known, and this would imply that the shortest teaching sequence for the full set of parity is of length 0.

Author response:
“**(Q3d)** Reviewer: *Appendix F: why would this not work as a teaching sequence for the entire set of parity?
Note that k samples suffice only if the value of k is known, and this would imply that the shortest teaching sequence for the full set of parity is of length 0.*”

That is not correct. We think you may have misinterpreted the definition of the full set of Parities. The statement you are referring to is the following,

> Page 24: “This would not work as a teaching sequence for the entire set of Parity functions Parity-n.”

For Parity-(n, k), the function value depends on exactly k bits and the full set of Parities for any n (Parity-n) contains all such functions (for all k). Quoting the description from the paper,

> Page 4: “The class Parity-n contains all $2^n$ possible Parity functions over $\{0, 1\}^n$. We denote the class of sparse parities with Parity-(n, k) which contain functions with $k$ relevant variables defined over $\{0, 1\}^n$.”

What you said would have been true if the full set of Parities would have had only one function which depends on all bits (Parity-(n, n)).

Reviewer comment (Theory):
I'm not sure it's fair to say that Transformers learn "two distinct algorithms" for learning with and without teaching sequences, since this seems to me as simply a matter of distribution shift: if the Transformer has only seen sequences where the first $t$ samples are from a teaching sequence, then there is a drastic distribution shift in the test time when there's no such samples. What am I missing here?

Author response:
“**(Q3f)** Reviewer: *I'm not sure it's fair to say that Transformers learn "two distinct algorithms" for learning with and without teaching sequences, since this seems to me as simply a matter of distribution shift: …*”

When we test Transformers trained for vanilla Conjunctions on examples with teaching sequences, there is a distribution shift as well but they still perform well (almost in the same way as without the distribution shift). Ignoring Transformers for now, one can consider two types of algorithms for learning Conjunctions. Algo (a): a general purpose algo such as an FFN trained with gradient descent. Algo (b): The second algorithm is dependent on teaching sequences and uses the teaching sequence to identify the target conjunction.

Note that the first algorithm can work with and without teaching sequences but it is natural to expect it to be suboptimal for inputs with teaching sequences. Similarly, the second algorithm will be optimal for inputs with a teaching sequence but will completely fail when the inputs do not have a teaching sequence. When we say that Transformers can learn two distinct algorithms, we mean it in this sense. From Figure 3 left and Table 1 (or Figure 6 top left), you can see that Transformers trained on vanilla Conjunctions behave like the first algorithm we described above. Similarly, from Figure 3 center-left and Figure 2 left, you can see that when trained with teaching sequences, they behave more like the second algorithm dependent on the teaching sequence. The emphasis here is on two different algorithms for the same class of functions (Conjunctions) where their effectiveness is dependent on the input distribution. If all aspects of the learning problem are identical then there is no reason to find two different algorithms.

Reviewer comment (Theory):
Section 5: the paper says real-value functions are not straightforward to evaluate "since LLMs receive and produce discrete values": I'm not sure how much this distinction is true or relevant, since LLMs do have embeddings which are vectors of real values.

Author response:
“**(Q4a)** Reviewer: *Section 5: the paper says real-value functions are not straightforward to evaluate "since LLMs receive and produce discrete values": I'm not sure how much this distinction is true or relevant, since LLMs do have embeddings which are vectors of real values.*”

We do not agree and the distinction is quite true. While training a Transformer of width (or d_model) $d$ on real-valued regression problems with $n$ dimensional inputs to learn functions of the form $f:  \mathbb{R}^{n} \rightarrow  \mathbb{R}$,  there is a trainable linear map $W:  \mathbb{R}^{n} \rightarrow  \mathbb{R}^{d}$ that transforms a real-valued vector to the same space as the width of the model. Pretrained LLMs do not have this linear map and even if we can provide real-valued vectors of the same dimension as the width (or d_model) of the model, it does not seem reasonable to draw conclusions from its performance since it has never been trained on those vectors. While training Transformers on Boolean functions in the meta-learning-like setup there is a similar linear map as well  $W:  \{0, 1\}^{n} \rightarrow \mathbb{R}^{d}$, however since the inputs are discrete and there are separate embeddings for the tokens ‘0’ and ‘1’ in the LLM, we can test the models by providing an n-dimensional Boolean input with n different tokens. This is the setup with ‘Direct evaluation’ in Section 5.

Reviewer comment (Novelty):
I am not sure which is the message of this paper compared to previous work. It has already been observed that transformers can "choose" between algorithms in [1], while we already knew that there are some tasks that transformers perform well and some others that they don't.

Author response:
Thank you for your thoughtful comments and time.

Regarding weakness and Question 3. 

We regret that the set of new insights in the paper was not clear. We first list the set of new insights below and then discuss each insight in some detail. We believe the key insights have not been explored in past works. For certain insights (such as learning distinct algorithms for the same task) we discuss the differences to past work (such as [Bai et al 2023] that you highlighted). If you disagree with any of the claims made regarding novelty in the discussion below, then it would be helpful if you could cite the relevant papers.

List of insights

1. Show concrete limitations of Transformers in in-context learning certain precise function classes which are known to be efficiently learnable.

2. Transformers can leverage more informative examples such as teaching sequences.

3. Attention-free architectures can perform such kind of in-context learning as well but they still lag behind on certain tasks. 

4. LLMs primarily pretrained on text data such as LLaMA-2 and GPT-4 can learn new functions in-context akin to baselines such as the nearest neighbour algorithm.

5. There are mechanisms encoded in the weights of pre-trained models such as GPT-2 which enable it to perform as well as baselines on learning classes like Conjunctions and also implement the nearest neighbour algorithm in a setup which is very close to the meta-learning-like setup. 




**(Insight 1)** While prior works have demonstrated that Transformers can learn various classes of real-valued functions by implementing gradient descent or by acting as a Bayes optimal predictor, our finding provides evidence that there are classes of functions where they perform poorly and some classes such as Parities where they fail to perform beyond chance-level accuracy. This can be of interest to researchers in the community and it can be further explored by future work to understand why this happens (as noted by Reviewer bf3N). While it is natural to expect any model to have limitations, it is important to note that Parities are learnable in polynomial time with algorithms such as Gaussian elimination and for the problem of learning Parity-(10, 2) we provide the models with input points which are even sufficient for FFN+GD to solve the task. Hence, the problem itself is not unsolvable and it was to some degree surprising to us that Transformers as well as other models fail to solve this task. 

**Why new?** We are not aware of any other precisely defined function classes where Transformers fail and which are known to be PAC-learnable in polynomial time.

**(Insight 2)** We show that various other architectures can match the performance of Transformers on various learning problems but there still exists a gap. Various architectures struggle on multiple tasks and Hyena is closest to Transformers in terms of performance falling behind only in the nearest neighbour problem. This could be of interest to the community (as noted by Reviewer bf3N) and could be useful for researchers working on developing new architectures that can match Transformers’ performance as LLMs. 

**Why new?** We are not aware of prior works that either compare the ability of recently proposed attention-free architectures to implement learning algorithms or works that even have a benchmark with a wide spectrum of performance for evaluating the ability of architectures to implement learning algorithms.

Reviewer comment (Experiments):
When the authors create the datasets, it seems like they always sample a function and $m$ Boolean inputs. At test time, is the prompt used also generated in the same way? Sampling one function and $m$ Boolean inputs? My question is related to whether the authors observe that these models can "generalize" to the case that the sequence length is not the one that the models were trained on.

Author response:
“**(Q1)** Reviewer: *When the authors create the datasets, it seems like they always sample a function and m Boolean inputs. At test time, is the prompt used also generated in the same way? Sampling one function an m Boolean inputs? My question is related to whether the authors observe that these models can ``generalize'' to the case that the sequence length is not the one that the models were trained on.*”

Yes during both training and evaluation, each prompt is created by sampling a function and $m$ input points. In some scenarios, the distribution over functions or input could be different but we have not conducted experiments to check if models can generalize to higher lengths during evaluation. We think most prior works have focused on this since Pretrained LLMs typically have a maximum context length (such as 4k, 8k, etc) and in-context learning is typically done within that.

Having said that, we are conducting some experiments based on your question to test the ability to generalize to higher lengths and will report the results within the next few days.

Reviewer comment (Reproducibility):
How exactly the authors prompted GPT4? Is it just a sequence of examples? It seems that in the supplemental material there is a folder prompts that contains some algorithms.

Author response:
“**(Q2)** Reviewer: *How exactly the authors prompted GPT4? Is it just a sequence of examples? It seems that in the supplemental material there is a folder prompts that contains some algorithms.*”

The details of the prompt are explained in Section G.3 and an example prompt is provided in Figure 17. Quoting a part from Section G.3 in the paper

> The main prompt (see Figure 17 for an example) comprises of an instruction in natural language followed by a series of $k$ example points ($x_i, y_i$), where $x_i \in \{0, 1\}^d$ is provided as a sequence of tokens $x_1, \ldots, x_d$, $x_j \in \{0, 1\}$ and $y_i \in \{0, 1\}$ is also provided as a single token. The model's goal is to predict the correct label $y_{k+1} \in \{0, 1\}$ for the query point $x_{k+1}$. For a particular function, we sample a sequence of $n$ points and prompt the model with $k < n$ points as in-context exemplars of the function. We call the model multiple times, increasing $k$ by a value of 5 every time.

Reviewer comment (Experiments):
It seems all the architectures considered (except LSTMs) have the number of hidden layers and latent dimensions in the same range. It would be good to include details of the number of parameters or compute used for each architecture so as to ensure that there is some normalizing factor among the architectures considered. This is to make sure that architectures like LSTMs and RetNets are not performing worse due to having a substantially smaller parameter count.

Author response:
For all architectures apart from LSTMs, the parameter counts are almost the same since in architectures such as RetNet, Hyena, and DSS, only the attention layer is replaced with a different mechanism. The feedforward block has the same number of parameters which dominates the parameter count and the replacements for attention with Retention/LongConv also have projection matrices. LSTMs typically have a smaller parameter count because deeper LSTMs are more difficult to train and do not perform well. For instance in Figure 14 (right), you can see the generalization error of LSTMs get poorer for a larger number of layers.

In the next version (in a few days), we will specify the largest model (in terms of parameters) we trained for each architecture. We hope that suffices.

Reviewer comment (Writing):
In the abstract, it is mentioned that Transformers can learn gradient-based learning algorithms. The prior works that I know of either claim that Transformers can represent such algorithms or that their performance matches gradient-based algorithms - none of these imply that the trained Transformers actually encode gradient-based learning algorithms. This line should be rephrased to avoid confusion.

Author response:
We think this is not entirely correct. While [Garg, et al.] showed that Transformers can match the performance of gradient-based algorithms and [ Akyürek, et al.] showed that Transformers can represent such algorithms, later [Von Oswald, et al.] showed quite convincingly that Transformers learn to encode gradient-based algorithms. In particular, they provided a much simpler construction to show how Transformers can implement a gradient-based algorithm for linear regression and then they inspected the weight/parameters of a trained Transformer model and found that they encode the same mechanism as their construction (apart from some scaling corrections) (See for instance Figure 9 in their paper). In other words, their results seem quite convincing to support their main claim that Transformers can learn gradient-based learning algorithms.

Reviewer comment (Novelty):
Here is a blog post that might be worth mentioning in the context of in-context learning abilities of real LLMs: https://www.alignmentforum.org/posts/c2RzFadrxkzyRAFXa/who-models-the-models-that-model-models-an-exploration-of

Author response:
Thank you for the suggestion. We were not aware of this and will cite this in our next version.