DistrAttention：一种在现代GPU上高效且灵活的自注意力机制

Transformer架构革新了深度学习，在自然语言处理、计算机视觉和时间序列预测等领域实现了最先进的性能。然而，其核心组件自注意力（self-attention）相对于输入序列长度具有二次时间复杂度，限制了Transformer的可扩展性。现有优化自注意力的方法要么舍弃了全上下文信息，要么缺乏灵活性。在本工作中，我们设计了DistrAttention，一种具备全上下文的高效且灵活的自注意力机制。DistrAttention通过在嵌入维度（通常称为d）对数据进行分组来实现此目的。我们采用轻量级采样与融合方法，利用局部敏感哈希对相似数据进行分组。进一步设计了分块分组框架以限制局部敏感哈希引入的误差。通过优化分块大小的选择，DistrAttention可以轻松集成到FlashAttention-2中，在现代GPU上获得高性能。我们通过大量实验评估了DistrAttention。结果显示，本方法在计算自注意力时比FlashAttention-2快37%。在ViT推理中，DistrAttention是所有近似自注意力机制中速度最快且最准确的。在Llama3-1B中，DistrAttention依然以仅1%的准确率损失实现了最低推理时间。