Looped Transformers are Better at Learning Learning Algorithms
Authors: Liu Yang
Kangwook Lee
Robert D Nowak
Dimitris Papailiopoulos
Venue: ICLR 2024 poster (2024)
Keywords: in-context learning
transformers
looped transformers
Primary area: unsupervised, self-supervised, semi-supervised, and supervised representation learning
TL;DR: We train a looped transformer from scratch to perform in-context learning of simple function classes. Empirical results indicate the looped transformer can match or outperform the standard transformer.

Abstract: Transformers have demonstrated effectiveness in in-context solving data-fitting problems from various (latent) models, as reported by Garg et al. (2022). However, the absence of an inherent iterative structure in the transformer architecture presents a challenge in emulating the iterative algorithms, which are commonly employed in traditional machine learning methods. To address this, we propose the utilization of looped transformer architecture and its associated training methodology, with the aim of incorporating iterative characteristics into the transformer architectures. Experimental results suggest that the looped transformer achieves performance comparable to the standard transformer in solving various data-fitting problems, while utilizing less than 10% of the parameter count.

Published as a conference paper at ICLR 2024
LOOPED TRANSFORMERS ARE BETTER AT LEARNING LEARNING ALGORITHMS
Liu Yang, Kangwook Lee, Robert D. Nowak & Dimitris Papailiopoulos University of Wisconsin, Madison, USA
kangwook llee, rddcwak)ewisc..ed edu, dimitris@papail ii
ABSTRACT
Transformers have demonstrated effectiveness in in-context solving data-fitting problems from various (latent) models, reported by Garg et al. (2022). However, the absence of an inherent iterative structure in the transformer architecture presents a challenge in emulating the iterative algorithms, which are commonly employed in traditional machine learning methods. To address this, we propose the utilization of looped transformer architecture and its associated training methodology, with the aim of incorporating iterative characteristics into the transformer architectures. Experimental results suggest that the looped transformer achieves performance comparable to the standard transformer in solving various data-fitting problems, while utilizing less than 10% of the parameter count
INTRODUCTION
Transformers (Vaswani et al. 2017: Brown et al. 2020: Devlin et al., 2019 have emerged as
the preferred model in the field of natural language processing (NLP) and other domains requiring sequence-to-sequeen modeling. Besides their state-of-art performance natural language processing tasks, large language models (L.M)) such as GPT-3 (Brown et al. (020) and P (Chowdhery et al. 2022) also exhibit the ability to learn in-context: they can adapt to various downstream tasks based on a brief prompt, thus bypassing the need for additional model fine-tuning. This intriguing ability of in-context learning has sparked interest in the research community, leading numerous studies (Min)
et al. 2022 Olsson et al. 2022 Li et al. 2023). However, the underlying mechanisms enabling these transformers to perform in context learning remain unclear.
In an effort to understand the in-context learning behavior of L.LMs, Garg et al. 2022) investigated the performance of transformers, when trained from scratch, in solving specific function class learning problems in in-context. Notably, transformers exhibited strong performance across all tasks, matching or even surpassing traditional solvers. Building this, /Aküürk et al. (2022) explored the transformer- based model's capability to address the linear regression learning problem, interpreting it as an implicit form of established learning algorithms. Their study included both theoretical and empirical perspectives to understand how transformers learn these functions. Subsequently, /von Oswald et al. (0022) demonstrated empirically that, when trained to predict the linear function output, a linear self-attention-one transformer inherently learns to perform single step of gradient descent to solve the linear regression task in-context. While the approach and foundational theory presented by von Oswald et al. (2022) are promising. there exists significant gap between the simplified architecture they examined and the standard decoder transformer used in practice. The challenge of training a standard decoder transformer from scratch, with only minor architectural modifications, to effectively replicate the learning algorithm remains an open question.
In traditional machine learning, iterative algorithms are commonly used to solve linear regression. However, the methodologies employed by standard transformers are not naturally structured for iterative computation. A looped transformer architecture, extensively studied in the literature such as Giannou et al. (2023), provides a promising avenue to bridge this gap. In addition to its inherent advantage of addressing problem-solving in an iterative manner, the looped transformer also breaks down tasks into simpler subtasks, potentially leading significant savings in model parameters.
To illustrate how task breakdown leads to saving parameters, let's look closely at using the trans- former model to solve linear regression task, specifically w min ||x y (Fig. Dp. To train a transformer this task, we input prompt sequence formatted as (æ1, w" x1..... ... Lk, w æx, Frst)) Here w represents the parameters of the linear regression model, {x1.... x are in-context sam- ples, and Besa is the test sample. The transformer can potentially try predict Yeest by approximating the ordinary least squares solution in a single forward pass. The computation of the matrix inverse, as required in the ordinary least squares solution (X X) XTy, y, is more difficult for transformers to
Our code is available at https::/gittubb.co com/ i ay ay transpefeerreeereme

---

