PaperlayerView exact revision on arXiv

Source-linked AI summary

Looped Transformers are Better at Learning Learning Algorithms

Liu Yang, Kangwook Lee, Robert Nowak, Dimitris Papailiopoulos

arXiv:2311.12424v3cs.LGcs.NE

TL;DR

Standard transformers lack an inherent iterative structure, despite iterative algorithms being common in data-fitting, and the mechanisms behind in-context learning remain unclear. The paper introduces a looped transformer and a training method for emulating iterative algorithms; on tested function classes, the looped model outperforms standard transformers particularly on sparse linear functions and decision trees.

  • Problem

    Standard transformers are not naturally structured for iterative computation, while iterative algorithms are commonly used for data-fitting; in-context learning mechanisms also remain unclear.

    Source: §1 · pp. 1–2

  • Method

    The paper introduces a looped transformer and a training methodology designed to emulate iterative algorithms, with input injected into each loop iteration.

    Source: §1 · p. 2, §4.1 · p. 5

  • Results

    Looped transformers trained from scratch learned in context across linear functions, sparse linear functions, decision trees, and 2-layer neural networks, outperforming standard transformers particularly on sparse linear functions and decision trees.

    Source: §1 · p. 2

  • Takeaways & Limitations

    Within the training distribution, the method approximates fixed-point solutions beyond the loop iterations used in training.

    Source: §7 · p. 11

    The trained model has an error floor and limited out-of-distribution performance, so it emulates the algorithm only for its training input distribution rather than identifying an algorithm that generalizes to any distribution.

    Source: §6 · p. 10

Abstract

from arXiv · show

Transformers have demonstrated effectiveness in in-context solving data-fitting problems from various (latent) models, as reported by Garg et al. 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.

1 INTRODUCTION

The paper asks whether looped transformers can emulate iterative learning algorithms more efficiently, and introduces a looped architecture and training method to address the mismatch between iterative algorithms and standard transformers. In tested function classes, looped transformers learn in context and outperform standard transformers particularly on sparse linear functions and decision trees.

1 INTRODUCTION

  • The paper frames its question against unclear mechanisms of in-context learning and an open question about training standard decoder transformers to replicate learning algorithms.Prior work on linear self-attention had shown an implicit gradient-descent step, but the cited passage identifies a gap between that simplified architecture and standard decoders.Source: §1 · pp. 1–2
  • Standard transformers lack a natural iterative structure, while looped architectures can break tasks into simpler subtasks and may reduce parameter needs.The paper motivates this contrast using linear regression, where gradient descent repeatedly applies matrix multiplication while least squares requires a matrix inverse.Source: §1 · pp. 1–2
  • The proposed training method aims to make looped transformers emulate iterative algorithms by studying their structural design and the loop iterations used during training.Its stated assumption is that performance should improve or converge as loop iterations increase when simulating a convergent algorithm.Source: §1 · p. 2
  • Looped transformers trained from scratch learned in context across four tested function classes and outperformed standard transformers particularly on sparse linear functions and decision trees.The tested classes were linear functions, sparse linear functions, decision trees, and 2-layer neural networks.Source: §1 · p. 2
  • For linear regression, a looped transformer trained for 30 iterations reaches a stable fixed-point solution beyond the trained iterations and uses 1/12 the parameters of a standard 12-layer transformer.The figure caption also reports performance matching the standard transformer and closely aligning with the least-squares solver.Source: Figure 2 · p. 3

2 RELATED WORKS

Related work covers transformer weight sharing and implicit models, neural networks for iterative reasoning tasks, and empirical and theoretical accounts of in-context learning for data-fitting problems.

2 RELATED WORKS

  • Weight-sharing transformers reuse weights across forward-pass layers or use a halting criterion, while implicit models repeatedly apply a function to find fixed points.Prior work also used recurrent transformer frameworks to adapt forward passes or extend context length.Source: §2 · p. 3
  • Deep implicit models use black-box solvers for fixed points and have been applied to math extrapolation, but training may require extensive tuning and remain unstable.The cited work also includes Jacobian regularization as a stabilization approach.Source: §2 · p. 3
  • Prior work applies recurrent neural networks to iterative math and reasoning tasks with finite states, while this paper studies continuous-output regression using a decoder with nonlinear attention.The paper situates its architecture and empirical focus in relation to Garg et al. (2022).Source: §2 · p. 3
  • Theoretical and empirical studies connect transformer in-context linear regression to implicit gradient updates, higher-order optimization, Bayesian estimation, and generalization error.For one-layer linear self-attention, cited work reports that one gradient-descent step is optimal under an ℓ2 prediction loss.Source: §2 · pp. 3–4

