?作者 | 李文

單位 | 中南大學(xué)

研究方向 | 自然語言處理

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

論文標題：

cosFormer: Rethinking Softmax in Attention

論文鏈接：

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

代碼鏈接：

https://github.com/OpenNLPLab/cosFormer

摘要

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

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

INTRODUCTION

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

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

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

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

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

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

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

OUR METHOD

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

標準Attention

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

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

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

線性復(fù)雜度的self-attention

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

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

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

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

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

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

softmax attention的分析

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

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

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

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

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

COSFORMR

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

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

，繼續(xù)進行變換得到下面的等式：

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

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

cos-Based Re-weighting Mechanism

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

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

根據(jù) Ptolemy 定理，我們可以進行拆分：

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

, 得到下面的公式：

推導(dǎo)過程如下：

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

與位置編碼的關(guān)系

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

代碼

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

作者的 Pytorch 實現(xiàn) Github：

https://github.com/OpenNLPLab/cosFormer

更多閱讀

#投 稿?通 道#

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

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

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

PaperWeekly 鼓勵高校實驗室或個人，在我們的平臺上分享各類優(yōu)質(zhì)內(nèi)容，可以是最新論文解讀，也可以是學(xué)術(shù)熱點剖析、科研心得或競賽經(jīng)驗講解等。我們的目的只有一個，讓知識真正流動起來。

📝?稿件基本要求：

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

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

? PaperWeekly 尊重原作者署名權(quán)，并將為每篇被采納的原創(chuàng)首發(fā)稿件，提供業(yè)內(nèi)具有競爭力稿酬，具體依據(jù)文章閱讀量和文章質(zhì)量階梯制結(jié)算

📬?投稿通道：

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

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

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

△長按添加PaperWeekly小編

🔍

現(xiàn)在，在「知乎」也能找到我們了

進入知乎首頁搜索「PaperWeekly」

點擊「關(guān)注」訂閱我們的專欄吧

·

總結(jié)

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

如果覺得生活随笔網(wǎng)站內(nèi)容還不錯，歡迎將生活随笔推薦給好友。