Looped Transformers are Better at Learning Learning Algorithms

Thank you for your feedback on our paper. We appreciate your comments and would like to address your concerns.

an extensive hyperparameter search for b and T: The authors make suggestions how this could be avoided, which should maybe be subject to further research.

Thank you for pointing this out. An extensive hyperparameter search for $b$ and $T$ is subject to the goal of finding the appropriate values for $b$ and $T$ such that the looped transformer:

• requires shortest training and inference time (related to the value of $b$),
• occupies least training memory (related to the value of $T$),
• is able to find a fixed-point solution beyond trained iterations.

From our understanding of looped transformer training presented in the paper, the larger $b$ is, the better the looped transformer is at finding a fixed-point solution beyond trained iterations. Consequently, if the main objective of employing a looped transformer is to achieve this fixed-point solution, the strategy would involve progressively increasing $b$ until the solution is attained. This is akin to determining the optimal number of layers in a standard transformer.

The search for $T$ mainly relates to balancing between achieving a fixed-point solution and avoiding training instability. Including regularization strategies can ease this instability, thus reducing the intensity of tuning for $T$.

How much time and resources does it take to find optimal values for b and T? If you consider the hyperparameter search is it still more useful to apply looped transformers instead of the standard transformers?

As discussed in the above section, we believe if the goal is to find the value of $b$ and $T$ such that the looped TF can find a fixed-point solution, the burden of hyper-paremter search will be reduced by including the regularization strategies. So we would think there is no harm applying the looped transformer over the standard transformer.

the structure of the paper can be improved, and minor details in the paper: ending the paper with a short conclusion instead of the related works.

Thanks for your suggestions. We have modified the draft accordingly.

What is the experimental setting for the plots shown in Figure 2? Why are the curves overlapping?

In figure 2, we train the vanilla TF and looped TF on linear regression problem, with problem dimension $d=20$. In the left figure, we demonstrate the performance of looped transformer on in-context samples $k=40$. For the overlapping curve, we test the looped TF and standard TF on the same set of test prompts, therefore they have very similar performance.

Thank you once again for your valuable feedback. We hope that our response effectively addresses your concerns.

We thank the reviewer for the thoughtful feedback and appreciate the opportunity to address your concerns.

experiments are only on simluated datasets: whether the proposed method is effective in dealing with real-world data remains to be validated. The classes of functions studied in the paper (including linear regression, decision tree, 2-layer ReLU NN, etc.) are ideal and relatively simple compared to the functions emerged in practical applications.

Indeed, we acknowledge that the datasets in our study are comparatively simplistic. This is due to the focus of this paper is on training a transformer to solve data-fitting problem with simple function classes.

To establish a connection with real-world applications and further validate our approach, we have conducted additional experiments using the following 10 datasets from OpenML, which we believe better represent practical scenarios:

Let $S =$ {$720, 725, 737, 803, 807, 816, 819, 847, 871, 42192$} be the set of all datasets. The datasets we examined have 0/1 binary labels. We trained the transformer and the looped transformer on 9 datasets and evaluated its in-context learning ability on the unseen test dataset. Both transformer and looped transformer have 256 embedding size, 8 heads. The tranformer has 12 layers, and looped transformer has 1 layer, trained to maximum loop iteration $b=30$, and $T=15$.

We uniformly sampled prompts from 9 datasets, where for each prompt, we first randomly selected a training set, then randomly selected $k+1$ samples from this training set, with $k$ being the number of in-context samples. During testing, we applied a similar approach for each test sample, selecting $k$ in-context samples from the test dataset, with care taken to exclude the test sample itself from these in-context pairs. The test accuracies are presented below:

As the result indicates, the looped transformer demonstrates comparable, and in some cases, better performance to the standard transformer in solving these OpenML datasets. We have updated the draft with this experiment results, and we believe these experiments offer a more realistic representation of practical applications.

In terms of further real-world applications, consider predicting user / customer behavior based on their historical data. Each customer represents a different context, and the optimal model form (linear, nonlinear, neural network, etc.) is often unknown. A transformer could learn an effective model for user data and then apply this learned model to new users. This is an area of ongoing investigation in our work, and we look forward to reporting our findings in the future.

Thank you for your valuable feedback.

relevance of looped transformers to in-context learning: It’s not entirely clear why this reduction in parameters is crucial for in-context learning, especially in cases involving linear functions, sparse linear functions, random decision trees, and 2-layer ReLU networks

We do not claim that a reduction in parameters is key to success in in-context learning. Instead, we aim to introduce a promising training method for a looped transformer, specifically tailored for in-context data fitting problems. The advantages of applying a looped transformer are twofold: 1) it reinforces the iterative nature commonly employed in solving data-fitting problems, and 2) it offers the benefit of saving parameters.

lacks in-depth mathematical insights: What is the mathematical insight of looping transformers for in-context learning?

From the perspective of the expressive power of looped transformer: it has been extensively explored. Bai et al., Theorem 3 demonstrated that any traditional $L$-layer neural network can be represented by a weight-tied, input-injected network, highlighting the expressiveness of such networks. Furthermore, Giannou et al. investigated how an encoder looped transformer can function as a universal computer. These studies provide solid theoretical foundations and constructions to demonstrate the significant expressiveness of transformers with looped architectures. Our work, instead, focuses on the empirical validation of this architecture, and presents a practical training method to effectively train a looped transformer architecture for in-context learning.

From the perspective of applying a looped transformer to in-context data fitting problems: the most common and effective training algorithms are iterative and recursive in nature, such as gradient descent. Our insight is that equipping transformer architectures with a recursive (looped) structure enables them to learn policies for in-context learning that leverage the advantages of recursion.

We hope our response address your concerns. We have updated the draft to include the discussion. Thanks again for your valuable insights and comments to our paper.

Hi reviewer gzpZ, we would like to inquire if our recent responses and revisions have adequately addressed your concerns regarding our ICLR submission. Could you please share your thoughts on the current rating of our paper? We value your feedback and appreciate your time and effort in this process.