[論文レビュー] Estimating Epistemic and Aleatoric Uncertainty with a Single Model

Matthew A. Chan, María J. Molina+1 more|arXiv (Cornell University)|Feb 5, 2024

Forecasting Techniques and ApplicationsDecision Sciences被引用数 3

ひとこと要約

本稿では、条件付き拡散モデルとハイパーネットワークを用いて、機械学習におけるエピステミック的不確実性とアレアトリック不確実性の両方を推定する単一モデルフレームワーク、Hyper-Diffusionを提案する。ハイパーネットワークから複数の重みセットをサンプリングし、拡散モデルを介して予測を生成することで、トレーニングコストの僅か数分の1でディープアンサンブルと同等の不確実性推定精度を達成する。

ABSTRACT

Estimating and disentangling epistemic uncertainty, uncertainty that is reducible with more training data, and aleatoric uncertainty, uncertainty that is inherent to the task at hand, is critically important when applying machine learning to high-stakes applications such as medical imaging and weather forecasting. Conditional diffusion models' breakthrough ability to accurately and efficiently sample from the posterior distribution of a dataset now makes uncertainty estimation conceptually straightforward: One need only train and sample from a large ensemble of diffusion models. Unfortunately, training such an ensemble becomes computationally intractable as the complexity of the model architecture grows. In this work we introduce a new approach to ensembling, hyper-diffusion models (HyperDM), which allows one to accurately estimate both epistemic and aleatoric uncertainty with a single model. Unlike existing single-model uncertainty methods like Monte-Carlo dropout and Bayesian neural networks, HyperDM offers prediction accuracy on par with, and in some cases superior to, multi-model ensembles. Furthermore, our proposed approach scales to modern network architectures such as Attention U-Net and yields more accurate uncertainty estimates compared to existing methods. We validate our method on two distinct real-world tasks: x-ray computed tomography reconstruction and weather temperature forecasting.

研究の動機と目的

• 医療画像診断や気象予測など、高リスクな機械学習応用分野における信頼性の高い不確実性評価の必要性に対応する。
• 特に生成的設定において、複雑なモデルに対してディープアンサンブルをトレーニングすることが計算的に非現実的であるという課題を克服する。
• 1つのトレーニング済みモデルを用いて、エピステミック不確実性（モデルの無知）とアレアトリック不確実性（データ固有のばらつき）を分離する。
• 高い予測精度を維持しながら、不確実性推定を可能にする、アンサンブル手法のスケーラブルな代替策を開発する。

提案手法

• 推論時にハイパーネットワークを用いて、条件付き拡散モデルの（擬似）アンサンブルとしてのネットワーク重みを生成する。
• 拡散モデルを用いて、各重みセットから複数の予測をサンプリングし、予測分布を形成する。
• アレアトリック不確実性を、拡散モデルの出力におけるサンプル予測の分散として推定する。
• エピステミック不確実性を、ハイパーネットワークが生成する異なる重みセット間の平均予測の分散として推定する。
• 重みセットが多様で、データと整合的であり、予測の不確実性を反映するように、ハイパーネットワークをエンドツーエンドでトレーニングする。
• 推論時に、重みセットの数（M）と各セットあたりの予測数（N）を調整することで、柔軟なサンプリングを可能にし、アンサンブルの挙動を模倣する。

Figure 1: Block diagram of a hyper-diffusion model. A hyper-network is optimized to generate a (pseudo-)ensemble of network weights for a diffusion model which then outputs a distribution of predictions. The mean “ensemble” prediction is formed by averaging over all samples from the predictive distr

実験結果

リサーチクエスチョン

• RQ1 複雑な逆問題において、1つのモデルがディープアンサンブルの不確実性推定性能をどれだけ正確に近似できるか？
• RQ2 提案されたハイパードイフュージョンフレームワークは、予測精度を損なわず、エピステミック的およびアレアトリック不確実性を正確に分離できるか？
• RQ3 サンプリングレート（M および N）が、分布外の状況におけるエピステミック的およびアレアトリック不確実性推定の信頼性にどのように影響するか？
• RQ4 CT再構成や気象予測などの現実世界の高リスクタスクにおいて、計算コストを抑えたまま、競争力のある性能を達成できるか？

主な発見

• Hyper-Diffusionは、推論時ドロップアウトによる性能劣化に苦しむMC-Dropoutよりもわずかに高いPSNRおよびSSIMスコアをx線CT再構成で達成し、DPS-UQと同等の性能を示した。
• 重みセットのサンプリング数（M）を増やすことでエピステミック不確実性の推定が向上し、分布外の特徴は高い不確実性、分布内領域は低い不確実性として正しく識別された。
• 1つの重みセットあたりの予測数（N）を増やすと、アレアトリック不確実性の推定がより一様かつ安定するようになり、N ≥ 8で不規則なピークが消失した。
• 本手法は正確な不確実性の分離を実現した：アレアトリック不確実性はデータ固有のばらつきを捉え、エピステミック不確実性はモデルの不確実性を反映した。
• Hyper-Diffusionは、CT再構成および気象予測の両タスクでディープアンサンブルと同等の性能を達成したが、トレーニングコストを著しく削減した。
• 本フレームワークは、複数の完全なモデルをトレーニングする計算的負担なしに、高リスク応用分野におけるスケーラブルな不確実性推定を可能にした。

Figure 2: Estimating aleatoric uncertainty with hyper-diffusion models. The distribution of sample variance for four hyper-diffusion models trained on datasets with varying noise levels is shown above. Notice that each distribution’s mean is approximately equal to the aleatoric variance $\sigma_{\et