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

arXiv 2502.20681 ↗First posted 28 Feb 2025cs.CL cs.AI cs.LGcited by 3

Not on the map. Why?

Abstract (arXiv)

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 could be conceptually attributed to 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 normalized ReLU self-attention 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 within this theoretical framework. A corollary further indicates that such a two-stage process is closely related to the spectral properties of attention weights.

Key facts. Cited by 3 papers in the map (data through 14 Oct 2026). Builds on 42 of its 54 references that are in the map. First posted on arXiv 28 Feb 2025.

Builds on

42 of 54 references.

References · 42

2018BERT: Pre-training of Deep Bidirectional Transformers for Language Understandingcited by 40,399
2020Language Models are Few-Shot Learnerscited by 32,419
2021Swin Transformer: Hierarchical Vision Transformer using Shifted Windowscited by 12,139
2021Masked Autoencoders Are Scalable Vision Learnerscited by 7,403
2021Emerging Properties in Self-Supervised Vision Transformerscited by 6,968
2022Locating and Editing Factual Associations in GPTcited by 2,691
2018Neural Tangent Kernel: Convergence and Generalization in Neural Networkscited by 2,621
2021Scaling Language Models: Methods, Analysis & Insights from Training Gophercited by 1,202
2018A Convergence Theory for Deep Learning via Over-Parameterizationcited by 787
2022What Can Transformers Learn In-Context? A Case Study of Simple Function Classescited by 715
2022GPT-NeoX-20B: An Open-Source Autoregressive Language Modelcited by 713
2022What learning algorithm is in-context learning? Investigations with linear modelscited by 684

30 more

2022Towards Reasoning in Large Language Models: A Surveycited by 596
2018Gradient Descent Finds Global Minima of Deep Neural Networkscited by 574
2023Trained Transformers Learn Linear Models In-Contextcited by 337
2023Transformers as Statisticians: Provable In-Context Learning with In-Context Algorithm Selectioncited by 298
2020Towards Understanding Ensemble, Knowledge Distillation and Self-Distillation in Deep Learningcited by 297
2018Learning Overparameterized Neural Networks via Stochastic Gradient Descent on Structured Datacited by 261
2019Towards Explaining the Regularization Effect of Initial Large Learning Rate in Training Neural Networkscited by 162
2023The Truth is in There: Improving Reasoning in Language Models with Layer-Selective Rank Reductioncited by 123
2021Toward Understanding the Feature Learning Process of Self-supervised Contrastive Learningcited by 119
2024The Evolution of Statistical Induction Heads: In-Context Learning Markov Chainscited by 118
2023In-Context Convergence of Transformerscited by 115
2023Scan and Snap: Understanding Training Dynamics and Token Composition in 1-layer Transformercited by 102
2023How Do Transformers Learn Topic Structure: Towards a Mechanistic Understandingcited by 86
2019How Much Over-parameterization Is Sufficient to Learn Deep ReLU Networks?cited by 84
2023A Theoretical Understanding of Shallow Vision Transformers: Learning, Generalization, and Sample Complexitycited by 77
2024Training Dynamics of Multi-Head Softmax Attention for In-Context Learning: Emergence, Convergence, and Optimalitycited by 74
2023Transformers Implement Functional Gradient Descent to Learn Non-Linear Functions In Contextcited by 63
2019Generating Sentences from Disentangled Syntactic and Semantic Spacescited by 59
2023A Study on ReLU and Softmax in Transformercited by 53
2023JoMA: Demystifying Multilayer Transformers via JOint Dynamics of MLP and Attentioncited by 52
2024Understanding Factual Recall in Transformers via Associative Memoriescited by 51
2024Transformers Learn Nonlinear Features In Context: Nonconvex Mean-field Dynamics on the Attention Landscapecited by 46
2023Replacing softmax with ReLU in Vision Transformerscited by 44
2019A Multi-Task Approach for Disentangling Syntax and Semantics in Sentence Representationscited by 37
2023On the Optimization and Generalization of Multi-head Attentioncited by 37
2021Disentangling Semantics and Syntax in Sentence Embeddings with Pre-trained Language Modelscited by 22
2024Unveil Benign Overfitting for Transformer in Vision: Training Dynamics, Convergence, and Generalizationcited by 17
2024On the Training Convergence of Transformers for In-Context Classification of Gaussian Mixturescited by 15
2024Learning and Transferring Sparse Contextual Bigrams with Linear Transformerscited by 9
2024On Mesa-Optimization in Autoregressively Trained Transformers: Emergence and Capabilitycited by 5

Built on by

Cited by year (2026 to 14 Oct)

Citing papers per year: 2026: 3.

3 papers.