任務通用是基礎模型研究的核心目標之一，同時也是深度學習研究通向高級智能的必經之路。近年來，得益于注意力機制的通用關鍵建模能力，Transformer在眾多領域中表現優異，逐漸呈現出通用架構的趨勢。但是隨著序列長度的增長，標準注意力機制的計算呈現二次復雜度，嚴重阻礙了其在長序列建模與大模型中的應用。

為此，來自清華大學軟件學院的團隊深入探索了這一關鍵問題，提出了任務通用的線性復雜度主干網絡Flowformer，在保持標準Transformer的通用性的同時，將其復雜度降至線性，論文被ICML 2022接受。

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

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

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

相比于標準Transformer，本文提出的Flowformer模型，具有以下特點：

• 線性復雜度，可以處理數千長度的輸入序列；
• 沒有引入新的歸納偏好，保持了原有注意力機制的通用建模能力；
• 任務通用，在長序列、視覺、自然語言、時間序列、強化學習五大任務上取得優秀效果。

1. 問題分析

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

2. 動機

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

Softmax中的競爭機制

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

網絡流中的競爭機制

我們關注到在圖論中的經典網絡流（Flow network）模型中，「守恒」（Conservation）是一個重要現象，即每個節點的流入量等于流出量。受到「固定資源情況下，必定引起競爭」的啟發，在本文中，我們試圖從網絡流視角重新分析經典注意力機制中的信息流動，并通過守恒性質將競爭引入注意力機制設計，以避免平凡注意力問題。

3. Flowformer

3.1 網絡流視角下的注意力機制

在注意力機制內部：信息流動可以表示為：從源（source，對應）基于學習到的流容量（flow capacity，對應注意力權重）匯聚至匯（sink，對應）。

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

3.2 Flow-Attention

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

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

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

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

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

（3）整體設計

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

4. 實驗

本文在標準數據集上進行了廣泛的實驗：

• 覆蓋了長序列、視覺、自然語言、時間序列、強化學習五大任務；
• 考察了標準（Normal）和自回歸任務（Causal）兩種注意力機制類型。
• 涵蓋了多種序列長度的輸入情況（20-4000）。
• 對比了各領域經典模型、主流深度模型、Transformer及其變體等多種基線方法。

如下表所示，Flowformer在五大任務上均表現優秀，驗證了模型的通用性。詳細實驗結果請見論文。

5. 分析

為了進一步說明Flowformer的工作原理，我們對ImageNet分類任務中的注意力（對應Flow-Attention中的）進行了可視化實驗，從中可以發現：

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

上述可視化表明，通過Flow-Attention將競爭引入注意力機制設計可以有效避免平凡注意力。更多可視化實驗可見論文。

6. 總結

本文提出的Flowformer通過將網絡流中的守恒原理引入設計，自然地將競爭機制引入到注意力計算中，有效避免了平凡注意力問題，在實現線性復雜度的同時，保持了標準Transformer的通用性。Flowformer在長序列、視覺、自然語言、時間序列、強化學習五大任務上取得優秀效果。此外，Flowformer中「無特殊歸納偏好」的設計理念也對通用基礎架構的研究具有一定的啟發性。在未來工作中，我們將進一步探索Flowformer在大規模預訓練上的潛力。