Published as a conference paper at ICLR 2024
2022). This is attributed to the increased number of layers and heads required in the inverse operation (Giannou et al. (223). Nevertheless, gradient descent offers alternative solution to linear regression, which requires only the matrix multiplication: X (Xw - y). but is applied repeatedly.
learn compared to matrix multiplication (Charton 2021 von Oswald et al.
Problem Sctup
ourput
82
Solve the lincar least square problem using decodcr-auly Transfomerr (TF). (x.]2.. each x, w'x,
Learned TF Salver
prompi
X y2 N2 Nea)
Trans former
This work:
Looped Transformer
Leurned TF Solver
closed-form solution, e.g.
iterative alporithm, C.D. gradicut descent:
Truditional Solver
(X'X)-XXy..
W4s1 + W, -1------ (X -N,
How can a transformer be trained to learn an iterative learning algorithm? Here we con- sider the task of training transformer to solve linear regression in context. The provided prompt (æ11, y1; 222, y/.... 3k,, Mes Weest) is fed into decoder transformer. The objective is to reduce the squared loss between the predicted thea based on this prompt, and the target value ((thea) Garg et al. (2022) demon- strated that a decoder transformer can learn to solve linear regression, which potentially involves learning the approximation of the least squares solution. In this study, we aim to train transformer to learn iterative learning algorithms. Our goal is to achieve performance par with standard transformers but with fewer parameters, To this end, introduce the looped transformer architecture and its accompanying training methodology.
Figure I:
Motivated by this observation, we ask the following question:
Can looped transformers emulate iterative learning algorithms more efficiently than standard, non-recursive transformers?
Within the specific function classes tested, the answer is positive. Our preliminary findings on using looped transformer to solve the linear regression task are illustrated in Fig. [ The remainder of
the paper is organized as follows. In Sec. 4 we develop training method for the looped transformer to emulate the desired performance of the iterative algorithm. Subsequently, in Sec. [ we compare
the empirical performance of the standard and the looped transformer and analyze the trade-off
between them. Our contributions and findings are outlined below:
Training methodology for Looped Transtormer. We propose a training methodology for looped transformers, aiming to effectively emulate iterative algorithms. The assumption for looped transformer simulating convergent algorithm is as loop iterations increase, the performance of the looped transformer should improve or converge. In aligmment with this assumption, we delve into the structural design of the looped transformer, as well as investigate the number of loop iterations required during training. These investigations lead to the formulation of our training method.
Performance of Looped Transformers for in-context Learning. Based our proposed training algorithm, empirical evidence demonstrates that looped transformer can be trained from scratch to in-context learn data generated from linear functions, sparse linear functions, decision trees, and 2-layer neural networks. Among the varied function classes examined, the looped transformer consistently outperforms the standard transformer, particularly when data is generated from sparse linear functions or decision trees. Our findings hint at the possibility that the looped transformer is more effective at in-context emulating learning algorithms, specifically for the learning tasks explored in this paper.
squared error
0.
06
Least squares
squared
Looped transformer Transformer
10-3
30
80
100
10 20 20
loop Nerrtions
in-conlext examples
Figure 2: The looped transformer can emulate iterutive learning algorithms, offering performance comparable to standard transformers with reduced parameters. We train looped transformer to solve lincar regression in-context. (Left): While trained for 30 loop iterations, the looped transformer during inference achieves stable fixed fixed-point solution beyond the trained loop iterations. (Right): The trained looped transformer matches the performance standard 12-layer transformer and closely aligns with the least squares solver, while using only 1/12 of the transformer's parameters.
2

---

Published as a conference paper at ICLR 2024
2 RELATED WORKS
Weight Sharing Models. Weight sharing is inherent to recurrent models such as RNN (Bengio al. 2003) and LSTM (Hochreiter & Schmidhuber 1997). These sequence models reuse the weights in the direction of sequence processing, enabling possibly infinite lengths of sequence inputs. The transformer model (Vaswani et al. 2017 Devlin et al.) 2019 typically has fixed context length. To accommodate a longer context without mcreasing parameters, several works (Dai et al. 2019; Wang et al. 2019; Hutchins et al. 2022) employed transformers
recurrent framework to enable an adaptive forward pass (Dai et al. (019)) or extend the context length (Hutchins et al. (022). Moreover, several works (Lan et al. 2019; Jaegle et al. 2021; Dehghani et al. 2018) also proposed weight sharing along the forward pass of the model. This weight sharing is either across a limited number of layers, or with a halting criterion to exit the recurrence if needed. At one extreme is the implicit model (Bai et al.) 2019), where the function is repeatedly applied an infinite number of times. The applications of these implicit models are addressed in the next paragraph.
Deep Implicit Model. Deep Implicit Models (Bai et al.) 2018: 2019; 2020: Winston Kolter 2020; Bai et al.] 2022) employ black-box solvers to find fixed-point solutions for implicit deep models. Later, |ai et al. (202T) proposed the Jacobian regularization to stabilize the training process. Nevertheless, this approach requires extensive hyperparameter tuning and still suffers from instability challenges. Implicit models have been demonstrated to solve math problems with extrapolation ability (Decugis et al. 2022). This is a special case of using the recurrent model to solve math tasks, as we will examine more closely in the next paragraph.
Using Transformer to Solve Math or Reasoning Tasks. Several works (Ongie et al. 2020
Doncevic et al. 2022; Zhang et al. 20222] Zhou et al. 2002bb Wei et al. 2022 Zhou et al.) 2022a
Bansal et al. [022] Charton 2022 Goel et al. 2022 /hhng et al. 2023b Dziri et al. 2023 Lee et al. 2023 Liu et al. (022) have explored the ability of deep neural networks, including recurrent models, in solving mathematical or reasoning tasks that involve iterative processes. These tasks involve finite states choices. Moreover, Zhang al. (2023b) studied how transformers can emulate the structural recursion problem, while Ongie et al. (020) investigated the use of recursive models in solving inverse problems in imaging. The regression problem, on the other hand, involves a continuous output space. Garg et al. (2022) investigated the transformer's 's ability to learn the continuous function classes. Akyürek et al. (2022) studied the phase transition of the solution obtained by transformers of varying depths when trained on linear regression and noisy linear regression problems. [Raventos] et al. (2023) studied the relationship between the number of tasks encountered during training and the performance of the transformer learning ridge regression problems in-context. Our work, however, utilizes the same architecture as in Garg et al. (2022)-a decoder model with non-linear attention and aims to provide empirical evidence and insights into the learning dynamics of this model.
Understanding How Transformer Learns In-Context to Solve Data-Fitting Problem. In ad- dition to approaching this problem from an empirical perspective, several studies have examined from a theoretical standpoint, focusing on the construction of transformers and the linear attention network (excluding softmax counterparts). They connected the learnability of the transformer for the linear regression task with the implicit gradient descent updates on query prompt predictions ( (von) Oswald et al. 2022 Ahn et al. 2023 Zhhng et al.] [023a] Mahankali et al.. 2023 Ding et al., 2023). von Oswald et al. (2022) empirically demonstrated that when minimizing the &g distance between predictions and true labels, one step of gradient descent is the optimal solution for a one-layer linear self-attention (LSA) transformer learn when solving the linear regression problem, while Ahn et al.)
(2023): y: Zhang et al. (2023a: 2024) and Mahankali et al. (2023) provided theoretical explanations for it. Beyond 1-layer LSA, Fu et al. (2023): Giannou et al. (2024) have demonstrated that multi-layer transformers can perform higher-order optimization. Moreover, another line of work (Muller et al.)
202T Xie al., 202T Akyūrek et al. 2022 Bai al. 2023) provided a theoretical understanding of the transformer's 's in in-context learnability through the lens of the Bayesian estimator. Li et al.
(2023) studied the in-context learning from the viewpoint of generalization error. Lin & Lee (2024) investigated the dual operating modes of in-context learning.
3 PROBLEM SETTING
Let F denote a class of functions defined on Rd. Let Dz denote probability distribution over F and let Dx denote a probability distribution on Rd. A random learning prompt P is generated as follows. A function f is sampled from D, and inputs {x1..... .x,) well the test sample Deest are sampled from Dx. The output of x; is computed by f(x;). The prompt is then =
(m), (11)) Lk,, I (xx)) ttsst) and the goal of leamiing system is predict (x Let M
3

---

Published as a conference paper at ICLR 2024
be learning system and let (P) denote its prediction (note it is not given f explicitly). The performance of M is measured by the squared error (( = (M se In this work, focus on transformer-based learning systems and compare them with other known learning systems depending on the tasks. Specifically, we examine the GPT-2 decoder model (Radford et al. 2019) with I. layers. By default, L = 12 for the unlooped transformer, following Garg et al. (2022) 's setting, and L = 1 for the looped transformer.
4 TRAINING ALGORITHM FOR LOOPED TRANSFORMER
In this section, we delve into the design choice for the algorithm-emulane looped transformer. For an Algorithm-emmlater looped transformer, we anticipate the following characteristics: 1) As loop iterations progress, the looped transformer should maintain or improve the performance; 2) the loop iterations have the potential to continue indefinitely without deterioration in performance.
Training Strategy. Building on the problem setting in Sec.B) let the prompt to the transformer be P- (z f21), .... Lk. fxx), Bteat), with fs denotes the prompt prefix with in-context samples P = ( f(x1). ... Dz, (2:)). xi+1). The output of the looped transformer after t looping iterations is
(P(P))0) = MM Ma Yo + Pi) + Pi). ......+PP' Pi)
iterations
where the transformer M is parameterized by 0, and Yo is zero tensor with the same shape as pi Then we train the transformer by minimizing the following expected loss:
min Ep
Σ (Y((P"))) T=0
bo
where we only measure the loss of the transformer over all prompt prefixes with loop iteration bo to b,
with bg = max(b T' O). This truncated loss is inspired by the truncated backpropagation through time (Hochreiter & Schmidhuber 1997 Hinton & Sutskever 2013) for computation saving.
Model Configuration and Parameter Count. For comparison with the unlooped transformer. we use a transformer identical to the one described in Garg et al. (022). except for the number of layers. Specifically, we employ a GPT-2 model with an embedding dimension of D 256 and = attention heads. The standard (unlooped) transformer has L = 12 layers, and the looped transformer has L = 1 layer. In terms of the number of parameters, the unlooped transformer comprises 9.48M parameters, and the looped transformer uses 0. 0.79M. We follow the same training strategy: train with Adam optimizer, learning rate 0.0001, no weight decay or other explicit regularization (such as gradient clip, data augmentation).
Key Factors for Finding a Fixed-point Solution. When deploying this looped architecture training, two key factors come into consideration: 1) the input injection in the looped architecture (Sec. 4.1). and 2) the maximum loop iterations during training (Sec. (42)). The subsequent sections will illustrate the impact of these two factors, using linear regression as the specific task for in-context learning.
4.1 INPUT INJECTION TO LOOPED TRANSFORMER
Reusing some notation, let P represent the inputs to the transformer model (-), and Y, be the output after applying M (-) for iterations. A looped transformer can be represented as Y = M(;; P), Vt. Several studies (Lan al. 2019 Dehghani et al., 2018) have investi- gated a specific variant termed the weight-tying form: Y.+1 = (Y.) with Yo- P, and / T < oo for some constant T.
eguanne arror
40
66
130
loop itorations
However, as t approaches infinity, the influence of the initial input Yo diminishes. and the solution becomes essentially random or un- predictable (Bai et al. 2019; Bansal et al., [002). To incorporate the input P into the solution of the looped transformer, we propose setting Y = MY, + P). but the choice of injecting is not limited to addition. In Fig. 3 we display the outcomes when training a looped transformer either with (default) input injection without (referred to as weight-tying). During inference, the weight-tying transformer will quickly diverge after the trained loop iterations.
Figure 3: Test error for the linear regression problem us- ing looped transformer, com- paring models with default in- put injection to those without. Without input injection, the transformer's 's performance de- teriorates beyond the trained loop iterations.
4.2 CHOICE OF LOOP ITERATIONS
To assess the looped transformer performance at the b-th loop iteration, we must execute this computation b times consecutively. This can be analogized to transformer architecture possessing b
4

---

Published as a conference paper at ICLR 2024
layers in effect, albeit with shared weights throughout these layers. Choosing the value of b requires balance. A higher b leads to longer training and inference time, whereas a lower b limits the model's potential. Furthermore, denoted in Eq. [ the loss is truncated, with only the outputs of I loop iterations contributing to the loss function.
Similarly, there's trade-off when choosing T: smaller T results in reduced memory usage but negatively impacts performance, whereas a larger T may cause the gradient to fluctuate, making the training process unstable. This section focuses on the optimal selection of b and T values (Fig. 4 for the linear regression task. We analyze the suitability of & values within the set {12, 20, 30, 40, 50}. and T values in {5, 10, 15) for the linear function.
10-1
[0
10
10-1
1.5
200
100
230
100
200
100g sterations
1oop iterations
loep iterations
Figure 4: Evaluation of the looped transformer on in-context learning linear functions with different 6 and T' during training (b and 7' are defined in Eq. [T.. The figure from left to right trained with T 5, 10, 15. and different colors present different values (denoted the legend). The solid lines of various colors depict how the looped transformer, trained with a specific value of b, performs as the loop iteration increases during inference. The corresponding dashed line represents the value of b.
Fig. 4 suggests that a lower b value might cause the looped transformer to diverge after the trained loop iterations, leading to a less robust fixed-point solution. On the other hand, higher value allows the model to locate a exed-point solution that avoids divergence after the trained loop iterations. However, exceeding a certain b value initiates decline in the convergence rate of the loop iteration.
similar trade-off applies the selection of T'; a lower I' may compromise performance, whereas a larger 7' could induce instability in training, a phenomenon documented in the recurrent model training literature [Jaeger] 2005). A looped transformer with 12 matches the effective depth of the unlooped transformer studied in Garg et al. (2022). Nevertheless, it fails to replicate the performance of the standard (unlooped) transformer.
An additional observation is that the looped transformer consistently discovers a fixed-poin solution that saturates prior to the trained iteration b. This saturation of the fixed-point solution occurs due to the loss objective, which requires the looped transformer to match the target within b steps. Consequently, selecting smaller value of b value expedites convergence to the fixed-point solution. However, beyond certain value of b. the convergence rate reaches saturation regardless of the increase in b. For instance, training with T 15, b = 40, 50 yields similar convergence rates.
Optimal Choice of and T. Our goal is to efficiently train the looped transformer for the in- distribution task. This requires determining the minimal b value that prevents divergence after training loop iterations. To minimize memory usage, we prefer a smaller T value. Guided by these criteria, we adopt b 20 and T 15 for the linear regression task.
5 EXPERIMENTAL RESULTS
5.1 EXPERIMENTAL SETUP
We focus the data-fitting problems generated by the linear model with problem dimension d 20, and in-context sample & = 40. The parameters are sampled from (0, Ta), and the in-context samples o; ~ (00 Ia) as well. When measuring the performance. we evaluate 1280 prompts and report the 90% confidence interval over 1000 bootstrap trials. To ensure the error is invariant to the problem dimension, we also do normalization the error as specified in Garg et al. (2022).
Scheduled Training. We follow the training curriculum in terms of d and k. as specified in Garg ct al. (2022). Besides, we also implement a schedule on the parameter b, where is progressively increased during the course of training. As consequence, the truncated loss window starts at [1, T] and then shifts over time to [b- - T, b.. It is worth noting that minimizing the mean squared error between the T] loop iteration and the target can be more challenging than for [b T.
Consequently, we avoid referring to this approach as "curriculum learning." as that term typically refers to starting with less difficult task and gradually moving to more challenging ones. Rather,
5

---

Published as a conference paper at ICLR 2024
13
1.5
12.5
squared crror
3.2
35.01
squared
7.3
3.5
1.0
Teep Roavers
3.0
0.5
10
26
30
20 36
10 In-context examples
20
45
In-context exmmples
in-context examples
(b) Noisy Linear Regression Figure 6: Evaluation of the transformer trained on linear functions, tested under various conditions: a) inputs with skewed covariance, b) noisy linear functions, and c) scaled inputs.
(a) Skewed Covariance
(c) Scaled with scaling=3
this strategy be seen method kick-start the training of the transformer model in loop structure. In general, the decision to use or not to use the scheduling does not significantly impact the outcome. However, in some instances, schedule for b. More details can be found in Appendix
observe superior results achieved by training with
5.2 LOOPED TRANSFORMER CAN IN-CONTEXT LEARN LINEAR REGRESSION
Looped Transformer Matches the Performance of Standard Transformer when Evaluating In- Distribution. As indicated in Fig. (right), the looped transformer trained with 20 and T 15 is able to match the performance of the standard transformer, almost matching the performance of the traditional optimal solver: the ordinary least squares.
Looped Transformer Learns Efficiently with Fewer Distinet In- Distribution Prompts. We further evaluate the sample complexity of transformer training. Specifically, we ask how many distinct prompts need to be seen during training for the transformer or the looped transformer to learn the function. To do so, instead of generating the linear function on the fly in each iteration, we generate the training dataset before training. Thus, during training, the transformer may encounter the same prompt multiple times.
12-
Tra excee
1.8
0.8.
0.2
1k distioct prompssfuatiion seen durlog troining
To isolate the impact of curriculum learning. during and focus on the linear regression task with problem dimen- sion d 10, in-context samples k 20. The results are presented in Fig. 5 Due to the fewer parameters in the looped transformer, is able to learn the function with fewer distinct prompts/functions pared to the standard transformer. More results regarding different problem dimensions are presented Appendix C
disable this strategy
Figure 5: Performance of trans- former on linear functions with = 10 and = 21, when trained with different numbers of distinct prompts/functions.
Looped Transformer Exhibits a Certain Inductive Bias when Evaluating
Recall that the transformer is trained with æ ~ (O, Ia). We now evaluate it on: a) inputs with skewed covariance, b) inputs with noise, and e) in-context example inputs and query inputs that lie in different orthants. The result is shown in Fig. 6 Rather than mastering the true algorithm applied to the linear regression problem irrespective shifts in the input distribution, the looped transformer exhibits an inductive bias favoring simpler solutions when compared to the unlooped transformer. This results in improved handling of (a) skewed covariance inputs.
In the task of (b) noisy linear regression, the looped transformer displays reduced peak in the double descent curve, similar to the effects of applying minor regularization to the base model for resolving noisy linear regression. [Bhattamishra al. (2022) also reported the simplicity bias of the random transformer towards low sensitivity in the boolean function base. However, this inductive bias towards simplicity can negatively impact performance when there is scaling shift in the input distribution. As depicted in Fig. (c), during the inference process, if we sample æ such that æ 3z, and (0, the looped transformer underperforms compared the unlooped transformer. More extrapolation results can be found in Appendix D
5.3 IMPACT OF MODEL ARCHITECTRRE VARIATIONS ON LOOPED TRANSFORMER PERFORMANCE
In this section, we explore the impact of varying the number of layers (L), heads (h), and embedding dimension (D) on the looped transformer. These experiments
trained with b = 30, T == 15.
Varying the Number of Layers. In Fig. 7 (left), we plot the squared error for the transformer with L layers, and the looped transformer with L layers, applying t loop iterations. From the figure, observe that as L increases, convergence achieved more rapidly. To match the performance of a standard transformer with T. layers, the looped transformer needs more than T. loop iterations, i.e. the

---

Published as a conference paper at ICLR 2024
10F
squared orror
10:11
D-53
10
20
20
40 60
20
20
100
loop Iterations
loap iterations
Figure 7: Let D be the transformer embedding dimension, L be the number of layers, and / be the number of heads. For linear functions with problem dimension d = 20, and in-context samples k 40, test the following: (left) the standard (unlooped) transformer and looped transformer with D 256 and 8, but varying I,, and (right) looped transformer with varying I. and D, h.
effective depth of the looped transformer exceeds that of the standard transformer. For instance, a looped transformer with single layer requires roughly 20 iterations in order to achieve an 8-layer transformer's 's performance. This suggests that the transformer and the looped transformer learn different representations at each layer (iteration).
To delve deeper into the learning dynamies of each layer in both the transformer and the looped transformer, we employ the model probing technique suggested by Akyürek et al. (2022); Alain & Bengio (0016). Reusing some notation, we represent the output of the standard transformer's s t-th layer or the looped transformer's loop iterations Y,. An MLP model is trained on Y, to learn target probes, specifically X y for the gradient component and the Wols for the optimal least squares solution. The mean squared error for model probing is illustrated in Fig. 8 While standard transformers initially capture representations relevant to X x y and Wals, and subsequently shift focus to other statistics, the looped transformer consistently refines its representations related to the target probes across its loop iterations. Details of the model probing are presented in Appendix F
Transformer
Looped Tran Transformer
X probe
probe
0.04
0.04
0.02
0.02
0.00
10
15
20
0.00
5
10
15
20
Layer/ Leeg Iteration
Layer / Loop Eterwtion
Figure 8: How do the transformer and the looped transformer encode information across layers/iteratios? We train 2-layer MLP probing model to recover the target probe from the transformer's output at 1-th layer/loop iteration. The dashed line indicates minimal probe error obtained in a control task, where the linear regression parameter w is fixed to be 1. Contrasting with the standard transformer, which decodes the target probe optimally around the 10-th layer, the looped transformer stcadily refines its representation, improving the decoding of the target probe with each subsequent loop iteration.
Varying the Number of Heads and Embedding Dimension. Fig. 7 (right) illustrates the squared error of the looped transformer for various combinations of I., D, and h. With I. = 1 and h 8 held constant, increasing D leads to improved convergence, saturating at D = 256. When holding D 256 and h = & constant and varying L, the looped transformer shows similar convergence performance. The primary differences are in the speed of convergence. Adjusting only h with fixed L. and D, the looped transformer's 's error decreases as h rises from to However, the error increases when h progresses from 8 to 32. Here, we follow the standard implementation of GPT-2 where the per-head embedding dimension is defined as D. This definition implies trade-off in the transformer between the number of heads and the per-head embedding dimension as la varies.
5.4 HIGHER COMPLEXITY FUNCTION CLASSES
In this section, we shift our focus to function classes of higher complexity: a) the sparse linear model with d = 20, & = 100, and non-sparse entry s = 3: b) the decision tree with depth= 4 and input dimension d 20; c) the 2-layer ReLU neural network with 100 hidden neurons and input dimension d = 20. For these functions, parameters are sampled from (0. Ia). and in-context samples 2; AO, Tz). This setup follows the methodology described in Garg et al. (2022), and is further detailed in Appendix[A] We also conduct experiments OpenML datasets, which better represents practical scenarios. Details of experiment setup and results for OpenML (Vanschoren et al. 2013) datasets are available in Appendix G
Optimal Loop Iterations in the Looped Transformer for Different Functions. In Scc. 4.2 we highlight that for the looped transformer trained on linear regression task, a larger value of 6 and an
httpp::/ thub. com www.kaapathy//aaaaeeeao
7

