JoMA: Demystifying Multilayer Transformers via JOint Dynamics of MLP and Attention (2310.00535v3)

Abstract: 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.

• The paper introduces the JoMA framework that mathematically integrates MLP and self-attention, shedding light on Transformer training dynamics.
• It reveals a two-phase convergence where attention moves from a sparse focus on salient tokens to a denser distribution over time.
• The study highlights implicit hierarchical learning in Transformers, offering insights for more efficient training and improved model design.

Analysis of JoMA: Joint Dynamics of Multilayer Transformers

The paper "JoMA: Demystifying Multilayer Transformers via JOint Dynamics of MLP and Attention" presents a novel framework, JoMA, which aims to mathematically elucidate the multifaceted training dynamics of multilayer Transformer architectures. Key to this paper is the integration of multi-layer perceptron (MLP) and self-attention mechanisms—the two primary components of the Transformer model—within a unified mathematical paradigm. By exploring the interactions between these components, the authors seek to enhance our understanding of how Transformers achieve their impressive capabilities.

The JoMA framework introduces an invariant representation that effectively eliminates the explicit need to model self-attention as a separate parameterized entity during training. This novel approach results in modified dynamics where attention is indirectly captured through the MLP layers alone. The theoretical underpinning suggests that during training, attention behavior initially becomes sparse, focusing on the most salient tokens before becoming denser to incorporate tokens with less pronounced salience. This behavior is akin to the inductive biases seen throughout machine learning, wherein more evident patterns are learned first before subtler ones are incorporated.

Key Insights and Findings

1. Linear and Nonlinear Dynamics:
   • In the context of a linear activation function within the MLP, the updates to the self-attention weights suggest a winner-take-all dynamic, whereby the system tends towards emphasizing the most prominent features.
   • With nonlinear activations, attention dynamics exhibit a two-phase convergence: more significant components are prioritized, followed by a gradual capture of minor components. This observation is crucial for understanding the temporal behavior of learned representations, especially in deeper Transformer layers.
2. Attention and Sparsity:
   • The framework predicts attention patterns that oscillate between sparse and dense distributions. Such dynamics were empirically validated using both synthetic and real-world data, including experiments with pre-trained models such as OPT and Pythia.
   • The observed attention sparsity, along with its "drop-and-bounce back" characteristic, aligns with the theoretical predictions of JoMA, highlighting its potential for explaining multistage learning processes that Transformers might utilize.
3. Hierarchical Learning:
   • The paper extends to explore how multi-layer Transformers can implicitly learn hierarchies in data distributions without explicit supervision. Using hierarchical binary latent tree (HBLT) models, the authors illustrate how the model effectively moves from learning direct associations within lower layers to forming more complex structures as depth increases.

Implications and Speculations

The JoMA framework opens pathways for a nuanced understanding of Transformers, potentially guiding future model architecture designs and optimization strategies. By detailing the dynamics of both linear and nonlinear scenarios, this work paves the way for improved model interpretability. Moreover, understanding when and how Transformers learn different information hierarchies could assist in designing models that are more efficient with training data.

Future advancements might focus on integrating these findings into the development of more efficient training algorithms. Additionally, exploring how embedding vector interactions play a role in these dynamics presents another frontier. Given the framework's reliance on assumption simplifications—such as orthogonality and independent dynamics—a promising area lies in relaxing these constraints to accommodate real-world complexities.

In conclusion, JoMA provides a detailed, mathematically rigorous attempt at demystifying how Transformers learn, revealing intricate dynamics that converge to enable robust and diverse context understanding in various tasks. This framework contributes significantly to our theoretical grasp of deep learning architectures, with potential ramifications for the broader field of AI research and application.