mlsanityData to 7 Oct 2026

JoMA: Demystifying Multilayer Transformers via JOint Dynamics of MLP and Attention

arXiv 2310.00535 ↗First posted 1 Oct 2023cs.LG cs.AI cs.CLcited by 52

Not on the map. Why?

Abstract (arXiv)

We propose Joint MLP/Attention (JoMA) dynamics, a novel mathematical framework to understand the training procedure of multilayer Transformer architectures. This is achieved by integrating out the self-attention layer in Transformers, producing a modified dynamics of MLP layers only. JoMA removes unrealistic assumptions in previous analysis (e.g., lack of residual connection) and predicts that the attention first becomes sparse (to learn salient tokens), then dense (to learn less salient tokens) in the presence of nonlinear activations, while in the linear case, it is consistent with existing works that show attention becomes sparse over time. We leverage JoMA to qualitatively explains how tokens are combined to form hierarchies in multilayer Transformers, when the input tokens are generated by a latent hierarchical generative model. Experiments on models trained from real-world dataset (Wikitext2/Wikitext103) and various pre-trained models (OPT, Pythia) verify our theoretical findings. Code can be found in https://github.com/facebookresearch/luckmatters/tree/yuandong3.

Key facts. Cited by 52 papers in the map (data through 7 Oct 2026). Builds on 58 of its 70 references that are in the map. First posted on arXiv 1 Oct 2023.

Builds on

58 of 70 references.

References · 58

  1. 2017Attention Is All You Needcited by 59,932
  2. 2018BERT: Pre-training of Deep Bidirectional Transformers for Language Understandingcited by 40,204
  3. 2020An Image is Worth 16x16 Words: Transformers for Image Recognition at Scalecited by 29,915
  4. 2022OPT: Open Pre-trained Transformer Language Modelscited by 3,464
  5. 2016Pointer Sentinel Mixture Modelscited by 2,992
  6. 2018Neural Tangent Kernel: Convergence and Generalization in Neural Networkscited by 2,597
  7. 2023Pythia: A Suite for Analyzing Large Language Models Across Training and Scalingcited by 1,734
  8. 2022In-context Learning and Induction Headscited by 930
  9. 2018A Convergence Theory for Deep Learning via Over-Parameterizationcited by 784
  10. 2022Transformers learn in-context by gradient descentcited by 714
  11. 2022What Can Transformers Learn In-Context? A Case Study of Simple Function Classescited by 700
  12. 2018On Lazy Training in Differentiable Programmingcited by 671
