虽然MCMC方法能够从贝叶斯神经网络权重w的后验分布p(w∣D)中提供渐近精确的样本,但它们在深度学习中常见的高维度和大数据集上往往表现不佳。每次迭代的计算成本以及收敛所需的迭代次数可能变得过高。这就是变分推断(VI)提供有吸引力替代方案的地方。VI不进行采样,而是将贝叶斯推断重新构建为优化问题,在易于处理的族中寻求一个最接近真实后验p(w∣D)的近似分布qϕ(w),这通常通过最小化库尔巴克-莱布勒(KL)散度KL[qϕ(w)∣∣p(w∣D)]来衡量。
正如我们在第三章中看到的那样,最小化此KL散度等同于最大化证据下界(ELBO):
L(ϕ)=Eqϕ(w)[logp(D∣w)]−KL[qϕ(w)∣∣p(w)]
这里,p(D∣w)是给定权重的数据的似然,p(w)是权重的先验分布,而qϕ(w)是由ϕ参数化的变分近似。第一项促使近似后验良好地解释数据,而第二项作为正则项,使近似保持与先验接近。
然而,由于非线性激活函数引入的权重与数据之间复杂的非共轭关系以及参数的庞大数量,将坐标上升变分推断(CAVI)等传统VI方法直接应用于BNN是困难的。
反向传播贝叶斯:使用梯度优化ELBO
Blundell等人(2015)提出了一种将VI应用于BNN的突破性技术,即神经网络中的权重不确定性,通常称为反向传播贝叶斯(BBB)。其核心思想是使用随机梯度上升法优化ELBO,以变分参数ϕ为目标,这与使用反向传播训练标准深度学习模型非常相似。
主要挑战在于计算ELBO的梯度,特别是期望项Eqϕ(w)[logp(D∣w)],因为期望是针对qϕ(w)取的,并且我们要微分的参数ϕ位于此分布内部。计算期望的梯度很棘手。
重参数化技巧
反向传播贝叶斯巧妙地绕过这个问题,使用了重参数化技巧。我们不直接从qϕ(w)中采样权重w,而是引入一个具有固定分布的辅助噪声变量ϵ(例如:标准高斯,ϵ∼N(0,I)),并定义一个确定性变换w=g(ϕ,ϵ),使得得到的w具有qϕ(w)的分布。
对于将qϕ(w)选择为对角高斯分布的常见情况,其中ϕ={μ,ρ}表示均值向量μ以及与标准差向量σ相关的参数ρ(通常σ=log(1+exp(ρ))以保证正数性),重参数化很简单:
w=μ+σ⊙ϵ=μ+log(1+exp(ρ))⊙ϵ,此处 ϵ∼N(0,I)
现在,ELBO中的期望可以根据ϵ的固定分布进行重写:
Eqϕ(w)[logp(D∣w)]=Eϵ∼p(ϵ)[logp(D∣g(ϕ,ϵ))]
现在可以使用链式法则将梯度∇ϕ推入期望内部,因为给定ϕ和ϵ,g(ϕ,ϵ)是确定性的:
∇ϕEqϕ(w)[logp(D∣w)]=Eϵ∼p(ϵ)[∇ϕlogp(D∣g(ϕ,ϵ))]
这个期望通常使用蒙特卡洛采样来近似。对于小批量数据Di,我们采样一个(或几个)ϵ值,计算相应的w=g(ϕ,ϵ),计算对数似然logp(Di∣w),然后通过确定性变换g和网络本身使用标准反向传播计算梯度∇ϕlogp(Di∣w)。
完整的目标函数和梯度
结合似然项和KL散度项,对于单个数据点(x,y)要最大化的目标函数(ELBO)变为:
L(ϕ)≈logp(y∣x,w)−KL[qϕ(w)∣∣p(w)]
其中w=g(ϕ,ϵ)是使用重参数化技巧采样的。如果qϕ(w)和p(w)选择得当(例如,都是高斯分布),KL散度项KL[qϕ(w)∣∣p(w)]通常可以解析计算。
训练中使用的损失函数通常是负ELBO,为整个数据集或小批量数据进行了缩放:
损失=−i=1∑N(Eqϕ(w)[logp(yi∣xi,w)]−N1KL[qϕ(w)∣∣p(w)])
在实践中,对于大小为M的小批量数据,损失近似为:
损失小批量≈−j=1∑Mlogp(yj∣xj,wj)+B1KL[qϕ(w)∣∣p(w)]
这里,wj=g(ϕ,ϵj)是为小批量数据中的每个数据点独立采样的(或者对整个批次使用一个单独的样本w),B是数据集中小批量的数量(适当缩放KL项)。该损失函数对ϕ(均值μ和方差参数ρ)的梯度通过反向传播计算,并用于使用Adam或RMSprop等优化器更新ϕ。
反向传播贝叶斯的训练过程。变分参数ϕ(均值μ和方差参数ρ)定义了近似后验qϕ(w)。在每一步中,通过重参数化采样噪声ϵ以生成网络权重w。前向传播计算小批量数据的预测输出。损失(负ELBO)结合了数据的对数似然和近似后验qϕ(w)与先验p(w)之间的KL散度。损失对ϕ的梯度通过反向传播计算,并使用优化器更新ϕ。
实践考量
- 先验的选择p(w): 标准高斯先验(N(0,σp2I))是常见的选择,当qϕ(w)也是高斯分布时,它能简化KL散度的计算。可以使用更复杂的先验,但可能需要对KL项进行数值估计。
- 初始化: 将均值μ初始化为与标准网络权重类似,并保持初始方差(由ρ控制)较小,通常有助于稳定性。
- 方差参数化: 使用ρ使得σ=log(1+exp(ρ))能够确保σ始终为正。
- 梯度方差: 梯度估计的随机性(由于采样ϵ)可能导致训练过程中出现高方差。例如在每一步中使用更多蒙特卡洛样本或方差减少方法(如局部重参数化技巧,它将重参数化应用于激活而不是权重)的技术有所帮助,尽管这可能会增加计算成本。
- 计算成本: 尽管比MCMC更具可扩展性,但BNN的VI计算成本仍然高于训练标准确定性神经网络。每次前向/反向传播都涉及权重采样和KL项计算,大致使参数数量翻倍(每个权重的μ和ρ),并增加了计算图的复杂性。
VI应用于BNN的优点与局限性
优点:
- 可扩展性: 使用小批量随机梯度优化,使其适用于MCMC不可行的大型数据集和复杂模型。
- 兼容性: 与现有深度学习框架和硬件加速(GPU/TPU)结合得相对较好。
- 单次优化: 在一次优化运行后提供变分参数ϕ的点估计,这与MCMC需要多个链和收敛检查不同。
局限性:
- 近似质量: 后验近似的准确性受到所选变分族qϕ(w)的限制。像平均场高斯这样简单的族可能无法捕捉真实后验中的复杂依赖关系。
- 优化困难: 优化ELBO可能具有挑战性,因为存在噪声梯度和潜在的局部最优。
- 方差低估: VI有时会低估后验分布的方差,可能导致过度自信的预测。
与蒙特卡洛Dropout(在后续章节讨论)相比,反向传播贝叶斯代表了一种更具原则性且更灵活的VI方法,它明确定义并优化权重的近似后验分布的参数。尽管通常比MC Dropout计算密集,但它允许对先验和变分近似有更多控制,如果优化成功且变分族足够,可能带来更准确的不确定性估计。