---

Published as a conference paper at ICLR 2024
appropriate T enhance its ability discover fixed-point solution that remains stable beyond the trained loop iterations. This observation consistent across various functions, but the optimal b and T values may vary. Fig. depicts the performance of looped transformer trained with T = 10 but varying b values on different functions.
Comparison with Fig. 4 indicates that tasks with higher complexity, like linear regression (which is more parameter-intensiie than sparse linear regression), necessitate increased and T values to achieve stable convergence during training. For instance, with T 10, linear regression diverges at b = 20, whereas sparse linear regression converges to a fixed point. The values of and ' required can be indicative of the complexity transformer faces when tackling specific task. Interestingly. while learning 2-layer neural network is more challenging (Chen et al., 2022), it requires fewer b and T for the looped transformer in comparison to seemingly simpler tasks, such as linear regression. Complete results of the looped transformer trained with various b and 7 are presented in Appendix B
50
12
20
30
59
10P
10
0.6-
10-1
0.4
10-1
120
150
200
10
100
150
200
0
10
1:00 150
200
loop Ieratiinn
loop Iterations
loop iterations
(a) Sparse Linear Regression.
(b) Decision Tree.
(c) 2-layer ReLU NN.
Figure 9: Evaluation of the looped transformeron in-context learning data generated from a) the sparse linear function, b) the random decision tree, and c) the 2-layer ReLU neural network with T = 10 and different b during training (b is defined in Eq. [T. The solid lines of various colors depict how the looped transformer, trained with a specific value of b. performs as the loop iteration increases during inference. The corresponding dashed line represents value of b.
Looped Transformer Matches hensive hyperparameter search, we identify optimal values for b and T, as detailed in Appendix B Specifically, for data generated from a) sparse linear functions, we use b 20 and T 10; b) decision trees, b = 70 and T = 15; and c) 2-layer ReLU NNs, 6 = 12 and T = 5. We then compare the performance of looped transformer against the standard transformer and other baseline solvers, as depicted in Fig. TO. Of all the tasks we examine, the looped transformer consistently matches the standard transformer, except for the sparse linear regression tasks, where the looped transformer outperforms the standard transformer, and even the Lasso (Tbbshirani 1996). By outperforming, we do not refer to the precision of the final solution but to the trend relative to the number of in-context samples. For instance, with 40 in- in-context samples, the Lasso achieves solution with an error of 1.32e-04, the Least Squares reaches 8.21e- the looped transformer achieves 6. 12e-04, and the standard transformer achieves 0.0017. However, with fewer than 10 in-context samples, the looped transformer has smaller error than the Lasso solution. Here we do a hyperparameter search for Lasso over a E {Le - 4, 1e 3, 0.01, 0.1, 1}. where a is the l parameter, and select the best a 0.01. We conjecture that this performance gain is a result of the inductive bias of the looped transformer, which favors simpler (i.c., sparser) solutions over their more complex counterparts. When tackling sparse problems, this inherent bias enhances the performance of the looped transformer.
Outperforms Standard Transformer. Through a compre-
Performance Analysis of Looped Transformer Across Varying Sparsity Levels in Sparse Linear Regression. The looped transformer demonstrates enhanced performance in the sparse linear regression task compared to the unlooped transformer. To further investigate this, we examine the task under varying levels of sparsity, which represent different task difficulties. To enable the looped transformer to handle all complexity levels, we choose parameters b 30 and T 15. As presented
3.00
0.00
6.00
mncenteat wammppex
53
22-coalent moumppee
(a) Sparse Linear Regression. Figure 10: Performance evaluation of the trained transformer on in-context learning data from: a) sparse linear functions, b) random decision trees, and c) l-aerer ReLU neural networks. The looped transformer matches or even surpasses the transformer's performance using only 1/12th of the parameters employed by the latter.
(b) Decision Tree.
(c) 2-1ayer ReLU NN.
8

