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

arXiv: https://arxiv.org/abs/2502.20681 · HTML: https://mlsanity.com/paper/2502.20681/

## 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 (most cited first, 30 of 42)

- [BERT: Pre-training of Deep Bidirectional Transformers for Language Understanding](https://mlsanity.com/paper/1810.04805/) (2018)
- [Language Models are Few-Shot Learners](https://mlsanity.com/paper/2005.14165/) (2020)
- [Swin Transformer: Hierarchical Vision Transformer using Shifted Windows](https://mlsanity.com/paper/2103.14030/) (2021)
- [Masked Autoencoders Are Scalable Vision Learners](https://mlsanity.com/paper/2111.06377/) (2021)
- [Emerging Properties in Self-Supervised Vision Transformers](https://mlsanity.com/paper/2104.14294/) (2021)
- [Locating and Editing Factual Associations in GPT](https://mlsanity.com/paper/2202.05262/) (2022)
- [Neural Tangent Kernel: Convergence and Generalization in Neural Networks](https://mlsanity.com/paper/1806.07572/) (2018)
- [Scaling Language Models: Methods, Analysis & Insights from Training Gopher](https://mlsanity.com/paper/2112.11446/) (2021)
- [A Convergence Theory for Deep Learning via Over-Parameterization](https://mlsanity.com/paper/1811.03962/) (2018)
- [What Can Transformers Learn In-Context? A Case Study of Simple Function Classes](https://mlsanity.com/paper/2208.01066/) (2022)
- [GPT-NeoX-20B: An Open-Source Autoregressive Language Model](https://mlsanity.com/paper/2204.06745/) (2022)
- [What learning algorithm is in-context learning? Investigations with linear models](https://mlsanity.com/paper/2211.15661/) (2022)
- [Towards Reasoning in Large Language Models: A Survey](https://mlsanity.com/paper/2212.10403/) (2022)
- [Gradient Descent Finds Global Minima of Deep Neural Networks](https://mlsanity.com/paper/1811.03804/) (2018)
- [Trained Transformers Learn Linear Models In-Context](https://mlsanity.com/paper/2306.09927/) (2023)
- [Transformers as Statisticians: Provable In-Context Learning with In-Context Algorithm Selection](https://mlsanity.com/paper/2306.04637/) (2023)
- [Towards Understanding Ensemble, Knowledge Distillation and Self-Distillation in Deep Learning](https://mlsanity.com/paper/2012.09816/) (2020)
- [Learning Overparameterized Neural Networks via Stochastic Gradient Descent on Structured Data](https://mlsanity.com/paper/1808.01204/) (2018)
- [Towards Explaining the Regularization Effect of Initial Large Learning Rate in Training Neural Networks](https://mlsanity.com/paper/1907.04595/) (2019)
- [The Truth is in There: Improving Reasoning in Language Models with Layer-Selective Rank Reduction](https://mlsanity.com/paper/2312.13558/) (2023)
- [Toward Understanding the Feature Learning Process of Self-supervised Contrastive Learning](https://mlsanity.com/paper/2105.15134/) (2021)
- [The Evolution of Statistical Induction Heads: In-Context Learning Markov Chains](https://mlsanity.com/paper/2402.11004/) (2024)
- [In-Context Convergence of Transformers](https://mlsanity.com/paper/2310.05249/) (2023)
- [Scan and Snap: Understanding Training Dynamics and Token Composition in 1-layer Transformer](https://mlsanity.com/paper/2305.16380/) (2023)
- [How Do Transformers Learn Topic Structure: Towards a Mechanistic Understanding](https://mlsanity.com/paper/2303.04245/) (2023)
- [How Much Over-parameterization Is Sufficient to Learn Deep ReLU Networks?](https://mlsanity.com/paper/1911.12360/) (2019)
- [A Theoretical Understanding of Shallow Vision Transformers: Learning, Generalization, and Sample Complexity](https://mlsanity.com/paper/2302.06015/) (2023)
- [Training Dynamics of Multi-Head Softmax Attention for In-Context Learning: Emergence, Convergence, and Optimality](https://mlsanity.com/paper/2402.19442/) (2024)
- [Transformers Implement Functional Gradient Descent to Learn Non-Linear Functions In Context](https://mlsanity.com/paper/2312.06528/) (2023)
- [Generating Sentences from Disentangled Syntactic and Semantic Spaces](https://mlsanity.com/paper/1907.05789/) (2019)

## Built on by (most cited first, 3 of 3)

- [LAMMI-Pathology: A Tool-Centric Bottom-Up LVLM-Agent Framework for Molecularly Informed Medical Intelligence in Pathology](https://mlsanity.com/paper/2602.18773/) (2026)
- [Hallucination is a Consequence of Space-Optimality: A Rate-Distortion Theorem for Membership Testing](https://mlsanity.com/paper/2602.00906/) (2026)
- [Effective Frontiers: A Unification of Neural Scaling Laws](https://mlsanity.com/paper/2602.02593/) (2026)

## 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.