Abstract
Recurrent neural networks (RNNs) and self-attention are both widely used sequence-mixing layers that maintain an internal memory. However, this memory is constructed using two orthogonal mechanisms: RNNs compress the entire past into a fixed-size state, whereas self-attention's state stores every past time step growing its state (the KV cache) linearly with the sequence length. This results in orthogonal strengths and weaknesses. Self-attention layers excel at retrieving information in the context but have large memory and computational costs, while RNNs are more efficient but degrade over longer contexts and underperform for precise recall tasks. Prior work combining these mechanisms has focused primarily on naively interleaving them to reduce computational cost without regard to their complementary mechanisms. We propose the Hybrid Associative Memory (HAM) layer, which combines self-attention and RNNs while leveraging their individual strengths: the RNN compresses the entire sequence, while attention supplements it *only* with information that is difficult for the RNN to predict, which is hence the most valuable information to explicitly store. HAM layers enable data-dependent growth of the KV cache, which can be precisely controlled by the user with a single, continuous threshold. We find that this fine-grained control of the KV cache growth rate has a smooth trade-off with loss and performance. Empirically, we show that our hybrid architecture offers strong, competitive performance relative to RNNs and Transformers even at substantially lower KV-cache usage.
Abstract (translated)
循环神经网络（RNN）和自注意力机制都是广泛使用的序列混合层，它们都维持着内部记忆。然而，这种记忆是通过两种正交的机制构建的：RNN将整个过去压缩为固定大小的状态，而自注意力的状态则存储每个过去的时间步，使其状态（KV缓存）随序列长度线性增长。这导致了互补的优缺点。自注意力层在从上下文中检索信息方面表现出色，但具有较大的内存和计算成本；而RNN更高效，但在更长上下文上性能会下降，且在精确回忆任务中表现不佳。先前结合这些机制的工作主要集中在简单地将它们交替使用，以减少计算成本，而忽略了它们互补的机制。我们提出了混合关联记忆（HAM）层，它结合了自注意力和RNN，同时利用它们各自的优势：RNN压缩整个序列，而注意力仅用RNN难以预测的信息来补充它，因此这是最值得显式存储的信息。HAM层实现了KV缓存的数据依赖增长，用户可以通过一个单一的连续阈值进行精确控制。我们发现，对KV缓存增长率的这种细粒度控制与损失和性能之间存在平滑的权衡。实验表明，即使在KV缓存使用率显著降低的情况下，我们的混合架构也提供了与RNN和Transformer相比强大且具有竞争力的性能。
URL
https://arxiv.org/abs/2603.22325
