We propose a scalable kinetic Langevin dynamics algorithm for sampling parameter spaces of big data and AI applications. Our scheme combines a symmetric forward/backward sweep over minibatches with a symmetric discretization of Langevin dynamics. For a particular Langevin splitting method (UBU), we show that the resulting Symmetric Minibatch Splitting-UBU (SMS-UBU) integrator has bias $O(h^2 d^{1/2})$ in dimension $d>0$ with stepsize $h>0$, despite only using one minibatch per iteration, thus providing excellent control of the sampling bias as a function of the stepsize. We apply the algorithm to explore local modes of the posterior distribution of Bayesian neural networks (BNNs) and evaluate the calibration performance of the posterior predictive probabilities for neural networks with convolutional neural network architectures for classification problems on three different datasets (Fashion-MNIST, Celeb-A and chest X-ray). Our results indicate that BNNs sampled with SMS-UBU can offer significantly better calibration performance compared to standard methods of training and stochastic weight averaging.
翻译:我们提出了一种可扩展的动力学朗之万算法,用于对大数据和人工智能应用中的参数空间进行采样。该方案结合了对小批量的对称前向/后向扫描与朗之万动力学的对称离散化。针对特定的朗之万分裂方法(UBU),我们证明了由此产生的对称小批量分割-UBU(SMS-UBU)积分器在维度 $d>0$ 和步长 $h>0$ 下具有 $O(h^2 d^{1/2})$ 的偏差,尽管每次迭代仅使用一个小批量,从而在步长函数方面提供了对采样偏差的极佳控制。我们将该算法应用于探索贝叶斯神经网络(BNNs)后验分布的局部模态,并评估了具有卷积神经网络架构的神经网络在三个不同数据集(Fashion-MNIST、Celeb-A 和胸部 X 射线)上分类问题的后验预测概率的校准性能。我们的结果表明,与标准的训练和随机权重平均方法相比,使用 SMS-UBU 采样的 BNNs 能够提供显著更好的校准性能。