[论文解读] Wide Bayesian neural networks have a simple weight posterior: theory and accelerated sampling

Jiri Hron, Roman Novak+2 more|arXiv (Cornell University)|Jun 15, 2022

Gaussian Processes and Bayesian Inference被引用 4

一句话总结

该论文提出了重加权（repriorisation），一种依赖于数据的重参数化方法，可将宽泛贝叶斯神经网络（BNNs）的后验分布进行变换，使其与标准正态先验之间的Kullback–Leibler（KL）散度随宽度增加而趋于零。通过实现马尔可夫链蒙特卡洛（MCMC）采样中更快的混合速度——有效样本量最高提升50倍——该方法显著加速了全连接网络和残差网络中后验推断的效率，尤其在更宽的网络中表现更优。

ABSTRACT

We introduce repriorisation, a data-dependent reparameterisation which transforms a Bayesian neural network (BNN) posterior to a distribution whose KL divergence to the BNN prior vanishes as layer widths grow. The repriorisation map acts directly on parameters, and its analytic simplicity complements the known neural network Gaussian process (NNGP) behaviour of wide BNNs in function space. Exploiting the repriorisation, we develop a Markov chain Monte Carlo (MCMC) posterior sampling algorithm which mixes faster the wider the BNN. This contrasts with the typically poor performance of MCMC in high dimensions. We observe up to 50x higher effective sample size relative to no reparametrisation for both fully-connected and residual networks. Improvements are achieved at all widths, with the margin between reparametrised and standard BNNs growing with layer width.

研究动机与目标

• 为解决贝叶斯神经网络（BNNs）在参数空间中理论理解有限且计算成本高的问题，特别是针对宽网络的情形。
• 通过表征权重空间中宽BNN后验行为，弥补现有函数空间结果（如NNGP极限）的不足，从而弥合对宽BNN理解的空白。
• 开发一种实用且高效的BNN MCMC采样方法，使其在宽度增加时仍能保持良好可扩展性，克服典型高维采样挑战。
• 通过将后验分布转化为随宽度增长而越来越接近标准正态分布的形式，实现BNN中更快、更高效的后验推断。

提出的方法

• 提出重加权（repriorisation），一种重参数化映射 $ T(\theta) $，将BNN权重 $ \theta $ 映射为新参数化形式 $ \theta = T(\theta) $，使得重参数化后后验与标准正态先验之间的KL散度随层宽增加而趋于零。
• 理论上证明（定理2.1）：当最小层宽 $ d_{\min} \to \infty $ 时，KL散度 $ \mathrm{KL}(\mathcal{N}(0,I_d) \| p(\phi|\mathcal{D})) \to 0 $，从而确立后验在参数空间中收敛于先验的性质。
• 提出一种计算高效的实现方式，利用Cholesky分解同时计算重参数化及其雅可比行列式，从而实现可处理的MCMC采样。
• 在重参数化后验上应用拉普拉斯蒙特卡洛（LMC）方法，利用后验对数密度梯度趋于标准正态分布梯度的特性，提升混合速度。
• 采用二阶泰勒展开与Hutchinson估计器，高效近似雅可比行列式的对数行列式，降低内存占用并实现可扩展的梯度计算。
• 利用JAX的 `vjp` 和 `jvp` 函数高效计算梯度，并采用自定义前向传播策略，避免反向传播过程中的内存溢出错误。

实验结果

研究问题

• RQ1 随着网络宽度增加，宽泛贝叶斯神经网络在参数空间中的后验分布行为如何？
• RQ2 能否设计一种重参数化方法，使得重参数化后后验与标准正态先验之间的KL散度在无限宽度极限下趋于零？
• RQ3 此类重参数化是否能显著提升BNN中MCMC采样的混合速度，特别是在高维权重空间中？
• RQ4 该方法能否在不同架构（包括残差网络）和不同数据集规模下实现高效实现与可扩展性？
• RQ5 在严格NNGP范式之外，重加权的优势在实际有限宽度设置下能持续多大程度？

主要发现

• 重加权使重参数化后BNN后验与标准正态先验之间的KL散度在无限宽度极限下趋近于零，为参数空间中后验的简化提供了理论基础。
• 与标准BNN相比，该方法在MCMC采样中可实现最高50倍的有效样本量（ESS）提升，且在所有宽度下均观察到改进，且随层宽增加而增强。
• 在CIFAR-10数据集上，对每层含1024个神经元的3隐层全连接网络，即使远离NNGP范式，仍观察到10倍的ESS提升。
• 在ResNet-20上，仅当顶层宽度与训练样本数量相当或更大时，才实现10倍的ESS提升，表明其性能具有范式依赖性。
• 通过基于Cholesky的雅可比行列式计算与结合Hutchinson估计器和JAX的 `vjp` 的内存优化梯度估计，显著提升了计算效率。
• 该方法在非高斯似然下仍具有效性，如在小规模实验中使用分类似然时，重参数化版本实现了约99%的接受率，且可使用比标准参数化更高的步长。