注意力是 Transformer 架構(gòu)的關(guān)鍵部分，負(fù)責(zé)將每個序列元素轉(zhuǎn)換為值的加權(quán)和。將查詢與所有鍵進(jìn)行點積，然后通過 softmax 函數(shù)歸一化，會得到每個鍵對應(yīng)的注意力權(quán)重。

盡管 SoftmaxAttn 中的 softmax 具有廣泛的用途和有效性，但它并非沒有局限性。例如，softmax 函數(shù)有時會導(dǎo)致注意力集中在少數(shù)幾個特征，而忽略了其他信息。

近來，一些研究探索了 Transformer 中 softmax 注意力的替代方案，例如 ReLU 和 sigmoid 激活函數(shù)。最近，來自蘋果的研究者重新審視了 sigmoid 注意力并進(jìn)行了深入的理論和實驗分析。

該研究證明：從理論上講，與 softmax 注意力相比，具有 sigmoid 注意力的 Transformer 是通用函數(shù)逼近器，并且受益于改進(jìn)的正則化。

• 論文地址：https://arxiv.org/pdf/2409.04431
• 項目地址：https://github.com/apple/ml-sigmoid-attention
• 論文標(biāo)題：Theory, Analysis, and Best Practices for Sigmoid Self-Attention

該研究還提出了一種硬件感知且內(nèi)存高效的 sigmoid 注意力實現(xiàn) ——FLASHSIGMOID。FLASHSIGMOID 在 H100 GPU 上的推理內(nèi)核速度比 FLASHATTENTION2 提高了 17%。

跨語言、視覺和語音的實驗表明，合理歸一化的 sigmoid 注意力與 softmax 注意力在廣泛的領(lǐng)域和規(guī)模上性能相當(dāng)，而之前的 sigmoid 注意力嘗試無法實現(xiàn)這一點。

此外，該研究還用 sigmoid 內(nèi)核擴(kuò)展了 FLASHATTENTION2，將內(nèi)核推理掛鐘時間減少了 17%，將現(xiàn)實世界推理時間減少了 8%。

論文作者 Jason Ramapuram 表示：如果想讓注意力快 18% 左右，你不妨試試 Sigmoid 注意力機(jī)制。他們用 Sigmoid 和基于序列長度的常量標(biāo)量偏置取代了注意力機(jī)制中的傳統(tǒng) softmax。

Sigmoid 注意力

假設(shè)為向量 n 的輸入序列，每個向量是 d 維。接著研究者定義了三個可學(xué)習(xí)權(quán)重矩陣、以及。這三個矩陣用于計算查詢，鍵，以及值。可以得到如下公式：

根據(jù)先前的研究，自注意力可以簡寫為：

其中 Softmax 函數(shù)將輸入矩陣的每一行進(jìn)行了歸一化。該研究將 Softmax 做了以下替換：

實際上，將作為輸出序列，可以得到：

將多個 SigmoidAttn 輸出進(jìn)行組合，得到多個頭的形式，如下所示：

Sigmoid 注意力理論基礎(chǔ)

該研究對 SigmoidAttn 進(jìn)行了分析，分析的目的主要有兩個：（1）證明當(dāng) SigmoidAttn 取代 SoftmaxAttn 時，Transformer 架構(gòu)仍然是一個通用函數(shù)逼近器；（2）通過計算 SigmoidAttn 的 Lipschitz 常數(shù)來恢復(fù)其規(guī)律性。

具有 Sigmoid 注意力的 Transformer 是通用逼近器嗎？

經(jīng)典 Transformer 可以將連續(xù)的序列到序列函數(shù)近似到任意精度，這一特性稱為通用近似特性 (UAP，Universal Approximation Property)。UAP 非常受歡迎，因為它證明了架構(gòu)的通用性和表示能力。由于 SigmoidAttn 修改了 Transformer 架構(gòu)，因此從理論上保證這種修改不會影響表示能力并保留 UAP 的性能至關(guān)重要。該研究通過以下定理提供此保證。

