← 论文海报合集← Paper Notes|
NeurIPS 2024 · 不确定性估计NeurIPS 2024 · Uncertainty Estimation

Estimating Epistemic and Aleatoric Uncertainty with a Single Model

HyperDM:单模型同时估计两类不确定性,媲美深度集成HyperDM: a single model estimates both kinds of uncertainty, rivaling deep ensembles
Matthew A. Chan  ·  Maria J. Molina  ·  Christopher A. Metzler  ·  University of Maryland, College Park

在医学影像与天气预测等高风险领域,区分可通过更多数据消减的 epistemic uncertainty(认知不确定性)与任务本身固有的 aleatoric uncertainty(偶然不确定性)至关重要。 本文提出 HyperDM——将 Bayesian hyper-network 与条件 diffusion model 融合,仅用单个模型便能生成 M×N 规模的预测分布,从而精确分解两类不确定性,预测质量媲美甚至优于 M 倍代价的 deep ensemble。In high-stakes domains such as medical imaging and weather forecasting, distinguishing epistemic uncertainty, which is reducible with more data from aleatoric uncertainty, which is inherent to the task itself is critically important. This paper proposes HyperDM — fusing a Bayesian hyper-network with a conditional diffusion model so that a single model can produce a predictive distribution of M×N samples, thereby precisely decomposing the two kinds of uncertainty, with prediction quality matching or even surpassing a deep ensemble that costs M times as much.

arXiv 2024-02 医学影像 · 天气预报Medical Imaging · Weather Forecasting 单一模型 · 集成近似Single Model · Ensemble Approximation 📄 arXiv:2402.03478 GitHub 代码GitHub Code
不确定性估计Uncertainty Estimation epistemic uncertainty aleatoric uncertainty diffusion model hyper-network Bayesian neural network 医学影像Medical Imaging 天气预测Weather Forecasting uncertainty decomposition ensemble

01 动机Motivation

Conditional diffusion model 已能高效地从数据后验分布中采样,这在理论上使不确定性估计变得简单:训练并采样一个大型 diffusion model 集成即可。 然而,随着模型架构复杂度增长,训练这样一个集成在计算上变得"intractable"(难以承受)。 现有单模型方案——如 Monte-Carlo dropout 和 Bayesian neural network(BNN)——要么损害预测精度,要么推理代价过高,且通常无法同时精确估计两类不确定性。Conditional diffusion models can already sample efficiently from the data posterior, which in principle makes uncertainty estimation simple: just train and sample a large ensemble of diffusion models. However, as model architectures grow more complex, training such an ensemble becomes computationally "intractable". Existing single-model alternatives — such as Monte-Carlo dropout and Bayesian neural networks (BNN) — either compromise predictive accuracy or incur prohibitive inference cost, and generally cannot estimate both kinds of uncertainty precisely at the same time.

"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."
Toy problem: AU and EU estimates
玩具实验验证。 (左)固定训练数据量,改变噪声方差 σ²η ∈ {0.01, 0.04, 0.16, 0.64}:HyperDM 估计的 AU 为 {0.02, 0.05, 0.21, 0.64},与真实值高度吻合。 (右)固定噪声,改变数据集规模 |D| ∈ {100, 200, 400, 800}:EU 估计值 {1.92×10⁻⁴, 2.20×10⁻⁵, 1.17×10⁻⁵, 1.83×10⁻⁶},随数据量增大单调递减,符合预期。Toy-problem validation. (Left) With the amount of training data fixed and the noise variance varied over σ²η ∈ {0.01, 0.04, 0.16, 0.64}: the AU estimated by HyperDM is {0.02, 0.05, 0.21, 0.64}, closely matching the ground truth. (Right) With the noise fixed and the dataset size varied over |D| ∈ {100, 200, 400, 800}: the EU estimates {1.92×10⁻⁴, 2.20×10⁻⁵, 1.17×10⁻⁵, 1.83×10⁻⁶} decrease monotonically as data grows, as expected.
训练开销(vs. M×倍 deep ensemble)Training cost (vs. an M× deep ensemble)
~3%额外训练时间(vs. MC-Dropout baseline)Extra training time (vs. MC-Dropout baseline)
DPS-UQ 训练开销倍数(vs. HyperDM)DPS-UQ training cost multiple (vs. HyperDM)
M×N总预测样本数(M 组权重,每组 N 次采样)Total predicted samples (M sets of weights, N samples each)

