?PaperWeekly 原創 ·?作者｜蘇劍林

單位｜追一科技

研究方向｜NLP、神經網絡

標準 Attention 的 復雜度可真是讓研究人員頭大。前段時間我們在文章 Performer：用隨機投影將 Attention 的復雜度線性化中介紹了 Google 的 Performer 模型，它通過隨機投影的方式將標準 Attention 轉化為線性 Attention。無獨有偶，前些天 Arxiv 上放出了 AAAI 2021 的一篇論文，里邊又提出了一種從另一個角度把標準 Attention 線性化的方案。

論文標題：

Nystr?mformer: A Nystr?m-Based Algorithm for Approximating Self-Attention

論文鏈接：

https://arxiv.org/abs/2102.03902

代碼鏈接：

https://github.com/mlpen/Nystromformer

該方案寫的是 Nystr?m-Based，顧名思義是利用了 Nystr?m 方法來近似標準 Attention 的。但是坦白說，在看到這篇論文之前，筆者也完全沒聽說過 Nystr?m 方法，而縱觀整篇論文，里邊也全是筆者一眼看上去感覺很茫然的矩陣分解推導，理解起來頗為困難。

不過有趣的是，盡管作者的推導很復雜，但筆者發現最終的結果可以通過一個相對來說更簡明的方式來理解，遂將筆者對 Nystr?mformer 的理解整理在此，供大家參考。

簡單的回顧

如果讀者對線性 Attention 還不是很了解，那么建議先通讀一下線性 Attention 的探索：Attention 必須有個 Softmax 嗎？和 Performer：用隨機投影將 Attention 的復雜度線性化。總的來說，線性 Attention 是通過矩陣乘法的結合律來降低 Attention 的復雜度。

1.1 標準Attention

標準的 Scaled-Dot Attention 寫成矩陣形式就是（有時候指數部分還會多個縮放因子，這里我們就不顯式寫出來了）：

這里 （對應 Self Attention）。此外，本文的所有 softmax，都是對矩陣的第二個維度做歸一化。

在上式中， 這一步必須要先算出來，然后才能算 softmax，它導致了我們不能使用矩陣乘法的結合律。而 是 個向量的內積，因此時間和空間復雜度都是 。

1.2 線性Attention

而線性 Attention 比較樸素的做法就是：

其中 是值域非負的激活函數。為了方便對比，上式還沒有顯式地寫出歸一化因子，只突出了主要計算量的部分。上式左端的復雜度依然是 的，由于矩陣乘法滿足結合律，我們可以先算后面兩個矩陣的乘法，這樣整體復雜度就降為 了。

上式是直接將 Attention 定義為兩個矩陣的乘法來利用乘法結合律的，也可以將標準 Attention（近似地）轉化為矩陣的乘法來利用結合律，如下一節提到的 Performer；此外，相乘矩陣也不一定是兩個，比如本文要介紹的 Nystr?mformer 就是將注意力表示為三個矩陣相乘的。

1.3 Performer

對于 Performer 來說，它是通過隨機投影來找到矩陣 使得 softmax 中的 ，這樣一來標準 Attention 就可以近似為上一節的線性 Attention 來算了，細節請看之前的文章 Performer：用隨機投影將 Attention 的復雜度線性化。

如果對 SVM 和核方法等比較熟悉的讀者可能會聯想到，這個做法其實就是核函數的思想，即低維空間中兩個向量的核函數可以映射為高維空間中兩個向量的內積。它也可以跟 LSH（Locality Sensitive Hashing）聯系起來。

Nystr?mformer

在這部分內容中，我們以一個簡單的雙重 softmax 形式的線性 Attention?為出發點，逐步尋找更加接近標準 Attention 的線性 Attention，從而得到 Nystr?mformer。

▲ Nystro?mformer結構示意圖。讀者可以讀完下面幾節后再來對照著理解這個圖。

2.1 雙重Softmax

在文章線性 Attention 的探索：Attention 必須有個 Softmax 嗎？中我們提到了一種比較有意思的線性 Attention，它使用了雙重 softmax 來構建 Attention 矩陣：