3 PROBLEM SETTING

The problem setting samples a function and examples to form a prompt, then evaluates a learning system by its prediction of the function at a test input.

3 PROBLEM SETTING

  • A function is sampled from a function class distribution, and example and test inputs are sampled from an input distribution.The function class consists of functions defined on R^d.Source: §3 · p. 4
  • The learning prompt contains input-output examples from the sampled function followed by a test input, and the system must predict the function's output there.The sampled function itself is not explicitly given to the learning system.Source: §3 · p. 4
  • Performance is measured by squared prediction error, with GPT-2 decoders as the transformer systems and 12 layers by default for unlooped models versus 1 for looped models.The layer settings follow the stated setup for comparison with Garg et al. (2022).Source: §3 · p. 4

4 TRAINING ALGORITHM FOR LOOPED TRANSFORMER

The section describes a looped-transformer training strategy intended to preserve performance across iterations, with input injection and loop-iteration choices shaping convergence and stability. The experiments report that b = 12 does not match the standard transformer’s performance, while larger b values can avoid post-training divergence.

4 TRAINING ALGORITHM FOR LOOPED TRANSFORMER

  • Training repeatedly applies the shared model to the prompt prefix, adds that prefix at each loop step, and minimizes loss over selected iterations and prompt prefixes.The truncated loss covers loop iterations from b0 = max(b − T, 0) through b.Source: §4 · p. 4
  • The one-layer looped transformer uses 0.79M parameters, compared with 9.48M for the 12-layer unlooped transformer.Both models use the stated GPT-2 configuration with embedding dimension 256 and 8 attention heads.Source: §4 · p. 4

4.1 INPUT INJECTION TO LOOPED TRANSFORMER

  • Input injection uses Yt+1 = M(Yt + P), reintroducing the input at each iteration; without injection, performance deteriorates beyond trained iterations.The paper also notes that injection methods need not be limited to addition.Source: §4.1 · p. 5, Figure 3 · p. 5

4.2 CHOICE OF LOOP ITERATIONS

  • For b = 12, the looped transformer matches the unlooped transformer's effective depth but not its performance; higher b avoids divergence after trained loop iterations, while beyond a certain b the convergence rate declines.Source: §4.2 · pp. 5–6
  • Lower b may lead to post-training divergence, whereas higher b can avoid it but may eventually reduce the convergence rate.For T = 15, b = 40 and b = 50 yield similar convergence rates.Source: §4.2 · pp. 5–6, Figure 4 · p. 5
  • The looped transformer can reach a fixed-point solution before iteration b because the loss requires matching the target within b steps.A smaller b can expedite convergence, but increasing b eventually stops improving the convergence rate.Source: §4.2 · p. 6
  • Choosing b trades off model potential against training and inference time, while T trades off memory use against performance and training stability; the selected values are b = 20 and T = 15.The authors choose the smallest b that prevents post-training divergence and prefer a smaller T to reduce memory use.Source: §4.2 · pp. 5–6

5 EXPERIMENTAL RESULTS

The experiments test looped transformers on linear regression and higher-complexity function classes, including in-distribution, shifted-distribution, and architecture-variation settings.

5.1 EXPERIMENTAL SETUP

  • The schedule progressively increases b and shifts the truncated loss window from [1, T] to [b −T, b]; its effect is usually small, though some cases improve.Minimizing mean squared error over [1, T] can be more challenging than over [b −T, b].Source: §5.1 · p. 6

5.2 LOOPED TRANSFORMER CAN IN-CONTEXT LEARN LINEAR REGRESSION

  • With b = 20 and T = 15, the looped transformer matches the standard transformer in-distribution and nearly matches ordinary least squares.Source: §5.2 · p. 6
  • The looped transformer learns linear regression with fewer distinct training prompts than the standard transformer in the reported sample-complexity experiment.The experiment uses problem dimension d = 10 and k = 20 in-context samples.Source: §5.2 · p. 7
  • The looped model's simplicity bias helps with skewed covariance and reduces noisy-regression double-descent peaks, but hurts under input-scaling shifts.In the scaling test, x = 3z with z ∼ N(0, I), and the looped model underperforms the unlooped transformer.Source: §5.2 · p. 7

