?作者 | 李文

單位 | 中南大學

研究方向 | 自然語言處理

最近讀了一篇 ICLR 2022 的論文，對于 attention-softmax 的計算復雜度有了一定的了解，下面是對論文的翻譯和筆記，如有理解錯誤，還請海涵指正。

論文標題：

cosFormer: Rethinking Softmax in Attention

論文鏈接：

https://arxiv.org/pdf/2202.08791.pdf

代碼鏈接：

https://github.com/OpenNLPLab/cosFormer

摘要

Transformer 在自然語言處理、計算機視覺和音頻處理，已經取得了巨大的成功。作為核心組件的 attention-softmax，可以有效捕捉長距離的特征，但是計算復雜度是序列長度的平方；核方法的提出就是為了減少其計算的時間復雜度，但是這是一個近似方法，受語料和任務場景影響，并且有一定的誤差。

這篇論文提出 cosformer 的一種方法，媲美標準 attention 甚效果更好。COFORMER 基于兩個關鍵 softmax注意力的屬性，第一個：注意矩陣的非負性， 第二個：非線性的權重轉換（re-weighting）方案，放大局部的注意力權重值。

INTRODUCTION

Figure 1：上圖中，y 軸表示性能，x 軸表示計算速度，圓圈的大小表示顯存占用，可以看到 cosFormer 在保持長距離任務下，與左上角的任務相比：不僅保持了性能，提升了計算速度，顯存占用也減少了。

transformer 及其變體在過去幾年對文本、音視頻任務產生了很大的影響，相較于 cnn 和 rnn，這一類模型可以學習長度不一的數據集，適用范圍更大，并且有利于捕捉其全局信息，刷爆了多個任務。點乘的 attention+softmax 機制是其捕捉長距離任務的關鍵，遺憾的是，其空間和時間復雜度，受到了序列長度的限制，尤其是在長距離任務下影響比較明顯一些。

為了解決這個問題，多種方法被提出，比如，2020 年提出的稀疏 attention 矩陣，低秩表達，核方法等等，這些方法都是為了解決標準 attention 的計算復雜度。然而這些線性時間復雜下的變體，通常存在一定的假設前提，并且是一個近似值，自帶約束，換句話說，這些方法可能并不能適應所有場景，尤其是在 GLUE benchmark 上，通常效果沒有標準（香草）的 attention 要好。另外不適用于 AR 語言模型的的 causal attention 注意力，也不適用 Linformer 和 BigBird 的交叉注意力。

softmax 是優化復雜度的主要難點所在（如果沒有 softmax，利用矩陣的乘法結合律（AB）C=A（BC），可以看下面的圖），我們（下文都指：作者們認為）很那去逼近它，因此作者提出疑問：“我們要得到一個關鍵的權重信息時，能否適用一個線性方法去替代 softmax”，這樣既保留了 attention 的關鍵特性，又可以利用結合律，先算 KV，通過深入挖掘。

我們找到了兩個關鍵信息：1、attention 中的權重元素是非負的，2：將權重進行 re-weigthing 的轉化是為了模型的穩定性，比如，（Katharopoulos et al., 2020）提出的線性 Transformer 使用了 的指數激活函數，來保證 Q、K 非負（這樣乘出來的 A 注意力矩陣也是非負的）， 然而，由于缺少一個 softmax 去 re-weigthing 權重, 模型表現不盡人意。

這篇 paper 中，我們提出了一個新的線性 Transformer 變體，命名為“COSFORMER”，滿足上面的兩個性質，在計算內積之前，我們先使用 relu 進行激活，強制將 K、Q 變成非負；另外我們使用了一個新的方法 將 K、Q 得到的 A 矩陣進行縮放 re-weigthing，來保障穩定性。這可以幫助模型關注局部上下文信息，通常局部上下文在 NLP 任務中比較重要。

由于 Ptolemy 定理的存在，attention 的計算可以分解，讓計算可以變成線性復雜度（先算 K、V），我們使用了 AR 和 AE 模型預訓練模型，在很多任務上進行了實驗，結果顯示 cosformer 使用更小的顯存，更快的速度就可以媲美標準的 Attention。并且刷新了 Long-Range Arena benchmark。

OUR METHOD

