arXiv:2501.14577v3 宣告类型: replace-cross 摘要：近年来，Transformer 成为序列建模架构中的一个基础构建块。然而，其核心在于自注意力机制，该机制的记忆消耗和计算成本随着序列长度 \(N\) 的平方增长，使得对于长序列而言变得极其昂贵。一种有前景的方法是 top-\(k\) 注意力，它只选择最相关的 \(k\) 个标记，并在显著降低空间和计算需求的同时，实现与普通的自注意力相当的性能。然而，因果掩码要求当前查询标记只能关注过去的标记，这阻止了现有的 top-\(k\) 注意力方法并行搜索最相关的标记，从而限制了训练效率。在这项工作中，我们提出了 ZETA，利用 \textbf{Z}-Order 曲线进行 \textbf{E}fficient \textbf{T}op-\(k\) \textbf{A}ttention，以实现对整个序列的过去标记的并行查询。虽然空间和时间复杂度为 \(\mathcal{O}(N \log N)\)。我们首先从理论上表明，键和查询维度的选择涉及到维度灾难和投影后相对距离保持之间的权衡。鉴于此见解，我们提议减少键和查询的维度，而与值的维度无关，同时利用 \(\textbf{Z}\)-Order 曲线将低维度的键和查询映射到 \(\emph{一}\) 维空间，这允许并行排序，从而大大提高了 top-\(k\) 标记选择的效率。实验结果表明，ZETA 在合成的 \textsc{Multi-Query Associative Recall} 任务中与标准注意力匹配，在 \textsc{Long Range Arena} 和 \textsc{WikiText-103} 语言建模任务中也优于注意力及其变体。