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

導讀：Transformer在自然語言處理、計算機視覺和音頻處理方面取得了巨大成功。作為其核心組成部分之一，Softmax Attention模塊能夠捕捉長距離的依賴關(guān)系，但由于Softmax算子關(guān)于序列長度的二次空間和時間復雜性，使其很難擴展。針對這點，研究者提出利用核方法以及稀疏注意力機制的方法來近似Softmax算子，從而降低時間空間復雜度。但是，由于誤差的存在，效果往往不盡如人意。

我們（商湯多模態(tài)研究組）認為，近似操作本身存在的誤差使得其效果很難超越Softmax Attention。我們的觀點是，與其近似Softmax，不如設(shè)計一種方式代替Softmax，并且同時降低時間空間復雜度。因此，本文提出了名為cosFormer的方法，在時間空間復雜度關(guān)于序列長度為線性復雜度的同時，其性能接近或者超越Softmax Attention，并在LRA benchmark上取得SOTA結(jié)果。我們的設(shè)計核心理念基于兩點，首先是注意力矩陣的非負性，其次是對局部注意力的放大（非極大值抑制）。

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

 

論文地址：https://arxiv.org/abs/2202.08791

部分開源代碼：https://github.com/OpenNLPLab/cosFormer

 

圖片表示各種transformer在LRA benchmark上的表現(xiàn)，其中y軸表示性能，x軸表示速度，圓圈大小表示內(nèi)存，我們提出的cosFormer取得了明顯的優(yōu)勢。

一、背景

 

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

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

1. 注意力矩陣的非負性

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

 

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

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

二、cosFormer

1、方法

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

 

2、實驗結(jié)果

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

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

論文地址：https://arxiv.org/abs/2202.08791

雷峰網(wǎng)(公眾號：雷峰網(wǎng))