Research map / Architectures and efficiency / Attention and transformers
JoMA: Demystifying Multilayer Transformers via Joint Dynamics of MLP and Attention
Read the paperCheck your idea against 108,967 papers
AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.
Proposes JoMA, a framework that integrates out self-attention to derive joint training dynamics of MLP and attention, showing attention first sparsifies then densifies under nonlinear activations and explaining hierarchical token composition.
Problem
Understanding how the basic Transformer block (self-attention plus MLP) learns, especially in multilayer Transformers, where existing theory is limited to 1-layer models, fixed MLPs, linear activations, or unrealistic assumptions such as no residual connections.
Approach
Integrate out the self-attention logits so that attention is expressed as a closed-form function of the lower MLP weights (first integral/invariant), yielding joint dynamics of MLP and attention with residual connections and nonlinear activations; combined with a hierarchical latent tree model, this explains how multilayer Transformers learn token hierarchies.
Key results
- With nonlinear activations, attention first becomes sparse (learning salient tokens) then denser, while linear case yields monotone sparsification consistent with Scan&Snap.
- Learned hidden MLP nodes correlate with latent variables with normalized correlation about 0.99 at the lowest layer and 0.55-0.81 at the next layer in synthetic experiments.
- Attention sparsity drop-and-bounce patterns are observed in training from scratch on Wikitext2/Wikitext103 and qualitatively in OPT-2.7B and Pythia models, with stable ranks of lower MLP layers showing salient drop-and-bounce patterns in top layers.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Use nonlinear MLP activation with self-attention instead of linear MLP activation | JoMA dynamics · Transformer training dynamics | attention entropy | Mixed |
| Increase generative-model complexity (larger Nch and latent counts) | 3-layer Transformer trained on HBLT data · HBLT latent-to-hidden-node alignment | NCorr | Hurt |
Surprising findings
- Under nonlinear activations, the attention entropy drops and then rebounds as non-salient tokens are learned later, and the attention patterns giving best validation loss coincide with the theoretical predictions.
Limitations
- Assumes orthonormal token embeddings (d sufficiently large), which may not hold when d is smaller than vocabulary size.
- Assumes stationary backpropagated gradients and fixed (untrained) embeddings and treats attention logits as free parameters rather than parameterized by embeddings.
- Softmax-case invariant relies on an additional approximation assumption that the mean attention is constant over time; multilayer analysis is qualitative.
Experimental setup
Datasets: Wikitext2 Wikitext103 synthetic HBLT data
Benchmarks: Wikitext2 Wikitext103
Metrics: attention entropy/sparsity stable rank normalized correlation between latents and hidden nodes validation loss
Baselines: Scan&Snap
Builds on: Scan&Snap linear attention dynamics multi-layer linear network dynamics
Closest papers by meaning
- On the Role of Attention Masks and LayerNorm in Transformers (NeurIPS 2024 · similarity 0.89)
- Transformers Learn Nonlinear Features In Context: Nonconvex Mean-field Dynamics on the Attention Landscape (ICML 2024 · similarity 0.88)
- Training-Induced Escape from Token Clustering in a Mean-Field Formulation of Transformers (NeurIPS 2026 · similarity 0.88)
- Transformers Learn Faster with Semantic Focus (NeurIPS 2025 · similarity 0.88)
- In-context Convergence of Transformers (ICML 2024 · similarity 0.87)
- Unveiling Induction Heads: Provable Training Dynamics and Feature Learning in Transformers (NeurIPS 2024 · similarity 0.87)
Writing a paper in this area? scime draft is a free LaTeX editor with AI included, real-time co-authors and a check of your references and venue rules before you submit. It also imports OpenAI Prism projects.