AI-Papers
- Attention機構を「動的パラメータを持つMLP」として数学的に再定式化し、明示的なAttention計算なしに線形時間でグローバル視覚モデリングを実現
- ImageNet分類でWeightFormer-Sが81.3%を達成し、DeiT-S（79.8%）を上回りながらFLOPs削減と高速推論を両立
- 高解像度1248×1248入力においてDeiT比7.7倍の処理速度とメモリ91%削減を実現し、スケーラビリティを実証
研究の背景
Vision Transformer（ViT）がコンピュータビジョンの主流となって以来、その中心にあるSelf-Attention機構のグローバルモデリング能力、つまり画像全体の文脈を一度に把握する力が、高精度の源泉と見なされてきました。しかしSelf-Attentionはトークン数Nに対してO(N²)の計算量を要するため、高解像度画像や長系列の処理では計算コストが急激に膨らみます。
この問題に対応するために、Linear AttentionやState Space Model（SSM）など線形時間の代替手法がこれまで多数提案されてきました。しかしその多くは、トークン間の「明示的な重み付き集約」こそがAttentionの強さの本質であるという前提のもとに設計されていました。清華大学LeapLabの研究チームは、この前提自体を問い直し、Attention機構の数学的な再定式化から出発する新しいアプローチを提案しています。
Attentionを動的MLPとして捉え直す
本研究の核心は、Self-Attentionの出力を単一クエリの視点で解釈し直すことにあります。従来、Attentionの出力はO = Softmax(QK⊤/√d)Vと書かれますが、あるトークンiに着目すると、その出力はK⊤を第1層の重み、Vを第2層の重み、Softmaxを活性化関数とする2層MLPの出力と数学的に等価であることが示されます。
さらに重要なのは、K⊤とVが入力全体Xから生成される点です。これをW = G(X)と表記すると、Attentionは「入力依存の動的パラメータを持つMLP」O = F_W(X)として書き直せます。従来の解釈が「トークン間の明示的な重み付き集約」を本質と見ていたのに対し、この新視点では「動的パラメータがグローバル文脈を暗黙的に圧縮する」ことが本質とされます。明示的なAttention行列の計算は不要であり、適切な動的パラメータ予測さえできれば同等のグローバルモデリングが達成できるという主張です。
WeightFormerの設計
この洞察をもとに提案されたのがWeightFormerです。入力Xから動的パラメータを予測してMLPに適用することで、Attention行列の明示的な計算を回避します。動的重み予測には2種類の戦略が採用されています。
一つ目は動的線形層戦略です。重みの基本形をW(X) = W₀ + ΔW(X)として、グローバル平均プーリング（GAP）を用いた単純な実装から出発し、最終的に双方向活性化（ΔW = W₁σ(W₂X⊤)σ(XW₃)W₄）を用いる高精度な形式に洗練させています。二つ目は動的深さ方向畳み込み戦略で、入力を固定解像度K×Kに圧縮して畳み込みカーネルを動的に生成します。効率性のため、これらの動的層は3ブロックごとにスパース配置されています。
有効受容野（ERF）の可視化実験では、静的なCNNや静的重み戦略と異なり、動的重み予測を用いたモデルが訓練後に画像全体をカバーする広域の受容野を獲得することが確認されています。この結果は、動的パラメータが明示的なAttentionなしにグローバル文脈を効果的に統合できることを視覚的に裏付けるものです。
実験結果
ImageNet-1K画像分類では、WeightFormer-S（2700万パラメータ、4.4G FLOPs）が81.3%のTop-1精度を達成しています。二次計算量を持つDeiT-S（79.8%、4.6G FLOPs）を精度・計算量ともに上回っており、線形時間モデルが二次計算量のTransformerを超えた初めての事例の一つとして注目されます。より大規模なWeightFormer-B（448×448入力）では83.4%を記録しています。
下流タスクでも同様の傾向が確認されています。COCO物体検出ではWeightFormer-TがBox AP 45.0を記録し、DeiT-T（44.4、FLOPs 106G）を精度・計算量ともに凌駕しています（FLOPs 77G）。ADE20KセマンティックセグメンテーションではWeightFormer-Sが45.6 mIoUを達成しています。また、高解像度1248×1248の入力においてDeiT比7.7倍の処理速度とGPUメモリ91%削減を実現しており、推論高速化の観点からも実用上の優位性が示されています。画像生成タスク（ImageNet条件付き）でもWeightFormer-B/2がFID 38.21を記録しており、生成モデルへの適用可能性も示されています。
まとめと今後の展望
WeightFormerは、Attention機構の本質を「トークン間の明示的な重み付き集約」から「動的パラメータによる暗黙的なグローバル文脈圧縮」へと再定義し、明示的なAttention計算なしに同等の性能を線形計算量で実現できることを実証しました。ImageNet・COCO・ADE20Kにわたる広範な実験がその有効性を支持しており、GitHubでコードも公開されています。
一方で、著者らも認める課題があります。現時点の評価はビジョンタスクに限定されており、NLPや音声など他ドメインへの汎化性は未検証です。動的パラメータ化がなぜ機能するのかの理論的理解も十分ではなく、入力条件付きの重み予測は勾配の流れを複雑にするため、訓練安定性の確保も引き続き検討が必要です。最適な重み生成機構の設計空間は広く、探索余地が多く残されています。
グローバルモデリングの実現に「明示的なAttentionが不可欠である」という常識に揺さぶりをかけたこの研究は、線形時間モデルの設計に新たな方向性を示しています。動的パラメータ化という枠組みをNLPや音声認識へ展開することや、より効率的な重み生成機構の探索によって、今後さらに幅広いドメインで実用的な線形時間モデルが登場することが期待されます。
