ALTA: Compiler-Based Analysis of Transformers

Peter Shaw,James Cohan,Jacob Eisenstein,Kenton Lee,Jonathan Berant,Kristina Toutanova
2024-10-24
Abstract:We propose a new programming language called ALTA and a compiler that can map ALTA programs to Transformer weights. ALTA is inspired by RASP, a language proposed by Weiss et al. (2021), and Tracr (Lindner et al., 2023), a compiler from RASP programs to Transformer weights. ALTA complements and extends this prior work, offering the ability to express loops and to compile programs to Universal Transformers, among other advantages. ALTA allows us to constructively show how Transformers can represent length-invariant algorithms for computing parity and addition, as well as a solution to the SCAN benchmark of compositional generalization tasks, without requiring intermediate scratchpad decoding steps. We also propose tools to analyze cases where the expressibility of an algorithm is established, but end-to-end training on a given training set fails to induce behavior consistent with the desired algorithm. To this end, we explore training from ALTA execution traces as a more fine-grained supervision signal. This enables additional experiments and theoretical analyses relating the learnability of various algorithms to data availability and modeling decisions, such as positional encodings. We make the ALTA framework -- language specification, symbolic interpreter, and weight compiler -- available to the community to enable further applications and insights.
Machine Learning,Artificial Intelligence,Computation and Language
What problem does this paper attempt to address?
The core problem that this paper attempts to solve is to understand and improve the ability of Transformer models in performing compositional generalization and "System 2" reasoning. Specifically, the authors propose a new framework - ALTA (A Language for Transformer Analysis) to compile interpretable symbolic programs into Transformer model weights. ### Main Problems and Background In recent years, there has been a great deal of discussion and controversy regarding whether Transformers can perform compositional generalization and complex reasoning tasks. Some studies have shown that certain types of Transformers perform poorly on specific evaluation tasks (such as Dziri et al., 2023; Qiu et al., 2023). Do these negative results reflect the variability of training methods or more fundamental architectural limitations? To better understand this, an interpretable framework is required to clarify whether Transformers can represent and learn solutions to various tasks. ### Goals of the ALTA Framework The ALTA framework aims to provide a new programming language and compiler to convert high - level programming logic into Transformer weights. In this way, it can: 1. **Demonstrate constructive expressive power**: Prove how Transformers can implement length - invariant algorithms, such as calculating parity and addition. 2. **Solve the SCAN benchmark**: Demonstrate how to use the Shift - Reduce parsing algorithm to solve the compositional generalization tasks in the SCAN (Lake & Baroni, 2018) benchmark. 3. **Analyze the reasons for training failure**: When it is known that an algorithm is expressible, but end - to - end training fails to induce the desired behavior, ALTA can help analyze the reasons and propose improvement plans. ### Key Contributions 1. **Support for dynamic control flow**: Compared with previous RASP and Tracr, ALTA supports loop structures and can implement certain forms of loops without relying on intermediate decoding steps. 2. **Compile to Universal Transformers**: ALTA can compile programs to Universal Transformers, thereby supporting conditional computation mechanisms and enhancing the model's expressive power. 3. **Sparse transformation rules**: ALTA represents the calculations of the MLP sub - layer as sparse transformation rules, reducing the number of required MLP hidden dimensions and improving the model's efficiency. ### Application Examples 1. **Constructive expressive power**: Demonstrates how Transformers can implement length - invariant parity and addition algorithms. 2. **Compositional generalization tasks**: Solves the compositional generalization tasks in the SCAN benchmark. 3. **Training analysis tool**: Provides a method for extracting intermediate supervision signals from ALTA execution trajectories to help diagnose and improve the model's training process. ### Theoretical Analysis The paper also introduces theoretical analysis tools to determine whether a program is minimal given a training set, that is, whether there are cases where some components can be removed without affecting the prediction of the training set. This analysis helps to understand whether the model parameters are fully specified during the training process and provides a theoretical basis for improving the model. ### Summary Through the ALTA framework, the authors not only demonstrate the potential of Transformers in expressing complex algorithms but also provide a set of tools and methods to help researchers better understand the learning behavior of Transformers, especially their performance in compositional generalization and reasoning tasks.