---

Published as a conference paper at ICLR 2024
in Fig. I1 the results show that across all sparsity levels of the weight vector for the sparse linear function, the looped transformer consistently outperforms the unlooped transformer, especially when the number of in-context samples is limited.
Figure 11: Performance of louped transformer in solving sparse linear functions with problem di- mension d = 20. We examine sparsity levels $ $= 1, 6, 9, 12, where only s entries are non-zero zero. Across all sparsity levels, the looped transformer consistently matches outperforms the unlooped transformer, achieving more accurate solutions with fewer in-context samples.
squared error
1.00
Logged Transfeee
0.75
0.50
16 20 50
40
in-context examples
DISCUSSION AND FUTURE WORK
Mathematical Insights of Looped Transformers. The expressive power of recurrent models has been widely investigated in the literature: Theorem in Bai al. (2019) demonstrated that any traditional L-yrayer neural network can be represented by weight-tied, input-injected network. Further, Giannou et al. 2023) explored the potential of an encoder looped transformer to function as a universal computer. Our work concentrates on empirically validating the looped architecture, offering a practical training method for a looped decoder transformer in in-context learning. Investigating the expressiveness of this transformer presents a promising future research direction.
Looped Transformer Emulates Tixed-Point Iteration Method Within Training Distribution. Through our designed training method, the looped transformer successfully approximates fixed -oint solutions beyond the iterations it was trained on, effectively emulating the fixed-point iteration method within the training distribution. Nevertheless, these solutions come with an inherent error floor and have limited performance on out-of-distributioon prompts. As a result, our trained looped transformer fails to identify an actual algorithm that generalizes to any input distribution, but only emulates the algorithm with the particular input distribution it is trained on. Overcoming this limitation, and enhancing the mansformer-derieee solver to match the generalization capabilities of conventional algorithms presents a promising avenue for future research.
Choice of and T and the Notion of Task Difficulty. As discussed in Sec. Sec.55.4] increasing the loop iterations (b) and truncated loss window size (7)) enables the looped transformer to find a fixed-point solution that does not diverge beyond the trained loop iterations. The optimal values for and T are determined through hyperparameter tuning, selecting the minimum parameters necessary to prevent divergence in the looped transformer. These optimal values of b and T reflect the complexity of a task from the transformer's 's perspective. A task that is inherently more complex will require larger values of and I to achieve convergence to a fixed-point solution, whereas simpler tasks will require smaller values. Thus, the parameters b and T could be used as metric to assess the difficulty level of tasks for transformers. Expanding on this, in practice where tasks exhibit varying degrees of complexity, an adaptive looping strategy could be beneficial. These strategies could include token-level adaptive loop iterations, as proposed by Dehghani et al. (2018), or in-context multi-task inference.
Trade-off. The computational cost of training a looped transformer is tied to the selected values of b and T. Specifically, b determines the transformer's forward pass time, since the looped transformer has assess b-th loop iteration's output by unrolling b times. T' influences the memory requirements during training, given that the loss is computed over T' iterations. If the goal is approximate fixed-point solution using the looped transformer, one could circumvent intensive tuning of b and T. Instead, opting for a larger b, a suitably chosen T, and incorporating regularization strategies such as mixed-precision training, gradient clipping, and weight decay may be beneficial. These regularization strategies could potentially alleviate the instability issues associated with T. We omit these regularization methods to maintain a fair comparison with the transformer training methodologies outlined in Garg et al. (2022).
Choice of b and T and their Cemorr-Coomutatioon
CONCLUSION
Motivated by the widespread application of iterative algorithms in solving data-fitting problems across various function classes, we analyzed the loop architecture choice and training strategy, then developed a training method for the looped transformer. This methodology enables the looped transformer to approximate fixed-point solutions beyond the initially trained loop iterations. We hope this work can contribute to the community in understanding and developing better training strategies for transformers in solving data- fata-fitting problems.

---

