线性Attention的探索：Attention必须有个Softmax吗？
?PaperWeekly 原創 ·?作者｜蘇劍林
單位｜追一科技
研究方向｜NLP、神經網絡
眾所周知，盡管基于 Attention 機制的 Transformer 類模型有著良好的并行性能，但它的空間和時間復雜度都是 級別的，n 是序列長度，所以當 n 比較大時 Transformer 模型的計算量難以承受。
近來，也有不少工作致力于降低 Transformer 模型的計算量，比如模型剪枝、量化、蒸餾等精簡技術，又或者修改 Attention 結構，使得其復雜度能降低到 甚至 。
前幾天筆者讀到了論文 Transformers are RNNs: Fast Autoregressive Transformers with Linear Attention?[1]?，了解到了線性化?Attention （Linear Attention）這個探索點，繼而閱讀了一些相關文獻，有一些不錯的收獲，最后將自己對線性化 Attention 的理解匯總在此文中。
Attention
當前最流行的 Attention 機制當屬 Scaled-Dot Attention [2] ，形式為：
這里的 ，簡單起見我們就沒顯式地寫出 Attention 的縮放因子了。
本文我們主要關心 Self Attention 場景，所以為了介紹上的方便統一設 ，一般長序列場景下都有 （BERT base 里邊 d=64）。
相關解讀可以參考筆者的一文讀懂「Attention is All You Need」| 附代碼實現，以及它的一些改進工作也可以參考突破瓶頸，打造更強大的 Transformer [3]、Google 新作 Synthesizer：我們還不夠了解自注意力，這里就不多深入介紹了。
1.1 摘掉Softmax
讀者也許想不到，制約 Attention 性能的關鍵因素，其實是定義里邊的 Softmax！事實上，簡單地推導一下就可以得到這個結論。
這一步我們得到一個 的矩陣，就是這一步決定了 Attention 的復雜度是 ；如果沒有 Softmax，那么就是三個矩陣連乘 ，而矩陣乘法是滿足結合率的，所以我們可以先算 ，得到一個 的矩陣，然后再用 左乘它，由于 ，所以這樣算大致的復雜度只是 （就是 左乘那一步占主導）。
也就是說，去掉 Softmax 的 Attention 的復雜度可以降到最理想的線性級別 ！這顯然就是我們的終極追求：Linear Attention，復雜度為線性級別的 Attention。所以，本文的主題就是探究摘掉 Softmax 后的線形 Attention。
1.2 一般的定義
問題是，直接去掉 Softmax 還能算是 Attention 嗎？它還能有標準的 Attention 的效果嗎？為了回答這個問題，我們先將 Scaled-Dot Attention 的定義（1）等價地改寫為（本文的向量都是列向量）。
所以，Scaled-Dot Attention 其實就是以 為權重對 做加權平均。所以我們可以提出一個 Attention 的一般化定義：
也就是把 換成 的一般函數 ，為了保留 Attention 的相似特性，我們要求 恒成立。也就是說，我們如果要定義新式的 Attention，那么要保留式（3）的形式，并且滿足 。
這種一般形式的 Attention 在 CV 中也被稱為 Non-Local 網絡，來自文章 Non-local Neural Networks [4]。
幾個例子
如果直接去掉 Softmax，那么就是 ，問題是內積無法保證非負性，所以這還不是一個合理的選擇。下面我們簡單介紹幾種可取的方案。
值得指出的是，下面介紹的這幾種 Linear Attention，前兩種只做了 CV 的實驗，第三種是筆者自己構思的，所以都還沒有在 NLP 任務上做過什么實驗，各位做模型改進的 NLPer 們就有實驗方向了。
2.1 核函數形式
一個自然的想法是：如果 的每個元素都是非負的，那么內積自然也就是非負的。為了完成這點，我們可以給 各自加個激活函數 ，即：
其中 是值域非負的激活函數。本文開頭提到的論文 Transformers are RNNs: Fast Autoregressive Transformers with Linear Attention [5]?選擇的是 。
非要講故事的話，式（4）可以聯想到“核方法（kernal method）”，尤其是 時 就相當于一個核函數，而 就是通過核函數所定義的內積。
這方面的思考可以參考論文 Transformer disp: An unified understanding for transformer’s attention via the lens of kernel [6]，此處不做過多延伸。
2.2 妙用Softmax
另一篇更早的文章 Efficient Attention: Attention with Linear Complexities?[7]?則給出了一個更有意思的選擇。它留意到在 中，，如果“ 在 d 那一維是歸一化的、并且 在 n 那一維是歸一化的”，那么 就是自動滿足歸一化了，所以它給出的選擇是：
其中 、 分別指在第一個（n）、第二個維度（d）進行 Softmax 運算。也就是說，這時候我們是各自給 加 Softmax，而不是 算完之后才加 Softmax。
其實可以證明這個形式也是式（4）的一個特例，此時對應于 ，讀者可以自行推導一下。
2.3 自己的構思
在這里，筆者給出自己的一種構思。這個構思的出發點不再是式（4），而是源于我們對原始定義（2）的近似。由泰勒展開我們有：
如果 ，那么就可以保證右端的非負性，而從可以讓 。到這里讀者可能已經想到了，想要保證 ，只需要分別對 做 歸一化。所以，筆者最終提出的方案就是：
這不同于形式（4），但理論上它應該是最接近原始的 Scaled-Dot Attention 了。
相關工作
通過修改 Attention 的形式來降低它的計算復雜度，相關的工作有很多，這里簡要列舉一些。
3.1 稀疏Attention
我們之前介紹過 OpenAI 的 Sparse Attention，通過“只保留小區域內的數值、強制讓大部分注意力為零”的方式，來減少 Attention 的計算量。經過特殊設計之后，Attention 矩陣的非 0 元素只有 個，所以理論上它也是一種線性級別的 Attention。類似的工作還有 Longformer。
但是很明顯，這種思路有兩個不足之處：
如何選擇要保留的注意力區域，這是人工主觀決定的，帶有很大的不智能性；
它需要從編程上進行特定的設計優化，才能得到一個高效的實現，所以它不容易推廣。
3.2 Reformer
Reformer 也是有代表性的改進工作，它將 Attention 的復雜度降到了 。
某種意義上來說，Reformer 也是稀疏 Attention 的一種，只不過它的稀疏 pattern 不是事先指定的，而是通過 LSH（Locality Sensitive Hashing）技術（近似地）快速地找到最大的若干個 Attention 值，然后只去計算那若干個值。
此外，Reformer 通過構造可逆形式的 FFN（Feedforward Network）替換掉原來的 FFN，然后重新設計反向傳播過程，從而降低了顯存占用量。
所以，相比前述稀疏 Attention，Reformer 解決了它的第一個缺點，但是依然有第二個缺點：實現起來復雜度高。要實現 LSH 形式的 Attention 比標準的 Attention 復雜多了，對可逆網絡重寫反向傳播過程對普通讀者來說更是遙不可及。
3.3 Linformer
跟本文所介紹的 Linear Attention 很相似的一個工作是 Facebook 最近放出來的 Linformer，它依然保留原始的 Scaled-Dot Attention 形式，但在進行 Attention 之前，用兩個 的矩陣 分別對 進行投影，即變為：
這樣一來， 就只是一個 的矩陣，而作者聲稱對于哪怕對于很大的序列長度 n，m 也可以保持為一個適中的常數，從而這種 Attention 也是線性的。
但是，筆者認為“對于超長序列 m 可以保持不變”這個結論是值得質疑的，原論文中對于長序列作者只做了 MLM 任務，而很明顯 MLM 并不那么需要長程依賴，所以這個實驗沒什么說服力。因此，Linformer 是不是真的 Linear，還有待商榷。
3.4 自回歸生成
Linformer 的另一個缺點是 這兩變直接把整個序列的信息給“糅合”起來了，所以它沒法簡單地把將來信息給 Mask 掉（Causal Masking），從而無法做語言模型、Seq2Seq 等自回歸生成任務，這也是剛才說的原作者只做了 MLM 任務的原因。
相比之下，本文介紹的幾種 Linear Attention 都能做到這一點。以式（3）和式（4）為例，如果要 Mask 掉未來信息，那么只需要把求和 改為 ：
實現上式有兩種方式：第一方式是設 以及 ，我們有：
這說明這種 Attention 可以作為一個 RNN 模型用遞歸的方式實現，它的空間復雜度最低，但是要串性計算，適合預測解碼時使用；第二種是直接將 做外積，得到一個 的矩陣，然后對 n 那一維執行 運算，這樣就一次性得到 了，它的速度最快，但空間占用最大，適合訓練時使用。
3.5 下采樣技術
從結果上來看，Linformer 的 就是將序列變短（下采樣）了，而將序列變短的一個最樸素的方法就是 Pooling 了，所以筆者之前也嘗試過把 Pooling 技術引入到 Transformer 中去。
近來也有類似的工作發出來，比如IBM的PoWER-BERT: Accelerating BERT Inference via Progressive Word-vector Elimination [8] 和 Google 的 Funnel-Transformer: Filtering out Sequential Redundancy for Efficient Language Processing [9] 。
除了 Pooling 之外，其實還有其他的下采樣技術，比如可以通過 stride > 1 的一維卷積來實現，基于這個思路，或許我們可以把 FFN 里邊的 Position-Wise 全連接換成 stride > 1 的一維卷積？總之這方面應該也能玩出很多花樣來，不過跟 Linformer 一樣，這樣糅合之后做自回歸生成就很難了。
文章小結
本文介紹了一些從結構上對 Attention 進行修改從而降低其計算復雜度的工作，其中最主要的 idea 是去掉標準 Attention 中的 Softmax，就可以使得 Attention 的復雜度退化為理想的 級別（Linear Attention）。
相比于其他類似的改進結構的工作，這種修改能在把復雜度降到 的同時，依然保留所有的 “token-token” 的注意力，同時還能保留用于做自回歸生成的可能性。
參考鏈接
[1] https://arxiv.org/abs/2006.16236
[2] https://arxiv.org/abs/1706.03762
[3] https://kexue.fm/archives/7325
[4] https://kexue.fm/archives/1711.07971
[5] https://arxiv.org/abs/2006.16236
[6] https://arxiv.org/abs/1908.11775
[7] https://arxiv.org/abs/1812.01243
[8] https://arxiv.org/abs/2001.08950
[9] https://arxiv.org/abs/2006.03236
更多閱讀
#投 稿?通 道#
?讓你的論文被更多人看到?
如何才能讓更多的優質內容以更短路徑到達讀者群體，縮短讀者尋找優質內容的成本呢？答案就是：你不認識的人。
總有一些你不認識的人，知道你想知道的東西。PaperWeekly 或許可以成為一座橋梁，促使不同背景、不同方向的學者和學術靈感相互碰撞，迸發出更多的可能性。?
PaperWeekly 鼓勵高校實驗室或個人，在我們的平臺上分享各類優質內容，可以是最新論文解讀，也可以是學習心得或技術干貨。我們的目的只有一個，讓知識真正流動起來。
?????來稿標準：
? 稿件確系個人原創作品，來稿需注明作者個人信息（姓名+學校/工作單位+學歷/職位+研究方向）?
? 如果文章并非首發，請在投稿時提醒并附上所有已發布鏈接?
? PaperWeekly 默認每篇文章都是首發，均會添加“原創”標志
?????投稿郵箱：
? 投稿郵箱：hr@paperweekly.site?
? 所有文章配圖，請單獨在附件中發送?
? 請留下即時聯系方式（微信或手機），以便我們在編輯發布時和作者溝通
????
現在，在「知乎」也能找到我們了
進入知乎首頁搜索「PaperWeekly」
點擊「關注」訂閱我們的專欄吧
關于PaperWeekly
PaperWeekly 是一個推薦、解讀、討論、報道人工智能前沿論文成果的學術平臺。如果你研究或從事 AI 領域，歡迎在公眾號后臺點擊「交流群」，小助手將把你帶入 PaperWeekly 的交流群里。
總結
以上是生活随笔為你收集整理的线性Attention的探索：Attention必须有个Softmax吗？的全部內容，希望文章能夠幫你解決所遇到的問題。
- 上一篇： 10 月 1~7 日全社会跨区域人员流动
- 下一篇： DSCC：苹果 OLED iPad Pr