ATP：通过关注顶部主键实现快速大语言模型服务

我们提出了一种具有线性复杂度的新注意力机制 ATP，它将注意力固定在顶部主键 (Top Principal keys) 上，而不是每个单独的 token 上。特别是，ATP 是由一个重要的观察结果驱动的，即输入序列通常是低秩的，也就是说，输入序列可以由少数几个主基表示。因此，ATP 不直接遍历所有输入 token，而是将输入转换为正交空间，并仅计算顶部主基（键）上的注意力。由于输入序列中观察到的低秩结构，ATP 能够用少数几个主键捕捉输入序列中的语义关系。此外，注意力复杂度从二次方降低到线性，而没有引起明显的性能下降。ATP 进一步降低了具有低秩输入的其他线性层的复杂度，与仅针对注意力模块的先前工作相比，带来了更多的加速。我们在各种模型（如 BERT 和 Llama）上的评估表明，ATP 在计算和内存复杂度远低于标准注意力机制的情况下，实现了相当的准确性。特别是，ATP 在使用仅 $1/2$ 的主键时几乎不损失准确性，而在使用 $1/4$ 的主键时仅产生约 $2\%$ 的准确性下降。