Looped Transformers with Layer Normalization Provably Learn the Power Method

• URL: http://arxiv.org/abs/2606.00605v1
• Date: Sat, 30 May 2026 08:05:27 GMT
• Title: Looped Transformers with Layer Normalization Provably Learn the Power Method
• Abstract summary: We study principal component prediction as a testbed for understanding the training dynamics of transformers with layer normalization (LN)<n>Our results provide the first theoretical analysis of the training dynamics of looped and single-layer transformers with LN, and shed light on the role of LN in transformer models.
• Score: 8.400141074200935
• License: http://creativecommons.org/licenses/by-sa/4.0/
• Abstract: Transformers have achieved remarkable success across a wide range of applications, and a growing body of work suggests that part of their strength comes from their ability to learn and execute algorithmic procedures. However, our understanding of how transformers learn such algorithms remains limited, especially in the presence of layer normalization (LN). In this work, we study principal component prediction as a concrete testbed for understanding the training dynamics of transformers with LN. We prove that a looped linear transformer with LN, trained by gradient descent, converges to a solution that implements the power method, with each self-attention layer performing one power iteration. Notably, the model is trained only for principal component prediction, rather than being explicitly supervised to implement the power method. Our finding thus reveals an "algorithmic implicit bias" of looped transformers with LN: principal-component prediction can in principle be achieved by many mechanisms, yet gradient descent selects one that realizes the power method. We further provide a concrete comparison between transformers with and without LN: even with layerwise guidance from power iterations, a transformer without LN cannot exactly learn the power method, whereas the corresponding transformer with LN can, leading to a provable performance gap in principal component prediction. Our results provide, to our knowledge, the first theoretical analysis of the training dynamics of looped and single-layer transformers with LN, and shed light on the role of LN in transformer models.

Related papers

• Transformers Efficiently Perform In-Context Logistic Regression via Normalized Gradient Descent [9.440916748352722]
  Transformer with softmax attention perform in-context learning on linear classification data.<n>We first construct a class of multi-layer transformers that can perform in-context logistic regression.<n>We show that our constructed transformer can be obtained through (i) training a single self-attention layer supervised by one-step gradient descent, and (ii) recurrently applying the trained layer to obtain a looped model.
  arXiv  Detail & Related papers  (2026-05-07T17:27:55Z)
• Weights to Code: Extracting Interpretable Algorithms from the Discrete Transformer [65.38883376379812]
  We propose the Discrete Transformer, an architecture engineered to bridge the gap between continuous representations and discrete symbolic logic.<n> Empirically, the Discrete Transformer not only achieves performance comparable to RNN-based baselines but crucially extends interpretability to continuous variable domains.
  arXiv  Detail & Related papers  (2026-01-09T12:49:41Z)
• From Shortcut to Induction Head: How Data Diversity Shapes Algorithm Selection in Transformers [67.02076505996284]
  We study how the choice of pretraining data distribution steers a shallow transformer toward one behavior or the other.<n>Our results shed light on the algorithmic biases of pretrained transformers and offer conceptual guidelines for data-driven control of their learned behaviors.
  arXiv  Detail & Related papers  (2025-12-21T08:10:26Z)
• Optimality and NP-Hardness of Transformers in Learning Markovian Dynamical Functions [32.71332125930795]
  Transformer architectures can solve unseen tasks based on input-output pairs in a given prompt due to in-context learning (ICL)<n>We investigate Markovian function learning through a structured ICL setup to reveal underlying optimization behaviors.
  arXiv  Detail & Related papers  (2025-10-21T13:42:48Z)
• Transformers Don't In-Context Learn Least Squares Regression [5.648229654902264]
  In-context learning (ICL) has emerged as a powerful capability of large pretrained transformers.<n>We study how transformers implement learning at inference time.<n>We highlight the role of the pretraining corpus in shaping ICL behaviour.
  arXiv  Detail & Related papers  (2025-07-13T01:09:26Z)
• On the Learn-to-Optimize Capabilities of Transformers in In-Context Sparse Recovery [15.164710897163099]
  We show that a K-layer Transformer can perform an L2O algorithm with a provable convergence rate linear in K.<n>Unlike the conventional L2O algorithms that require the measurement matrix involved in training to match that in testing, the trained Transformer is able to solve sparse recovery problems generated with different measurement matrices.
  arXiv  Detail & Related papers  (2024-10-17T19:18:28Z)
• Can Looped Transformers Learn to Implement Multi-step Gradient Descent for In-context Learning? [69.4145579827826]
  We show a fast flow on the regression loss despite the gradient non-ity algorithms for our convergence landscape. This is the first theoretical analysis for multi-layer Transformer in this setting.
  arXiv  Detail & Related papers  (2024-10-10T18:29:05Z)
• How Do Nonlinear Transformers Learn and Generalize in In-Context Learning? [82.51626700527837]
  Transformer-based large language models displayed impressive in-context learning capabilities, where a pre-trained model can handle new tasks without fine-tuning. We analyze how the mechanics of how Transformer to achieve ICL contribute to the technical challenges of the training problems in Transformers.
  arXiv  Detail & Related papers  (2024-02-23T21:07:20Z)
• Transformers as Statisticians: Provable In-Context Learning with In-Context Algorithm Selection [88.23337313766353]
  This work first provides a comprehensive statistical theory for transformers to perform ICL. We show that transformers can implement a broad class of standard machine learning algorithms in context. A emphsingle transformer can adaptively select different base ICL algorithms.
  arXiv  Detail & Related papers  (2023-06-07T17:59:31Z)
• XAI for Transformers: Better Explanations through Conservative Propagation [60.67748036747221]
  We show that the gradient in a Transformer reflects the function only locally, and thus fails to reliably identify the contribution of input features to the prediction. Our proposal can be seen as a proper extension of the well-established LRP method to Transformers.
  arXiv  Detail & Related papers  (2022-02-15T10:47:11Z)
• N-ODE Transformer: A Depth-Adaptive Variant of the Transformer Using Neural Ordinary Differential Equations [1.2183405753834562]
  We use neural ordinary differential equations to formulate a variant of the Transformer that is depth-adaptive in the sense that an input-dependent number of time steps is taken by the ordinary differential equation solver. We consider the simple problem of determining the parity of a binary sequence, for which the standard Transformer has known limitations. We find, however, that the depth-adaptivity of the N-ODE Transformer does not provide a remedy for the inherently nonlocal nature of the parity problem.
  arXiv  Detail & Related papers  (2020-10-22T00:48:24Z)

This list is automatically generated from the titles and abstracts of the papers in this site.

This site does not guarantee the quality of this site (including all information) and is not responsible for any consequences.