局部和全局注意力的双重维度

仅解码器 Transformer 计算对先前词元的 KV 缓存的注意力。键（和值）通常用相同的维度表示，无论其与预测目标的距离如何。然而，在自然语言中，下一个单词受前一个标记的影响最为强烈。我们假设本地和远程词元对表示能力提出了不对称的要求：本地词元对于预测即时输出更为关键，因此需要更丰富的表示，而远程词元主要用作远程记忆，较低维的表示可能就足够了。我们将这个想法形式化为距离自适应表示（DAR），在受控设置中实现，在本地上下文窗口内保留全维表示，同时将降维表示（例如原始维度的 1/4）分配给该窗口之外的标记。在多个预训练尺度（70M 到 410M 参数）以及 1B 尺度模型上的持续监督 微调 中，这种方法与全维基线的性能非常匹配。相反，统一降低所有词元位置的维度会导致性能更差。这些结果挑战了常见的假设，即键和值维度在词元位置上应该是统一的。我们的研究结果提出了设计注意力架构的新方向，该架构可以跨序列自适应地分配表示能力，从而在推理过程中进一步减少 KV 缓存。