Beyond Point Estimates: Benchmarking Uncertainty Quantification Methods on the AION-1 Astronomical Foundation Model

技術要約：AION-1における不確実性定量化のベンチマーク

問題提起
AION-1のような天文学的サーベイのための基盤モデルは、銀河の特性推定といったダウンストリームの回帰タスクに対して強力な学習済み表現を提供する。しかし、点予測だけでは科学的推論には不十分であり、誤差を伝播させ、分布外（out-of-distribution）の入力を特定し、堅牢な集団レベルの推論を確保するために、信頼できる不確実性定量化（UQ）が必要となる。UQは極めて重要であるが、高次元の埋め込みを持つ大規模なブラックボックス型の基盤モデルにおいては、その性質が十分に理解されていない。既存の手法は、形式的な被覆保証を欠いているか、あるいは複雑で不均質な天体データにおける局所的な予測の困難さに適応できないことが多い。

手法
著者らは、凍結されたAION-1基盤モデル（具体的には3億パラメータを持つAION-Bバリアント）に適用された7つのUQ手法の直接比較を行う。本研究は、Legacy Surveyの測光・撮像およびDESIスペクトルから得られる5つの物理量、すなわち赤方偏移（$z$）、星質量（$\log M_\star$）、星形成世代の年齢（$t_{age}$）、ガス相金属量（$\log Z$）、および比星形成率（$\log sSFR$）の回帰に焦点を当てている。ラベルはPROVABGS SEDフィッティングの後験分布から導出されている。

実験セットアップは以下の通りである：

1. 基本アーキテクチャ: 3層のMLPが768次元のAION-1埋め込みをスカラー予測へと写像する。エンコーダの重みは凍結されており、回帰ヘッドのみが訓練される。
2. 評価されたUQ手法:
   • 共形手法（分布に依存しない保証）:
     • VanillaSplit: 全域的な分位点を用いる標準的なスプリット共形予測（一定幅の区間）。
     • MADSplit: 区間の幅を局所的な困難さに適応させる、MAD正規化されたバリアント。
     • CQR (Conformalized Quantile Regression): 分位点ネットワークを訓練した後、共形補正を行う。
     • LVD (Locally Valid and Discriminative): 低次元化された特徴空間上でカーネルを学習し、条件付き非適合度分位点を推定する。LCAion（生のAION-1埋め込み上）とLCNet（256次元のMLP隠れ活性化上）の2つの空間でテストされる。
     • LVD-MAD: スケール適応型の区間のために、学習された残差大きさ予測器を拡張したLVD（LCMADAionおよびLCMADNet）。
   • 非共形ベースライン:
     • Deep Ensembles (DE): 5つの独立に初期化されたMLP。
     • MC Dropout: ドロップアウトを有効にした確率的なフォワードパス。

評価指標には、周辺被覆率（marginal coverage）、裾の被覆率（tail coverage: 上限/下限10%）、最悪予測ビン被覆率（worst-predicted bin coverage: 最も困難な10%の点）、および効率性（平均予測区間幅）が含まれる。

主な結果

• 周辺被覆率: すべての共形手法（LVDバリアント、VanillaSplit、MADSplit、CQR）は、すべての特性において、目標とする公称90%の±1パーセントポイント以内の周辺被覆率を達成した。対照的に、非共形ベースラインは較正に失敗した。Deep Ensemblesは系統的に過小被覆（12–38%）となり、MC Dropoutは過剰な幅を持つ区間とともに強い過剰被覆（57–99%）を示した。
• 裾および最悪予測の性能:
  • 最悪予測ビンにおいて、CQRは最も高い被覆率（44–66%）を達成し、次いでLVD-MADバリアント（22–42%）となった。
  • VanillaSplitは、最も困難な領域において被覆率が約0%に崩壊した。
  • 生のAION-1埋め込み上で動作するLVDバリアント（LCAion, LCMADAion）は、圧縮されたMLP活性化を使用するものよりも一般に優れた性能を示し、これは768次元の空間が局所的な較正に有益な構造を保持していることを示唆している。
  • MADバリアント（LCMADNet, LCMADAion）は、非MAD版を一貫して上回り、不均一分散性に対するスケール適応型正規化の重要性を強調している。
• 効率性: 適切に較正された手法の中で、LVDバリアントは被覆率を維持しながら最もタイトな予測区間を生成した。CQRは、周辺的な妥当性を保証するために広い区間を生成する。

主な貢献

1. 体系的なベンチマーク: 本論文は、凍結されたAION-1埋め込みに対する5つの異なる銀河特性回帰タスクにおける、7つのUQ手法の初の直接比較を提供する。
2. 局所的妥当性の実証: 本研究は、共形手法が周辺的な妥当性を保証する一方で、LVDフレームワーク（特に基盤モデルの埋め込み上で動作する場合）が、有限サンプルにおける局所的な妥当性を提供することを実証している。これにより、予測区間は単に全域的な平均に依存するのではなく、各銀河の予測の特定の困難さに適応することが可能になる。
3. 埋め込み空間の分析: 結果は、圧縮された中間特徴（256次元）よりも生の高次元埋め込み（768次元）を利用することによって、特に裾の領域における局所的な較正の堅牢性が向上することを示している。

意義と主張
著者らは、共形予測、特にLVDフレームワークが、天文学における不確実性を考慮した推論のための推奨されるUQアプローチであると結論付けている。分布に依存しない保証を欠く非共形ベースラインや、条件付きの誤較正に陥る可能性のある標準的な共形手法とは異なり、LVDは分布に依存しない保証と局所的な適応性の間の理にかなったバランスを提供する。これは、個々の予測の信頼性が集団レベルの統計と同様に重要な科学的推論において極めて重要である。論文では、AION-Bが使用されたが、埋め込みにおけるLVDの優位性は、より大規模なモデル（AION-LおよびAION-XL）におけるさらなる調査の価値がある旨を述べている。