ユニバーサルミッション！清華大学、線形複雑性を実現するバックボーンネットワークFlowformerを提案｜ICML2022

タスクの一般性は、基本モデル研究の中心的な目標の 1 つであり、ディープラーニング研究が高度なインテリジェンスへと進む唯一の方法でもあります。近年、アテンションメカニズムの汎用キーモデリング機能のおかげで、Transformer は多くの分野で優れたパフォーマンスを発揮し、徐々に汎用アーキテクチャのトレンドを示してきました。しかし、シーケンスの長さが長くなると、標準的なアテンション メカニズムの計算は 2 次的な複雑さを示し、長いシーケンスのモデリングや大規模なモデルへの適用が著しく妨げられます。

この目的のために、清華大学ソフトウェア学院のチームはこの重要な問題を詳細に調査し、標準のTransformerの汎用性を維持しながら複雑さを線形に削減するタスクユニバーサル線形複雑性バックボーンネットワークFlowformerを提案しました。この論文はICML 2022に採択されました。

著者リスト: Wu Haixu、Wu Jialong、Xu Jiehui、Wang Jianmin、Long Mingsheng

リンク: https://arxiv.org/pdf/2202.06258.pdf

コード: https://github.com/thuml/Flowformer

標準的なTransformerと比較して、本論文で提案するFlowformerモデルには次の特徴があります。

• 線形複雑度、数千の長さの入力シーケンスを処理できます。
• 新たな帰納的選好は導入されず、元の注意メカニズムの一般的なモデリング機能は維持されます。
• 汎用性が高く、長いシーケンス、ビジョン、自然言語、時系列、強化学習の5 つの主要タスクで優れた結果を達成しています。

1. 問題分析

標準のアテンション機構の入力には、querys()、keys()、values() の 3 つの部分が含まれており、次のように計算されます。ここで、はアテンション重み行列であり、最終的な計算結果は重み付けされた融合結果です。上記のプロセスの計算の複雑さは次のとおりです。古典的なアルゴリズムにおける複数の行列の乗算の問題については多くの研究が行われてきたことに注意してください。特に、アテンション メカニズムでは、行列乗算の結合法則を使用して、元の 2 次複雑度を線形に削減するなどの最適化を実現できます。しかし、注意機構の機能により、連想法則を直接適用することは不可能です。したがって、注意機構内の機能をどのように削除するかが、線形複雑性を実現するための鍵となります。しかし、最近の多くの研究では、関数が些細な注意学習を回避する上で重要な役割を果たすことが実証されています。要約すると、我々は以下の目標を達成するモデル設計を期待しています：(1) 機能を削除する、(2) 些細な注意を避ける、(3) モデルの一般性を維持する。

2. 動機

目標（１）については、これまでの研究では、カーネル法を用いて関数を置き換える、つまり、注意計算（非線形関数）を近似することが多いが、それを直接削除すると、注意は些細なものになってしまう。このため、目標（2）に関しては、これまでの研究では、cosFormerの局所性仮定など、モデルの汎用性を制限し、目標（3）を達成できないようないくつかの帰納的バイアスを導入する必要がありました。

ソフトマックスにおける競争メカニズム

上記の目的を達成するために、 の基本的な特性から始めます。当初は、「勝者総取り」の最大演算を微分可能な形式に拡張することが提案されていたことに注意してください。したがって、固有の「競争」メカニズムのおかげで、トークン間の注意の重みを区別することができ、些細な注意の問題を回避できます。上記の考慮に基づいて、カーネル法の分解によって引き起こされる些細な注意の問題を回避するために、注意メカニズムの設計に競争メカニズムを導入することを試みます。

ネットワークフローにおける競争メカニズム

グラフ理論の古典的なネットワークフローモデルでは、 「保存」が重要な現象であること、つまり各ノードの流入が流出に等しいことに気付きました。 「固定されたリソースの下では競争は避けられない」という考えに触発され、本稿では、ネットワークフローの観点から古典的な注意メカニズムの情報フローを再分析し、保存特性を通じて注意メカニズムの設計に競争を導入して、自明な注意問題を回避することを試みます。

3. フローフォーマー

3.1 ネットワークフローの観点から見た注目メカニズム

注意メカニズムの内部では、情報フローは、学習されたフロー容量（注意重みに対応するフロー容量）に基づいて、ソース（ソース、対応）からシンク（シンク、対応）へのフローとして表すことができます。

注意メカニズムの外側では、ソース (v) 情報はネットワークの前の層から取得され、シンク (R) 情報も次のフィードフォワード層に提供されます。

3.2 フローアテンション