結(jié)果表明，即使使用 SigmoidAttn，一系列 transformer 塊也可以實現(xiàn)上下文映射。

Sigmoid 注意力的正則性

與神經(jīng)網(wǎng)絡(luò)中的任何層一樣，SigmoidAttn 的正則性值得研究，因為它可以深入了解相應(yīng)網(wǎng)絡(luò)的魯棒性及其優(yōu)化的難易程度。

SigmoidAttn 正則性定理為：

結(jié)果證明，SigmoidAttn 的局部 Lipschitz 常數(shù)遠(yuǎn)低于 SoftmaxAttn 的最差局部 Lipschitz 常數(shù)。

FLASHSIGMOID：硬件感知實現(xiàn)

現(xiàn)代架構(gòu)上的注意力計算往往會受到內(nèi)存訪問 IO 的限制。FLASHATTENTION 和 FLASHATTENTION2 通過優(yōu)化 GPU 內(nèi)存層次結(jié)構(gòu)利用率來加速注意力計算。得益于這些方法提供的速度提升，該研究開發(fā)了 SigmoidAttn 的硬件感知實現(xiàn) ——FLASHSIGMOID，采用了三個核心思路：

• Tiling：注意力分而治之的方法：與 FLASHATTENTION 和 FLASHATTENTION2 類似，F(xiàn)LASHSIGMOID 并行處理輸入部分以計算塊中的注意力輸出，有效地組合部分結(jié)果以生成最終的注意力輸出。 
• 內(nèi)核融合：與 FLASHATTENTION 和 FLASHATTENTION2 一樣，F(xiàn)LASHSIGMOID 將 SigmoidAttn 的前向和后向傳遞的計算步驟實現(xiàn)為單個 GPU 內(nèi)核，通過避免高帶寬內(nèi)存 (HBM) 上的中間激活具體化，最大限度地減少內(nèi)存訪問并提高內(nèi)存效率。
• 激活重計算：sigmoid 注意力的向后傳遞需要 sigmoid 激活矩陣，如果在 GPU HBM 上具體化，則會導(dǎo)致執(zhí)行速度變慢和內(nèi)存效率低下。FLASHSIGMOID 通過僅保留查詢、鍵和值張量來解決這個問題，以便在向后傳遞期間重新計算 sigmoid 激活矩陣。盡管增加了 FLOPs，但事實證明，與具體化和保留注意力矩陣的替代方法相比，這種方法在掛鐘時間上更快，并且內(nèi)存效率更高。

實驗

為了實驗驗證 SigmoidAttn，該研究在多個領(lǐng)域進(jìn)行了評估：使用視覺 transformer 進(jìn)行監(jiān)督圖像分類、使用 SimCLR 進(jìn)行自監(jiān)督圖像表示學(xué)習(xí)、BYOL（Bootstrap Your Own Latent）和掩碼自動編碼器 (MAE) 以及自動語音識別 (ASR) 和自回歸語言建模 (LM)。

該研究還在 TED-LIUM v3 上驗證了 ASR 的序列長度泛化，在所有這些領(lǐng)域和算法中，該研究證明 SigmoidAttn 的性能與 SoftmaxAttn 相當(dāng)（圖 2 和 21），同時提供訓(xùn)練和推理加速。

該研究得出以下觀察結(jié)果：

SigmoidAttn 對于沒有偏置的視覺任務(wù)是有效的（MAE 除外），但依賴于 LayerScale 以無超參數(shù)的方式匹配基線 SoftmaxAttn（圖 9-a）的性能。除非另有說明，否則為 SoftmaxAttn 呈現(xiàn)的所有結(jié)果也公平地添加了 LayerScale。

LM 和 ASR 對初始范數(shù)較為敏感。需要通過 (a) 相對位置嵌入進(jìn)行調(diào)整；(b) 適當(dāng)初始化 b 以實現(xiàn)相同效果 —— 允許使用任何位置嵌入。

感興趣的讀者可以閱讀論文原文，了解更多研究內(nèi)容。