【人工知能・深層学習】論文紹介：MetaMixer Is All You Need

論文紹介 MetaMixer Is All You Need 2024/12/14 立教大学人工知能科学研究科 瀧雅人研究室D2 安木 駿介

論文概要 論文タイトル • MetaMixer Is All You Need 原論文 Submitted on 2024/06/04 http://arxiv.org/abs/2406.02021 貢献1 • VisionモデルにおけるFFNのクエリキーバリューメカニズムの詳細な分析を提供 • このメカニズムを抽象化したMixerアーキテクチャ“MetaMixer”を提案 • MetaMixerは、TransformerとConvNeXtのそれぞれの主要ミキサーモジュールを内包する 貢献2 • Transformer由来のミキサーモジュール (Self Attention(SA)ミキサー) が非効率である問題に対処するた め、ラージカーネル畳み込みを使用してSAミキサーを模倣したFFN設計を提案＝MetaMixer内の FFNified attention 貢献3 • 様々なモダリティ（画像、点群、時系列）とタスクにおいて、最先端スコアを達成。 MetaMixerフレームワークを通してミキサーの設計空間を包括的に分析。

• http://arxiv.org/abs/2406.02021

事前知識

事前知識：FFNはkey-valueメモリー FFNは、Self-Attentionのようにkey-valueメモリーとして動作する FFNの2つのprojection weights (𝑊1 と𝑊2 )は、 Self-AttentionにおけるKey・Valueと同様に動作する 正確には以下のような若干の違いがある • FFNがチャンネル軸に沿って適用される • ソフトマックスの代わりにGELU関数を使用する • KeyとValueは動的に計算されない FFNの手順をkey-valueメモリーの観点で解釈すると・・・ 1. Keyにあたる𝑊1 は学習サンプルのパターンを検出。これらのKeyは入力の特定のパターンと相関 2. 𝑐 = 𝐺𝐸𝐿𝑈(𝑋𝑊1𝑇 )は各メモリの非正規化係数を表し、キーにより検出された各パターンの重要性と解釈できる • ※ 𝑐 = 𝐺𝐸𝐿𝑈(𝑋𝑊1𝑇 )を係数と呼ぶ 3. Valueにあたる𝑊2 は対応するメモリ係数を掛け合わせて統合される（係数の加重による分布シフト）

事前知識： FFNはkey-valueメモリー FFNは、Self-Attentionのようにkey-valueメモリーとして動作する FFNのKeyとValueは訓練中に学習され、 モデルのパラメータの一部となる。 Query・Key・Valueは、 入力シーケンスに基づいて動的に計算される →入力ごとに動的に計算されるのでは なく、訓練データからパターンをキャ プチャして保存する永続的なメモリセ ルとして機能する FFNの2つのprojection weights (𝑊1 と𝑊2 )は、Self-AttentionにおけるKey・Valueと同様に動作する 正確には以下のような若干の違いがある • FFNがチャンネル軸に沿って適用される • ソフトマックスの代わりにGELU関数を使用する • KeyとValueは動的に計算されない FFNの手順をkey-valueメモリーの観点で解釈すると・・・ 1. Keyにあたる𝑊1 は学習サンプルのパターンを検出。これらのKeyは入力の特定のパターンと相関 2. 𝑐 = 𝐺𝐸𝐿𝑈(𝑋𝑊1𝑇 )は各メモリの非正規化係数を表し、キーにより検出された各パターンの重要性と解釈できる 参考：Transformer Feed-Forward Layers Are Key-Value Memories • ※ 𝑐 = 𝐺𝐸𝐿𝑈(𝑋𝑊1𝑇 )を係数と呼ぶ 3. Valueにあたる𝑊2 は対応するメモリ係数を掛け合わせて統合される（係数の加重による分布シフト）

• https://arxiv.org/abs/2012.14913

事前知識：PoolFormer PoolFormerは、TokenMixerとして単純にPoolingを使用するため、FFNがモデルの支配的なコンポーネント →FFNの能力を分析する上で望ましいモデル ＝FFN 参考：MetaFormer Is Actually What You Need for Vision Transformerの成功は、Attentionに基づくTokenMixerモジュールがその能力に最も貢献していると考えられていたが、 この研究ではTokenMixerモジュールではなく、Transformerの一般的なアーキテクチャが、モデルの性能にとってより 重要であるという仮説を立て、AttentionをただのPoolingにしても競争力のある性能を得られることを示した。

• https://arxiv.org/abs/2111.11418