這一節主要講 COSFORMER 詳細的技術實現，最關鍵的點就是如何替換掉標準 Attention 中，非線性的并且不可拆分的 softmax 函數，使用一個非線性的可以拆分成線性操作進行替換；我們的方法適用范圍比較廣，causal 和 cross attentions 都支持，并且是線性的時間空間復雜度（說的都是跟輸入序列長度 N 成線性），在長文本上表現較好。

標準Attention

標準的 Attention 通常是給定一個輸入長度為? 的序列 ，然后通過? embedding 層表示為 的矩陣形式，然后經過一個 encoder 的 self-attention 模塊 ：，公式如下：

其中 F 是一個包含殘差的前向連接，A 表示的是 attention 計算結果，A 的計算時間和空間復雜度為 ，所以對于長文本而言，這個公式就會存在一個瓶頸。 中有三個重要的權重 ，用于計 、、，，這里使用 表示一個矩陣的第 行，用 來表示 A 函數的輸出，那么有如下公式成立：

其中 表示的是一個 token 的向量，來自于多個 的加權和， 表示點乘，可以理解為相似度，計算每一個 的時間復雜度為 O（N），這里就可以看到時間復雜度與輸入序列長度 成平方關系。如果 ，就是 softmax 了。

線性復雜度的self-attention

根據上面的公式，我們可以使用一個相似的方法，去計算 attention 矩陣，我們的目的是希望，、 的計算式可以拆分的成線性的，而不是跟 softmax 一樣不可拆分，所以對 進行如下表示：

其中 是將 和 映射到另外一個隱空間,我們重寫上一節的公式得到了：

因為矩陣乘法的結合律，我們就可以得到如下公式是成立的：

這樣就解決了計算復雜度的問題，可以參見下面的圖：

上圖已經很明顯的標出了其復雜度，當 ，時間復雜度與輸入序列長度的關系顯而易見。那么之前的解法有哪些呢，綜上所以，核心點是尋找一個可以拆分的 方法，一種方法是希望找到 softmax 的一個近似值，比如 RFA 和 Performer。

另外一種是進行替換，比如 linear transformer 直接替換了 softmax，先使用 將 、 的分量變成非負的，然后直接 得到非負的 Attention 矩陣。但是在我們的實驗中，我們發現這種方法對于長距離任務而言，效果并不好（見上圖），本論文也是提出了另外的方法替換 softmax 實現了更好的性能。

softmax attention的分析

在標準的 self-attention 中，，實際上是在 矩陣的行上面做了一個 re-weighting，換句話說，是對一個 token 和其他 token 的關聯關系權重做了一個 re-weighting，然而拋開之前的論文和經驗，softmax-attention 的關鍵和必要特征是什么呢？（作者的意思應該是，這里要拋開之前的解釋，只談這個矩陣的特征，我對這句話的理解是深度學習本身也不需要解釋，即使解釋不了，也可以擬合得到好的效果，歡迎指正）。

這里很容易得到兩個重要的特征，1、 矩陣是非負的 2、提供了一個非線性的 re-weighting 機制，擬合權重的分布，增加訓練穩定性。為了驗證上面的猜想，設計了一些實驗，實驗結果在下面表格中。

上圖中 表示對 、 都乘以單位陣，值沒有任何變化， 對 、 激活后負數會變小，正數不變。 對 、 進行激活，這樣所有分量就非負了， 表示標準的 attention 權重計算，上面的實驗是依次增加我們總結的關鍵點。

首先驗證非負的的重要性，上圖中 表示單位陣，即不對輸入進行任何轉化，對比 和其他中實驗的結果可以看到非負是比較重要額，其次驗證 re-weighting 的重要性，其中 softmax 是唯一進行了 re-weighting 的，同其他三個進行比較也可以看到其區別。

對比 和 可以看到，softmax 可以讓模型更快收斂，對于很多下游任務更好（圖里面有兩個都更好一些），這可能是因為 softmax 會放大數值，模型會更加關注一些敏感的模式。

COSFORMR

針對之前的觀察，我們提出了 COSFORMER，保留了 Attention 矩陣的非負性和 re-weighting 權重。COSFORMER 包含了兩個重要的模塊：線性的映射核函數 ，用于保證非負性；cos-Based re-weighting 機制，線性核函數 呼應了最開始的 ，定義如下：

