Exploring the Space of Key-Value-Query Models with Intention
Marta Garnelo Wojciech Marian Czarnecki
Attention-based models have been a key element of many recent breakthroughs
in deep learning. Two key components of Attention are the structure of its
input (which consists of keys, values and queries) and the computations by
which these three are combined. In this paper we explore the space of models
that share said input structure but are not restricted to the computations of
Attention. We refer to this space as Keys-Values-Queries (KVQ) Space. Our goal
is to determine whether there are any other stackable models in KVQ Space that
Attention cannot efficiently approximate, which we can implement with our
current deep learning toolbox and that solve problems that are interesting to
the community. Maybe surprisingly, the solution to the standard least squares
problem satisfies these properties. A neural network module that is able to
compute this solution not only enriches the set of computations that a neural
network can represent but is also provably a strict generalisation of Linear
Attention. Even more surprisingly the computational complexity of this module
is exactly the same as that of Attention, making it a suitable drop in
replacement. With this novel connection between classical machine learning
(least squares) and modern deep learning (Attention) established we justify a
variation of our model which generalises regular Attention in the same way.
Both new modules are put to the test an a wide spectrum of tasks ranging from
few-shot learning to policy distillation that confirm their real-worlds
applicability.