Disentangling Feature Structure: A Mathematically Provable Two-Stage Training Dynamics in Transformers

• URL: http://arxiv.org/abs/2502.20681v2
• Date: Sat, 11 Oct 2025 04:45:15 GMT
• Title: Disentangling Feature Structure: A Mathematically Provable Two-Stage Training Dynamics in Transformers
• Authors: Zixuan Gong, Shijia Li, Yong Liu, Jiaye Teng,
• Abstract summary: This paper theoretically demonstrates how the two-stage training dynamics potentially occur in transformers.<n>To our best knowledge, this is the first rigorous result regarding a feature-level two-stage optimization process in transformers.
• Score: 18.662154648423087
• License: http://arxiv.org/licenses/nonexclusive-distrib/1.0/
• Abstract: Transformers may exhibit two-stage training dynamics during the real-world training process. For instance, when training GPT-2 on the Counterfact dataset, the answers progress from syntactically incorrect to syntactically correct to semantically correct. However, existing theoretical analyses hardly account for this feature-level two-stage phenomenon, which originates from the disentangled two-type features like syntax and semantics. In this paper, we theoretically demonstrate how the two-stage training dynamics potentially occur in transformers. Specifically, we analyze the feature learning dynamics induced by the aforementioned disentangled two-type feature structure, grounding our analysis in a simplified yet illustrative setting that comprises a normalized ReLU self-attention layer and structured data. Such disentanglement of feature structure is general in practice, e.g., natural languages contain syntax and semantics, and proteins contain primary and secondary structures. To our best knowledge, this is the first rigorous result regarding a feature-level two-stage optimization process in transformers. Additionally, a corollary indicates that such a two-stage process is closely related to the spectral properties of the attention weights, which accords well with our empirical findings.

Related papers

• Understanding the Staged Dynamics of Transformers in Learning Latent Structure [5.944972519558522]
  We train a small decoder-only transformer on three task variants.<n>We show that the model acquires capabilities in discrete stages.<n>We also identify a crucial asymmetry, where the model can compose fundamental rules robustly, but struggles to decompose complex examples to discover the fundamental rules.
  arXiv  Detail & Related papers  (2025-11-24T17:20:42Z)
• The Kinetics of Reasoning: How Chain-of-Thought Shapes Learning in Transformers? [25.29458951592086]
  Chain-of-thought (CoT) supervision can substantially improve transformer performance.<n>We investigate these learning dynamics through the lens of grokking by pretraining transformers on symbolic reasoning tasks.
  arXiv  Detail & Related papers  (2025-10-28T20:14:26Z)
• How do Transformers Learn Implicit Reasoning? [67.02072851088637]
  We study how implicit multi-hop reasoning emerges by training transformers from scratch in a controlled symbolic environment.<n>We find that training with atomic triples is not necessary but accelerates learning, and that second-hop generalization relies on query-level exposure to specific compositional structures.
  arXiv  Detail & Related papers  (2025-05-29T17:02:49Z)
• How Transformers Learn Regular Language Recognition: A Theoretical Study on Training Dynamics and Implicit Bias [48.9399496805422]
  We focus on two representative tasks in the category of regular language recognition, known as even pairs' and parity check'<n>Our goal is to explore how a one-layer transformer, consisting of an attention layer followed by a linear layer, learns to solve these tasks.
  arXiv  Detail & Related papers  (2025-05-02T00:07:35Z)
• Exact Learning Dynamics of In-Context Learning in Linear Transformers and Its Application to Non-Linear Transformers [1.7034813545878589]
  Transformer models exhibit remarkable in-context learning (ICL) Our work offers an exact dynamical model for ICL and theoretically grounded tools for analyzing complex transformer training.
  arXiv  Detail & Related papers  (2025-04-17T13:05:33Z)
• Interpreting Affine Recurrence Learning in GPT-style Transformers [54.01174470722201]
  In-context learning allows GPT-style transformers to generalize during inference without modifying their weights. This paper focuses specifically on their ability to learn and predict affine recurrences as an ICL task. We analyze the model's internal operations using both empirical and theoretical approaches.
  arXiv  Detail & Related papers  (2024-10-22T21:30:01Z)
• How Transformers Get Rich: Approximation and Dynamics Analysis [11.789846138681359]
  We provide both approximation and dynamics analyses of how transformers implement induction heads.<n>In the em approximation analysis, we formalize both standard and generalized induction head mechanisms.<n>For the em dynamics analysis, we study the training dynamics on a synthetic mixed target, composed of a 4-gram and an in-context 2-gram component.
  arXiv  Detail & Related papers  (2024-10-15T10:22:27Z)
• Training Dynamics of Transformers to Recognize Word Co-occurrence via Gradient Flow Analysis [97.54180451650122]
  We study the dynamics of training a shallow transformer on a task of recognizing co-occurrence of two designated words. We analyze the gradient flow dynamics of simultaneously training three attention matrices and a linear layer. We prove a novel property of the gradient flow, termed textitautomatic balancing of gradients, which enables the loss values of different samples to decrease almost at the same rate and further facilitates the proof of near minimum training loss.
  arXiv  Detail & Related papers  (2024-10-12T17:50:58Z)
• Non-asymptotic Convergence of Training Transformers for Next-token Prediction [48.9399496805422]
  Transformers have achieved extraordinary success in modern machine learning due to their excellent ability to handle sequential data. This paper provides a fine-grained non-asymptotic analysis of the training dynamics of a one-layer transformer. We show that the trained transformer presents non-token prediction ability with dataset shift.
  arXiv  Detail & Related papers  (2024-09-25T20:22:06Z)
• Strengthening Structural Inductive Biases by Pre-training to Perform Syntactic Transformations [75.14793516745374]
  We propose to strengthen the structural inductive bias of a Transformer by intermediate pre-training. Our experiments confirm that this helps with few-shot learning of syntactic tasks such as chunking. Our analysis shows that the intermediate pre-training leads to attention heads that keep track of which syntactic transformation needs to be applied to which token.
  arXiv  Detail & Related papers  (2024-07-05T14:29:44Z)
• How Do Transformers Learn In-Context Beyond Simple Functions? A Case Study on Learning with Representations [98.7450564309923]
  This paper takes initial steps on understanding in-context learning (ICL) in more complex scenarios, by studying learning with representations. We construct synthetic in-context learning problems with a compositional structure, where the label depends on the input through a possibly complex but fixed representation function. We show theoretically the existence of transformers that approximately implement such algorithms with mild depth and size.
  arXiv  Detail & Related papers  (2023-10-16T17:40:49Z)
• In-Context Convergence of Transformers [63.04956160537308]
  We study the learning dynamics of a one-layer transformer with softmax attention trained via gradient descent. For data with imbalanced features, we show that the learning dynamics take a stage-wise convergence process.
  arXiv  Detail & Related papers  (2023-10-08T17:55:33Z)
• Transformers learn in-context by gradient descent [58.24152335931036]
  Training Transformers on auto-regressive objectives is closely related to gradient-based meta-learning formulations. We show how trained Transformers become mesa-optimizers i.e. learn models by gradient descent in their forward pass.
  arXiv  Detail & Related papers  (2022-12-15T09:21:21Z)

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.