事前知識：Transformer系モデルの活性化マップは疎な傾向にある ここで疎（スパース）とは、活性化マップ（ReLU後のMLPの中間出力）の非ゼロ割合が少ない （例：T5-Base→3.0%、ViT-B16→6.3%）ことを指す。 【活性化マップの非ゼロ割合】 ViTでもReLUを使用していることに注意。ViT、T5どちらにおいても、浅い層の時 点で非ゼロ割合が少なく、層が進むにつれてさらに非ゼロ割合は小さくなる 参考 ReLUは負値も0に丸めるため、特に疎な活性化 マップを生成する。 ・The Lazy Neuron Phenomenon: On Emergence of Activation Sparsity in Transformers ・ReLU Strikes Back: Exploiting Activation Sparsity in Large Language Models

• https://arxiv.org/abs/2210.06313
• https://arxiv.org/abs/2310.04564

事前知識：DepthwiseとPointwiseの畳み込み まずは通常の畳み込み K x K x InputD x OutputD バイアス + + 出力特徴マップ 入力特徴マップ 重みフィルター

事前知識：DepthwiseとPointwiseの畳み込み DepthwiseとPointwiseの畳み込み Depthwise K x K x InputD Pointwise InputDxOutputD

事前知識：DepthwiseとPointwiseの畳み込み DepthwiseとPointwiseの畳み込み Depthwise K x K x InputD Pointwise InputDxOutputD

コンピュータビジョンモデルにおいても FFNはキーバリューメモリとして機能するのか？ NLP領域でのFFNは十分に研究されてきたが、ビジョンモデルでは不十分 まずはFFNが支配的なコンポーネントであるPoolFormerで分析してみる

CVモデルでもFFNはキーバリューメモリとして機能する：理由① PoolFormerの活性化値は、他のTransformerで見られるものと類似した傾向にある 【Fig2より】 【比較用先行研究】 先行研究の観察と同様に、最終段階のレイヤは10%未満の活性化係数を持つ 各値が分類の出力予測に対する分布として解釈できることを考慮すると、出力層付近の極端な変化は直感的な結果 ほとんどFFNのみで構成されるPoolFormerが、Self-Attentionを含むモデルと同じような活性化値推移を示すということは FFNのみでもキーバリューメモリのような機能が見られる一つの間接的な証拠

CVモデルでもFFNはキーバリューメモリとして機能する：理由② メモリ係数マップの可視化から、projection weight 𝑊1 がFFNのキーのように動作していることが読み取れる 【Fig2より：最終層で最も活性化されたキーに対応する係数マップ。括弧内の数字は係数の平均値】 各クラスに特化されたキーは一貫してターゲットオブジェクトと相関しており、特徴的な視覚的特徴をとらえる潜在的役割があることを示唆 →FFNのみでもキーバリューメモリのような機能が見られるもう一つの間接的な証拠 【FFNのメモリ係数（再掲）】 メモリ係数 再掲： FFNの手順をkey-valueメモリーの観点で解釈すると・・・ 1. Keyにあたる𝑊1 は学習サンプルのパターンを検出。これらのKeyは入力の特定のパターンと相関 2. 𝑐 = 𝐺𝐸𝐿𝑈(𝑋𝑊1𝑇 )は各メモリの非正規化係数を表し、キーにより検出された各パターンの重要性と解釈できる 3. Valueにあたる𝑊2 は対応するメモリ係数を掛け合わせて統合される

CVモデルでもFFNはキーバリューメモリとして機能する：理由②+α 同じキーは類似の異なるクラスの画像の関連領域で一貫した相関性を示す 同じキーは同一クラスの画像の意味的に一貫 した領域で相関性を示す さらに一部のキーは70~80%ほどの多くのクラスと強い相関を示し、 例えばあるキーは対象物体の周辺と相関し、形状認識に役立つ機能を持つと思われる

ConvNeXtブロックはQKVメカニズムとして 解釈できる

ConvNeXtブロックはQKVメカニズムとして解釈できる DepthwiseConvがクエリprojectionとして使用される場合、ConvNeXtブロックは、 Self-Attentionと同様に、QKVのメカニズムを介して動作すると解釈することもできる Self-AttentionやConvNeXtブロックのような最も成功しているミキサーの基本であるQKVフレームワークは、 ミキサー設計の出発点として不可欠であるという仮説が立つ→MetaMixerの提案へ