02 方法Method

HyperDM 的核心思想是将 Bayesian hyper-network(BHN)conditional diffusion model(DM) 串联,用 BHN 采样 DM 的权重以模拟权重后验 p(φ|D),再用 DM 对每组权重多次采样以模拟似然函数 p(x|y,φ)。 整个框架仅需训练一个 BHN,推理时从中抽取 M 组权重,每组生成 N 个预测,共得 M×N 个样本,再用全方差分解公式(Law of Total Variance)将总不确定性分解为 AU 与 EU。The core idea of HyperDM is to chain a Bayesian hyper-network (BHN) with a conditional diffusion model (DM): the BHN samples the weights of the DM to model the weight posterior p(φ|D), and the DM then samples repeatedly under each set of weights to model the likelihood p(x|y,φ). The whole framework requires training only one BHN; at inference it draws M sets of weights, generates N predictions from each for M×N samples in total, and then applies the Law of Total Variance to decompose the total uncertainty into AU and EU.

HyperDM framework overview
HyperDM 总体框架。 (a) Bayesian hyper-network h_θ 从随机噪声 z∼N(0,σ²) 生成 M 组 diffusion model 权重 φ_i。 (b) 每组权重 φ_i 驱动 DM 随机生成 N 个预测,共得 M×N 个样本 x̂_{i,j}。 (c) 集成预测聚合为最终输出,并分别计算 epistemic uncertainty 图(EU = Var over M)与 aleatoric uncertainty 图(AU = E over M of Var over N)。Overall HyperDM framework. (a) The Bayesian hyper-network h_θ maps random noise z∼N(0,σ²) to M sets of diffusion model weights φ_i. (b) Each set of weights φ_i drives the DM to randomly generate N predictions, giving M×N samples x̂_{i,j} in total. (c) The ensemble predictions are aggregated into the final output, and the epistemic uncertainty map (EU = Var over M) and the aleatoric uncertainty map (AU = E over M of Var over N) are computed separately.

Implicit Likelihood — Aleatoric Uncertainty

条件 diffusion model 通过学习反转逐步将数据分布变换为高斯分布的前向扩散过程,使得我们可以从隐式条件分布 q(x|y) 中采样。 对固定权重 φ_i,反复采样 N 次所得方差即为该权重下的 aleatoric 不确定性估计: ÂU = E_{i∈M}[Var_{j∈N}[X̂]]。 由于 DM 不对预测分布做高斯假设,相比 BNN 的混合高斯近似更为灵活。Conditional diffusion models learn to reverse the forward diffusion process that gradually turns the data distribution into a Gaussian, which lets us sample from the implicit conditional distribution q(x|y). For a fixed set of weights φ_i, the variance obtained by sampling N times is the aleatoric uncertainty estimate under those weights: ÂU = E_{i∈M}[Var_{j∈N}[X̂]]. Because the DM makes no Gaussian assumption about the predictive distribution, it is more flexible than the mixture-of-Gaussians approximation used by a BNN.

Implicit Posterior — Epistemic Uncertainty

