Attention-Only Transformers and Implementing MLPs with Attention Heads (2309.08593v1)

Abstract: The transformer architecture is widely used in machine learning models and consists of two alternating sublayers: attention heads and MLPs. We prove that an MLP neuron can be implemented by a masked attention head with internal dimension 1 so long as the MLP's activation function comes from a restricted class including SiLU and close approximations of ReLU and GeLU. This allows one to convert an MLP-and-attention transformer into an attention-only transformer at the cost of greatly increasing the number of attention heads. We also prove that attention heads can perform the components of an MLP (linear transformations and activation functions) separately. Finally, we prove that attention heads can encode arbitrary masking patterns in their weight matrices to within arbitrarily small error.

• The paper demonstrates that any MLP layer can be simulated using masked attention heads with specific activation functions like SiLU and approximations of ReLU and GeLU.
• It introduces a method to convert conventional transformers into attention-only architectures by replacing MLP layers with a greater number of attention heads.
• The study highlights practical challenges in computational efficiency and training complexity while suggesting pathways for improved model interpretability.

The paper "Attention-Only Transformers and Implementing MLPs with Attention Heads" by Robert Huben and Valerie Morris explores the theoretical possibility of implementing Multilayer Perceptron (MLP) layers within transformer architectures entirely using attention mechanisms. This paper emerges within the context of improving the interpretability of transformers by converting them into attention-only variants, which could then be subjected to interpretability techniques well-suited for attention mechanisms.

Key Contributions:

1. MLP Representation with Attention Heads:
   • The authors demonstrate that a neuron within an MLP can be simulated using a masked attention head with an internal dimension of 1, given that the activation function belongs to a specific class, including SiLU and close approximations of ReLU and GeLU.
   • Using Theorem 1, they establish that any MLP layer, under these activation functions, can be expressed as a sum of masked attention heads. This involves a drastic increase in the number of attention heads required to match the MLP’s functionality.
2. Conversion of Transformers:
   • The paper introduces a method to convert a conventional transformer (with alternating attention and MLP layers) into an attention-only transformer. This is facilitated through Theorem 2, where they advocate for models that apply exclusively attention heads in lieu of MLP sublayers.
3. Component Functionality with Attention:
   • It is proven that attention heads can separately perform MLP components: linear transformations and activation functions. Theorems 3 and 4 address the capability of attention heads to implement row-wise linear operations and to apply generalized SiLU functions, respectively.
4. Encoding Masking Patterns in Weight Matrices:
   • The investigation is extended to encoding arbitrary masking patterns into the weight matrices of attention heads. Theorem 5 discusses a technique for embedding mask patterns directly into weight matrices and achieving functionality equivalent to original masking through parameter adjustments.
   • This approach, however, introduces practical challenges in training and could impact computational efficiency, given the substantial increase in required attention heads and potential inefficiencies in operations.

Practical Implications and Limitations:

• The substitution of MLP layers with attention heads is accompanied by a considerable increase in computational overhead, due to the need for many more attention heads to replicate MLP functionality.
• While theoretically viable, the method could negatively influence both training and inference efficiency, raising questions about practical implementation in current transformer architectures.
• Additionally, the process introduces significant complexities in terms of proscribed masking patterns, which complicate model implementations and could conflict with regularization techniques.

Future Directions:

• Despite practical limitations, this work suggests potential pathways for interpretable models, leveraging the expressiveness of attention-only architectures.
• Further exploration is warranted to assess whether attention-only architectures, as envisioned by the authors, would match or surpass traditional transformers in performance metrics without unrealistically impeding computational efficiency.
• The prospects for the integration of interpretability techniques that effectively manage the additional complexity introduced by this architecture change remain promising areas for future research.

Overall, this paper provides a significant theoretical expansion in understanding the capabilities of transformer architectures and offers insights into potential avenues for advancing interpretability in these systems.