任務(wù)通用是基礎(chǔ)模型研究的核心目標(biāo)之一，同時(shí)也是深度學(xué)習(xí)研究通向高級(jí)智能的必經(jīng)之路。近年來，得益于注意力機(jī)制的通用關(guān)鍵建模能力，Transformer在眾多領(lǐng)域中表現(xiàn)優(yōu)異，逐漸呈現(xiàn)出通用架構(gòu)的趨勢。但是隨著序列長度的增長，標(biāo)準(zhǔn)注意力機(jī)制的計(jì)算呈現(xiàn)二次復(fù)雜度，嚴(yán)重阻礙了其在長序列建模與大模型中的應(yīng)用。

為此，來自清華大學(xué)軟件學(xué)院的團(tuán)隊(duì)深入探索了這一關(guān)鍵問題，提出了任務(wù)通用的線性復(fù)雜度主干網(wǎng)絡(luò)Flowformer，在保持標(biāo)準(zhǔn)Transformer的通用性的同時(shí)，將其復(fù)雜度降至線性，論文被ICML 2022接受。

作者列表：吳海旭，吳佳龍，徐介暉，王建民，龍明盛

鏈接：https://arxiv.org/pdf/2202.06258.pdf

代碼：https://github.com/thuml/Flowformer

相比于標(biāo)準(zhǔn)Transformer，本文提出的Flowformer模型，具有以下特點(diǎn)：

• 線性復(fù)雜度，可以處理數(shù)千長度的輸入序列；
• 沒有引入新的歸納偏好，保持了原有注意力機(jī)制的通用建模能力；
• 任務(wù)通用，在長序列、視覺、自然語言、時(shí)間序列、強(qiáng)化學(xué)習(xí)五大任務(wù)上取得優(yōu)秀效果。

1. 問題分析

標(biāo)準(zhǔn)的注意力機(jī)制輸入包含queries（），keys（）和values（）三部分，，其計(jì)算方式如下：其中為注意力權(quán)重矩陣，最終計(jì)算結(jié)果為將進(jìn)行加權(quán)融合所得，上述過程計(jì)算復(fù)雜度為。注意到，對于多項(xiàng)矩陣的連乘問題，在經(jīng)典算法中已有較多研究。特別地，對于注意力機(jī)制，我們可以利用矩陣乘法的結(jié)合律來實(shí)現(xiàn)優(yōu)化，如，即可將原本的二次復(fù)雜度降至線性。但是注意力機(jī)制中的函數(shù)使得無法直接應(yīng)用結(jié)合律。因此，如何移除注意力機(jī)制中的函數(shù)是實(shí)現(xiàn)線性復(fù)雜度的關(guān)鍵。但是，近期的眾多工作證明，函數(shù)在避免平凡注意力學(xué)習(xí)上起到了關(guān)鍵性作用。綜上，我們期待一種模型設(shè)計(jì)方案，實(shí)現(xiàn)以下目標(biāo)：（1）移除函數(shù)；（2）避免平凡注意力；（3）保持模型的通用性。

2. 動(dòng)機(jī)

針對目標(biāo)（1），在之前的工作中，往往使用核方法來替代函數(shù)，即通過近似注意力計(jì)算（為非線性函數(shù)），但直接去掉會(huì)造成平凡注意力。為此，針對目標(biāo)（2），之前工作不得不引入一些歸納偏好，這限制了模型的通用性，因此不滿足目標(biāo)（3），比如cosFormer中的局部性假設(shè)等。

Softmax中的競爭機(jī)制

為滿足上述目標(biāo)，我們從的基本性質(zhì)出發(fā)進(jìn)行分析。我們注意到，最初被提出是用于：將「贏者通吃」的取極大值操作擴(kuò)展為可微分形式。因此，得益于其內(nèi)在的「競爭」機(jī)制，它可以使各個(gè)token之間的注意力權(quán)重差異化，從而避免了平凡的注意力的問題。基于以上考慮，我們試圖將競爭機(jī)制引入注意力機(jī)制設(shè)計(jì)，從而避免核方法分解帶來平凡注意力問題。

網(wǎng)絡(luò)流中的競爭機(jī)制

我們關(guān)注到在圖論中的經(jīng)典網(wǎng)絡(luò)流（Flow network）模型中，「守恒」（Conservation）是一個(gè)重要現(xiàn)象，即每個(gè)節(jié)點(diǎn)的流入量等于流出量。受到「固定資源情況下，必定引起競爭」的啟發(fā)，在本文中，我們試圖從網(wǎng)絡(luò)流視角重新分析經(jīng)典注意力機(jī)制中的信息流動(dòng)，并通過守恒性質(zhì)將競爭引入注意力機(jī)制設(shè)計(jì)，以避免平凡注意力問題。

3. Flowformer

3.1 網(wǎng)絡(luò)流視角下的注意力機(jī)制

在注意力機(jī)制內(nèi)部：信息流動(dòng)可以表示為：從源（source，對應(yīng)）基于學(xué)習(xí)到的流容量（flow capacity，對應(yīng)注意力權(quán)重）匯聚至匯（sink，對應(yīng)）。

在注意力機(jī)制外部，源（v）的信息來自于上一層網(wǎng)絡(luò)，匯（R）的信息也將提供給下面的前饋層。

3.2 Flow-Attention

