Sigmoid自注意力的理论、分析与最佳实践
机器学习
2025-01-23 v2
摘要
注意力机制是Transformer架构的关键组成部分。它是一种序列到序列的映射，将每个序列元素转换为其价值向量的加权和。这些权重通常通过查询和键之间点积的softmax函数获得。近期研究探索了softmax注意力在Transformer中的替代方案，如ReLU和Sigmoid激活函数。本文重新审视Sigmoid注意力，进行深入的理论和实证分析。理论上，我们证明了包含Sigmoid注意力的Transformer是通用函数逼近器，并且相对于softmax注意力在正则性方面具有改进优势。通过详细的实证分析，我们识别出在训练早期阶段大初始注意力范数稳定化是一个关键因素，这对于成功训练包含Sigmoid注意力的模型至关重要，显著优于先前的尝试。我们还引入了FLASHSIGMOID，一种面向硬件且高效的Sigmoid注意力实现，在H100 GPU上相较于FLASHATTENTION2实现了17%的推理内核加速。实验在语言、视觉和语音领域显示，经过适当归一化的Sigmoid注意力在广泛的领域和规模范围内，匹配了softmax注意力的强大性能，而此前尝试的Sigmoid注意力无法完全实现这一点。我们的工作统一了先前的工作，确立了Sigmoid注意力的最佳实践，作为Transformer中softmax替代方案。
引用
@article{arxiv.2409.04431,
title = {Theory, Analysis, and Best Practices for Sigmoid Self-Attention},
author = {Jason Ramapuram and Federico Danieli and Eeshan Dhekane and Floris Weers and Dan Busbridge and Pierre Ablin and Tatiana Likhomanenko and Jagrit Digani and Zijin Gu and Amitis Shidani and Russ Webb},
journal= {arXiv preprint arXiv:2409.04431},
year = {2025}
}