BHN h_θ 将随机向量 z∈ℝ⁸ 映射为 DM 的权重 φ,从而隐式地近似权重后验 p(φ|D)。 由于存在多组权重都能合理解释训练数据(尤其在数据稀少时),BHN 不会塌缩到单一模式。 M 组权重预测均值之间的方差即为 epistemic 不确定性: ÊU = Var_{i∈M}[E_{j∈N}[X̂]]。 相比 BNN,HyperDM 训练快 2× 以上(toy 实验:30 s vs. 70 s),推理速度快一个数量级(生成 10,000 预测:0.7 s vs. 8.7 s)。The BHN h_θ maps a random vector z∈ℝ⁸ to the weights φ of the DM, thereby implicitly approximating the weight posterior p(φ|D). Since several sets of weights can plausibly explain the training data (especially when data are scarce), the BHN does not collapse to a single mode. The variance among the predictive means of the M sets of weights is the epistemic uncertainty: ÊU = Var_{i∈M}[E_{j∈N}[X̂]]. Compared with a BNN, HyperDM trains more than 2× faster (toy experiment: 30 s vs. 70 s) and infers an order of magnitude faster (generating 10,000 predictions: 0.7 s vs. 8.7 s).

03 实验Experiments

在两个真实高风险任务上验证:CT 图像重建(LUNA16 数据集,1,200 张肺部 CT,稀疏 Radon 变换 45 视角,Gaussian 噪声 σ²=0.16)和天气温度预测(ERA5 数据集,1,240 张地面气温图,6 小时步长,2009–2018 年 1 月)。 基线:MC-Dropout(单 DM + dropout 采样)、DPS-UQ(M 成员 deep ensemble)。 采样设置统一为 M=10,N=100,共 1000 个预测。在单张 NVIDIA RTX A6000 上训练,400 个 epoch,Adam 优化器,学习率 1×10⁻⁴。Validated on two real high-stakes tasks: CT image reconstruction (LUNA16 dataset, 1,200 lung CT slices, sparse Radon transform with 45 views, Gaussian noise σ²=0.16) and weather temperature forecasting (ERA5 dataset, 1,240 surface temperature maps, 6-hour steps, January of 2009–2018). Baselines: MC-Dropout (a single DM with dropout sampling) and DPS-UQ (an M-member deep ensemble). The sampling setting is fixed at M=10, N=100, for 1000 predictions in total. Training runs on a single NVIDIA RTX A6000 for 400 epochs with the Adam optimizer at a learning rate of 1×10⁻⁴.

预测质量对比(LUNA16 CT 重建 & ERA5 天气预测)Prediction quality comparison (LUNA16 CT reconstruction & ERA5 weather forecasting)

方法Method LUNA16 SSIM ↑ LUNA16 PSNR (dB) ↑ LUNA16 CRPS ↓ ERA5 SSIM ↑ ERA5 PSNR (dB) ↑ ERA5 CRPS ↓
MC-Dropout 0.7730.250.023 0.9331.340.034
DPS-UQ 0.8934.950.01 0.9432.830.013
HyperDM(本文)HyperDM (ours) 0.8735.160.01 0.9533.150.012

在 LUNA16 上,HyperDM 的 PSNR(35.16 dB)优于 DPS-UQ(34.95 dB),而训练开销仅比 MC-Dropout 多约 3%,DPS-UQ 则需约 的训练时间。 在 ERA5 上,HyperDM 的 PSNR(33.15 dB)比 DPS-UQ 高约 1%,CRPS(0.012)也略优。On LUNA16, the PSNR of HyperDM (35.16 dB) is better than that of DPS-UQ (34.95 dB), while its training cost is only about 3% higher than MC-Dropout, whereas DPS-UQ requires about the training time. On ERA5, the PSNR of HyperDM (33.15 dB) is about 1% higher than that of DPS-UQ, and its CRPS (0.012) is slightly better as well.

