CVPR 2024 | ABNN：将大型预训练模型无缝转换为贝叶斯神经网络，提高模型综合安全性

作为现代人工智能技术的基石，深度神经网络（DNN）已经在各种感知和认知任务中进行应用，例如目标检测、自然语言处理等等。随着DNN的广泛使用，模型本身的安全性和可靠性也越来越受到研究者和用户的重视。如何对DNN模型中固有的不确定性（uncertainty）进行量化并提高稳定性在风险性较高的应用中（例如自动驾驶、医疗诊断和工业异常检测等）至关重要。

本文介绍一篇来自NVIDIA、valeo.ai和新加坡国立大学等单位合作完成的论文，本文提出了一种全新的自适应贝叶斯神经网络策略（ABNN），可以以post-hoc的形式来将预训练好的DNN无缝转换为BNN（贝叶斯神经网络天然具备实现不确定性估计的能力）。ABNN既保留了大型DNN的鲁棒预测能力，同时也可以通过添加简单的BNN适应层和一些轻量的微调步骤来提高模型的不确定性量化能力。

论文题目： Make Me a BNN: A Simple Strategy for Estimating Bayesian Uncertainty from Pre-trained Models 论文链接： https://arxiv.org/abs/2312.15297

一、引言

在一些安全性要求较高的场景（例如自动驾驶和医疗诊断）中部署DNN，需要我们全面了解DNN的局限性和脆弱性，尤其是关注一些超越其原始预测范围的异常情况。先前的工作大多使用贝叶斯方法对神经网络进行不确定性估计，虽然BNN具有坚实的理论基础，可以在给定训练数据集的情况下估计模型参数的后验分布。但是BNN很难对大型的DNN起效（训练不稳定、计算效率低下），特别是对于各种复杂的计算机视觉网络。

深度集成（Deep Ensembles）[1]是一种高效的替代方案，但是需要训练多个独立的神经网络，并在推理时对它们的输出进行平均来近似估计模型的不确定性，计算成本较高，因而很难应用在资源受限的场景中。为了解决上述问题，本文提出了一种事后策略（post-hoc）将预训练好的DNN转换为BNN，并且仅在网络的归一化层上添加一个简单的适应模块进行微调，作者将这种方式称为Adaptable-BNN（ABNN）。ABNN可以以资源高效的形式估计预训练模型局部最小值周围的后验分布，同时实现鲁棒的不确定性估计。上图展示了本文方法的不确定性估计性能与训练时间之间的权衡情况。

如上图所示，在使用ABNN对DNN进行转化后，模型在接受分布外样本时可以产生出多样化的不确定性预测来实现更加全面的决策，而普通的DNN只能做出最高置信度的错误预测。

二、本文方法

2.1 贝叶斯框架下的BNN

传统的DNN训练通常遵循最大似然估计（maximum likelihood estimate，MLE）方式，给定训练数据集，希望模型学习到参数使模型对训练数据的联合概率分布最大化，这可以通过最小化负对数似然损失函数 来实现。而在贝叶斯框架下，不仅需要考虑数据的似然项，还需要结合模型参数 的先验分布 ，并基于最大后验概率估计（maximum a posteriori，MAP）来进行优化，这等价于最小化 损失函数:

通常情况下，参数 的先验分布 会被设置为高斯分布，这就形成了损失函数中常见的 L2 正则化项。如果我们将MAP的思路应用到DNN中，就可以得到一个具有最大后验概率的单一确定性模型，可以将其表示为 。

与传统的确定性DNN不同，贝叶斯神经网络（BNN）并不是仅学习单一的 参数，而是对参数的后验概率分布 进行建模，具体来说，在给定新的输入时，BNN的输出是通过整合模型所有可能参数集合的期望得到的，可以表示如下：

其中 表示整个参数空间，相较于DNN，BNN的主要优势在于能够直接量化模型预测的不确定性，得到更加可靠、完整的输出概率分布 。

2.2 Adaptable-BNN（ABNN）

为了在现有DNN中实现更加高效、稳定的贝叶斯建模，本文引入了一种称为自适应BNN的贝叶斯建模方式，以提高DNN在实际应用中的不确定性量化能力。ABNN的整体框架图如下图所示，其创新之处在于巧妙的设计了一个Bayesian Normalization Layer（贝叶斯归一化层BNL）来将预训练的确定性 DNN 无缝转化为具有贝叶斯特性的网络，并且只需微调少量参数即可完成训练，避免了从头训练 BNN 所面临的不稳定性。

本文提出的BNL可以灵活的替换DNN中的归一化层（如批归一化、层归一化、实例归一化等），替换过程中还需要加入一定的高斯噪声来将确定性DNN转换为BNN，BNL的数学表达式可以表示如下：

其中， 表示第 层的预激活输出， 和 分别是 的经验均值和方差，通过标准归一化计算得到。 和 是 BNL 的可学习参数，与标准归一化层相比，BNL 引入了额外的随机噪声项，通过乘以可学习的缩放因子来赋予 BNL 贝叶斯特性。这使得原本确定性的 DNN 能够捕捉模型参数的不确定性，从而得到更可靠的不确定性估计。

2.3 ABNN的训练和推理

在对ABNN实现归一化层的结构转换后，只需要对BNL的参数 和 进行微调即可。但是在微调过程中需要避免模型陷入到局部最小值处，为了缓解这个问题，作者对原有损失函数引入了一个小的扰动，以防止训练崩溃并增强模型的多样性。这个扰动需要对交叉熵损失中类权重进行修改，可以定义为：

其中， 代表模型在训练开始时初始化的类相关随机权重。随后可以将该扰动项与原有的MAP损失进行结合：

通过这种受约束的微调过程，ABNN 能够在保留 DNN 主要预测能力的基础上，显著增强其不确定性建模能力，ABNN的具体训练过程如下图所示。

