基於2大理念改進蕭簫 轉載自 商湯AI量子位 | 公眾號 QbitAITransformer很受歡迎 但它架構上的不少問題依舊令人頭疼。典例之一就是其中的Softmax Attention模塊 商湯最新論文登上ICLR：給注意力機製Softmax找個“平替” - 趣味新聞網

蕭簫 轉載自 商湯AI

量子位 | 公眾號 QbitAI

Transformer 很受歡迎，但它架構上的不少問題依舊令人頭疼。

典例之一就是其中的 Softmax Attention 模塊，雖然能捕捉長距離依賴關係，但由於Softmax算子關於序列長度的二次空間和時間復雜性，導緻難以擴展。

雖然也有用核方法、稀疏注意力機製等來近似Softmax算子，以降低時間空間復雜度，但近似操作本身存在的誤差使得其效果很難超越Softmax Attention。

為此，商湯多模態研究組想到，與其近似Softmax，為何不重新設計一種方式“平替”Softmax？

他們提齣瞭一種叫做 cosFormer 的新方法，論文目前已經登上ICLR 2022。

一方麵，cosFormer在時間空間復雜度關於序列長度為綫性復雜度的同時，其性能接近或者超越Softmax Attention；

另一方麵，它也在LRA benchmark上取得瞭SOTA，其中y軸錶示性能，x軸錶示速度，圓圈大小錶示內存。

一起來看看。

此前的Softmax有什麼問題？

Softmax Attention

Softmax Attention的計算方式是這樣的：

綫性Attention

通過分析我們發現，性能瓶頸的主要原因是exp操作，如果相似度函數可以錶示為

那麼

根據矩陣運算的結閤律：

上式可以變換為：

經過計算後可以得到該方法的時間復雜度為，即關於序列長度是一次的。

Softmax Attention和綫性Attention的計算方式可以用下圖概括：

所以接下來的會介紹的選擇，以及核心的reweighting操作。

Softmax的兩大性質

我們經過分析以及實驗，歸納齣Softmax Attention中比較重要的性質，這兩個性質可以指導我們的模型設計：

注意力矩陣的非負性局部注意力的放大（非極大值抑製）

對於第一點，我們有如下實驗進行驗證（模型結構為RoBERTa）：

這裏Loss錶示驗證集損失（越低越好），其餘指標均為準確率（越高越好），可以看到，當保證瞭注意力矩陣的非負性之後，可以達到較好的效果。基於該實驗，我們選擇為ReLU函數。

對於第二點，我們的方式是在注意力矩陣中引入先驗locality信息，觀察Softmax注意力矩陣，如下圖所示，我們發現其注意力矩陣的權重在對角綫附近很集中：

所以我們的方法需要在加瞭reweighting操作後也更加集中在對角綫附近。注意並非所有的有類似權重的函數均適用，這個reweighting的函數需要跟前麵的QK一樣可以拆分成兩個矩陣的乘法的形式。

至此，就可以引入我們的cosFormer瞭。

cosFormer如何超越Softmax？

我們的方法基於綫性Attention，首先給齣符號定義：

根據之前的分析，我們選擇瞭：

可得：

為瞭進行reweighting操作，並且同時保證綫性Attention的計算方式依然成立，我們選擇瞭cos函數：

展開可得：

為瞭便於展示，我們把它記作：

最終得到：

上式和綫性Attention的計算方式一緻，經過分析不難得齣時間復雜度依然是O(N)。

具體性能究竟有多好？

我們在單嚮模型，雙嚮模型以及LRA benchmark上測試瞭我們的方法，均取得瞭非常不錯的效果。

單嚮語言模型，指標錶示睏惑度（越低越好）：

雙嚮語言模型，指標錶示準確率（越高越好）：

LRA benchmark：

性能實驗，指標錶示準確率（越高越好）：

內存速度實驗，指標錶示速度（越高越好，如果內存溢齣，則標記為叉）：

目前代碼已開源。