論文「ENSEMBLING PRUNED ATTENTION HEADS FOR UNCERTAINTY-AWARE EFFICIENT TRANSFORMERS」の技術的サマリー

本論文は、ICLR 2026 にて発表された研究で、深層学習モデルにおける不確実性推定（Uncertainty Quantification: UQ）の効率化と精度向上を目的とした新しいフレームワーク**「Hydra Ensembles」**を提案しています。

以下に、問題定義、手法、主要な貢献、実験結果、および意義について詳細をまとめます。

---

1. 背景と問題定義

背景

深層ニューラルネットワーク（DNN）は視覚や言語タスクで卓越した性能を示しますが、安全クリティカルな領域（医療、自動運転など）では、過信した誤った予測を行うリスクがあります。これを防ぐために、モデルの予測に対する「不確実性」を定量化する UQ が不可欠です。

既存手法の課題

• Deep Ensembles（深層アンサンブル）: 複数の独立して訓練されたモデルを組み合わせる手法で、UQ 性能のゴールドスタンダードですが、計算コスト、メモリ使用量、推論時間が非常に高く、大規模なファウンデーションモデル（CLIP, BERT など）へのスケーリングが困難です。
• 既存の効率的な代替手法（MC Dropout, BatchEnsemble など）: 計算コストは削減されますが、アンサンブルメンバーの独立性が低く、UQ 性能が Deep Ensembles に劣る傾向があります。
• プルーニング（剪定）: 既存のモデルから不要な重みを削除して効率化する手法は一般的ですが、「単純なプルーニングは精度を維持できても、不確実性の較正（Calibration）という重要な問題が指摘されていました。特にノイズのあるデータや分布外（OOD）データにおいて、プルーニングされたモデルは信頼性の低い不確実性推定を行うことが示唆されています。

2. 提案手法：Hydra Ensembles

著者らは、大規模トランスフォーマーモデルにおいて、Deep Ensembles の多様性を維持しつつ、単一モデルに近い推論コストで動作するアンサンブル手法として**「Hydra Ensembles」**を提案しました。

核心的なアイデア

1. アテンションヘッドレベルの構造化プルーニング:

   • 単一の事前学習済みトランスフォーマーモデルから出発し、各コピーに対して異なる「アテンションヘッド」のサブセットを剪定（削除）します。
   • これにより、元のモデルの重みを再利用しつつ、構造的に異なる多様なサブネットワーク（アンサンブルメンバー）を生成します。
   • 従来の「重み」レベルのプルーニングではなく、「機能単位（回路/Circuit）」であるアテンションヘッドを単位とすることで、モデルの多様性を確保します。

2. グループ化された全結合層（GFC）:

   • 生成された複数のプルーニング済みサブネットワークを、単一のトランスフォーマーモデルにマージします。
   • **Fused Multi-Head Attention **(MHA): 各メンバーが保持する異なるアテンションヘッドを、Grouped Fully-Connected (GFC) 層を用いて統合し、一度のフォワードパスで並列処理します。
   • Merged MLP: 各メンバーの MLP 層の重みを平均化してマージします。
   • この構造により、アンサンブル全体を単一のモデルとして推論でき、Deep Ensembles のような複数回の推論パスを不要にします。

3. メンバー生成戦略:

   • Taylor 法: 勾配に基づく古典的な重要性スコアでヘッドを剪定（検証セットなしの場合）。
   • Circuit 法（Headmap）: 不確実性推定や OOD 検出に特化した「回路（Circuit）」を特定し、それに基づいてヘッドを剪定（検証セットがある場合）。これにより、不確実性推定に有害なヘッドを意図的に除外できます。

4. 微調整（Fine-tuning）:

   • 剪定後のモデルをターゲットタスクで微調整することで、精度を回復させます（ゼロショット設定では微調整を行いません）。

3. 主要な貢献

1. プルーニングと不確実性の関係の解明:

   • 理論的および実証的に、単純なプルーニング（特にノイズのあるデータ条件下）がモデルの較正を損ない、不確実性推定を不安定にすることを示しました。
   • 逆に、適切な構造（アテンションヘッド）の選択と多様性の確保が、堅牢な UQ に不可欠であることを証明しました。

2. Hydra Ensembles の提案:

   • トランスフォーマーベースの大規模モデル向けに設計された、最初のプルーニング・アンサンブルフレームワークです。
   • 事前学習済みモデルを再利用し、ゼロから訓練する必要がないため、トレーニングコストと推論コストの両方を大幅に削減します。

3. SOTA 性能の達成:

   • 画像分類、テキスト分類、ゼロショット分類のタスクにおいて、Deep Ensembles と同等かそれ以上の UQ 性能（較正、OOD 検出）を達成しつつ、推論コストを大幅に削減しました。

4. 実験結果

評価タスク

• 教師あり画像分類: ImageNet-1K, CIFAR-100 (ViT-B/16)
• 教師ありテキスト分類: SST-2 (BERT-base)
• ゼロショット画像分類: ImageNet-1K OODベンチマーク (OpenCLIP-ViT)

主な結果

• 精度と較正:
  • ImageNet-1K および CIFAR-100 において、Hydra Ensembles は Deep Ensembles と同等の Top-1 精度を維持しつつ、OOD 検出性能（AUROC, FPR95, AUPR）を改善しました。
  • 特に CIFAR-100 では、Deep Ensembles に比べて OOD 平均 AUROC が +3.4% 向上しました。
• 推論コスト:
  • BF16 精度での推論において、Hydra Ensembles の推論時間は単一モデルの約 1.07 倍のみでした。
  • 対照的に、Deep Ensembles は単一モデルの約 3 倍の時間を要しました。
  • メモリ使用量も Deep Ensembles よりも大幅に削減されています。
• ゼロショット分類（ImageNet-1K OOD）:
  • 追加の訓練なし（ゼロショット）で、既存の SOTA 手法（ViLU など）を上回る性能を達成しました。
  • ViLU と比較して、AUROC が +1.3% 向上、FPR95 が -3.5% 改善、AUPR が +4.0% 向上しました。

考察

• 多様性の源泉: 実験により、アンサンブルメンバー間の多様性は「プルーニングのシード（どのヘッドを削除するか）」によって主に決定されることが示されました。
• MLP マージの影響: MLP 層を平均化してマージしても、アテンションヘッドの多様性が保たれているため、アンサンブルの性能や較正は劣化しませんでした。

5. 意義と結論

本論文は、大規模トランスフォーマーモデルにおける不確実性推定の「効率性」と「精度」のトレードオフを打破する重要なステップです。

• 実用性: 大規模モデルを安全クリティカルな用途に適用する際、Deep Ensembles のような重すぎる手法に依存せず、Hydra Ensembles を用いることで、高い信頼性と低いコストを両立できます。
• 理論的洞察: 「プルーニングは必ずしも UQ に有害ではないが、その方法論（どの部分をどう剪定するか）が重要である」という知見は、今後のモデル圧縮と信頼性向上の研究に指針を与えます。
• スケーラビリティ: 事前学習済みモデルをそのまま活用できるため、CLIP や BERT などの巨大モデルに対しても適用可能であり、将来的な大規模 AI システムの信頼性確保に寄与します。

要約すれば、Hydra Ensemblesは、トランスフォーマーの「頭（アテンションヘッド）」を巧みに剪定・再構成することで、「単一モデルの速さ」で「アンサンブルの信頼性」を実現する画期的な手法です。