5.3 IMPACT OF MODEL ARCHITECTURE VARIATIONS ON LOOPED TRANSFORMER PERFORMANCE

  • A one-layer looped transformer needs roughly 20 iterations to match an 8-layer standard transformer's performance, so its effective depth can exceed its nominal depth.Source: §5.3 · p. 7, Figure 7 · p. 8
  • Probing finds that the standard transformer initially represents target statistics and later shifts away, whereas the looped model steadily refines target-probe representations across iterations.The probes target the gradient component X^T y and the optimal least-squares solution w_ols.Source: §5.3 · p. 8, Figure 8 · p. 8
  • With L = 1 and h = 8, convergence improves as D rises to 256; at fixed L and D, increasing heads helps through 8 but worsens by 32.Under the stated GPT-2-style definition, varying h trades off the number of heads against per-head embedding dimension.Source: §5.3 · p. 8

5.4 HIGHER COMPLEXITY FUNCTION CLASSES

  • 1/12th the parameters suffices to match or surpass the standard transformer on higher-complexity tasks; sparse-regression gains persist across sparsity levels, especially with few samples.For sparse regression, the reported advantage over Lasso concerns the error trend across in-context sample counts, not the precision of the final solution.Source: §5.4 · pp. 9–10, Figure 10 · p. 10

6 DISCUSSION AND FUTURE WORK

The discussion reports that looped transformers can emulate fixed-point behavior beyond their trained iteration count within the training distribution, but not generalize as an algorithm to other input distributions. It also considers how loop settings relate to task difficulty and training cost.

6 DISCUSSION AND FUTURE WORK

  • At d = 20, the looped transformer matches or outperforms the unlooped transformer across the tested sparsity levels, with more accurate solutions using fewer in-context samples.The tested sparsity levels are s = 1, 3, 6, 9, and 12.Source: Figure 11 · p. 10
  • The work empirically validates a practical training method for a looped decoder transformer, while further study of its expressiveness remains future work.The paper distinguishes this empirical focus from prior work on recurrent-model expressiveness and looped-transformer universal computation.Source: §6 · p. 10
  • Fixed-point solutions extend beyond trained loop iterations within the training distribution, but performance is limited on out-of-distribution prompts.The authors say the model emulates the algorithm for its training input distribution rather than identifying an algorithm that generalizes to any distribution.Source: §6 · p. 10
  • More complex tasks require larger b and T for convergence, although the low b needed by the 2-layer ReLU network remains unexplained.The authors conjecture that b indicates task difficulty for the looped transformer.Source: §6 · pp. 10–11
  • Training cost rises with loop count b, which sets forward-pass time, and loss-window size T, which determines training memory requirements.The authors suggest larger b with a suitable T and regularization as an alternative to intensive tuning.Source: §6 · p. 11

7 CONCLUSION

The conclusion presents a training method that lets looped transformers approximate fixed-point solutions beyond their initially trained iterations. The appendix specifies the input and function-class settings used for the experiments.

7 CONCLUSION

  • The training method enables looped transformers to approximate fixed-point solutions beyond the loop iterations used in training.The authors frame the work as analysis of loop architecture and training strategy for data-fitting problems.Source: §7 · p. 11

A DETAILED SETUP FOR FUNCTION CLASSES

B CHOICE OF LOOP ITERATIONS AND ITS EFFECT ON IN-CONTEXT PERFORMANCE

The appendix examines how loop count b and loss-window size T affect in-context performance across several function classes. Smaller T is associated with higher mean squared error, while increasing b improves performance until saturation; the needed settings vary by task.

B CHOICE OF LOOP ITERATIONS AND ITS EFFECT ON IN-CONTEXT PERFORMANCE

C EFFECT OF NUMBER OF DISTINCT PROMPTS / FUNCTIONS SEEN DURING TRAINING

The section extends the distinct-prompt experiments across three problem dimensions without curriculum learning. Across dimensions, the looped transformer demonstrates better sample efficiency than the vanilla transformer.

C EFFECT OF NUMBER OF DISTINCT PROMPTS / FUNCTIONS SEEN DURING TRAINING

D EXTRAPOLATION BEHAVIOR OF LOOPED TRANSFORMER AND UNLOOPED TRANSFORMER

The section evaluates looped and standard transformers on linear-regression inputs and prompts that differ from training, including scaled prompts and function parameters. Across these settings, the looped model is often comparable to standard transformers with fewer layers, and in some cases matches or exceeds deeper ones.