ConvNeXtブロックはQKVメカニズムとして解釈できる DepthwiseConvがクエリprojectionとして使用される場合、ConvNeXtブロックは、 Self-Attentionと同様に、QKVのメカニズムを介して動作すると解釈することもできる Self-AttentionやConvNeXtブロックのような最も成功しているミキサーの基本であるQKVフレームワークは、 ミキサー設計の出発点として不可欠であるという仮説が立つ→MetaMixerの提案へ

抽象化アーキテクチャMetaMixerと 具体化アーキテクチャFFNetの提案

提案の全体像：MetaMixer・FFNification・FFNet MetaMixer FFNification FFNet ： サブオペレーションが事前定義されていないQKVメカニズムの抽象化フレームワーク ： Self-Attentionの非効率なサブオペレーションをFFNのサブオペレーションに置き換え、 Self-AttentionとFFNの効率ギャップを埋めること ： MetaMixerファミリーの具体的なアーキテクチャとして論文中で更に提案されているモデル

提案の理解：Self-AttentionのFFNification Self-Attentionは以下の思想に基づきFFNificationされ、FFNified Attentionと呼ばれるブロックに置き換えられる 【 Self-AttentionをFFNificationする際の設計思想】 ■を■に置き換える (1) KeyとValueの生成 Self-AttentionはLinear Projection（動的計算）によってKとVを生成する FFNはKとVとして静的な重みを採用する。 →そのため、簡単で効率的な実装のために静的重みスキームを採用。 (2) 活性化関数(係数生成) Attention Softmaxは、指数計算とトークン長に渡る総和計算を含むためハイコスト で、並列化に課題がある。 そのため、SoftmaxをGELUに置き換え、要素ごとの活性化計算で効率化。 (3) Q・K・Vの相互作用 Self-AttentionにおけるQ-KとCoef-Vの相互作用は、1つのトークンと他の全ての トークンの間の関係を計算するため、トークン長に対して2次的なコスト増加を招く。 →これらの高価な計算を、単純な局所的集約である畳み込みで代用（Kは7以上）

提案の理解：Depthwise畳み込みによる相互作用の置き換えと利点 Q-KとCoef-Vの2つの相互作用はDepthwiseConvで置き換えられ(右図)、効率性と帰納的バイアスからの恩恵を得る FFNificationプロセスにおける最重要ポイント ↓ Q-KとCoef-Vの二か所の相互作用に Depthwise Convを活用すること 【畳み込み相互作用の活用の利点】 ・Self-AttentionはVisionドメインでは特に非効率。代替演算としてDWConvは効率的 ・局所性や位置間の重み共有といった帰納的バイアスからの恩恵 →モデルの複雑さを軽減し、並進的に等価な表現(※)の学習をサポート ※参考：Translation Invariance & Equivariance in Convolutional Neural Networks 【畳み込みにおける補足的な工夫】 ・受容野拡大のため、主に7×7のカーネルサイズを利用 ・re-parameterization trickを活用し、推論コストを増加させることなく性能を向上

• https://blog.paperspace.com/pooling-and-translation-invariance-in-convolutional-neural-networks/

提案の理解：活性化の再考 Softmaxによるexponential normalizationの出力は小数値に集中する ↓ Depthwise畳み込みと組み合わせるのには適していない （局所受容野による畳み込みが十分な文脈集約を達成できない可能性があるため） ↓ この問題を回避できるためGELUを採用

提案の理解：FFNet FFNified AttentionとConvNeXtブロックを、それぞれトークンミキサーとチャンネルミキサー として採用するFFNet（Fast-Forward Network）を提案 チャンネルの Mixing 空間的特徴の Mixing 空間的な パターンの検出 チャンネルの パターン検出 空間的なパターン の重要性を評価 係数の加重による 空間的な分布シフト チャンネルパターン の重要性を評価 係数の加重による チャンネルの分布シフト FFNified attentionとConvNeXtブロックはMixing次元の点で対称であるため（QKVメカニズムの動作軸が異なるため）、 トークンミキサーとチャンネルミキサーとして相補的な役割が期待される

実験

アーキテクチャの詳細

画像分類 ImageNet-1Kの分類結果

Segmentation・Object Detection 【2D Semantic Segmentation】 【COCO object detection and instance segmentation results】 論文中のその他のタスク（Super-Resolutionや3D semantic segmentation、Time-Series Forecasting）も 含めて最先端あるいはそれと競合的な結果を達成している

その他・議論