負責將 映射到 ， 是一個可以拆分的的相似度計算方法。本文定義了 ，一方面保證 非負，另外又是非線性的轉換。因此我們可以得到 的表示如下：

，繼續進行變換得到下面的等式：

這樣子就已經讓時間復雜度降下來了，換句話說就是本來標準的 self-Attention： = [（向量·向量=標量）]*向量 ，換成了一個 Q、K、V 三個矩陣點乘，這個滿足結合律，可以交換計算的先后順序，沒懂可以看上面的圖，Linear Transfomer 就是這么干的，做完這個還不夠，還需要對權重進行 re-weighting，加速收斂，增強穩定性。

而 re-weighting 設計的目標也是必須可以拆分的函數，作者在設計這個函數的時候，首先分析了 softmax 起到的作用，然后構思了一個可以拆分的 cos-Based Re-weighting。

cos-Based Re-weighting Mechanism

進行非線性的 re-weighting 去擬合分布，可以加速收斂，穩定訓練過程。我們同樣發現，由于 其收縮放大的屬性，可以懲罰遠距離的 token 權重，增大近距離的局部上下文權重。可以參考下面的兩張圖：

針對這兩點，提出了 cos-Based Re-weighting 方法，

根據 Ptolemy 定理，我們可以進行拆分：

其中 表示位置，，M 是一個超參，代碼中通常用的就是 seq_length（causal attention 對應的是 query 和 key 中最長的）， 表示的就是距離,距離越遠， 越大， 上面的值就小；令

, 得到下面的公式：

推導過程如下：

最終，在不損失標準 Tansformer 優勢的情況下，我們的到 COSFORMER 公式如下：

與位置編碼的關系

上面的公式中，我們可以看到公式中存在 這樣的位置關系，cosformer 提出了一種新的想法，可以解釋為相對位置偏差。相較于 Rotary Position Embedding（2021 RoFormer），參考蘇神：Transformer升級之路：博采眾長的旋轉式位置編碼，膜拜下蘇神。相較于 Stochastic Positional Encoding（2021），它是利用采樣策略近似 softmax，并且使用了相對位置編碼。

代碼

代碼作者已經實現了，自己大概看了下，長距離下對計算速度是有提升的，總結起來就是論文的第一張圖。有興趣可以自己看下，cosFormer 論文內的偽代碼貼在下面了。

作者的 Pytorch 實現 Github：

https://github.com/OpenNLPLab/cosFormer

更多閱讀

#投 稿?通 道#

?讓你的文字被更多人看到?

如何才能讓更多的優質內容以更短路徑到達讀者群體，縮短讀者尋找優質內容的成本呢？答案就是：你不認識的人。

總有一些你不認識的人，知道你想知道的東西。PaperWeekly 或許可以成為一座橋梁，促使不同背景、不同方向的學者和學術靈感相互碰撞，迸發出更多的可能性。?

PaperWeekly 鼓勵高校實驗室或個人，在我們的平臺上分享各類優質內容，可以是最新論文解讀，也可以是學術熱點剖析、科研心得或競賽經驗講解等。我們的目的只有一個，讓知識真正流動起來。

📝?稿件基本要求：

? 文章確系個人原創作品，未曾在公開渠道發表，如為其他平臺已發表或待發表的文章，請明確標注?

? 稿件建議以?markdown?格式撰寫，文中配圖以附件形式發送，要求圖片清晰，無版權問題

? PaperWeekly 尊重原作者署名權，并將為每篇被采納的原創首發稿件，提供業內具有競爭力稿酬，具體依據文章閱讀量和文章質量階梯制結算

📬?投稿通道：

? 投稿郵箱：hr@paperweekly.site?

? 來稿請備注即時聯系方式（微信），以便我們在稿件選用的第一時間聯系作者

? 您也可以直接添加小編微信（pwbot02）快速投稿，備注：姓名-投稿

△長按添加PaperWeekly小編

🔍

現在，在「知乎」也能找到我們了

進入知乎首頁搜索「PaperWeekly」

點擊「關注」訂閱我們的專欄吧

·

總結

以上是生活随笔為你收集整理的ICLR 2022 | 在注意力中重新思考Softmax，商汤提出cosFormer实现多项SOTA的全部內容，希望文章能夠幫你解決所遇到的問題。

如果覺得生活随笔網站內容還不錯，歡迎將生活随笔推薦給好友。