Learning Linear Attention in Polynomial Time

TL;DR

This paper proves that single-layer Transformers with linear attention can be learned in polynomial time using strong, agnostic PAC learning, bridging the gap between their theoretical expressivity and practical learnability.

Contribution

It introduces the first polynomial-time learnability results for linear attention Transformers, connecting them to linear predictors in an RKHS and demonstrating their practical learnability.

Findings

Linear attention models are polynomial-time learnable.

Empirical validation on tasks like automata and key-value learning.

Examples include associative memories and finite automata.

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.