键值表示：具有可扩展块循环压缩记忆的 Transformer

召回提出了一个困难的选择：Transformer 具有线性增长的记忆，会减慢每个连续标记的速度，而线性 RNN 通常具有固定成本，但召回有限。我们提出了键值均值（“KVM”），这是一种新颖的注意力块循环，可以适应固定大小或增长状态。为强大的 Transformer 基线配备固定大小的 KVM 注意力层，可以生成强大的 $O(N)$ 分块 RNN，同时仅添加少量新参数。我们使用可增长的 KVM 缓存训练 Transformer，并证明它在长上下文测试中具有竞争力，仅需要次二次预填充时间和次线性状态增长。 KVM 可通过标准操作实现，无需自定义内核，并支持分块并行化训练和预填充。它在一个统一的包中提供了传统 Transformer（可扩展上下文记忆、分块并行训练和预填充）和 RNN 的许多优点。它可以在每一层上使用，节省 KV 缓存记忆，并允许在 $O(N)$ 和 $O(N^2)$ 之间连续选择预填充时间复杂度。我们在 https://github.com/featherless-ai/KVM-paper 上发布了我们的代码，并在 Apache 2.0 许可证下在 https://huggingface.co/collections/featherless-ai/kvm-paper 上发布了经过训练的模型。