DistrAttention：现代GPU上高效灵活的自注意力机制
机器学习
2025-07-24 v1 人工智能
摘要
Transformer架构彻底改变了深度学习，在自然语言处理、计算机视觉和时间序列预测等领域取得了最先进的性能。然而，其核心组件自注意力机制相对于输入序列长度具有二次时间复杂度，这阻碍了Transformer的可扩展性。现有的优化自注意力的方法要么丢弃了全上下文信息，要么缺乏灵活性。在这项工作中，我们设计了DistrAttention，一种具有全上下文的高效且灵活的自注意力机制。DistrAttention通过在嵌入维度（通常称为）上对数据进行分组来实现这一点。我们通过一种轻量级的采样和融合方法来实现DistrAttention，该方法利用局部敏感哈希对相似数据进行分组。进一步设计了一种分块分组框架，以限制局部敏感哈希引入的误差。通过优化块大小的选择，DistrAttention可以轻松集成到FlashAttention-2中，从而在现代GPU上获得高性能。我们通过大量实验评估了DistrAttention。结果表明，在计算自注意力时，我们的方法比FlashAttention-2快37%。在ViT推理中，DistrAttention是近似自注意力机制中速度最快且最准确的。在Llama3-1B中，DistrAttention仍然实现了最低的推理时间，仅损失1%的准确率。
引用
@article{arxiv.2507.17245,
title = {DistrAttention: An Efficient and Flexible Self-Attention Mechanism on Modern GPUs},
author = {Haolin Jin and Mengbai Xiao and Yuan Yuan and Xiao Zhang and Dongxiao Yu and Guanghui Zhang and Haoliang Wang},
journal= {arXiv preprint arXiv:2507.17245},
year = {2025}
}