Sigmoid 自注意力的理論、分析與最佳實踐

主要貢獻者

注意力是變壓器架構中的一個關鍵部分。它是一種序列到序列的映射，將每個序列元素轉換為加權值的總和。這些權重通常是通過鍵（keys）和查詢（queries）之間的點積（dot products）計算出來的 softmax。最近的研究探索了變壓器中 softmax 注意力的替代方案，例如 ReLU 和 sigmoid 激活函數。在這項工作中，我們重新檢視了 sigmoid 注意力，並進行了深入的理論和實證分析。在理論上，我們證明了使用 sigmoid 注意力的變壓器是通用函數逼近器，並且相比於 softmax 注意力，具有更好的規則性。通過詳細的實證分析，我們發現，在訓練的早期階段，大的初始注意力範數的穩定性是成功訓練使用 sigmoid 注意力模型的關鍵因素，超越了之前的嘗試。我們還介紹了 FLASHSIGMOID，這是一種硬體友好且內存高效的 sigmoid 注意力實現，能在 H100 GPU 上比 FLASHATTENTION2 提高 17% 的推理速度。跨越語言、視覺和語音的實驗顯示，適當正規化的 sigmoid 注意力在廣泛的領域和規模上能夠匹配 softmax 注意力的強大表現，而之前的 sigmoid 注意力嘗試未能完全實現。我們的工作統一了先前的研究，並確立了將 sigmoid 注意力作為變壓器中 softmax 替代方案的最佳實踐。