CT OOD uncertainty maps
CT 重建:OOD 检测。通过向脊柱人工插入金属植入物构造 OOD 测量。DPS-UQ 与 HyperDM 的 epistemic uncertainty 图均成功高亮了异常区域;MC-Dropout 未能在 EU 图中识别 OOD 像素。CT reconstruction: OOD detection.An OOD measurement is constructed by artificially inserting a metal implant into the spine. The epistemic uncertainty maps of both DPS-UQ and HyperDM successfully highlight the anomalous region; MC-Dropout fails to identify the OOD pixels in its EU map.
Weather forecasting OOD uncertainty maps
天气预测:OOD 检测。在加拿大东北部人工插入一个热点构造 OOD 测量。HyperDM 的 EU 图能更精准地定位异常热点;MC-Dropout 反而错误地将美国中部标记为 OOD 区域。所有方法的 AU 图均显示海洋低于北美大陆,与水的高热容物理先验一致。Weather forecasting: OOD detection.An OOD measurement is constructed by artificially inserting a hot spot in northeastern Canada. The EU map of HyperDM localizes the anomalous hot spot more precisely; MC-Dropout instead wrongly marks the central United States as an OOD region. The AU maps of all methods show the ocean to be lower than the North American landmass, consistent with the physical prior of the high heat capacity of water.

Ablation 研究Ablation studies

论文通过改变集成规模 M 和采样次数 N 进行消融实验: (1)增大 M 能改善 out-of-distribution 检测能力(更准确的 EU 估计); (2)增大 N 能平滑 aleatoric uncertainty 估计中的不规则性,使 AU 图更稳定。 聚合方式上,取集成均值(mean)在 SSIM / CRPS 综合指标上优于 median 和 mode。The paper ablates by varying the ensemble size M and the number of samples N: (1) increasing M improves out-of-distribution detection (more accurate EU estimates); (2) increasing N smooths the irregularities in the aleatoric uncertainty estimate, making the AU map more stable. As for aggregation, taking the ensemble mean outperforms median and mode on the combined SSIM / CRPS metrics.

04 局限性Limitations

说明:以下局限性均为论文作者在"Limitations and Future Work"一节中明确陈述。Note: All limitations below are explicitly stated by the authors in the "Limitations and Future Work" section of the paper.
Diffusion model 推理速度慢Diffusion model inference is slow

"As a consequence of their iterative denoising process, inference on DMs is slow compared to inference on classical neural network architectures." 尽管近年加速采样策略(如 DDIM 的 few-step 采样、consistency model 的 single-step 采样)已大幅缓解此问题,但仍是实际部署的障碍。Although recent accelerated sampling strategies (such as the few-step sampling of DDIM and the single-step sampling of consistency models) have greatly mitigated this issue, it remains an obstacle to practical deployment.

Hyper-network 参数量随主网络规模线性增长Hyper-network parameter count grows linearly with the size of the primary network

"Hyper-networks suffer from a scalability problem in that their number of parameters scales with the number of primary network parameters. This stems from the fact that the dimensionality of the hyper-network's output layer is (in most cases) proportional to the number of parameters in the primary network." 对于现代大规模网络(数十亿参数),BHN 本身的存储和训练代价可能无法承受。LoRA 等参数高效权重生成策略是潜在解决方向。For modern large-scale networks (billions of parameters), the storage and training cost of the BHN itself may be unaffordable. Parameter-efficient weight-generation strategies such as LoRA are a potential direction.

实验规模受限(128×128 图像,M≤10)Limited experimental scale (128×128 images, M≤10)

当前实验受限于单张 NVIDIA RTX A6000,图像分辨率为 128×128,集成规模仅测试至 M=10。The current experiments are limited by a single NVIDIA RTX A6000: the image resolution is 128×128 and the ensemble size is tested only up to M=10. "Due to the high computational costs required to train M>10 member deep ensembles, we limited baselines to ten-member ensembles for this experiment." 在更高分辨率或更大集成规模下的性能尚未得到充分验证(此条为从实验设置推断)。Performance at higher resolutions or with larger ensembles has not yet been thoroughly validated (this item is inferred from the experimental setup).