ZETA：利用 Z 序曲线实现高效 Top-k 注意力
机器学习
2025-04-02 v3 人工智能
摘要
近年来，Transformer 已成为序列建模架构的基本构建模块。然而，其核心是自注意力机制，该机制的内存和计算成本随序列长度 呈二次增长，导致长序列处理时变得昂贵。一种有前景的做法是采用 Top- 注意力，仅选择最相关的 个 token，即可在显著降低空间和计算需求的同时保持与标准自注意力相当的性能。然而，因因果掩码要求当前查询 token 只能注意过去的 token，现有的 Top- 注意力方法难以高效地并行搜索最相关的 token，从而限制了训练效率。本文提出 ZETA (Z-Order Curves for Efficient Top- Attention)，通过 Z 序曲线实现对整个序列的并行查询过去的 token。我们首先理论地表明，键和查询维度的选择在“维度灾难”与保持相对距离的投影之间存在权衡。基于此洞见，我们提出减少键和查询的维度（而非值），并利用 -order 曲线将低维键和查询映射到一维空间，从而实现并行排序，大幅提高 Top- token 选择的效率。实验结果表明，ZETA 在合成的 \textsc{Multi-Query Associative Recall} 任务上匹配标准注意力，在 \textsc{Long Range Arena} 和 \textsc{WikiText-103} 语言建模任务中超越注意力及其变体。
引用
@article{arxiv.2501.14577,
title = {ZETA: Leveraging Z-order Curves for Efficient Top-k Attention},
author = {Qiuhao Zeng and Jerry Huang and Peng Lu and Gezheng Xu and Boxing Chen and Charles Ling and Boyu Wang},
journal= {arXiv preprint arXiv:2501.14577},
year = {2025}
}
备注
25 pages, 4 figures, accepted in International Conference on Learning Representations (ICLR) 2025