cosFormer：重新思考注意力機(jī)制中的Softmax

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

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

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

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

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

Part 1

背景

1. Softmax Attention

為了引出我們的方法，對(duì)Softmax Attention的計(jì)算方式進(jìn)行一定的推廣：

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

2. 線(xiàn)性 Attention

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

那么：

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

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

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

Softmax Attention和線(xiàn)性Attention的計(jì)算方式可以用下圖概括：

所以接下來(lái)將介紹的選擇，以及核心的reweighting操作。

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

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

1. 注意力矩陣的非負(fù)性

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

對(duì)于第一點(diǎn)，我們有如下實(shí)驗(yàn)進(jìn)行驗(yàn)證（模型結(jié)構(gòu)為RoBERTa）：

這里L(fēng)oss表示驗(yàn)證集損失（越低越好），其余指標(biāo)均為準(zhǔn)確率（越高越好）?？梢钥吹?，當(dāng)保證了注意力矩陣的非負(fù)性之后，可以達(dá)到較好的效果?；谠搶?shí)驗(yàn)，我們選擇為ReLU函數(shù)。

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

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

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

Part 2

cosFormer

1. 方法

我們的方法基于線(xiàn)性Attention，首先給出符號(hào)定義：

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

可得：

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

展開(kāi)可得：

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

最終得到：

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

2. 實(shí)驗(yàn)結(jié)果

我們?cè)趩蜗蚰Ｐ汀㈦p向模型以及LRA benchmark上測(cè)試了我們的方法，均取得了非常不錯(cuò)的效果。

單向語(yǔ)言模型，指標(biāo)表示困惑度（越低越好）：

雙向語(yǔ)言模型，指標(biāo)表示準(zhǔn)確率（越高越好）：

LRA benchmark：

1）性能實(shí)驗(yàn)，指標(biāo)表示準(zhǔn)確率（越高越好）：

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

審核編輯 ：李倩