38 more
  1. 2022What learning algorithm is in-context learning? Investigations with linear modelscited by 671
  2. 2018Gradient Descent Provably Optimizes Over-parameterized Neural Networkscited by 655
  3. 2018Gradient Descent Finds Global Minima of Deep Neural Networkscited by 572
  4. 2019Fine-Grained Analysis of Optimization and Generalization for Overparameterized Two-Layer Neural Networkscited by 562
  5. 2018A Mean Field View of the Landscape of Two-Layers Neural Networkscited by 463
  6. 2018On the Global Convergence of Gradient Descent for Over-parameterized Models using Optimal Transportcited by 422
  7. 2019Are Transformers universal approximators of sequence-to-sequence functions?cited by 404
  8. 2023Trained Transformers Learn Linear Models In-Contextcited by 331
  9. 2023Transformers as Statisticians: Provable In-Context Learning with In-Context Algorithm Selectioncited by 291
  10. 2022Tensor Programs V: Tuning Large Neural Networks via Zero-Shot Hyperparameter Transfercited by 272
  11. 2018Learning Overparameterized Neural Networks via Stochastic Gradient Descent on Structured Datacited by 260
  12. 2022Exploring Length Generalization in Large Language Modelscited by 204
  13. 2022Hidden Progress in Deep Learning: SGD Learns Parities Near the Computational Limitcited by 184
  14. 2018A Convergence Analysis of Gradient Descent for Deep Linear Neural Networkscited by 180
  15. 2019Towards moderate overparameterization: global convergence guarantees for training shallow neural networkscited by 176
  16. 2023Birth of a Transformer: A Memory Viewpointcited by 154
  17. 2021Inductive Biases and Variable Creation in Self-Attention Mechanismscited by 144
  18. 2020Infinite attention: NNGP and NTK for deep attention networkscited by 117
  19. 2021Self-Attention Networks Can Process Bounded Hierarchical Languagescited by 105
  20. 2023Scan and Snap: Understanding Training Dynamics and Token Composition in 1-layer Transformercited by 101
  21. 2022Vision Transformers provably learn spatial structurecited by 92
  22. 2020On the Computational Power of Transformers and its Implications in Sequence Modelingcited by 85
  23. 2023How Do Transformers Learn Topic Structure: Towards a Mechanistic Understandingcited by 85
  24. 2023A Theoretical Understanding of Shallow Vision Transformers: Learning, Generalization, and Sample Complexitycited by 77
  25. 2020A Rigorous Framework for the Mean Field Limit of Multilayer Neural Networkscited by 68
  26. 2023Transformers as Support Vector Machinescited by 67
  27. 2020Understanding Self-supervised Learning with Dual Deep Networkscited by 62
  28. 2017Globally Optimal Gradient Descent for a ConvNet with Gaussian Inputscited by 61
  29. 2023Max-Margin Token Selection in Attention Mechanismcited by 57
  30. 2020A Mean-field Analysis of Deep ResNet and Beyond: Towards Provable Optimization Via Overparameterization From Depthcited by 56
  31. 2017Learning ReLUs via Gradient Descentcited by 54
  32. 2017An Analytical Formula of Population Gradient for two-layered ReLU network and its Applications in Convergence and Critical Point Analysiscited by 52
  33. 2023Transformers learn through gradual rank increasecited by 52
  34. 2023On the Role of Attention in Prompt-tuningcited by 49
  35. 2023Do Transformers Parse while Predicting the Masked Word?cited by 48
  36. 2018Gradient descent with identity initialization efficiently learns positive definite linear transformations by deep residual networkscited by 44
  37. 2020Modeling from Features: a Mean-field Framework for Over-parameterized Deep Neural Networkscited by 44
  38. 2021On the Expressive Power of Self-Attention Matricescited by 44

50 most cited of 58 All 58 (JSON)

Built on by

Cited by year (2026 to 7 Oct)

Citing papers per year: 2023: 1; 2024: 35; 2025: 12; 2026: 4.

52 papers.

Sort:

By area:

  1. 2024GaLore: Memory-Efficient LLM Training by Gradient Low-Rank Projectioncited by 447
  2. 2024How Transformers Learn Causal Structure with Gradient Descentcited by 114
  3. 2024A Primer on the Inner Workings of Transformer-based Language Modelscited by 83
  4. 2024Found in the Middle: How Language Models Use Long Contexts Better via Plug-and-Play Positional Encodingcited by 72
  5. 2024On the Societal Impact of Open Foundation Modelscited by 64
  6. 2024Active-Dormant Attention Heads: Mechanistically Demystifying Extreme-Token Phenomena in LLMscited by 54
  7. 2024Why Larger Language Models Do In-context Learning Differently?cited by 44
  8. 2024An Information-Theoretic Analysis of In-Context Learningcited by 39
  9. 2024How Do Nonlinear Transformers Learn and Generalize in In-Context Learning?cited by 35
  10. 2025Transformers Learn to Implement Multi-step Gradient Descent with Chain of Thoughtcited by 32
  11. 2024Implicit Bias and Fast Convergence Rates for Self-attentioncited by 31
  12. 2024Linear Transformers are Versatile In-Context Learnerscited by 28
