Enhancing Layer Interaction Using Key-Correlated Layer Attention

技術要約：Key-Correlated Layer Attentionによるレイヤー間相互作用の強化

問題提起
深層ニューラルネットワーク（DNN）の設計における近年の進展では、レイヤーの出力をシーケンス・トークンとして扱うことで、レイヤー間の相互作用を強化するレイヤー・アテンション（Layer Attention）が活用されています。しかし、標準的なレイヤー・アテンションは、ネットワークの深さ $D$ に対して二次的な計算複雑性（$O(D^2)$）を伴い、全前層の出力をキャッシュする必要があるため、大幅なメモリ増大を招きます。

これらの課題に対処するため、先行研究では再帰的レイヤー・アテンション（Recurrent Layer Attention: RLA）や、MRLA-Lのような線形近似が導入されました。RLAは、Key ($K$) および Value ($V$) の表現をキャッシュすることで計算オーバーヘッドを削減しますが、一度生成されると初期レイヤーの特徴が固定されてしまう静的なメカニズムとして機能するため、動的な情報の更新が制限されます。さらに、MRLA-Lのような線形近似は、アテンション出力を線形近似することで複雑さを軽減しようと試みています。著者らは、このアプローチによって初期レイヤーの寄与が指数関数的に減衰すること（減衰係数 $\beta_l$ の分析により示されている通り）を指摘しており、これにより効果的な長距離のクロスレイヤー依存関係を確立することが困難になっています。加えて、既存の線形化手法は、過去のレイヤー出力のキャッシュに伴うメモリのボトルネックや、階層的なDNN構造に線形アテンションを適用する際の特有の課題に対処できていないことが多くあります。

手法：Key-Correlated Layer Attention (KCLA)
提案手法である**Key-Correlated Layer Attention (KCLA)**は、動的な更新と長距離依存関係のモデリングを維持しつつ、線形計算複雑性を実現するように設計された新しい線形化レイヤー・アテンション手法です。本手法は、以下の3つの核心的な洞察とメカニズムに基づいています。

1. Keyの相関性の観察： 実証的な分析により、標準的なレイヤー・アテンションにおける異なるレイヤー間のKey表現は、高いコサイン類似性を示すことが明らかになりました。これを利用して、KCLAはValue ($V$) の線形変換を因数分解し、Key表現が方向的に相関していると近似することで、標準的なレイヤー・アテンションの出力を近似します。これにより、$K$ および $V$ 行列の全履歴を保存することなく、レイヤー出力を再帰的な状態（$U^t$）へと集約することが可能になります。
2. 動的なSoftmax復元： アテンションを線形化する際の重大な課題は、アテンション・スコアの正規化と動的な分布調整を可能にするsoftmax操作の喪失です。KCLAは、この能力を復元するために修正されたマルチヘッド・アテンション・メカニズムを導入しています。
   • 対数一様数列から導出された**異なるスケーリング係数（温度パラメータ）**を持つ複数のヘッドを採用しています。
   • 学習可能な線形変換 $\Phi$ を用いて、Queryのノルム（$\|Q^t\|$）と相関角（$\cos \theta$）をヘッド固有の重みへとマッピングします。
   • これにより、モデルは各レイヤーにおいてアテンション分布を動的に調整でき、フルサイズの $N \times N$ 行列を計算することなく、softmaxにおける温度パラメータの役割を効果的にシミュレートできます。
3. クロスステージ適応： ステージごとに特徴マップのサイズやチャネル次元が変化するDNNに対応するため、KCLAは特定の適応手順（アルゴリズム1）を実装しています。これには、蓄積された特徴の正規化、ダウンサンプリングのためのストライドを伴う $1 \times 1$ 畳み込みの適用、およびステージ間の次元の不一致を埋めるための正規化因子の再計算が含まれます。

主な貢献

• 観察と証明： 著者らは、標準的なレイヤー・アテンションにおける異なるレイヤー間のKey表現が方向において高度に相関していることを示し、線形近似の理論的根拠を提供しました。
• 新しい線形化手法： KCLAは、RLAの静的な性質やMRLA-Lの指数関数的な減衰を回避する、線形化レイヤー・アテンション手法として提案されました。これは、線形計算複雑性およびネットワークの深さに依存しない固定された空間複雑性を実現します。
• 動的なSoftmax近似： マルチヘッドのスケーリングと学習可能な変換を用いることで、softmax操作を近似する新しい手法を導入し、各レイヤーにおける動的かつ適応的なアテンション分布の調整を可能にしました。

実験結果
著者らは、画像分類、物体検出、医療画像セグメンテーションの3つの基本的なビジョン・タスクにおいてKCLAを評価しました。

• 画像分類： CIFAR-100において、KCLAはResNetの深さが20、56、110のいずれの場合でも最高の性能を達成し、MRLA-L、DLA-L、および標準的なレイヤー・アテンションを上回りました。ImageNet-1Kでは、ResNet-50およびResNet-101に統合されたKCLAは、ベースラインに対してそれぞれ+1.7%および+1.4%のTop-1精度向上を実現しました。また、MRLA-Lに対してわずかながらも一貫した優位性を維持しつつ、動的なDLA-Lよりも大幅に少ないパラメータ数を使用しました。
• 物体検出： Faster R-CNNおよびMask R-CNNを用いたCOCO2017データセットにおいて、KCLAは競争力のある性能を示しました。ResNet-101を用いた場合、より複雑なDLA-Lアーキテクチャに対し、2.2M少ないパラメータを使用しながら、いくつかの指標（AP$_{75}$, AP$_S$, AP$_M$）で上回りました。
• 医療画像セグメンテーション： ISIC2017、ISIC2018、およびKvasir-SEGデータセットにおいて、KCLAはMRLA-LおよびベースラインのResNet-UNetを一貫して上回り、最小限のパラメータ・オーバーヘッドで優れたmIoUおよびDSCスコアを達成しました。
• 複雑性分析： メモリとFLOPsが深さに対して二次的に増加するMRLBとは異なり、KCLAは計算負荷とメモリコストの両方において線形な増加を示します。

意義と主張
本論文は、KCLAが線形アテンションの効率性と標準的なアテンションの表現力の間の溝を埋めることに成功したと主張しています。Keyの相関に基づき標準的なレイヤー・アテンション機構を直接近似することで、KCLAはRLAの静的な特徴表現の問題や、MRLA-Lの長距離依存関係の劣化を回避しています。

著者らは、提案手法が効率的な長距離のクロスレイヤー依存関係をモデリングする能力を維持し、効率と性能の望ましいバランスを提供することを強調しています。また、softmax関数を復元するための手法が、これらの結果を達成するための重要な要因であると述べています。現在の実装は事前設定された温度パラメータに依存していますが、著者らは、完全な自己適応型・動的温度マルチヘッド線形レイヤー・アテンションの設計が、将来の最適化に向けた方向性であるとしています。再現性を促進するため、コードは公開されています。