cosFormer：重新思考注意力機制中的Softmax

導(dǎo)讀：Transformer在自然語言處理、計算機視覺和音頻處理方面取得了巨大成功。作為其核心組成部分之一，Softmax Attention模塊能夠捕捉長距離的依賴關(guān)系，但由于Softmax算子關(guān)于序列長度的二次空間和時間復(fù)雜性，使其很難擴展。

針對這點，研究者提出利用核方法以及稀疏注意力機制的方法來近似Softmax算子，從而降低時間空間復(fù)雜度。但是，由于誤差的存在，效果往往不盡如人意。

商湯多模態(tài)研究組認為，近似操作本身存在的誤差使得其效果很難超越Softmax Attention。我們的觀點是，與其近似Softmax，不如設(shè)計一種方式代替Softmax，并且同時降低時間空間復(fù)雜度。

因此，本文提出了名為cosFormer的方法，在時間空間復(fù)雜度關(guān)于序列長度為線性復(fù)雜度的同時，其性能接近或者超越Softmax Attention，并在LRA benchmark上取得SOTA結(jié)果。我們的設(shè)計核心理念基于兩點，首先是注意力矩陣的非負性，其次是對局部注意力的放大（非極大值抑制）。

本文主要介紹已收錄于ICLR 2022的一篇文章 cosFormer : Rethinking Softmax in Attention。

Part 1

背景

1. Softmax Attention

為了引出我們的方法，對Softmax Attention的計算方式進行一定的推廣：

其中表示相似度計算函數(shù)，如果，上式即變?yōu)镾oftmax Attention（不考慮除以的縮放操作）。注意到計算的時間復(fù)雜度為，的時間復(fù)雜度為，所以總時間復(fù)雜度為，即關(guān)于序列長度是二次的。

2. 線性 Attention

通過分析我們發(fā)現(xiàn)，性能瓶頸的主要原因是操作，如果相似度函數(shù)可以表示為：

那么：

根據(jù)矩陣運算的結(jié)合律：

上式可以變換為（編者修正：下方公式未變換，請參照論文）：

經(jīng)過計算后可以得到該方法的時間復(fù)雜度為，即關(guān)于序列長度是一次的。

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

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

3. Softmax 的兩大性質(zhì)

我們經(jīng)過分析以及實驗，歸納出Softmax Attention中比較重要的性質(zhì)，這兩個性質(zhì)可以指導(dǎo)我們的模型設(shè)計：

1. 注意力矩陣的非負性

2. 局部注意力的放大（非極大值抑制）

對于第一點，我們有如下實驗進行驗證（模型結(jié)構(gòu)為RoBERTa）：

這里L(fēng)oss表示驗證集損失（越低越好），其余指標均為準確率（越高越好）?？梢钥吹?，當保證了注意力矩陣的非負性之后，可以達到較好的效果。基于該實驗，我們選擇為ReLU函數(shù)。

對于第二點，我們的方式是在注意力矩陣中引入先驗locality信息，觀察Softmax注意力矩陣，如下圖所示，我們發(fā)現(xiàn)其注意力矩陣的權(quán)重在對角線附近很集中：

所以我們的方法需要在加了reweighting操作后也更加集中在對角線附近。注意并非所有的有類似權(quán)重的函數(shù)均適用，這個reweighting的函數(shù)需要跟前面的QK一樣可以拆分成兩個矩陣的乘法的形式。

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

Part 2

cosFormer

1. 方法

我們的方法基于線性Attention，首先給出符號定義：

根據(jù)之前的分析，我們選擇了：

可得：

為了進行reweighting操作，并且同時保證線性Attention的計算方式依然成立，我們選擇了cos函數(shù)：

展開可得：

為了便于展示，我們把它記作：

最終得到：

上式和線性Attention的計算方式一致，經(jīng)過分析不難得出時間復(fù)雜度依然是。

2. 實驗結(jié)果

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

單向語言模型，指標表示困惑度（越低越好）：

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

LRA benchmark：

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

2）內(nèi)存速度實驗，指標表示速度（越高越好，如果內(nèi)存溢出，則標記為叉）：

審核編輯 ：李倩