可以證明這樣構造出來的 Attention 矩陣自動滿足歸一化要求，不得不說這是一種簡單漂亮的線性 Attention 方案。

不過，直接對 做 softmax 似乎有點奇怪，總感覺沒有經過相似度（內積）對比就直接 softmax 會有哪里不對勁。為了解決這個問題，Nystr?mformer 先分別將 視為 n 個 d 維向量，然后聚成m類來得到 m 個聚類中心構成的矩陣 ，這時候我們可以通過下述公式來定義 Attention：

具體的聚類過程我們稍后再來討論。現在，softmax 的對象是內積的結果，具有比較鮮明的物理意義，因此可以認為上式比前面的式?(3)?更為合理。如果我們選定一個比較小的 m，那么上式右端的復雜度只是線性地依賴于 n，因此它也是一個線性 Attention。

2.2 向標準靠近

純粹從改進式?(3)?的角度來看，式 (4) 已經達到目標了，不過 Nystr?mformer 并不局限于此，它還希望改進后的結果與標準 Attention 更加接近。

為此，觀察到式?(4)?的注意力矩陣 是一個 的矩陣乘以一個 的矩陣，為了微調結果，又不至于增加過多的復雜度，我們可以考慮在中間插入一個 的矩陣 ：

如何選擇 呢？一個合理的要求是當 m=n 時應當完全等價于標準 Attention，此時 ，推出：

對于一般的 m， 恰好是一個 矩陣，因此選它作為 至少在矩陣運算上是合理的，而根據 m=n 時的特殊情況我們則“大膽地”推測選它作為 能讓新的 Attention 機制更接近標準 Attention，因此 Nystr?mformer 最終選擇的是：

作為 Attention 矩陣，它是三個小矩陣的乘積，因此通過矩陣乘法的結合律就能轉化為線性 Attention。

不過，還有一個理論上的小細節需要補充一下，那就是上式涉及到矩陣的求逆，而 未必是可逆的。當然，從實踐上來看，一個實數的方陣不可逆的概率幾乎為零（不可逆意味著行列式嚴格等于 0，從概率上來看不等于 0 自然比等于 0 的概率大得多）。

因此這種情況在具體實驗中可以不考慮，但理論上還是得完善的。這個其實也簡單，如果是不可逆的矩陣，那就換成“偽逆”就好（記號為 ），它對任意矩陣都存在，并且當矩陣可逆時偽逆跟逆相等。

因此，最終的 Nystr?mformer 的 Attention 矩陣形式為：

2.3 迭代求逆陣

從理論上看，式 (8) 已經達到目標了，不過落實到實踐上還需要處理好一些細節問題，比如上述偽逆怎么求。偽逆又叫廣義逆、Moore-Penrose 逆等，標準的求法是通過 SVD 來求，設矩陣 的 SVD 分解為 ，那么它的偽逆為：

其中對角陣 的偽逆 等于將它對角線所有非零值取倒數所得到的新對角陣。SVD 的求法雖然理論上比較簡單易懂，但計算量還是比較大的，而且也不容易求梯度，因此并不是實現偽逆的理想方式。

Nystr?mformer采用了迭代求逆的近似方法。具體來說，它采用了論文《Chebyshev-type methods and preconditioning techniques》[1] 提供的迭代算法：

若初始矩陣 滿足 ，那么對于下述迭代格式：

成立 。

這里的 可以是任意一種矩陣范數，滿足條件的一個比較簡單的初始值可以是：

在 Nystr?mformer 論文中，作者直接用上述初始值和迭代格式進行迭代，將迭代 6 次的結果來代替 。迭代 6 次看上去很多，但事實上論文所選取的 m 比較小（論文寫的是 64），迭代過程中又只涉及到矩陣乘法，因此迭代計算量不會太大，而且只有乘法的話求梯度就很輕松了。這樣求偽逆的問題就算是解決了，論文將這個迭代過程簡寫為?pINV。

2.4 池化當聚類

