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

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

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

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

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

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

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

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

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

Sigmoid 注意力

假設(shè)為向量 n 的輸入序列，每個(gè)向量是 d 維。接著研究者定義了三個(gè)可學(xué)習(xí)權(quán)重矩陣、以及。這三個(gè)矩陣用于計(jì)算查詢，鍵，以及值?？梢缘玫饺缦鹿剑?/span>

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

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

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

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

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

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

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

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

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

Sigmoid 注意力的正則性

與神經(jīng)網(wǎng)絡(luò)中的任何層一樣，SigmoidAttn 的正則性值得研究，因?yàn)樗梢陨钊肓私庀鄳?yīng)網(wǎng)絡(luò)的魯棒性及其優(yōu)化的難易程度。

SigmoidAttn 正則性定理為：

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

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

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

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

實(shí)驗(yàn)

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

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

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

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

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

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