注意力是 Transformer 架構的關鍵部分，負責將每個序列元素轉換為值的加權和。將查詢與所有鍵進行點積，然後透過 softmax 函式歸一化，會得到每個鍵對應的注意力權重。
儘管 SoftmaxAttn 中的 softmax 具有廣泛的用途和有效性，但它並非沒有侷限性。例如，softmax 函式有時會導致注意力集中在少數幾個特徵，而忽略了其他資訊。
近來，一些研究探索了 Transformer 中 softmax 注意力的替代方案，例如 ReLU 和 sigmoid 啟用函式。最近，來自蘋果的研究者重新審視了 sigmoid 注意力並進行了深入的理論和實驗分析。
該研究證明：從理論上講，與 softmax 注意力相比，具有 sigmoid 注意力的 Transformer 是通用函式逼近器，並且受益於改進的正則化。
論文地址：https://arxiv.org/pdf/2409.04431
專案地址：https://github.com/apple/ml-sigmoid-attention
論文標題：Theory, Analysis, and Best Practices for Sigmoid Self-Attention
該研究還提出了一種硬體感知且記憶體高效的 sigmoid 注意力實現 ——FLASHSIGMOID。FLASHSIGMOID 在 H100 GPU 上的推理核心速度比 FLASHATTENTION2 提高了 17%。
跨語言、視覺和語音的實驗表明，合理歸一化的 sigmoid 注意力與 softmax 注意力在廣泛的領域和規模上效能相當，而之前的 sigmoid 注意力嘗試無法實現這一點。
此外，該研究還用 sigmoid 核心擴充套件了 FLASHATTENTION2，將核心推理掛鐘時間減少了 17%，將現實世界推理時間減少了 8%。
論文作者 Jason Ramapuram 表示：如果想讓注意力快 18% 左右，你不妨試試 Sigmoid 注意力機制。他們用 Sigmoid 和基於序列長度的常量標量偏置取代了注意力機制中的傳統 softmax。
Sigmoid 注意力
假設為向量 n 的輸入序列，每個向量是 d 維。接著研究者定義了三個可學習權重矩陣
、
以及
。這三個矩陣用於計算查詢
，鍵
，以及值
。可以得到如下公式：
根據先前的研究，自注意力可以簡寫為：
其中 Softmax 函式將輸入矩陣的每一行進行了歸一化。該研究將 Softmax 做了以下替換：
實際上，將作為輸出序列，可以得到：
將多個 SigmoidAttn 輸出進行組合，得到多個頭的形式，如下所示：
Sigmoid 注意力理論基礎
該研究對 SigmoidAttn 進行了分析，分析的目的主要有兩個：（1）證明當 SigmoidAttn 取代 SoftmaxAttn 時，Transformer 架構仍然是一個通用函式逼近器；（2）透過計算 SigmoidAttn 的 Lipschitz 常數來恢復其規律性。
具有 Sigmoid 注意力的 Transformer 是通用逼近器嗎？
經典 Transformer 可以將連續的序列到序列函式近似到任意精度，這一特性稱為通用近似特性 (UAP，Universal Approximation Property)。UAP 非常受歡迎，因為它證明了架構的通用性和表示能力。由於 SigmoidAttn 修改了 Transformer 架構，因此從理論上保證這種修改不會影響表示能力並保留 UAP 的效能至關重要。該研究透過以下定理提供此保證。
結果表明，即使使用 SigmoidAttn，一系列 transformer 塊也可以實現上下文對映。
Sigmoid 注意力的正則性
與神經網路中的任何層一樣，SigmoidAttn 的正則性值得研究，因為它可以深入瞭解相應網路的魯棒性及其最佳化的難易程度。
SigmoidAttn 正則性定理為：
結果證明，SigmoidAttn 的區域性 Lipschitz 常數遠低於 SoftmaxAttn 的最差區域性 Lipschitz 常數。
FLASHSIGMOID：硬體感知實現
現代架構上的注意力計算往往會受到記憶體訪問 IO 的限制。FLASHATTENTION 和 FLASHATTENTION2 透過最佳化 GPU 記憶體層次結構利用率來加速注意力計算。得益於這些方法提供的速度提升，該研究開發了 SigmoidAttn 的硬體感知實現 ——FLASHSIGMOID，採用了三個核心思路：
Tiling：注意力分而治之的方法：與 FLASHATTENTION 和 FLASHATTENTION2 類似，FLASHSIGMOID 並行處理輸入部分以計算塊中的注意力輸出，有效地組合部分結果以生成最終的注意力輸出。
核心融合：與 FLASHATTENTION 和 FLASHATTENTION2 一樣，FLASHSIGMOID 將 SigmoidAttn 的前向和後向傳遞的計算步驟實現為單個 GPU 核心，透過避免高頻寬記憶體 (HBM) 上的中間啟用具體化，最大限度地減少記憶體訪問並提高記憶體效率。
啟用重計算：sigmoid 注意力的向後傳遞需要 sigmoid 啟用矩陣，如果在 GPU HBM 上具體化，則會導致執行速度變慢和記憶體效率低下。FLASHSIGMOID 透過僅保留查詢、鍵和值張量來解決這個問題，以便在向後傳遞期間重新計算 sigmoid 啟用矩陣。儘管增加了 FLOPs，但事實證明，與具體化和保留注意力矩陣的替代方法相比，這種方法在掛鐘時間上更快，並且記憶體效率更高。
實驗
為了實驗驗證 SigmoidAttn，該研究在多個領域進行了評估：使用視覺 transformer 進行監督影像分類、使用 SimCLR 進行自監督影像表示學習、BYOL（Bootstrap Your Own Latent）和掩碼自動編碼器 (MAE) 以及自動語音識別 (ASR) 和自迴歸語言建模 (LM)。
該研究還在 TED-LIUM v3 上驗證了 ASR 的序列長度泛化，在所有這些領域和演算法中，該研究證明 SigmoidAttn 的效能與 SoftmaxAttn 相當（圖 2 和 21），同時提供訓練和推理加速。
該研究得出以下觀察結果：
SigmoidAttn 對於沒有偏置的視覺任務是有效的（MAE 除外），但依賴於 LayerScale 以無超引數的方式匹配基線 SoftmaxAttn（圖 9-a）的效能。除非另有說明，否則為 SoftmaxAttn 呈現的所有結果也公平地新增了 LayerScale。
LM 和 ASR 對初始範數較為敏感。需要透過 (a) 相對位置嵌入進行調整；(b) 適當初始化 b 以實現相同效果 —— 允許使用任何位置嵌入。
感興趣的讀者可以閱讀論文原文，瞭解更多研究內容。