還需要解決的另一個問題是聚類方法的選擇，比較直接的想法自然就是直接套用 K-Means 了。然而，同前面求偽逆所面臨的問題一樣，在設計模型時不僅要考慮前向計算，還需要考慮反向傳播的求梯度，直接套用 K-Means 涉及到 操作，無法求出有意義的梯度，需要將它“軟化”才能嵌入到模型中。

這一系列操作下來，其實就相當于膠囊網絡的“動態路由”過程，細節我們在再來一頓賀歲宴：從 K-Means 到 Capsule 討論過。這個方案的主要問題是 K-Means 是一個迭代過程，需要迭代幾次才能保證效果，這導致計算量明顯加大，不是特別理想。

Nystr?mformer 選了一個非常簡單的方案：假設序列長度 n 是 m 的整數倍（如果不是，padding 零向量），那么將 的每 n/m 個向量求平均作為 的每個向量。這個操作叫做?Adaptive Average Pooling（原論文稱為 Segment-Means，簡稱 sMEANS），即是一種平均池化方法，通過自適應窗口大小使得平均池化后的特征矩陣具有固定的形狀。

Nystr?mformer 的實驗表明，不需要比較復雜的聚類方法，就這樣使用簡單的自適應池化就可以取得非常有競爭力的效果了，而且只需要選擇 m=64，跟映射前的 d 是一般大小，這比 Performer 要選擇比 d 大幾倍的 m 要好得多了。

不過，自適應池化的一個明顯缺點是會“糅合”每一個區間的信息，導致它不能防止未來信息泄漏而不能做自回歸生成（語言模型或者 Seq2Seq 的解碼器），這基本是任何帶有 Pooling 技術的模型的缺點。

實驗與分析

這里我們匯總一下 Nystr?mformer 的實驗結果，并且分享一下筆者對它的一些看法和思考。

3.1 性能與效果

可能受限于算力，原論文做的實驗不算特別豐富，主要是將 small 和 base 版本的 BERT 里邊的標準 Attention 替換為 Nystr?mformer 進行對比實驗，實驗結果主要是下面兩個圖。

其中一個是預訓練效果圖，其中比較有意思的是 Nystr?mformer 在 MLM 任務上的效果比標準 Attention 還要優；另外是在下游任務上的微調效果，顯示出跟標準 Attention（即 BERT）比還是有競爭力的。

▲ Nystro?mformer在預訓練任務（MLM和SOP）上的效果

▲?Nystro?mformer在下游任務的微調效果

不過，原論文并沒有比較 Nystr?mformer 跟同類模型的效果差異，只是提供下面的一張復雜度對比圖，因此無法更好地突出 Nystr?mformer 的競爭力：

▲?不同模型的時間和空間復雜度對比圖

3.2 個人的思考

總的來說，Nystr?mformer 對標準 Attention 進行近似線性化的思路還是比較新穎的，值得學習與參考。不過偽逆部分的處理總感覺有點不大自然，這部分可能是未來的一個改進點，如果可以做到不用近似，那就比較完美了。還有，如何定量地估計 Nystr?mformer 與標準 Attention 的誤差，也是一個值得思考的理論問題。

從實驗上來看，Nystr?mformer 跟標準 Attention 相比還是顯得有競爭力的，尤其是 MLM 的結果比標準 Attention 還好，顯示了 Nystr?mformer 的潛力。此外，前面說到包含了?Pooling?導致不能做自回歸生成是 Nystr?mformer 的一個顯著缺點，不知道有沒有辦法可以彌補，反正筆者目前是沒有想到好的方向。

跟 Performer 相比，Nystr?mformer 去除了線性化過程中的隨機性，因為 Performer 是通過隨機投影來達到線性化的，這必然會帶來隨機性，對于某些有強迫癥的讀者來說，這個隨機性可能是難以接受的存在，而 Nystr?mformer 則不存在這種隨機性，因此也算是一個亮點。

3.3 Nystr?m方法

可能有些讀者還是想學習一下 Nystr?m 方法，這里稍微補充一下。要理解 Nystr?m 方法，需要先簡單認識一下矩陣的 CUR 分解。