Published as a conference paper at ICLR 2024
REFERENCES
Kwangjun Ahn, Xiang Cheng, Hadi Daneshmand, and Suvrit Sra. Transformers learn to implement preconditioned gradient descent for in-context learning, ArXiv, abs/2306.0022 00297, 2023. URL
https://app semant wwm.mnncccchhlarr.. org/Corpus ID: : 258999480
Ekin Akyürek, Dale Schaurmans, Jacob Andreas, Tengyu Ma, and Denny Zhou. What learning algorithm is in-context learning? investigations with linear models. ArXiv, abs/2211..156.11 2022.
URL htwp::
Guillaume Alain and Yoshua Bengio. Understanding intermeiite layers using linear classifier probes. ArXiv, abs/1610.01644, 2016. URL htt.tttts:///ppl sempp: swwant icscho ar ort.oorrrrruu..
9794990
Shaojie Bai, Zico Kolter, and Vladlen Koltun. Trellis networks for sequence modeling. ArXiv, abs/1810.06682, 2018.
Shaojie Bai, J. Zico Kolter, and Vladlen Koltun. Deep equilibrium models. ArXiv, abs/1909.01377, 2019.
Shaojie Bai, Vladlen Koltun, and J. Zico Kolter. Multiscale deep equilibrium models. ArXiv, abs/2006.08656, 2020.
Shaojie Bai, Vladlen Koltun, and Zico Kolter. Stabilizing equilibrium models by jacobian regular- ization. In International Conference on Machine Learning. 2021.
Shaojie Bai, Vladlen Koltun, and J. Zico Kolter. Neural deep equilibrium solvers. In International Conference on Learning Representations, 2022.
Yu Bai, Fan Chen, Haiquan Wang. Caiming Xiong, and Song Mei. Transformers as statisticians: Provable co-context learning with in-context algorithm selection. ArXiv, abs/2336.046311 2023. URL https::pp://pppt semant csccollall.ooo ar wwwp/org/rrrrrrooo ID: 259095794
Arpit Bansal, Avi Schwarzschild, Eitan Borgnia, Zeyad Ali Sami Emam, Furong Huang, Micah Goldblum, and Tom Goldstein. End-to-end algorithm synthesis with recurrent networks: Logical extrapolation without overthinking. ArXiv, abs/2202.05826, 2022.
Yoshua Bengio, Réjean Ducharme, Pascal Vincent, and Christian Janvin. A neural probabilistic language model. J. Mach. Learn. Res., 3:1 137-1155,, 2003.
S. Bhattamishra, Arkil Patel, Varun Kanade, and Phil Blunsom. Simplicity bias in transformers and their ability learn sparse boolean functions. ArXiv, abs/2211.123166 2022.
Tom B. Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared Kaplan, Prafulla Dhariwal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, Sandhini Agarwal, Ariel Herbert-Voss Gretchen Krueger, T. J. Henighan, Rewon Child, Aditya Ramesh, Daniel M. Ziegler, Jeff Wu, Clemens Winter, Christopher Hesse, Mark Chen, Eric Sigler, Mateusz Litwin, Scott Gray, Benjamin Chess, Jack Clark, Christopher Berner, Sam McCandlish, Alec Radford, Ilya Sutskever, and Dario Amodei. Language models are few-shot learners. ArXiv, abs/2005. 14165. 2020.
Franccois Charton. Linear algebra with transformers. Trans. Mach. Learn. Res., 2022, 2021.
Franccois Charton. What is my math transformer doing? three results interpretability and generalization. ArXiv, abs/221 1001700, 2022.
Sitan Chen, Aravind Gollakota, Adam R. Klivans, and Raghu Meka. Hardness of noise-free learning for two-hidden-layer neural networks. ArXiv, abs/2202.05255. 2022. URL https: :/aps:///ppp /api : 246706042
semant cwccolcccccooooo..cccccocccccc ar wwtp:aaro.rrooooooooeeeeeeaapeee
Aakanksha Chowdhery, Sharan Narang, Jacob Devlin, Maarten Bosma, Gaurav Mishra, Adam Roberts, Paul Barham, Hyung Won Chung, Charles Sutton, Sebastian Gehrmann, Parker Schuh, Kensen Shi, Sasha Tsvyashchenko, Joshua Maynez, Abhishek Rao, Parker Barnes, Yi Tay, Noam M. Shazeer, Vinodkumar Prabhakaran, Emily Reif, Nan Du, Benton C. Hutchinson, Reiner Pope, James Bradbury, Jacob Austin, Michael Isard, Guy Gur-Ari, Pengcheng Yin, Toju Duke, Anselm
10

---

Published as a conference paper at ICLR 2024
Levskaya, Sanjay Ghemawat, Sunipa Dev, Henryk Michalewski, Xavier Garcia, Vedant Misra, Kevin Robinson, Liam Fedus, Denny Zhou, Daphne Ippolito, David Luan, Hyeontaek Lim, Barret Zoph, Alexander Spiridonov, Ryan Sepassi, David Dohan, Shivani Agrawal, Mark Omernick, Andrew M. Dai, Thanumalayan Sankaranarayana Pillai, Marie Pellat, Aitor Lewkowycz, Erica Moreira, Rewon Child, Oleksandr Polozov, Katherine Lee, Zongwei Zhou, Xuezhi Wang. Brennan Sacta, Mark Díaz, Orhan Firat, Michele Catasta, Jason Wei, Kathleen S. Meier-Hellstern, Douglas Eck, Jeff Dean, Slav Petrov, and Noah Fiedel. Palm: Scaling language modeling with pathways. URL httph: api semanticscholo
J. Mach. Learn, Res., 24:240:1 1:::4:::120::1111 240: 2022.
http:/Coopul:::2:::9:
Zihang Dai, Zhilin Yang, Yiming Yang. Jaime G. Carbonell, Quoe V. Le, and Ruslan Salakhutdinov. Transformer-xl: Attentive language models beyond a fixed-length context. ArXiv, 0e2//990102860 2019.
Juliette Decugis, Max Emerling, Ashwin Ganesh, Alicia Y. Tsai, and Laurent El Ghaoui. On the abilities of mathematical extrapolation with implicit models. 2022.
Mostafa Dehghani, Stephan Gouws, Oriol Vinyals, Jakob Uszkoreit, and Lukasz Kaiser. Universal transformers. ArXiv, abs/1807.0333 190.0,0.9, 2018.
Jacob Devlin, Ming-Wer Wei Chang, Kenton Lee, and Kristina Toutanova. Bert: Pre-training of deep bidirectional transformers for language understanding. ArXiv, abs/1810.048055 2019.
Nan Ding, Tomer Levinboim, Jialin Wu, Sebastian Goodman, and Radu Soricut. Causallm is not optimal for in-context learning. 2023. URL Lat.s///o...seemnicccccoa.aaaa.
Corpus CorpusID: 10:66877222
Danimir T. Doncevic, Alexander Mitsos, Yu Guo, Qianxiao Li, Felix Dietrich, Manuel Dahmen, and Ioannis G. Kevrekidis. A recursively recurrent neural network (r2n2) architecture for learning iterative algorithms. ArXiv, abs/2211 1.121.1.2222 1386, 2022.
Nouha Dziri, Ximing Lu, Melanie Sclar, Xiang Lorraine Li, Liwei Jian, Bill Yuchen Lin, Peter West, Chandra Bhagavatula, Ronan Le Bras, Jena D. Hwang, Soumya Sanyal, Sean Welleck, Xiang Ren, Allyson Ettinger. Zaïd Harchaoui, and Yejin Choi. Faith and fate: Limits of transformers compositionaliiyy ArXiv, abs/2305.18654 2023. URL https:::///// /apppppp.../////
lrrg/Corpus 25:966673229979
Deqing Fu, Tian-Qi Chen, Robin Jia, and Vatsal Sharan. Transformers learn higher-order optimization methods for in-context learning: A study with linear models. arXiv preprint 10. 2023.
Shivam Garg, Dimitris Tsipras, Percy Liang, and Gregory Valiant. What can transformers learn in-context? case study of simple function classes. ArXiv, abs/2208.01066, 2022.
Angeliki Giannou, Shashank Rajput, Jy yong Sohn, Kangwook Lee, Jason Lee, and Dimitris Papailiopoulos. Looped transformers programmable computers. ArXiv, abs/2301.13196, 2023.
Angeliki Giannou, Liu Yang, Tianhao Wang, Dimitris Papailiopoulos, and Jason D Lee. How well can transformers emulate in-context newton's 's method? arXiv preprint arXiv:2403.03113, 2024.
Surbhi Goel, Sham M. Kakade, Adam Tauman Kalai, and Cyril Zhang, Recurrent convolutional neural networks learn succinct learning algorithms. ArXiv, abs/2209.00735, 2022.
Geoffrey E. Hinton and llya Sutskever. Training recurrent neural networks. 2013.
Sepp Hochreiter and Jürgen Schmidhuber. Long short-term memory. 1780- 1997.
Neural Computation, 9:
DeLesley S. Hutchins, Imanol Schlag. Yuhuai Wu, Ethan Dyer, and Behnam Neyshabur. Block- 2022.
recurrent transformers. ArXiv,
Herbert Jaeger. tutorial on training recurrent neural networks, covering bppt rtrl, ekf and the " echo state network approach semantic scholar. 2005.
11

