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.

翻译：基于注意力的模型是深度学习近期多项突破的关键要素。注意力的两个核心组成部分包括其输入结构（由键、值和查询构成）以及这三者组合的计算方式。本文探索了共享上述输入结构但不受注意力计算约束的模型空间，我们将此空间称为键值查询空间。目标是确定该空间中是否存在注意力无法高效近似的其他可堆叠模型，这些模型可利用现有深度学习工具实现，并能解决学界关注的难题。出乎意料的是，标准最小二乘问题的解恰好满足这些特性。能够计算该解的神经网络模块不仅丰富了网络可表征的计算类型，且被严格证明是线性注意力的广义化形式。更令人惊讶的是，该模块的计算复杂度与注意力机制完全一致，使其成为理想的即插即用替代品。通过建立经典机器学习（最小二乘法）与现代深度学习（注意力机制）之间的这一创新性关联，我们验证了在该思路下推广常规注意力的变体模型。两个新模块在从少样本学习到策略蒸馏的广泛任务中经受检验，证实了它们的实际应用价值。