38 more
  1. 2024Towards a Theoretical Understanding of the 'Reversal Curse' via Training Dynamicscited by 27
  2. 2024In-Context Learning with Transformers: Softmax Attention Adapts to Function Lipschitznesscited by 22
  3. 2024Transformers Provably Learn Sparse Token Selection While Fully-Connected Nets Cannotcited by 21
  4. 2024Bypassing the Exponential Dependency: Looped Transformers Efficiently Learn In-context by Multi-step Gradient Descentcited by 21
  5. 2024SWAN: SGD with Normalization and Whitening Enables Stateless LLM Trainingcited by 20
  6. 2024Self-attention Networks Localize When QK-eigenspectrum Concentratescited by 19
  7. 2024Towards Tracing Trustworthiness Dynamics: Revisiting Pre-training Period of Large Language Modelscited by 18
  8. 2024On the Surprising Effectiveness of Attention Transfer for Vision Transformerscited by 16
  9. 2025Emergence of Superposition: Unveiling the Training Dynamics of Chain of Continuous Thoughtcited by 16
  10. 2024Implicit Optimization Bias of Next-Token Prediction in Linear Modelscited by 14
  11. 2024Do LLMs dream of elephants (when told not to)? Latent concept association and associative memory in transformerscited by 14
  12. 2024Training Nonlinear Transformers for Chain-of-Thought Inference: A Theoretical Generalization Analysiscited by 14
  13. 2024Non-asymptotic Convergence of Training Transformers for Next-token Predictioncited by 13
  14. 2024Next-token prediction capacity: general upper bounds and a lower bound for transformerscited by 11
  15. 2025Gating is Weighting: Understanding Gated Linear Attention through In-context Learningcited by 11
  16. 2024Understanding and Minimising Outlier Features in Neural Network Trainingcited by 10
  17. 2025How Transformers Learn Regular Language Recognition: A Theoretical Study on Training Dynamics and Implicit Biascited by 10
  18. 2025Generalization or Hallucination? Understanding Out-of-Context Reasoning in Transformerscited by 10
  19. 2024Optimizing Attention with Mirror Descent: Generalized Max-Margin Token Selectioncited by 9
  20. 2025From Condensation to Rank Collapse: A Two-Stage Analysis of Transformer Training Dynamicscited by 9
  21. 2024Learning on Transformers is Provable Low-Rank and Sparse: A One-layer Analysiscited by 8
  22. 2024Training Dynamics of Transformers to Recognize Word Co-occurrence via Gradient Flow Analysiscited by 8
  23. 2025Taming Transformer Without Using Learning Rate Warmupcited by 8
  24. 2023The Expressibility of Polynomial based Attention Schemecited by 7
  25. 2025Transformer learns the cross-task prior and regularization for in-context learningcited by 7
  26. 2024On the Learn-to-Optimize Capabilities of Transformers in In-Context Sparse Recoverycited by 6
  27. 2024Exploring Gradient Subspaces: Addressing and Overcoming LoRA's Limitations in Federated Fine-Tuning of Large Language Modelscited by 6
  28. 2024Provably Neural Active Learning Succeeds via Prioritizing Perplexing Samplescited by 5
  29. 2024On the Power of Convolution Augmented Transformercited by 5
  30. 2026Large Vision-Language Models Get Lost in Attentioncited by 5
  31. 2024Enhancing In-Context Learning Performance with just SVD-Based Weight Pruning: A Theoretical Perspectivecited by 4
  32. 2025Disentangling Feature Structure: A Mathematically Provable Two-Stage Training Dynamics in Transformerscited by 3
  33. 2026Unlabeled Data Can Provably Enhance In-Context Learning of Transformerscited by 2
  34. 2026Breaking the Reversal Curse in Autoregressive Language Models via Identity Bridgecited by 2
  35. 2024A Theoretical Survey on Foundation Modelscited by 1
  36. 2024Sequential Compression Layers for Efficient Federated Learning in Foundational Modelscited by 1
  37. 2025How Transformers Learn In-Context Recall Tasks? Optimality, Training Dynamics and Generalizationcited by 1
  38. 2025Towards Understanding Transformers in Learning Random Walkscited by 1

50 most cited of 52 All 52 (JSON)