而在推理阶段，ABNN 利用采样的方式来量化模型的不确定性。具体过程如下:

1. 对于一个给定的输入 ，ABNN 首先将其输入到多个 ABNN 模型中,每个模型对应一组不同的参数 。

2. 对于每个 ABNN 模型，ABNN 还会独立采样多次不同的噪声 ，并将其带入 BNL 计算得到预测输出 。

3. 最终，ABNN 将所有 ABNN 模型在不同噪声下的预测结果进行平均，得到综合的不确定性估计:

这种通过采样多个参数配置和噪声项来平均预测的方式，使 ABNN 能够有效捕捉模型参数的不确定性，从而得到更加可靠的输出概率分布。论文指出，这种推理过程的计算开销相对于标准 BNN 和深度集成(Deep Ensembles)方法都有显著降低。

三、实验效果

本文的实验主要在图像分类和语义分割两大任务上进行，前者主要使用 CIFAR-10、CIFAR-100 和 ImageNet 三个经典的图像分类数据集，使用的网络架构包括 ResNet-50 和 WideResNet28-10。后者主要在StreetHazards、BDD-Anomaly 和 MUAD等数据集上进行，网络架构为 DeepLabv3+。作者选取了现有流行的BNN方法作为对比baseline，包括BatchEnsemble[2]，Deep Ensembles[1]和Laplace[3]等。

下表展示了ABNN与其他对比方法在CIFAR-100数据集上的性能对比，可以看出，ABNN 在保持高分类准确率的同时，在不确定性指标(如 ECE、AUPR、AUC、FPR95)方面也取得了非常出色的表现，与深度集成(Deep Ensembles)相当。

在更大规模的ImageNet数据集上，ABNN的性能始终可以超越Laplace和单模型基线，如下表所示。当模型架构从ResNet家族切换到ViT时，ABNN仍然表现出具有强竞争力的性能，这表明ABNN具有一定的架构鲁棒性。

除了图像分类，作者还在语义分割任务上评估了 ABNN 的性能，下表展示了本文方法在StreetHazards、BDD-Anomaly和MUAD数据集上的性能。

从表中可以看出，ABNN在语义分割任务上的性能也非常出色，在 mIoU(平均交并比)指标上，ABNN 虽然略低于最优的Deep Ensembles，但优于大部分对比方法。在不确定性指标如 AUPR、AUC 和 FPR95 上，ABNN 都取得了出色的表现，与Deep Ensembles相当。值得一提的是，ABNN 的计算成本明显低于Deep Ensembles，这在实际部署中会带来很大优势。

除了常规的性能评估，作者还对ABNN的多样性进行了分析，如上图所示，作者绘制了ABNN特征的T-SNE降维可视化效果，该实验使用ResNet50在CIFAR-10上对三个ABNN进行微调，可以看出，ABNN可以从单个DNN开始，不断扩张特征边界使其偏离局部最小值并探索不同的模式。

四、总结

本文为神经网络不确定性量化研究引入了一种新颖的视角，利用预先训练的确定性模型的优势，作者提出了一种简单可自适应的贝叶斯神经网络ABNN，ABNN可以通过较少的修改性策略来将确定性DNN转变为贝叶斯网络，从而提高网络的稳定性。ABNN不仅保留了DNN的主要预测特性，仅需要插入轻量的BNN自适应层和少量的fine-tuning就可以获得对现实数据进行不确定性量化的能力。作者通过大量的实验证明，ABNN 能够在计算效率、训练稳定性和不确定性建模能力等多个方面优于标准 BNN 和其他后验概率近似方法。这些优势使 ABNN 成为一种高度实用和可扩展的贝叶斯网络建模方法，在安全应用领域具有广阔的应用前景。

参考资料

[1] Balaji Lakshminarayanan, Alexander Pritzel, and Charles Blundell. Simple and scalable predictive uncertainty estimation using deep ensembles. In NeurIPS, 2017.

[2] Yeming Wen, Dustin Tran, and Jimmy Ba. BatchEnsemble: an alternative approach to efficient ensemble and lifelong learning. In ICLR, 2019.

[3] Erik Daxberger, Agustinus Kristiadi, Alexander Immer, Runa Eschenhagen, Matthias Bauer, and Philipp Hen- nig. Laplace redux–effortless Bayesian deep learning. In NeurIPS, 2021.

llustration From IconScout By TanahAir Studio

-The End-

扫码观看！

本周上新！

“AI技术流”原创投稿计划

TechBeat是由将门创投建立的AI学习社区（

www.techbeat.net

投稿内容

// 最新技术解读/系统性知识分享 //

// 前沿资讯解说/心得经历讲述 //

投稿须知

稿件需要为原创文章，并标明作者信息。

我们会选择部分在深度技术解析及科研心得方向，对用户启发更大的文章，做原创性内容奖励

投稿方式

发送邮件到

chenhongyuan@thejiangmen.com

或添加工作人员微信（chemn493）投稿，沟通投稿详情；还可以关注“将门创投”公众号，后台回复“投稿”二字，获得投稿说明。

关于我“门”

将门是一家以专注于数智核心科技领域的新型创投机构，也是北京市标杆型孵化器。 公司致力于通过连接技术与商业，发掘和培育具有全球影响力的科技创新企业，推动企业创新发展与产业升级。

将门成立于2015年底，创始团队由微软创投在中国的创始团队原班人马构建而成，曾为微软优选和深度孵化了126家创新的技术型创业公司。

如果您是技术领域的初创企业，不仅想获得投资，还希望获得一系列持续性、有价值的投后服务，欢迎发送或者推荐项目给我“门”:

bp@thejiangmen.com

点击右上角，把文章分享到朋友圈