Learning Linear Attention in Polynomial Time (2410.10101v2)

Abstract: Previous research has explored the computational expressivity of Transformer models in simulating Boolean circuits or Turing machines. However, the learnability of these simulators from observational data has remained an open question. Our study addresses this gap by providing the first polynomial-time learnability results (specifically strong, agnostic PAC learning) for single-layer Transformers with linear attention. We show that linear attention may be viewed as a linear predictor in a suitably defined RKHS. As a consequence, the problem of learning any linear transformer may be converted into the problem of learning an ordinary linear predictor in an expanded feature space, and any such predictor may be converted back into a multiheaded linear transformer. Moving to generalization, we show how to efficiently identify training datasets for which every empirical risk minimizer is equivalent (up to trivial symmetries) to the linear Transformer that generated the data, thereby guaranteeing the learned model will correctly generalize across all inputs. Finally, we provide examples of computations expressible via linear attention and therefore polynomial-time learnable, including associative memories, finite automata, and a class of Universal Turing Machine (UTMs) with polynomially bounded computation histories. We empirically validate our theoretical findings on three tasks: learning random linear attention networks, key--value associations, and learning to execute finite automata. Our findings bridge a critical gap between theoretical expressivity and learnability of Transformers, and show that flexible and general models of computation are efficiently learnable.

• The paper demonstrates that multi-head linear attention architectures (MHLAs) can be learned in polynomial time by framing training as learning a linear predictor in an expanded feature space.
• It introduces a certifiable "second moment condition" that ensures functional equivalence of empirical risk minimizers for MHLAs, facilitating out-of-distribution generalization.
• MHLAs are shown to be expressive enough to simulate universal Turing machines with polynomially bounded computation histories, with empirical tests validating efficient learnability.

An Examination of Learning Linear Attention in Polynomial Time

The paper "Learning Linear Attention in Polynomial Time" addresses a critical gap in understanding the learnability of simplified Transformer models, particularly focusing on the polynomial-time learnability of single-layer Transformers with linear attention. This work investigates the theoretical aspects of learning Boolean circuits or simulating Turing machines using Transformers, providing strong agnostic PAC learning results within a reproducing kernel Hilbert space (RKHS).

Key Contributions

The paper offers several novel contributions to the field of computational learning theory applied to Transformers:

1. Polynomial-Time Learnability Arguably Established: The authors rigorously demonstrate that multi-head linear attention architectures (MHLAs) can be learned in polynomial time. They establish that training such models can be effectively converted into a problem of learning an ordinary linear predictor in an expanded feature space.
2. Certifiable Identifiability: A significant advancement is the introduction of a certifiable condition, termed "second moment condition," which ensures that empirical risk minimizers for MHLA are functionally equivalent across all data inputs. This facilitates out-of-distribution generalization.
3. Expressiveness of MHLAs: The expressiveness of MHLAs is highlighted through their capability to simulate a class of universal Turing machines with polynomially bounded computation histories. This is not only an intriguing theoretical finding but shows practical implications by proposing a method to learn models like associative memories and finite automata.
4. Empirical Validation: A suite of experiments supports the theoretical findings, emphasizing the role of head count over layers in optimization, the utility of certifiable identifiability in predicting generalization, and the efficient learnability of universal finite automata executors.

Implications and Discussion

The implications of this work are both theoretical and practical. Theoretically, it tightens the case for the potential efficiency of Transformers in terms of sample complexity and computational resources, when configured with linear attention. The empirical findings underscore these claims by showing tangible performance gains using over-parameterization (particularly in terms of adding heads rather than layers).

The paper posits avenues for future exploration, such as understanding the limits of identifiability and exploring the practical benefits of their identifiability certificate in real-world datasets. Moreover, the authors highlight the ability to simulate universal computation models efficiently as a point of convergence between theoretical computer science and machine learning, setting a stage for further exploration of whether such theoretical constructs can inspire better performing architectures in practice.

Future Developments

Given the findings of this paper, several potential future research directions arise. One significant direction is applying these learnability principles to enhance the architecture and training procedures of contemporary large-scale transformers. Understanding the balance between model architectural complexity and computational efficiency could yield improvements in both performance and sustainability of learning systems.

Additionally, exploring more complex data distributions beyond the assumptions set in this paper could provide more insights into the generalizability and robustness of linear attention models. Practical implementations that leverage the polynomial-time algorithms could further solidify these learning concepts as fundamental principles in building efficient and scalable machine learning systems.

In summary, this paper bridges a crucial gap between theory and practice in machine learning by presenting insights into the learnability of simple attention mechanisms used in Transformers. It provides a mathematical underpinning for efficiently learning these methods, offering a structured pathway to developing robust, computationally efficient algorithms with guaranteed performance bounds.