上記の観察に基づいて、流入と流出の両方の観点から注意メカニズムと外部ネットワーク間の相互作用を制御し、それによってソースとシンク内でそれぞれ競争を引き起こして些細な注意を回避することで、「固定リソース」を実現できます。一般性を損なうことなく、注意メカニズムと外部ネットワーク間の相互作用の量をデフォルト値の 1 に設定します。

（１）シンクへの流入量の保存（R）：

保存前にシンクに流入する情報量は次のようになることは難しくありません。各シンクに流入する情報量をユニット1に固定するために、正規化として情報フロー（アテンションウェイト）の計算を導入します。正規化後、シンクに流入する情報量は次のようになります。

このとき、シンク流入量の保存により、情報源（V）間には自然な競争関係があります。このとき、各情報源が提供する情報量（V）を計算すると、競争下での各情報源が提供する情報量が得られ、これは各情報源の重要性も表します。

（２）情報源からの流出量の保存（V）：前の過程と同様に、保存前では、情報源から流出する情報量は次の通りである。各ソースから流れ出る情報量をユニット1に固定するために、正規化として情報フロー（注目重み）の計算を導入します。正規化後、j番目の情報源の流出情報量は次のようになります。このとき、ソース流出の保存により、シンク（）間には自然な競争関係があります。このとき、各シンク（）が受け取る情報量を計算すると、競争下で最終結果において各結果が受け取る必要のある情報量が得られます。

（３）全体設計

上記の結果に基づいて、私たちは次の Flow-Attention メカニズムを設計しました。これは、競争、集約、割り当ての 3 つの部分から構成されます。競争は、重要な情報を強調するための競争メカニズムを導入します。集約は、行列の結合法則に基づいて線形複雑性を実装します。割り当ては、競争メカニズムを導入することで、次のレイヤーに渡される情報量を制御します。上記のプロセスのすべての操作には線形の複雑さがあります。同時に、Flow-Attention の設計は、ネットワーク フローの保存原理と情報フローの再統合のみに依存しているため、新しい帰納的設定は導入されず、モデルの汎用性が保証されます。標準の Transformer の二次複雑度の Attention を Flow-Attention に置き換えると、Flowformer が得られます。

4. 実験

私たちは標準的なデータセットに対して広範な実験を実施します。

• 長いシーケンス、ビジョン、自然言語、時系列、強化学習の 5 つの主要タスクをカバーします。
• 標準（Normal）と自己回帰タスク（Causal）の 2 種類の注意メカニズムが調査されました。
• さまざまなシーケンス長の入力状況 (20 ～ 4000) をカバーします。
• さまざまな分野の古典的なモデル、主流のディープ モデル、Transformer とそのバリエーションなど、さまざまなベースライン メソッドが比較されます。

下の表に示すように、Flowformer は 5 つのタスクすべてで優れたパフォーマンスを発揮し、モデルの汎用性を証明しています。詳細な実験結果については論文をご覧ください。

5. 分析

Flowformer の動作原理をさらに説明するために、ImageNet 分類タスク (Flow-Attention に対応) における注意に関する視覚化実験を実施しました。その結果、次のことがわかりました。

• 分解に Linear Transformer などのカーネル法のみを使用すると、モデルが混乱し、重要な領域を効果的に捕捉できなくなります。
• 従来の Transformer と Flowformer はどちらも画像の主要な位置を正確にキャプチャできますが、計算の複雑さの点で後者の方が有利です。
• cosFormer は、言語タスクに顕著な効果をもたらす 1 次元の局所性仮定を注意メカニズムに導入します。ただし、画像（2D データを 1D シーケンスに展開）では、局所性の仮定を 2 次元に拡張しなければ、視覚タスクに適応させることはできません。これは、「新しい誘導的な設定を導入しない」という Flowformer の設計アプローチの利点も裏付けています。

上記の視覚化は、Flow-Attention を通じて注意メカニズムの設計に競争を導入することで、些細な注意を効果的に回避できることを示しています。さらなる視覚化実験については論文をご覧ください。

6. 結論

本論文で提案するFlowformerは、ネットワークフローの保存原理を設計に導入し、競争メカニズムをアテンション計算に自然に導入することで、単純なアテンション問題を効果的に回避し、標準的なTransformerの汎用性を維持しながら線形複雑性を実現します。 Flowformer は、長いシーケンス、ビジョン、自然言語、時系列、強化学習の5 つの主要タスクで優れた結果を達成しました。さらに、Flowformer の「特別な誘導的好みはない」という設計コンセプトは、一般的なインフラストラクチャの研究にも刺激を与えます。今後の研究では、大規模な事前トレーニングにおける Flowformer の可能性をさらに探求していきます。