bde: 一个通过MILE实现贝叶斯深度集成的Python包

# 用于通过 MILE 实现贝叶斯深度集成的 Python 包
来源：https://arxiv.org/html/2605.14146
Vyron ArvanitisAngelos Aslanidis慕尼黑大学物理系，慕尼黑，德国同等贡献作者。Emanuel Sommer慕尼黑大学统计系，慕尼黑，德国慕尼黑机器学习中心，慕尼黑，德国同等贡献作者。David Rügamer慕尼黑大学统计系，慕尼黑，德国慕尼黑机器学习中心，慕尼黑，德国

## 摘要

bde(https://github.com/scikit-learn-contrib/bde) 是一个 Python 包，旨在将最先进的基于采样的贝叶斯深度学习 (BDL) 带给实践者和研究人员。该包结合了 JAX 和 blackjax(Bradbury 等, 2018 (https://arxiv.org/html/2605.14146#bib.bib5); Cabezas 等, 2024 (https://arxiv.org/html/2605.14146#bib.bib4)) 的高速与高性能能力，以及 scikit-learn(Pedregosa 等, 2011 (https://arxiv.org/html/2605.14146#bib.bib9)) 的用户友好 API。它针对表格监督学习任务，包括分布回归和（多类）分类，为贝叶斯深度集成 (BDEs)(Sommer 等, 2024 (https://arxiv.org/html/2605.14146#bib.bib2)) 提供无缝接口，并特别实现了微正则朗之万集成 (MILE)(Sommer 等, 2025 (https://arxiv.org/html/2605.14146#bib.bib1))。

bde 的工作流程实现了 MILE 的两阶段 BDE 推理过程。首先，它使用正则化经验风险最小化（以负对数似然作为损失），通过 AdamW 优化器(Loshchilov 和 Hutter, 2019 (https://arxiv.org/html/2605.14146#bib.bib8)) 优化 n_members 个独立配置灵活的前馈神经网络实例。其次，它过渡到采样阶段，使用微正则朗之万蒙特卡洛(Robnik 等, 2023 (https://arxiv.org/html/2605.14146#bib.bib6); Robnik 和 Seljak, 2024 (https://arxiv.org/html/2605.14146#bib.bib7))，并附带专门针对贝叶斯神经网络调整的调优阶段。这种结合被称为 MILE(Sommer 等, 2025 (https://arxiv.org/html/2605.14146#bib.bib1))。本质上，神经网络集成的优化首先找到不同的高似然模式，然后通过这些模式进行采样，探索局部后验结构。这个过程生成一组样本（模型）集成，构成隐式后验近似。

## 软件设计

由于集成成员之间的优化和采样是独立的，bde 利用 JAX 的并行化和即时编译，在 CPU、GPU 和 TPU 上高效扩展。给定新的测试数据，该包近似后验预测，通过统一接口提供点预测、置信区间、覆盖估计和其他不确定性度量。

## 领域现状

可靠的不确定性量化 (UQ) 越来越被视为现代机器学习系统的关键组成部分，而 BDL 为实现这一目标提供了原则性框架(Papamarkou 等, 2024 (https://arxiv.org/html/2605.14146#bib.bib3))。尽管有多个库支持诸如变分推理或经典贝叶斯建模等基于优化的方法，但用于贝叶斯神经网络中基于采样的推理的易用工具仍然稀缺。现有的概率编程框架提供了 MCMC，但需要大量手动配置才能在神经网络模型上实现竞争性能。

## 需求陈述

bde 通过提供首个用户友好的 MILE 实现来弥补上述差距——这是一种混合采样技术，已被证明能为贝叶斯神经网络提供强大的预测准确性和校准不确定度(Sommer 等, 2025 (https://arxiv.org/html/2605.14146#bib.bib1))。通过完全的 scikit-learn 兼容性，该包能够无缝集成到现有机器学习工作流程中，使用户无需具备 MCMC 动力学、初始化策略或 JAX 内部机制的专门知识，即可获得原则性的贝叶斯不确定性估计。

通过自动编排优化、采样、并行化和预测推理，bde 为将基于采样的 BDL 方法应用于表格监督学习任务提供了一种快速、可重复且实用的解决方案。

## 研究影响声明

bde 弥合了高性能 MCMC 研究与实际数据科学之间的差距。通过为表格数据提供精心实现的 MILE，它使研究人员和实践者都能轻松应用 BDE。其被纳入 scikit-learn-contrib 组织，确保遵循严格的软件标准和 API 一致性，使其成为表格机器学习中可重复 UQ 的可靠社区工具。

## 使用示例

以下示例展示了使用 bde 仅用几行代码完成带有 UQ 的回归任务。假设训练输入已经预处理和归一化。工作流程指定集成和采样超参数，拟合模型，并获取后验预测量，包括预测矩、置信区间和原始集成输出。

frombdeimportBdeRegressor

regressor=BdeRegressor(

n_members=8,

hidden_layers=[16,16],

activation="relu",

epochs=1000,

validation_split=0.15,

lr=1e-3,

weight_decay=1e-4,

patience=20,

warmup_steps=5000,

n_samples=200,

n_thinning=10,

desired_energy_var_start=0.5,

desired_energy_var_end=0.1,

)

regressor.fit(x=X_train,y=y_train)

means,sigmas=regressor.predict(X_test,mean_and_std=True)

means,intervals=regressor.predict(X_test,credible_intervals=[0.1,0.9])

raw=regressor.predict(X_test,raw=True)

分类任务类似地使用 BdeClassifier 完成。

## 回归基准测试

我们在 airfoil(Dua 和 Graff, 2017 (https://arxiv.org/html/2605.14146#bib.bib10)) 和 bikesharing(Fanaee-T, Hadi, 2013 (https://arxiv.org/html/2605.14146#bib.bib11)) 数据集上提供了一个小型 bde 基准测试。我们报告了平均预测性能（RMSE）、UQ 度量（分布回归和均值回归公式中的 NLL），结果为 5 次独立运行的均值±\pm 标准差。结果显示，BDE 在 UQ 方面具有即开即用的竞争力，尤其是在其原生的分布回归能力方面。

所有模型使用 10 个 CPU 核心，无需额外调优。对于 BDE，我们从具有四个宽度为 16 的隐藏层的前馈神经网络中生成了 1000 个后验样本。尽管由于 MCMC 采样的迭代性质，高度优化的 BDE 通常比基于优化的竞争对手产生更高的计算成本，但这一投入使得构建灵活的非参数后验近似成为可能。如上所示，这种权衡产生了强大的性能和严格的认知 UQ，例如通过校准的置信区间。所有实验配置都随发布包一起提供，以确保可重复性。

## AI 使用声明

在软件开发过程中，通过 GitHub Copilot 使用了生成式 AI 进行本地、行级或块级的代码自动补全，使用了 Claude Sonnet 3.7 和 4 模型。此外，还使用了 Codex 协助包文档的初稿以及一些小的重构，以确保与 scikit-learn API 的兼容性。在构思、架构或设计决策、代码审查或测试策略方面未使用任何 AI 工具。所有 AI 生成的建议均经过作者严格审查、必要时修改并完全验证，作者承担全部责任。

## 致谢

DR 的研究由德国研究基金会（DFG, German Research Foundation）资助——编号 578966082。

## 参考文献

- J. Bradbury, R. Frostig, P. Hawkins, M. J. Johnson, C. Leary, D. Maclaurin, G. Necula, A. Paszke, J. VanderPlas, S. Wanderman-Milne, 和 Q. Zhang (2018)JAX: composable transformations of Python+NumPy programs.外部链接:Link (http://github.com/google/jax)引用自:摘要 (https://arxiv.org/html/2605.14146#Sx1.p1.1).
- A. Cabezas, A. Corenflos, J. Lao, 和 R. Louf (2024)BlackJAX: Composable Bayesian inference in JAX.外部链接:2402.10797引用自:摘要 (https://arxiv.org/html/2605.14146#Sx1.p1.1).
- D. Dua 和 C. Graff (2017)UCI machine learning repository.加州大学欧文分校，信息与计算机科学学院.外部链接:Link (http://archive.ics.uci.edu/ml)引用自:回归基准测试 (https://arxiv.org/html/2605.14146#Sx7.p1.1).
- Fanaee-T, Hadi (2013)Bike Sharing Dataset.备注:UCI Machine Learning Repository引用自:回归基准测试 (https://arxiv.org/html/2605.14146#Sx7.p1.1).
- I. Loshchilov 和 F. Hutter (2019)Decoupled weight decay regularization.在国际学习表征会议中,引用自:摘要 (https://arxiv.org/html/2605.14146#Sx1.p2.1).
- T. Papamarkou, M. Skoularidou, K. Palla, L. Aitchison, J. Arbel, D. Dunson, M. Filippone, V. Fortuin, P. Hennig, J. M. Hernández-Lobato, A. Hubin, A. Immer, T. Karaletsos, M. E. Khan, A. Kristiadi, Y. Li, S. Mandt, C. Nemeth, M. A. Osborne, T. G. J. Rudner, D. Rügamer, Y. W. Teh, M. Welling, A. G. Wilson, 和 R. Zhang (2024)Position: Bayesian Deep Learning is Needed in the Age of Large-Scale AI.在第41届国际机器学习会议论文集,引用自:领域现状 (https://arxiv.org/html/2605.14146#Sx3.p1.1).
- F. Pedregosa, G. Varoquaux, A. Gramfort, V. Michel, B. Thirion, O. Grisel, M. Blondel, P. Prettenhofer, R. Weiss, V. Dubourg, J. Vanderplas, A. Passos, D. Cournapeau, M. Brucher, M. Perrot, 和 E. Duchesnay (2011)Scikit-learn: machine learning in Python.机器学习研究杂志12,页码 2825–2830.引用自:摘要 (https://arxiv.org/html/2605.14146#Sx1.p1.1).
- J. Robnik, G. B. De Luca, E. Silverstein, 和 U. Seljak (2023)Microcanonical hamiltonian monte carlo.机器学习研究杂志24(1),页码 14696–14729.引用自:摘要 (https://arxiv.org/html/2605.14146#Sx1.p2.1).
- J. Robnik 和 U. Seljak (2024)Fluctuation without dissipation: microcanonical langevin monte carlo.在近似贝叶斯推理进展研讨会,页码 111–126.引用自:摘要 (https://arxiv.org/html/2605.14146#Sx1.p2.1).
- E. Sommer, J. Robnik, G. Nozadze, U. Seljak, 和 D. Rügamer (2025)Microcanonical langevin ensembles: advancing the sampling of bayesian neural networks.在第十三届国际学习表征会议,引用自:摘要 (https://arxiv.org/html/2605.14146#Sx1.p1.1),摘要 (https://arxiv.org/html/2605.14146#Sx1.p2.1),需求陈述 (https://arxiv.org/html/2605.14146#Sx4.p1.1).
- E. Sommer, L. Wimmer, T. Papamarkou, L. Bothmann, B. Bischl, 和 D. Rügamer (2024)Connecting the dots: is mode-connectedness the key to feasible sample-based inference in bayesian neural networks?.在第41届国际机器学习会议论文集,引用自:摘要 (https://arxiv.org/html/2605.14146#Sx1.p1.1).