D EXTRAPOLATION BEHAVIOR OF LOOPED TRANSFORMER AND UNLOOPED TRANSFORMER

  • Both models are trained on linear regression with d = 20, k = 40, and inputs and weights sampled from N(0, I_d), then tested on six out-of-distribution input or query conditions.Those conditions include skewed covariance, a d/2-dimensional subspace, noise, orthogonal or repeated queries, and different orthants.Source: D EXTRAPOLATION BEHAVIOR OF LOOPED… · p. 18
  • On skewed-covariance, d/2-dimensional-subspace, and different-orthant prompts, the looped transformer can match or outperform a standard transformer with up to 20 layers.For linear regression, loop counts b < 20 were excluded because the looped transformer cannot find a fixed-point solution.Source: D EXTRAPOLATION BEHAVIOR OF LOOPED… · p. 18, Figure 15 · p. 19
  • On noisy linear regression, the looped transformer resembles a 4-layer standard transformer, while the standard model’s peak falls when d in-context samples are presented.The looped transformer’s behavior is especially similar when b is larger.Source: D EXTRAPOLATION BEHAVIOR OF LOOPED… · p. 18, Figure 15 · p. 19
  • When the query is orthogonal to the in-context samples or repeats one of them, the looped transformer performs similarly to the standard transformer.The evaluation also includes skewed covariance, lower-dimensional inputs, noise, and different orthants.Source: D EXTRAPOLATION BEHAVIOR OF LOOPED… · p. 18, Figure 15 · p. 19
  • When the prompt x is scaled by 2, the looped transformer performs similarly to a standard transformer with L = 4.The section also evaluates scaling the linear-function parameter w during inference.Source: D EXTRAPOLATION BEHAVIOR OF LOOPED… · p. 18, Figure 16 · p. 19

E EFFECTS OF SCHEDULING

The loop-iteration schedule has little effect on most tested function classes but improves decision-tree performance, so the experiments use scheduling. The probing experiments examine whether layer or iteration outputs encode gradient and least-squares-solution information.

E EFFECTS OF SCHEDULING

  • Scheduling leaves performance largely unchanged on linear, sparse-linear, and 2-layer ReLU tasks, but improves decision-tree performance; the experiments therefore use scheduling.The schedule gradually increases the maximum loop count b and shifts the truncated loss window during training.Source: E EFFECTS OF SCHEDULING · p. 20

F DETAILS OF MODEL PROBING

  • Probes test whether each layer or loop-iteration output contains the gradient component X^T y or the optimal least-squares solution w_ols.The output Y_t has shape k × D, and each target probe has shape d × 1.Source: F DETAILS OF MODEL PROBING · p. 20
  • A learned sequence-weighted representation is passed to a two-linear-layer MLP with a ReLU between layers to predict each target probe.Separate sequence weights and MLPs are trained for each layer output and in-context length; probe error is mean squared error normalized by problem dimension d.Source: F DETAILS OF MODEL PROBING · p. 21
  • The probe control fixes w at 1 during training, then samples w from N(0, I_d) at inference; the gap from model-probing results indicates use of in-context information.The reported probe setup uses D = 256, d = 20, and 64 hidden neurons.Source: F DETAILS OF MODEL PROBING · p. 21

G OPENML EXPERIMENT

The OpenML experiments train on nine datasets and evaluate in-context learning on an unseen tenth dataset. The looped transformer has comparable, and sometimes better, test accuracy than the standard transformer on these datasets.

G OPENML EXPERIMENT

  • The experiment uses 10 OpenML datasets with binary labels and numerical input features, training on nine datasets and testing on the remaining unseen dataset.The dataset collection is identified by the set S = {720, 725, 737, 803, 807, 816, 819, 847, 871, 42192}.Source: G OPENML EXPERIMENT · p. 21, Table 1 · p. 21
  • The standard transformer has 12 layers, while the looped transformer has one layer and is trained to a maximum of b = 30 iterations with window size T = 15.Both models use embedding size 256 and 8 attention heads.Source: G OPENML EXPERIMENT · p. 21
  • Training prompts sample k + 1 examples from a randomly selected training set; test prompts draw k in-context examples from the test dataset while excluding the query sample.The reported comparison is based on test accuracies.Source: G OPENML EXPERIMENT · pp. 21–22
  • On the unseen OpenML test datasets, the looped transformer achieves comparable and sometimes better test accuracy than the standard transformer.The comparison uses test accuracy across different test dataset IDs.Source: G OPENML EXPERIMENT · p. 22, Table 2 · p. 22
Loading 2311.12424v3…