揭秘大型生成网络:扩散模型中的认知不确定性估计

Shedding Light on Large Generative Networks: Estimating Epistemic Uncertainty in Diffusion Models

摘要

本文介绍了一种名为“扩散集合捕获不确定性”(DECU)的创新框架,旨在估计生成扩散模型中的认知不确定性。这些模型因其庞大的参数数量(超过1亿)和高维图像空间的操作而面临传统不确定性估计方法的计算挑战。DECU框架通过引入一种新颖的方法,即在预训练参数的静态集合中训练条件扩散模型的集合,显著减少了计算负担和需要训练的参数数量。此外,DECU利用成对距离估计器(PaiDEs)在高维空间中准确测量模型输出和权重之间的互信息,从而有效地估计认知不确定性。实验结果表明,DECU框架在ImageNet数据集上能够有效地捕捉认知不确定性,特别是在采样不足的图像类别中。

原理

DECU框架的核心在于其高效的集合训练技术和成对距离估计器(PaiDEs)的应用。首先,DECU通过在网络的子集中训练扩散模型的集合,利用Rombach等人[2022]的预训练网络,避免了传统方法中训练每个模型所需的巨大计算资源。其次,PaiDEs是一种非样本基础的方法,已被证明在估计高维回归任务中模型输出和权重之间的互信息方面非常有效。互信息是测量认知不确定性的一个公认指标,PaiDEs通过评估集合组件之间的分布距离来捕捉这种互信息。

流程

DECU的工作流程包括以下几个关键步骤:首先,使用预训练的权重初始化集合中的每个组件,这些权重在整个训练过程中保持静态。其次,仅训练条件标签嵌入层,这显著减少了需要训练的参数数量和训练时间。训练完成后,图像生成过程包括从随机噪声开始,选择一个集合组件,通过马尔可夫链前进到分支点,然后分支到不同的马尔可夫链,每个链对应一个不同的组件,最终生成图像。图2展示了这一流程,其中每个组件的编码器、解码器和UNet是共享的,而嵌入层接受类别标签作为输入。

应用

DECU框架的应用前景广泛,特别是在需要高精度图像生成的领域,如医学图像生成和自动驾驶系统。这些领域对不确定性的准确估计至关重要,因为错误的预测可能导致灾难性的后果。DECU不仅能够提供对生成图像的认知不确定性的深入理解,还可以帮助构建符合国际安全标准的自动化系统解决方案。随着大型扩散模型在自动化系统中的集成越来越广泛,DECU框架为理解和评估这些黑箱模型提供了一个重要的工具。