Paper: Looped Transformers are Better at Learning Learning Algorithms (ICLR-2024)

Reviewer comment (Experiments):
The application of the looped transformer (in order to match the performance of a standard transformer) requires 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.

Author response:
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$.

Reviewer comment (Presentation):
The structure of the paper can be improved, e.g., by ending the paper with a short conclusion instead of the related works.

Author response:
Thanks for your suggestions. We have modified the draft accordingly.

Reviewer comment (Experiments):
What is the experimental setting for the plots shown in Figure 2?

Author response:
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.

Reviewer comment (Evaluation):
Why are the curves for the looped transformer and the standard Transformer in Figure 2 exactly overlapping?

Author response:
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.

Reviewer comment (Experiments):
How much time and resources does it take to find optimal values for b and T?

Author response:
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.

Reviewer comment (Experiments):
If you consider the hyperparameter search is it still more useful to apply looped transformers instead of the standard transformers?

Author response:
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.

Reviewer comment (Experiments):
Since all experiments are made on simulated datasets, whether the proposed method is effective in dealing with real-world data remains to be validated.

Author response:
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](https://www.openml.org/), which we believe better represent practical scenarios:


| dataset id | Name | # numerical features | # instances |
| -------- | -------- | -------- | -------- |
| 720     | abalone     | 7     | 4177 |
| 725     | bank8FM     | 8     | 8192 |
| 737     | space_ga     | 8     | 3107 |
| 803     | delta_ailerons     | 5     | 7129 |
| 807     | kin8nm     | 8     | 8192 |
| 816     | puma8NH     | 8     | 8192 |
| 819     | delta_elevators     | 6     | 9517 |
| 847     | wind     | 14     | 6574 |
| 871     | pollen     | 5     | 3848 |
| 42192     | compas-two-years     | 7     | 3848 |

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: 

| test dataset id | train dataset ids | vanilla TF | looped TF |
| -------- | -------- | -------- | -------- |
| 720 | S \ {720} | 0.626 $\pm$ 0.008  | **0.662 $\pm$ 0.008** |
| 725 | S \ {725} | 0.511 $\pm$ 0.007    | 0.504 $\pm$ 0.008 |
| 737 | S \ {737} | 0.656 $\pm$ 0.006   | **0.72 $\pm$ 0.01** |
| 803 | S \ {803} | 0.394 $\pm$ 0.01    | 0.40 $\pm$ 0.01 |
| 807 | S \ {807} | 0.405 $\pm$ 0.004   | 0.416 $\pm$ 0.005 |
| 816 | S \ {816} | 0.463 $\pm$ 0.004   | 0.462 $\pm$ 0.004 |
| 819 | S \ {819} | 0.483 $\pm$ 0.005   | **0.568 $\pm$ 0.01** |
| 847 | S \ {847} | 0.668 $\pm$ 0.007   | **0.757 $\pm$ 0.006**|
| 871 | S \ {871} | **0.532 $\pm$ 0.004**   | 0.51 $\pm$ 0.005 |
| 42192 | S \ {42192} | 0.65 $\pm$ 0.005  | 0.65 $\pm$ 0.008 |

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.

Reviewer comment (Theory):
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.

Author response:
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](https://www.openml.org/), which we believe better represent practical scenarios:


| dataset id | Name | # numerical features | # instances |
| -------- | -------- | -------- | -------- |
| 720     | abalone     | 7     | 4177 |
| 725     | bank8FM     | 8     | 8192 |
| 737     | space_ga     | 8     | 3107 |
| 803     | delta_ailerons     | 5     | 7129 |
| 807     | kin8nm     | 8     | 8192 |
| 816     | puma8NH     | 8     | 8192 |
| 819     | delta_elevators     | 6     | 9517 |
| 847     | wind     | 14     | 6574 |
| 871     | pollen     | 5     | 3848 |
| 42192     | compas-two-years     | 7     | 3848 |

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: 

| test dataset id | train dataset ids | vanilla TF | looped TF |
| -------- | -------- | -------- | -------- |
| 720 | S \ {720} | 0.626 $\pm$ 0.008  | **0.662 $\pm$ 0.008** |
| 725 | S \ {725} | 0.511 $\pm$ 0.007    | 0.504 $\pm$ 0.008 |
| 737 | S \ {737} | 0.656 $\pm$ 0.006   | **0.72 $\pm$ 0.01** |
| 803 | S \ {803} | 0.394 $\pm$ 0.01    | 0.40 $\pm$ 0.01 |
| 807 | S \ {807} | 0.405 $\pm$ 0.004   | 0.416 $\pm$ 0.005 |
| 816 | S \ {816} | 0.463 $\pm$ 0.004   | 0.462 $\pm$ 0.004 |
| 819 | S \ {819} | 0.483 $\pm$ 0.005   | **0.568 $\pm$ 0.01** |
| 847 | S \ {847} | 0.668 $\pm$ 0.007   | **0.757 $\pm$ 0.006**|
| 871 | S \ {871} | **0.532 $\pm$ 0.004**   | 0.51 $\pm$ 0.005 |
| 42192 | S \ {42192} | 0.65 $\pm$ 0.005  | 0.65 $\pm$ 0.008 |

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.

Reviewer comment (Reproducibility):
Can you provide more details about the probability distribution over the used classes of functions, especially for decision trees and 2-layer ReLU NN?

Author response:
> Can you provide more details about the probability distribution over the used classes of functions, especially for decision trees and 2-layer ReLU NN?

We generate function classes following the approach outlined by [Garg et al.](https://arxiv.org/abs/2208.01066):

1. Linear function class: We sample $w \sim \mathcal{N}(0, I)$. 
2. Sparse linear function: We sample $w \sim \mathcal{N}(0, I)$, and randomly select $s=3$ coordinate to be non-zero.
3. Decision tree: We utilize a binary decision tree with a depth of 4. For each non-leaf node, the decision to move left or right is based on the sign of one randomly selected coordinate of the input $x$. This coordinate is chosen uniformly at random. The leaf nodes are initialized following a normal distribution $\mathcal{N}(0, I)$. The evaluation of an input $x$ involves traversing from the root to a leaf node, moving left for a negative coordinate and right for a positive one, with the output being the value of the reached leaf node.
4. 2-layer ReLU neural network: We initialize the weights of the network following $\mathcal{N}(0, I)$, and set the hidden layer dimension to 100.

We have modified the draft to make it more clear as well.

Thank you once again for your valuable feedback, and we hope that our response adequately addresses your concerns.