大家可能都聽說過矩陣的 SVD 分解，格式為 ，其中 是正交矩陣而 是對角矩陣。要注意 是正交矩陣意味著它們是稠密的，那么當 很大的時候 SVD 的計算成本和儲存成本都很大（哪怕是做了近似）。

現在假設 很大但很稀疏，那么它的 SVD 分解比原始矩陣還不劃算得多得多。為此，CUR 分解應運而生，它希望從原矩陣中選擇 k 列組成矩陣 、選擇 k 行組成矩陣 ，并插入一個 的矩陣 ，使得：

由于 都是原句子的一部分，因此也繼承了稀疏性。關于 CUR 分解，讀者還可以參考斯坦福的 CS246 課程的《Dimensionality Reduction》[2]一節。跟 SVD 不同的是，CUR 分解在筆者看來更多的是一種分解思想而不是具體的分解算法，它有不同的實現方式，比如 Nystr?m 方法也算是其中一種，分解形式為

其中 和 選出來的列矩陣和行矩陣，這里為了方便描述，假設了經過排列后選出來的行列均排在矩陣前面。Nystr?mformer 其實也沒有直接用 Nystr?m 方法（事實上也直接套用不了，原論文有描述），而是借鑒了 Nystr?m 方法的分解思想而已。

關于 Nystr?m 方法，原論文主要引用的是《Improving CUR Matrix Decomposition and the Nystr?m Approximation via Adaptive Sampling》[3]，但并不推薦新手讀這篇論文，而推薦讀《Matrix Compression using the Nystro ?m Method》[4] 和《Using the Nystr?m Method to Speed Up Kernel Machines》[5]。

要特別說明的是，對于 CUR 分解和 Nystr?m 方法，筆者也是新學的，可能有理解不當的地方，請讀者自行甄別理解，也歡迎熟悉相關理論的讀者交流指正。

來一個小結

本文介紹了提升 Transformer 效率的一個新工作 Nystr?mformer，它借鑒了 Nystr?m 方法的思想來構建一個能逼近標準 Attention 的線性 Attention，類似思想的工作還有 Performer，兩者相比各有自己的優缺點，都是值得學習的工作。

本文分享了筆者自己對 Nystr?mformer 的理解，竊認為這種途徑更加易懂一些，如有謬誤，肯請讀者指正。

參考文獻

[1] https://www.researchgate.net/publication/220562466_Chebyshev-type_methods_and_preconditioning_techniques

[2] https://web.stanford.edu/class/cs246/slides/06-dim_red.pdf

[3] https://arxiv.org/abs/1303.4207

[4] https://arxiv.org/abs/1305.0203

[5] https://www.researchgate.net/publication/49459305_Using_the_Nystroem_Method_to_Speed_Up_Kernel_Machines

更多閱讀

#投 稿?通 道#

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

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

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

PaperWeekly 鼓勵高校實驗室或個人，在我們的平臺上分享各類優質內容，可以是最新論文解讀，也可以是學習心得或技術干貨。我們的目的只有一個，讓知識真正流動起來。

?????來稿標準：

? 稿件確系個人原創作品，來稿需注明作者個人信息（姓名+學校/工作單位+學歷/職位+研究方向）?

? 如果文章并非首發，請在投稿時提醒并附上所有已發布鏈接?

? PaperWeekly 默認每篇文章都是首發，均會添加“原創”標志

?????投稿郵箱：

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

? 所有文章配圖，請單獨在附件中發送?

? 請留下即時聯系方式（微信或手機），以便我們在編輯發布時和作者溝通

????

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

進入知乎首頁搜索「PaperWeekly」

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

關于PaperWeekly

PaperWeekly 是一個推薦、解讀、討論、報道人工智能前沿論文成果的學術平臺。如果你研究或從事 AI 領域，歡迎在公眾號后臺點擊「交流群」，小助手將把你帶入 PaperWeekly 的交流群里。

總結

以上是生活随笔為你收集整理的提升Transformer效率又有新招？基于矩阵分解的线性化Attention方案的全部內容，希望文章能夠幫你解決所遇到的問題。

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