---

Published as a conference paper at ICLR 2024
Andrew Jaegle, Felix Gimeno, Andrew Brock, Andrew Zisserman, Oriol Vinyals, and João Carreira. Perceiver: General perception with iterative attention. In International Conference on Machine Learning, 2021.
Zhenzhong Lan, Mingda Chen, Sebastian Goodman, Kevin Gimpel, Piyush Sharma, and Radu Soricut. Albert: A lite bert for self-superviser learning of language representations. ArXiv, abs/1909. 1942, 1994, 2019.
Nayoung Lee, Kartik Sreenivasan, Jason D Lee, Kangwook Lee, and Dimitris Papailiopoulos. Teaching arithmetic to small transformers. arXiv preprint arXiv:2307.0338/ 2023.
Yingcong Li, Muhammed Emrullah Iidiz, Dimitris Papailiopoulos, and Samet Oymak. Transformers as algorithms: Generalization and stability in in-context learning. 2023.
Ziqian Lin and Kangwook Lee. Dual operating modes of in-context learning. arXiv preprint arXiv:2402.1888 18819, 2024.
Bingbin Liu, Jordan T. Ash, Surbhi Goel, Akshay Krishnamurthy, and Cyril Zhang. Transformers learn shortcuts to automata. ArXiv, abs/2210.107499 2022.
Arvind V. Mahankali, Tatsunori Hashimoto, and Tengyu Ma. One step of gradient descent is provably the optimal in-context learner with one layer of linear self-attention. ArXiv, abs/2307.033 2023.
Sewon Min, Xinxi Lyu, Ari Holtzman, Mikel Artetxe, Mike Lewis, Hannaneh Hajishirzi, and Luke Zettlemoyer. Rethinking the role of demonstrations: What makes in-context learning work? In Conference Empirical Methods in Natural Language Processing, 2022.
Samuel Muller, Noah Hollmann, Sebastian Pineda Arango, Josif Grabocka, and Frank Hutter. Transformers can do bayesian inference. ArXiv, abs/21 10510 2021.
Catherine Olsson, Nelson Elhage, Neel Nanda, Nicholas Joseph, Nova DasSarma, T. J. Henighan, Benjamin Mann, Amanda Askell, Yuntao Bai, Anna Chen, Tom Conerly, Dawn Drain, Deep Ganguli, Zac Hatficld-oodds, Danny Hernandez, Scott Johnston, Andy Jones, John Kernion, Liane Lovitt, Kamal Ndousse, Dario Amodei, Tom B. Brown, Jack Clark, Jared Kaplan, Sam McCandlish, and Christopher Olah. In-context learning and induction heads. ArXiv, abs/2209. 2022.
Greg Ongie, Ajil Jalal, Christopher A. Metzler, Richard Baraniuk, Alexandros G. Dimakis, and Rebecca M. Willett. Deep learning techniques for inverse problems in imaging. IEEE Journal Selected Areas in Information Theory, 1:39-56, 2020.
Alec Radford, Jeff Wu, Rewon Child, David Luan, Dario Amodei, and Ilya Sutskever. Language models are unsupervised multitask learners. 2019.
Allan Raventos, Mansheej Paul, Feng Chen, and Surya Ganguli. The effects of pretraining task diversity on in-context learning of ridge regression. In ICLR 2023 Workshop on Mathematical and Empirical Understanding of Foundation Models, 2023.
Robert Tibshirani. Regression shrinkage and selection via the lasso. Journal of the royal statistical society series b-methodological, 58:267-288, 1996.
Joaquin Vanschoren, Jan N. van Rijn, Bemnd Bischl, and Luis Torgo Openml: networked science in machine learning. SIGKDD Explorations, 15(2):49 60, 2013, doi: 10.1145/264119022641] URL htRl:/ttp://doo acm wwwp///./......0
90. 264119
Ashish Vaswani, Noam M. Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N. Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention is all you need. In NIPS, 2017.
Johannes Oswald, Eyvind Niklasson, E. Randazzo, João Sacramento, Alexander Mordvintsev, Andrey Zhmoginov, and Max Vladymyrov. Transformers learn in-context by gradient descent. ArXiv, abs/2212.07677, 2022.
Z. Wang, Yao Ma, Zitao Liu, and Jiliang Tang. R-transformer: Recurrent neural network enhanced transformer. ArXiv, abs/1907.05572 2019..
12

---

Published as a conference paper at ICLR 2024
Jason Wei, Xuezhi Wang, Dale Schuurmans, Maarten Bosma, Ed Huai hsin Chi, F. Xia, Quoe Le, and Denny Zhou. Chain of thought prompting elicits reasoning in large language models. ArXiv, abs/2201.11903, 2022.
Ezra Winston and J. Zico Kolter. Monotone operator equilibrium networks. ArXiv, abs/2006.08591 2020.
Sang Michael Xie, Aditi Raghunathan, Percy Liang, and Tengyu Ma. An explanation of in-context learning implicit bayesian inference. ArXiv, abs/21 1.02080, 2021.
Ruiqi Zhang, Spencer Frei, and Peter Bartlett. Trained transformers learn linear models in-context. ArXiv, abs/2306.09927, 2023a.
Ruiqi Zhang, Jingfeng Wu, and Peter L Bartlett. In-context learning of linear transformer block: Benefits of the mlp component and one-step gd initialization. arXiv preprint 244V::2402.149511 2024.
Shizhuo Dylan Zhang, Curt Tigges, Stella Rose Biderman, Maxim Raginsky, and Talia Ringer. Can transformers learn solve problems recursively? 2023b.
Yi Zhang, Arturs Backurs, Sébastien Bubeck, Ronen Eldan, Suriya Gunasekar, and Tal Wagner. Unveiling transformers with lego: synthetic reasoning task. ArXiv, abs/2206.04301, 2022.
Denny Zhou, Nathanael Scharli, Le Hou, Jason Wei, Nathan Scales, Xuezhi Wang, Dale Schuurmans, Olivier Bousquet, Quoc Le, and Ed Huai hsin Chi. Least-to-most prompting enables complex reasoning in large language models. ArXiv, abs/2205.10625 2022a.
Hattie Zhou, Azade Nova, H. Larochelle, Aaron C. Courville, Behnam Neyshabur, and Hanie Sedghi. Teaching algorithmic reasoning via in-context learning. ArXiv, abs/2211.09066 2022b.
A DETAILED SETUP FOR FUNCTION CLASSES
For completeness, we detail the setup of the function classes here, adhering to the settings in Garg al. (2022). Inputs x ~ (0, I) are consistent across all function classes.
Linear Function. For linear function with dimension d, we sample the parameter of linear function w ~ Rd following
Sparse Linear Function. We follow the same distribution setting as linear function, and we uniformly select out of d dimension to be nonzero.
We employ a binary decision tree with a depth of 4. At each non-leaf node, the decision branch left or right is based on the sign of randomly selected coordinate of the input æ. This coordinate is chosen uniformly at random. The leaf nodes are initialized following normal distribution N NO,, I). The evaluation of an input æ involves traversing from the root leaf node, moving left for negative coordinate and right for positive one, with the output being the value of reached leaf node.
Decision Tree.
2-layer ReLU NN.
We initialize the weights of the network following (N(0,, I), and the hidden layer dimension to be 100.
B CHOICE OF LooP ITERATIONS AND ITS EFFECT ON IN-CONTEXT PERFORMANCE
the looped transformer performance when trained with different b and T for the following function classes: linear functions with problem dimension = 20 and in co-contex samples lo = 40, b) sparse linear functions with problem dimension = 20, in-context samples k = 40, and non-sparse entry 3, e) decision tree with problem dimension
In this section, we present the result (Fig. 12
13

---

Published as a conference paper at ICLR 2024
d = 20, depth to be 4, and in-context samples k = 100, well d) 2-layer ReLU NN with problem dimension = 20, 100 hidden neurons, and in-context samples l: = 100.
Recall that the target function aim to minimize during the training of looped transformer is:
min Ep 9
Σ
b bu
t=b
where pi denotes the prompt prefix with in-context samples P (x1 f(z)),... 2.. f x(+1), (Pi())
the output of the looped transformer with pa- rameter 0 after looping iterations. bo max(b T, 7,0). The choice of b affects how many loop iterations the looped transformer needs to unroll during the training, and T represents the truncated window size for the loss calculation.
For linear regression, sparse linear regression, and 2-layer ReLU NN tasks, we investigate the different values of b in {12, 20, 30, 40, 50), and for the more challenging task decision tree, we search over b values in {12, 50, 70, 100). For all tasks, we investigate T values in {5, 10, 15). Larger T values such as T = 20 will result in training instability if trained without any stabilization technique such as weight decay, mixed-precisioe
or gradient norm clip. But once employed with those techniques, it is safe to train with large 7 well. In Fig. T2 and Fig. [3] present the result of the looped transformer, when trained with different and 7, the performance on different tasks, with respect to the loop iterations, and with respect the in-context samples.
For each task, the results are consistent: training with smaller T will yield a larger mean squared error in the in-context learning solution. When T is held constant, larger b will enable the looped transformer to find seidtpoint solution. As b increases, the performance improves until it reaches a point of saturation.
Across tasks, we can see a clear pattern of the choice of and T related to the level of difficulty of the task. Sparse linear regression, which has fewer parameters its function class, requires fewer b and 7 to find converged solution compared to linear regression. The decision tree task requires the learning of 2depth+
number of parameters (in our case, 31 parameters). Thus it requires a larger b to find a fixed-point solution. Counterintuitively, the 2-layer ReLU neural network, which contains in total 21 x 100 number of parameters, only requires The reason why data generated from 2-layer ReLU neural network only requires small value of b to find the fixed-point solution remains unclear. We conjecture that the value of b is an indication of the task's 's difficulty for the looped transformer to solve.
12 to find fixed-point solution.
C EFFECT OF NUMBER OF DISTINCT PROMPTS / FUNCTIONS SEEN DURING TRAINING
In Section 5 we investigated the impact of the number of distinct prompts during transformer training for problem dimension d = 10 and in-context samples & = 20. In this section, we extend the results to include (1) d = 5, = 10, and (2) d = 20, k = 40. It is important to note that all experiments were conducted without the use of curriculum learning for d and k. For the looped transformer, set b = 30, and T' = 15, and we don't 't apply scheduling in the loop iterations as well. For d = 5. investigated learning rate in {0.0001). for d 10 with learning rate (0.00001, 0.0001}, and for d = 20, we explored learning rate {0.00001, 0.0001, 0.001}. selecting the result with the best performance. The result is presented in Fig. 14 Across different problem dimensions, the looped transformer demonstrated better sample efficiency compared to the vanilla transformer.
D EXTRAPOLATION BEHAVIOR OF LOOPED TRANSFORMER AND UNLOOPED TRANSFORMER
We follow the setup from Garg et al. (2022) both the transformer and the looped transformer are trained to solve linear regression tasks in-context. During training, the inputs
on testing the out-of-distribution behavior. Recall that
sampled from .V Ia), and the linear function weights w (O Ia). Here d is 20, and the number of in content samples is l = 40. We now evaluate these models on: a) inputs with skewed covariance; b) inputs that lie in d/ d/2-dimensional subspace; c) inputs with noise:
14