基于上述觀察，我們可以通過分別從流入和流出兩個(gè)角度，控制注意力機(jī)制與外部網(wǎng)絡(luò)的交互，來實(shí)現(xiàn)「固定資源」，從而分別引起源和匯內(nèi)部的競爭，以避免平凡注意力。不失一般性，我們將注意力機(jī)制與外部網(wǎng)絡(luò)的交互信息量設(shè)置為默認(rèn)值1.

（1）匯（R）的流入守恒：

不難得到，未經(jīng)過守恒之前，對于第個(gè)匯，其流入的信息量為：。為了固定每個(gè)匯流入的信息量為單位1，我們將作為歸一化引入信息流（注意力權(quán)重）的計(jì)算。經(jīng)過歸一化之后，第個(gè)匯的流入信息量為：

此時(shí)，由于匯的流入守恒，各個(gè)源（V）之間存在天然的競爭關(guān)系，我們計(jì)算此時(shí)每個(gè)源（V）給出的信息量，即可得到：競爭情況下，每個(gè)源所提供的信息量，這也代表著每個(gè)源的重要性。

（2）源（V）的流出守恒：與前述過程類似，未經(jīng)過守恒之前，對于第個(gè)源，其流出的信息量為：。為了固定每個(gè)源流出的信息量為單位1，我們將作為歸一化引入信息流（注意力權(quán)重）的計(jì)算。經(jīng)過歸一化之后，第j個(gè)源的流出信息量為：。此時(shí)，由于源的流出守恒，各個(gè)匯（）之間存在天然的競爭關(guān)系，我們計(jì)算此時(shí)每個(gè)匯（）接受的信息量，即可得到：競爭情況下，每個(gè)結(jié)果所需要最終所接受的信息量。

（3）整體設(shè)計(jì)

基于上述結(jié)果，我們設(shè)計(jì)如下Flow-Attention機(jī)制，具體包含競爭（Competition）、聚合（Aggregation）、分配（Allocation）三部分：其中Competition將競爭機(jī)制引入中，突出重要的信息；Aggregation基于矩陣結(jié)合律實(shí)現(xiàn)線性復(fù)雜度；Allocation通過將競爭機(jī)制引入，控制傳遞到下一層的信息量。上述過程中的所有操作均為線性復(fù)雜度。同時(shí)，F(xiàn)low-Attention的設(shè)計(jì)僅僅依賴于網(wǎng)絡(luò)流中的守恒原理，對信息流的重新整合，因此并沒有引入新的歸納偏好，保證了模型的通用性。將標(biāo)準(zhǔn)Transformer中的二次復(fù)雜度Attention替換為Flow-Attention，即得到了Flowformer。

4. 實(shí)驗(yàn)

本文在標(biāo)準(zhǔn)數(shù)據(jù)集上進(jìn)行了廣泛的實(shí)驗(yàn)：

• 覆蓋了長序列、視覺、自然語言、時(shí)間序列、強(qiáng)化學(xué)習(xí)五大任務(wù)；
• 考察了標(biāo)準(zhǔn)（Normal）和自回歸任務(wù)（Causal）兩種注意力機(jī)制類型。
• 涵蓋了多種序列長度的輸入情況（20-4000）。
• 對比了各領(lǐng)域經(jīng)典模型、主流深度模型、Transformer及其變體等多種基線方法。

如下表所示，F(xiàn)lowformer在五大任務(wù)上均表現(xiàn)優(yōu)秀，驗(yàn)證了模型的通用性。詳細(xì)實(shí)驗(yàn)結(jié)果請見論文。

5. 分析

為了進(jìn)一步說明Flowformer的工作原理，我們對ImageNet分類任務(wù)中的注意力（對應(yīng)Flow-Attention中的）進(jìn)行了可視化實(shí)驗(yàn)，從中可以發(fā)現(xiàn)：

• 如果僅僅使用核方法進(jìn)行分解，如Linear Transformer，會(huì)造成模型注意力分散，無法有效捕捉到關(guān)鍵區(qū)域；
• 經(jīng)典Transformer和Flowformer均可以準(zhǔn)確捕捉到圖像的關(guān)鍵位置，但是后者在計(jì)算復(fù)雜度上具有優(yōu)勢；
• cosFormer在注意力機(jī)制中引入一維局部性假設(shè)，在語言任務(wù)上效果突出。但是在圖像（將2D數(shù)據(jù)展開成1D序列）中，如果不將局部性假設(shè)擴(kuò)展至二維，則無法適配視覺任務(wù)。這也印證了Flowformer中「沒有引入新的歸納偏好」設(shè)計(jì)方式的優(yōu)勢。

上述可視化表明，通過Flow-Attention將競爭引入注意力機(jī)制設(shè)計(jì)可以有效避免平凡注意力。更多可視化實(shí)驗(yàn)可見論文。

6. 總結(jié)

本文提出的Flowformer通過將網(wǎng)絡(luò)流中的守恒原理引入設(shè)計(jì)，自然地將競爭機(jī)制引入到注意力計(jì)算中，有效避免了平凡注意力問題，在實(shí)現(xiàn)線性復(fù)雜度的同時(shí)，保持了標(biāo)準(zhǔn)Transformer的通用性。Flowformer在長序列、視覺、自然語言、時(shí)間序列、強(qiáng)化學(xué)習(xí)五大任務(wù)上取得優(yōu)秀效果。此外，F(xiàn)lowformer中「無特殊歸納偏好」的設(shè)計(jì)理念也對通用基礎(chǔ)架構(gòu)的研究具有一定的啟發(fā)性。在未來工作中，我們將進(jìn)一步探索Flowformer在大規(guī)模預(yù)訓(xùn)練上的潛力。