Learning Linear Attention in Polynomial Time

Fuente: arXiv
Salvato in:
Dettagli Bibliografici
Autori principali: Yau, Morris, Akyürek, Ekin, Mao, Jiayuan, Tenenbaum, Joshua B., Jegelka, Stefanie, Andreas, Jacob
Natura: Preprint
Pubblicazione: 2024
Soggetti:
Accesso online:
Tags: Aggiungi Tag
Nessun Tag, puoi essere il primo ad aggiungerne!!
_version_ 1866911228564078592
author Yau, Morris
Akyürek, Ekin
Mao, Jiayuan
Tenenbaum, Joshua B.
Jegelka, Stefanie
Andreas, Jacob
author_facet Yau, Morris
Akyürek, Ekin
Mao, Jiayuan
Tenenbaum, Joshua B.
Jegelka, Stefanie
Andreas, Jacob
contents 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.
format Preprint
id arxiv_https___arxiv_org_abs_2410_10101
institution arXiv
publishDate 2024
record_format arxiv
spellingShingle Learning Linear Attention in Polynomial Time
Yau, Morris
Akyürek, Ekin
Mao, Jiayuan
Tenenbaum, Joshua B.
Jegelka, Stefanie
Andreas, Jacob
Machine Learning
Artificial Intelligence
Computation and Language
Data Structures and Algorithms
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.
title Learning Linear Attention in Polynomial Time
topic Machine Learning
Artificial Intelligence
Computation and Language
Data Structures and Algorithms
url https://arxiv.org/abs/2410.10101