arXiv:2409.04431cs.LG2024-09被引 73

sigmoid注意力可替代softmax,性能相当且更稳定。

Theory, Analysis, and Best Practices for Sigmoid Self-Attention

  • 用sigmoid替代softmax,理论证明其具备通用逼近能力。
  • 正确归一化后,性能媲美softmax,在多领域表现优异。
  • 提出硬件优化实现,推理速度提升17%,适合部署落地。

注意力是Transformer的核心组件,通常通过softmax计算查询与键的点积权重。近期研究探索了ReLU、sigmoid等替代激活函数。本文深入分析sigmoid注意力,理论上证明其为通用函数逼近器,且相比softmax具有更好正则性。实证发现,训练初期对大初始注意力范数进行稳定是成功关键,显著优于以往尝试。我们提出FLASHSIGMOID,一种面向硬件、内存高效的sigmoid注意力实现,在H100 GPU上比FLASHATTENTION2快17%。跨语言、视觉、语音多个任务实验表明,经适当归一化后,sigmoid注意力性能可匹敌softmax,而此前方法未能达成此效果。本工作统一现有成果,确立了sigmoid注意力作为softmax即插即用替代品的最佳实践。

原文摘要 · Abstract (English)

Attention is a key part of the transformer architecture. It is a sequence-to-sequence mapping that transforms each sequence element into a weighted sum of values. The weights are typically obtained as the softmax of dot products between keys and queries. Recent work has explored alternatives to softmax attention in transformers, such as ReLU and sigmoid activations. In this work, we revisit sigmoid attention and conduct an in-depth theoretical and empirical analysis. Theoretically, we prove that transformers with sigmoid attention are universal function approximators and benefit from improved regularity compared to softmax attention. Through detailed empirical analysis, we identify stabilization of large initial attention norms during the early stages of training as a crucial factor for the successful training of models with sigmoid attention, outperforming prior attempts. We also introduce FLASHSIGMOID, a hardware-aware and memory-efficient implementation of sigmoid attention yielding a 17% inference kernel speed-up over FLASHATTENTION2 on H100 GPUs. Experiments across language, vision, and speech show that properly normalized sigmoid attention matches the strong performance of softmax attention on a wide range of domains and scales, which previous attempts at sigmoid attention were unable to fully achieve. Our work unifies prior art and establishes best practices for sigmoid attention as a drop-in softmax replacement in transformers.

注意力机制Transformer模型优化高效推理

Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。