---

Published as a conference paper at ICLR 2024
12
20
10
squared error
10-1
10-1
10
10-2
10-
10-1
10-3
10aa
10
loop iterations
loop iterations
loop iterations
1.0
1.0
1.0
0.6
0.6
0.6
0.4
0.0
20
0.0
0.0
in-context samples
in-context samples
in-context samples
(a) Linear Regression
12
20
50
10
10
10'
squared orror
10-1
10-1
10-1
10-2
10-3
10-1
10-3
10-
10-2
50
100
150
200
50
100
150
200
50
100
150
200
loop Iterations
loop iterations
loop iterations
0.8
0.8
0.8
0.4
0.4
0.4
0.0
10
20
30
0.0
10
20
0.0
10
in-context samples
in-context samples
in-context samples
(b) Sparse Linear Regression
We evaluate the looped transformer on in-context learning of linear functions, and b) sparse linear functions with different and during training (b and T defined Eq. [ For each block figures, from left right with T = 5, 10, 15. (top) Performance with respect to different loop iterations, (bottom) performance with respect to different numbers of in-context samples.
Figure 12:
15

---

Published as a conference paper at ICLR 2024
12
50
70
100
squared error
10-1
10
10
10a
10
loop iterations
loop iterations
loop iterations
1.0
1.0
1.0
0.6
0.6
0.6
Squared
0.2
0.2
0.2
in-context samples
in-context samples
In-context samples
(a) Decision Tree
12
20
30
40
50
1.0
1.0-
0.6
0.8
0.6
0.4
0.2
50
100
150
200
50
100
150
200
0
50
100
150
200
loop iterations
loop iterations
loop iterations
squared error
1.00
1.00
1.00
0.50
0.50
0.50
0.25
0.00
20
60
60
160
0.00
20
40
60
80
100
0.00
20
40
60 80
100
in-context samples
in-context samples
in-context samples
(b) 2-yayer ReLU NN
Figure 13: We evaluate the looped transformer on in-context learning of a) random decision tree, and b) 2-layer ReLU neural network with different and T during training (b and 7' is defined in Eq. IT. For each block of figures, from left right is with T 5, 15. (top) Performance with respect to different loop iterations, (bottom) performance with respect different numbers of in-context samples.
1.2
1.2
1.0
Leapie Carat
L4876 Lanalo
Lesppc Tratabrt
0.8
0.8
0.8
60.4
0..4
80.4
1k
24
Sk
10te
25%
4llk
80k
1k
Sk
1ile
20k
Atlk
20kc
1k
5k
1ife
20k
46%
20k
distact pommpts/incctoon seen during tratning
dustact promppsstmncttoss seen during training
duttint poompts.tmccttoos sean daring training
(a) Problem dimension
(b) Problem dimension
(c) Problem dimension 20
Figure 14: Performance of the standard transformer and the looped transformer linear regression problem with problem dimension d' = 5, 10, 20, when trained with different number of distinct prompts/functions.
16

---

Published as a conference paper at ICLR 2024
1,5
1.5
squared error
squared error
1.0
0.5
0.5
0.0
10
20
30
10
20
30
10
20
20
40
in-context examples
in-context examples
in in-context examples
(a) Skewed Covariance
(b)
d/2mdiinnsioonl Subspace
(c) Noisy Output
1.5
1.5
1.5
squared error
squared arror
1.0
1.0
Tranthment
Loiped Transbrmen biid
0.5
Leaed
trankeeen kelp Laaged
6.9
0.0
10
20
30
10
20
30
10 2b
20
40
In-Context excamples
In-context examples
In-context examples
(d) Orthogomal Query
(e) Query Matches context Example
(f) Different Orthants
Figure 15: Evaluating the transformer of the different number of layers and looped transformer, with the number of layer to be L = 1, trained to maximum b loop iterations during training. The transformer is trained on the linear regression task, but evaluated on different distr--ditt
tasks.
d) the query is orthogonal to the in- in-context example; e) the query is equal to one of the in-context samples; and f) the in-context example inputs and query inputs lie in different orthants. The results are presented in Fig. 15
From the figure, it is evident that the looped transformer, some cuitof-distribution cases, exhibits performance similar to the standard transformer with fewer layers. For instance, in the noisy linear regression case, the transformer with a number layers = exhibits reduced peak when in-context samples are presented. This trend is similar to the looped transformer, especially when the b value is larger. certain other cases, the looped transformer demonstrates performance akin to the standard transformer with a higher number of layers. For instance, when the prompts (a) with skewed covariance, (b) in d/ 12-dimensional subspace, or (f) the in-context samples and the prompts lie in different orthants, for the standard transformer, the larger L is, the better the performance is. Moreover, the looped transformer can match or even outperform the standard transformer with as many L 20 number of layers. Recall that for the linear regression task, when < 20, the looped transformer cannot find fixed-point solution. Therefore, we didn't 't include the looped transformer with 6 < 20 in the comparison. In other cases, the looped transformer performs similarly to the standard transformer. This is evident in instances such as (d) when the in-context samples and the query lie in orthogonal space, and (e) when the query matches one of the in-context samples.
We further evaluated the out- out-of-distributinn performance of the looped transformer by either scaling the prompt æ or scaling the parameter the linear function w during inference. The results are depicted in Fig. 16 As illustrated in the figure, the looped transformer demonstrates performance comparable to that of the standard transformer with fewer layers. For instance, when x is scaled by 2, the looped transformer performs similarly to the standard transformer with L 4.
E EFFECTS OF SCHEDULING
In Sec. 5.1 we discuss the training schedule for the loop iterations we will gradually increase b during the course of training. As a result, the truncated loss window starts at [1, and then gradually shifts to - T, b). In this section, we investigate the impact of this scheduling during training. Specifically, we compare the looped transformer training the four function classes: function; b) sparse linear function; c) decision tree; d) 2-layer ReLU neural network, when trained with and without the scheduling. When training, in Sec. 3.4 Results are presented Fig. [7 As indicated in the figure, most cases, the looped transformer's performance does not differ when trained with without the scheduling. However, for
linear
adhere the selected b and T as indicated
17

---

Published as a conference paper at ICLR 2024
Transformer 1=1
Transforme: L=12
Looped Transformer b=
Looped Transformer b=
Least Squares
Transformer 1md
Transforme: L=20
Looged Transformer bm30
Louped Transformer b=100
Squardd
aquassed
contenl esamples
In content cromples
In-oontect vxmmples
conDoCt Champiess
(a) scaled x, to right scale 1/3, 1/2, 2, and 3.
4.15
1.10
in.conteni esamples
23
In content crampies
contoxt cxmmples
CE-CODOO exampies
(b) scaled y, to right scale is 1/3, 1/2, 2, and 3.
Figure 16: Evaluating the transformer of the different number of layers J., and looped transformer, with the number layer to be trained to maximum loop iterations during training. The transformer trained on the linear regression task, but evaluated on scaled prompt ær or scaled linear function parameter w.
w/ scheduling
scheduling
1.0
1.0
squared error
squared error
0.6
0.6
0.4
0.4
0.0
0.0
10
10
in in-context samples
in-context samples
(a) Linear Regression
(b) Sparse Linear Regression
1.0
1.0
squared error
Squared error
0.6
0.6
0.4
0.4
0.2
0.2
20
40
60
80
100
0
20
40 60 80
100
in in-context samples
in-context samples
(c) Decision Tree
(d) 2-layer ReLU NN.
Figure 17: Performance of the looped transformer with respect to varying number of in-context samples. We compared the trained looped transformer in solving (a) linear regression, (b) sparse linear regression, (c) decision tree, and (d) 2-layer ReL.U NN, when trained with (solid line) without (dashed line) scheduling. Whether scheduling is applied over has impacts on the performance when training the looped transformer on data generated from the decision tree, with the application of scheduling resulting in slightly better performance.
the decision tree task, the looped transformer trained with scheduling exhibits better performance compared to when trained without scheduling. Therefore, throughout our experiments, we opt to use scheduling for the loop iterations.
18

---

Published as a conference paper at ICLR 2024
F DETAILS OF MODEL PROBING
In Sec. 5.3
analyze the output of either the standard transformer's t-th layer the looped transformer's s t-th loop iteration by training a probing model on them. In this section, we present the details of the model probing experiments.
Reusing the notation, we denote the output at the I-th layer or loop iteration as Y, Y2 is of shape k x D, where k represents the sequence length, and 1) represents the embedding dimension of the transformer. In our experiments, we investigate two target probes: (a) X T y, representing the gradient component, and Wots, representing the optimal least square solution. Given the problem dimension is d, the target probes v are of shape d x 1. For simplicity, we omit the batch size.
Before feeding the representation Y, into the probing model, we first compute weighted sum over the sequence dimension. Given E R we perform the following operations:
softmax
(2)
Z, = Yo
With Zz E RIXD we employ a multi-layer perceptron model (MLP) SMLP to probe the output. The MLP model consists of two linear layers, with ReLU in between, with dhidden denote the number of hidden neurons in MLP:
i (Z,)
When training the s and fMLP, we aim to minimize the mean squared error between and v. For each layer's output, and each in-context length, we train a distinct and fMPP to probe the target. When presenting the results, we compute the average loss over various context lengths. Additionally, to normalize the error, we divide the mean squared error by the problem dimension d.
In our experiments, we set D = 256, d 20, and duidden = 64. We optimize the probing model with Adam optimizer and learning rate 0.001, which found to be the best among the set {0.01, 0.001, 0.0001}.
Control Task Configuration. To highlight the significance the learned probing model, we follow the setting in Akyürek et al. (2022), and train the probing model with the linear function weights w fixed at 1. In this case, the probing model will only learn the statistics of the input X, eliminating the need for in-context learning. Then during inference, we randomly sample w ~ _N Ia) and record the resulting probing error. The minimal probing error for the control tasks are shown as dashed lines in Fig. From the figure, the gap between the control task and the model probing results indicates the probing model utilizes the in-context information obtained in the output of each layer/loop iteration, suggesting that the probing error is significant.
G OPENML EXPERIMENT
To establish connection with real-world applications and further validate our approach. we have conducted additional experiments using 10 datasets from OpenML (Vanschoren et al. 2013), which we believe better represent practical scenarios. The details of the used datasets are presented in T
Let (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 datasets and evaluated its in-context learning ability on the unseen test dataset. Both transformer and looped transformer have 256 embedding size, heads. The tranformer has 12 layers, and looped transformer has layer, trained to maximum loop iteration = 30, and window size = 15.
During training, uniformly sampled prompts from 9 datasets, where for each prompt, we first randomly selected a training set, then randomly selected + samples from this training set, with k being the number of in-context samples. During testing, applied similar approach for each test sample, selecting à in-context samples from the test dataset, with care taken to exclude the test sample itself from these in-context pairs. The test accuracies presented in Table 2
19

---

Published as a conference paper at ICLR 2024
Dataset ID
Name
# Numerical
Features
Instances
720
abalone
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 compas-two-years
5
3848
42192
3848
Table 1: OpenML dataset used in the experiments. We only numerical features for the input features.
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 believe these experiments offer more realistic representation of practical applications.
Test Dataset ID
Train Dataset IDs
Transformer
Looped Transformer
720
S (720)
0.626 + 0.008
0.662 + 0.008
725
S 725
0.511 0.007
0.504 + 0.008
737
(737
0.656 + 0.006
0.72 + 0.01
803
(803)
0.394 + 0.01
0.40 = 0.01
807
(807
0.405 + 0.004
0.416 +: 0.005
816
S (816)
0.463 0.004
0.462 0.004
819
(89
0.483 + 0.005
0.568 0.568 0.01
847
(847))
0.668 + 0.007
0.757 0.006
871
(871)
0.532 + 0.004
0.51 0.005
42192
(2242192)
0.65 + 0.005
0.65 + 0.008
Table 2: The test accuracy for transformer and looped transformer, for different test dataset id. Looped transformer demonstrates comparable, and some cases, better performance compared to vanilla transformer.
20


Reviews:
[{"review_id": "rhZiAs0utI", "reviewer": "Reviewer_gzpZ", "summary": "This paper proposes to use looped transformers to solve in-context learning, which achieves comparable performance to the standard transformer, but utilizes less than 10% of the parameters.", "questions": "What is the mathematical insight of looping transformers for in-context learning?", "limitations": "", "rating": 5, "confidence": 4, "soundness": 2, "presentation": 2, "contribution": 2, "strengths": "It is interesting to see looped transformers work well for in-context learning.\n\nThe paper provides a thorough evaluation and ablations of looped transformers for in-context learning. \n\nThe paper is well written.", "weaknesses": "My primary concern lies in the relevance of looped transformers to in-context learning. It appears that their main advantage is in reducing the number of parameters. However, 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. I find that the paper lacks in-depth mathematical insights or a thorough exploration of the practical implications that would help address this concern.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "n6QINWM5G3", "reviewer": "Reviewer_xBYm", "summary": "This paper proposes a training methodology for looped transformers to effectively emulate iterative algorithms and provides empirical evidence that demonstrate the advantages of looped transformer on in-context learning. However, since all experiments are made on simulated datasets, whether the proposed method is effective in dealing with real-world data remains to be validated.", "questions": "1. Can you provide more details about the probability distribution over the used classes of functions, especially for decision trees and 2-layer ReLU NN?", "limitations": "", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "(1) The paper uses the looped transformer to emulate iterative learning algorithms and presents a novel methodology to train the looped transformer under reasonable assumptions.\n(2) The paper provides a wide range of evaluation and detailed ablation studies of the proposed method on simulated datasets and demonstrates its superior performance compared to standard, non-recursive transformers.", "weaknesses": "(1) Since all experiments are made on simulated datasets, whether the proposed method is effective in dealing with real-world data remains to be validated.\n(2) 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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "KOYZjo7FEV", "reviewer": "Reviewer_QsDA", "summary": "In the paper the authors examine the applicability of looped transformers to the task of learning linear regression in context (giving sampled examples). They compare the performance of a looped transformer with a standard 12 layer transformer and a least squares solver. They show that for some settings the looped transformer matches or even outperforms the performance of a standard transformer on the task of linear regression, while incorporating significantly less parameters.", "questions": "- What is the experimental setting for the plots shown in Figure 2? \n- Why are the curves for the looped transformer and the standard Transformer in Figure 2 exactly overlapping?\n- How much time and resources does it take to find optimal values for b and T?\n- If you consider the hyperparameter search is it still more useful to apply looped transformers instead of the standard transformers?", "limitations": "", "rating": 8, "confidence": 3, "soundness": 3, "presentation": 2, "contribution": 3, "strengths": "- It is an interesting approach to use looped transformers for the task of linear regression in-context. It seems to work well with significantly less parameters.\n- The authors examine different settings, like using input injection, choosing the number of iterations, altering the number of layers, heads and the dimension of the embeddings.\n- The experimental evaluation is convincing.", "weaknesses": "- 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.\n- The structure of the paper can be improved, e.g., by ending the paper with a short conclusion instead of the related works.\n\nMinor details:\n- \"It is worth noting\" instead of \"it's worth noting\"\n- \"use the scheduling does not significantly impact the outcome\" instead of \"use the scheduling doesn’t significantly impact the outcome\"\n- The Figures are too small\n- Fig. 2: The curve for the least squares solver is difficult to see\n- Fig. 2 description: (left) and (right) should not be written after the punctuation in the sentence but before.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (poster)