（議論用の再掲）提案の理解：活性化の再考 Softmaxによるexponential normalizationの出力は小数値に集中する ↓ Depthwise畳み込みと組み合わせるのには適していない （局所受容野による畳み込みが十分な文脈集約を達成できない可能性があるため） ↓ この問題を回避できるためGELUを採用 なぜだろう？ SoftmaxかGELUかでは左図の NormalizedかUnnormalizedかのような 違いが出るのはわかる。 畳み込みは赤いBOXのように局所特徴集 約的な手法なので、小さい値ばかりで活 性化値のメリハリが無いと畳み込み後の 値にあまり違いが出なさそうに思う。 そのような理解で正しいのか？

有効受容野（ERF） FFNetはConvNeXtと同じカーネルサイズのラージカーネルCNNだが その有効受容野はConvNeXtよりもはるかに大きい 【 ERFが “ConvNeXt ＜ FFNet” となる理由】 最終的な予測のためにより広い範囲のピクセルを考慮（重要視）すべく、同じ カーネルサイズの畳み込みをより効果的に利用しているためだと考えられる ConvNeXtに無くてFFNetにある機能、すなわち“トークンミキサーがラージカーネルにより周辺コンテキスト を集約する機能＝空間的なFFN”が、長距離依存の獲得を促進している可能性がある もしカーネルサイズが特別大きいわけではないがERFが大きいCNNモデルがある場合（NFNetとかそんな気がする）、 その要因としてモデルが“空間的なFNN”あるいはそれに近いプロセスを含む可能性が考えられるかも？

有効受容野（ERF）の補足 大きいERFの例 小さいERFの例 有効受容野＝任意のユニットの出力に、一定以上寄与する領域

有効受容野（ERF）の補足 参考：有効受容野を大きくさせるとされている既知のパーツは以下の通り 拡張畳み込み （Dilated Conv） サブサンプリング ラージカーネル （プーリング） 有効受容野の 拡大

トークンミキサーとチャンネルミキサーの補完的なメカニズム 右図より、「トークンミキサーがラージカーネルにより周辺コンテキストを集約し、次にチャンネル ミキサーが重要な領域に焦点を当てる」という補完的なメカニズムが機能しているように読み取れる 【FFNet-3の係数マップ】 チャンネルの Mixing 空間的特徴の Mixing 空間的な パターンの検出 チャンネルの パターン検出 空間的なパターン の重要性を評価 係数の加重による 空間的な分布シフト チャンネルパターン の重要性を評価 係数の加重による チャンネルの分布シフト トークンミキサーでは、Keyはターゲットオブジェクトとその周辺コンテキストの両方と相関しているが、 チャンネルミキサーでは、ターゲットオブジェクトに焦点が集中している トークンミキサーとチャンネルミキサーの補完的なメカニズムの働きを示唆 FFNetでCAMがどのような結果になるかも気になる

FFNの特徴の冗長性に関する懸念 典型的なチャンネル拡張率を採用したFFNは、特徴の冗長性の問題を抱えている FFNetはFFN中心のネットワークなので、アーキテクチャ設計のための特徴冗長性解析を行うべきでは？ 【EfficientViT論文におけるSwin-Tの特徴冗長性解析】 Swinは典型的なViTと同様、 Taylor structured pruningという枝刈り法で、重要でないチャネルを Attentionに対するFFNの 削除し、最も重要なチャネルを残して精度を最大限に維持させた結果 チャンネル拡張率は4倍 【EfficientViTにおけるFFNチャンネル比率調整の効果】 チャンネル比率の調整あり→ 空間的な パターンの検出 FFNチャンネル数 チャンネル比率の調整なし→ Attentionチャンネル数 EfficienctViTでの解析によれば、典型的なチャンネル構成を採用すると Q・K・FFNで特徴の冗長性が生じる。特にFFNの冗長性は顕著に見える。 ※参考：CVPR2023の方のEfficientViT: Memory Efficient Vision Transformer with Cascaded Group Attention Depthwise Convの有無 → 余談。EfficientViTにおいて もトークンミキサーとして DWConvが採用されており、 FFNetはこれに触発されてい ると思われる。実際性能も向 上している。 この結果を受け、EfficienctViTではチャンネル比率の調整（例えばFFNの拡張 率4倍を2倍にし、Vに対してQ・Kの次元を小さくし、ステージごとのチャン ネル増加を2倍以下にする工夫）により効率性と性能を向上させている FFNetではFFNの拡張率4倍を3倍に設定しているが、その効果測定は行われて いない。FFNetはFFN中心のネットワークなので、このあたりにも気を遣った 方がよさそう。

• https://arxiv.org/abs/2305.07027

以上