LAYA: Layer-wise Attention Aggregation for Interpretable Depth-Aware Neural Networks

技術要約: LAYA (Layer-wise Attention Aggregation)

問題提起
深層ニューラルネットワーク（DNN）は通常、最終隠れ層（$h_L$）によって生成される表現のみに依存して予測を行います。この設計は、最も深い埋め込みが、先行する変換を通じて抽出されたすべての関連する意味情報を完全に包含しているということを暗黙のうちに仮定しています。しかし、経験的な証拠によれば、中間層には低レベルのパターンから高レベルの抽象化に至るまで、豊かで相補的な情報が含まれていることが示唆されています。標準的なアーキテクチャは、これらの中間表現を出力段階で破棄することで、判別可能な情報を失う可能性があるだけでなく、異なる抽象化レベルの最終決定への寄与をも不明瞭にしてしまいます。さらに、これらの寄与を解釈するための既存の手法は、推論後の外部分析を必要とする「事後的（post-hoc）」なものが大半であり、モデルの予測メカニズム自体に組み込まれたものではありません。

手法
本論文では、標準的な分類器ヘッドの置き換えとして設計された、軽量でアーキテクチャに依存しない出力モジュールであるLAYA (Layer-wise Attention Aggregator) を提案します。LAYAは、最終的な埋め込みのみを投影するのではなく、すべての隠れ層 $\{h_i\}_{i=1}^L$ からの中間表現を動的に集約します。

そのメカニズムは以下の通りです：

1. 投影 (Projection): 各隠れ層の表現 $h_i$ は、学習可能なアダプター $g_i(\cdot)$ を介して、共通の潜在空間（次元 $d^*$）へと写像され、結果として $z_i = g_i(h_i)$ が得られます。
2. スコアリング (Scoring): これらの投影された表現は、オプションとして共有マッピング $\psi(\cdot)$（恒等写像または軽量なMLP）によって変換され、連結されます。スコアリングMLPはこの連結された情報を処理し、非正規化された関連性スコア $s(x) \in \mathbb{R}^L$ を生成します。
3. アテンション (Attention): スコアは、温度スケールされたソフトマックス関数を用いて、入力依存のアテンション係数 $\alpha_i(x)$ に変換されます：
   $$\alpha_i(x) = \frac{\exp(s_i(x)/\tau)}{\sum_{j=1}^L \exp(s_j(x)/\tau)}$$
   ここで、$\tau$ は分布の鋭さを制御します。
4. 集約と予測 (Aggregation & Prediction): 最終的な集約表現は、加重和 $h_{agg} = \sum_{i=1}^L \alpha_i(x) z_i$ として計算され、その後、タスク固有の出力ヘッドに渡されて分類が行われます。

極めて重要な点は、アテンション重み $\alpha_i(x)$ が推論プロセスの一部としてフォワードパス中に生成されるため、追加のバックワード計算や外部の説明モデルを必要とせずに、**内在的な層属性スコア（intrinsic layer-attribution scores）**として機能することです。

主な貢献

1. 出力層の再考: 本研究は、中間表現を無視する従来の出力ヘッドの限界を指摘し、深さ（depth）を考慮した集約への転換を提案します。
2. 内在的な解釈可能性: LAYAは、モデルの決定プロセスが本質的に解釈可能となるメカニズムを提供します。学習されたアテンション係数は、各層の最終予測への寄与を明示的に定量化するため、層レベルの属性特定においてIntegrated GradientsやGrad-CAMのような事後的な説明ツールを必要としません。
3. アーキテクチャへの非依存性: このモジュールは、標準的な出力ヘッドのドロップイン・リプレイスメント（置き換え）として機能し、バックボーンネットワーク（例：CNN、Transformer、MLP）を完全に変更することなく使用できます。
4. 包括的な評価: 本論文では、3つの異なる表現階層（Fashion-MNISTを用いた多層パーセプトロン（MLP）、CIFAR-10を用いた畳み込みニューラルネットワーク（CNN）、および芸術的スタイル分類タスクにおける凍結されたVision Transformer（ViT））にわたる厳格な評価を提示しています。

実験結果

• 予測性能: LAYAは、標準的なベースライン（LastLayer）、静的な連結（Concat）、およびグローバルな混合（ScalarMix）と比較して、同等、あるいは場合によってはより優れた予測性能を達成しました。
  • Fashion-MNIST および CIFAR-10 において、LAYAはテストされたすべてのヘッドの中で最高の平均精度を記録しました。
  • Best Artworks データセットでは、静的な連結が最も高い性能を示しましたが、LAYAは標準的なLastLayerベースラインを大幅に上回り、凍結された学習済みバックボーンにおいても中間表現が価値のある情報を含んでいることを証明しました。
• 説明の忠実度 (Faithfulness): 属性スコアと「1つの層を除外した際の影響（leave-one-layer-out perturbation effects）」との間のスペアマン相関を用いた Faithfulness、および Deletion@1（最も重要な層を削除した際の信頼度低下）を用いた定量的分析により、LAYAの内在的スコアが非常に高い忠実度を持つことが示されました。
  • Fashion-MNISTにおいて、LAYAの忠実度スコア（0.8575）は、勾配ベースの事後的（post-hoc）な手法（約0.89）に迫り、静的なグローバル重みを大幅に上回りました。
  • Best Artworksデータセットにおいて、LAYAは最高の Deletion@1 スコア（0.5787）を達成しました。これは、LAYAが最も重要であると特定した層が、実際に予測に対して最も影響力を持っていたことを示しています。
• 解釈のパターン: 定性的分析により、タスク依存的な深さの利用傾向が明らかになりました。例えば、CIFAR-10ではアテンションはほぼ最終層に集中していましたが、Best Artworksでは、アテンションは特定の中間層および後半のViTブロック（例：ブロック7, 9, 11, 12）に疎に分布しており、これは芸術的スタイルによって変化していました。

意義と主張
本論文は、LAYAが標準的な出力段階における階層的表現の「崩壊（collapse）」に対する、原理的な代替案を提供すると位置づけています。その主要な意義は、分類精度を普遍的に圧倒することにあるのではなく、競争力のある予測性能と、直接検査可能な入力依存の層属性（layer attribution）を両立させている点にあります。

著者は、LAYAが以下の特性を持つと主張しています：

• 内在的な解釈可能性を提供し、外部ツールなしで、特定の抽象化レベルへのモデルの依存関係を透明にします。
• 構造化されたタスク依存のパターンを明らかにし、異なるタスクや、同一タスク内の異なるクラスが、異なる組み合わせの層の深さに依存していることを示します。
• アテンション誘導型のモデル圧縮（アテンションが無視できる層のプルーニング）、適応的な早期終了（early-exit）メカニズム、およびアテンションの変化に基づく診断分析といった、将来的な応用への基礎を提供します。

結論として、本研究の焦点はコンピュータビジョンにありますが、そのフレームワークは他のシーケンスアーキテクチャにも適用可能であり、Transformerのようなモデルにおいて、深さとともに情報がいかに洗練されていくかについて、新たな知見を提供する可能性があるとしています。