arXivDaily arXiv每日学术速递 周一至周五更新
arXiv 2409.04431cs.LG
收藏

Sigmoid 自注意力的理论、分析与最佳实践

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

Apple(苹果公司)

Jason Ramapuram, Federico Danieli, Eeshan Dhekane, Floris Weers, Dan Busbridge, Pierre Ablin, Tatiana Likhomanenko, Jagrit Digani, Zijin Gu, Amitis Shidani, Russ Webb

AI总结:

本文通过理论证明和实验分析,确立了 sigmoid 注意力作为 Transformer 中 softmax 直接替代品的最佳实践,并提出了硬件加速实现 FLASHSIGMOID,在多种任务上达到与 softmax 相当的性能。

摘要:

注意力是 Transformer 架构的关键组成部分。它是一种序列到序列的映射,将每个序列元素转换为值的加权和。权重通常通过键和查询之间点积的 softmax 获得。近期工作探索了 Transformer 中 softmax 注意力的替代方案,例如 ReLU 和 sigmoid 激活函数。在本工作中,我们重新审视 sigmoid 注意力,并进行了深入的理论和实证分析。理论上,我们证明了具有 sigmoid 注意力的 Transformer 是通用函数逼近器,并且与 softmax 注意力相比,具有更好的正则性。通过详细的实证分析,我们确定在训练早期阶段稳定较大的初始注意力范数是成功训练具有 sigmoid 注意力模型的关键因素,其性能优于先前的尝试。我们还引入了 FLASHSIGMOID,一种硬件感知且内存高效的 sigmoid 注意力实现,在 H100 GPU 上比 FLASHATTENTION2 实现了 17% 的推理内核加速。跨语言、视觉和语音的实验表明,适当归一化的 sigmoid 注意力在广泛领域和规模上匹配了 softmax 注意力的强大性能,而先前尝试的 sigmoid 注意力未能完全实现这一点。我们的工作统一了先前的研究,并确立了 sigmoid 注意力作为 Transformer 中 softmax 的直接替代品的最佳实践。

英文摘要:

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.

机构、总结、中文摘要均由AI分析而来,可能存在错误,请以论文正文为准。

↑