SVI 第四部分:技巧与窍门¶
本教程之前的三篇 SVI 教程(第一部分、第二部分 和 第三部分)介绍了使用 Pyro 进行变分推断的各种步骤。在这些教程中,我们定义了模型和指导(即变分分布),设置了变分目标(特别是 ELBO),并构建了优化器(pyro.optim)。所有这些机制的作用是将贝叶斯推断转化为一个随机优化问题。
所有这些都非常有用,但为了达到我们的最终目标——学习模型参数、推断近似后验、使用后验预测分布进行预测等等——我们需要成功解决这个优化问题。根据具体问题的细节——例如潜在空间的维度、是否存在离散潜在变量等——这可能容易也可能困难。在本教程中,我们将介绍一些我们认为对 Pyro 用户进行变分推断普遍有用的技巧和窍门。ELBO 不收敛?!遇到 NaNs?!请看下方寻找可能的解决方案!
Pyro 论坛¶
如果阅读本教程后仍然在优化方面遇到困难,请随时在我们的论坛上提问!
1. 从小学习率开始¶
虽然较大的学习率可能适用于某些问题,但通常最好从 \(10^{-3}\) 或 \(10^{-4}\) 这样的小学习率开始
optimizer = pyro.optim.Adam({"lr": 0.001})
这是因为 ELBO 梯度是随机的,并且可能具有高方差,因此较大的学习率会迅速导致模型/指导参数空间进入数值不稳定或不受欢迎的区域。
在使用较小的学习率实现稳定的 ELBO 优化后,可以尝试较大的学习率。这通常是个好主意,因为过小的学习率可能导致优化效果差。特别是,小的学习率可能导致陷入 ELBO 的较差局部最优解。
2. 默认使用 Adam 或 ClippedAdam¶
在进行随机变分推断时,默认使用 Adam 或 ClippedAdam。注意 ClippedAdam
只是 Adam
的便捷扩展,提供了内置的学习率衰减和梯度裁剪支持。
这些优化算法在变分推断中表现良好的基本原因是,当优化问题具有高度随机性时,它们通过每参数动量提供的平滑通常是必不可少的。请注意,在 SVI 中,随机性可能来自潜在变量采样、数据子采样或两者兼有。
除了调整学习率外,在某些情况下可能还需要调整控制 Adam
使用的动量的超参数对 betas
。特别对于高度随机的模型,使用较高的 \(\beta_1\) 值可能更合理
betas = (0.95, 0.999)
而不是
betas = (0.90, 0.999)
3. 考虑使用衰减学习率¶
虽然在优化初期,当你离最优解较远并希望采取较大的梯度步长时,适度大的学习率很有用,但后期通常最好使用较小的学习率,这样就不会在最优解附近过度震荡而无法收敛。一种方法是使用 Pyro 提供的学习率调度器。例如用法请参阅此处的代码片段。另一种便捷方法是使用 ClippedAdam 优化器,它通过 lrd
参数提供内置的学习率衰减支持
num_steps = 1000
initial_lr = 0.001
gamma = 0.1 # final learning rate will be gamma * initial_lr
lrd = gamma ** (1 / num_steps)
optim = pyro.optim.ClippedAdam({'lr': initial_lr, 'lrd': lrd})
4. 确保你的模型和指导分布具有相同的支持域¶
假设你的 model
中有一个具有约束支持域的分布,例如 LogNormal 分布,其支持域在正实数轴上
def model():
pyro.sample("x", dist.LogNormal(0.0, 1.0))
那么你需要确保 guide
中相应的 sample
站点具有相同的支持域
def good_guide():
loc = pyro.param("loc", torch.tensor(0.0))
pyro.sample("x", dist.LogNormal(loc, 1.0))
如果未能做到这一点,并使用例如以下不可接受的指导
def bad_guide():
loc = pyro.param("loc", torch.tensor(0.0))
# Normal may sample x < 0
pyro.sample("x", dist.Normal(loc, 1.0))
你很可能很快会遇到 NaNs。这是因为 LogNormal 分布在满足 x<0
的样本 x
处的 log_prob
是未定义的,而 bad_guide
很可能会产生这样的样本。
5. 约束需要约束的参数¶
类似地,你需要确保用于实例化分布的参数是有效的;否则你很快会遇到 NaNs。例如,Normal 分布的 scale
参数需要是正的。因此,以下 bad_guide
是有问题的
def bad_guide():
scale = pyro.param("scale", torch.tensor(1.0))
pyro.sample("x", dist.Normal(0.0, scale))
而以下 good_guide
正确地使用了约束来确保正性
from pyro.distributions import constraints
def good_guide():
scale = pyro.param("scale", torch.tensor(0.05),
constraint=constraints.positive)
pyro.sample("x", dist.Normal(0.0, scale))
6. 如果构建自定义指导时遇到困难,请使用 AutoGuide¶
为了使模型/指导对能够实现稳定的优化,需要满足一些条件,其中一些我们在上面已经介绍过。有时很难诊断数值不稳定或收敛不良的原因。其中一个原因是,根本问题可能出现在多个不同的地方:模型中、指导中,或者优化算法或超参数的选择中。
有时问题实际上出在你的模型中,即使你认为问题出在指导中。反之,有时问题出在你的指导中,即使你认为问题出在模型或其他地方。出于这些原因,在尝试识别潜在问题时,减少可变因素的数量会有所帮助。一个便捷的方法是将你的自定义指导替换为 pyro.infer.AutoGuide。
例如,如果模型中的所有潜在变量都是连续的,可以尝试使用 pyro.infer.AutoNormal 指导。或者,可以使用 MAP 推断代替完整的变分推断。更多详情请参阅 MLE/MAP 教程。一旦 MAP 推断工作正常,就有充分的理由相信你的模型设置正确(至少在基本数值稳定性方面)。如果你想获得近似后验分布,现在可以接着进行完整的 SVI。事实上,一个自然的操作顺序可能使用以下一系列日益灵活的 autoguides
AutoDelta → AutoNormal → AutoLowRankMultivariateNormal
如果你发现需要更灵活的指导,或者希望更好地控制指导的具体定义方式,此时可以着手构建自定义指导。一种方法是利用 easy guides,它在完全自定义指导的控制度和 autoguide 的自动化之间取得了平衡。
另请注意,autoguides 提供多种初始化策略,在某些情况下可能需要对这些策略进行实验以获得良好的优化性能。控制初始化行为的一种方法是使用 init_loc_fn
。关于 init_loc_fn
的示例用法,包括 easy guide API 的示例用法,请参阅此处。
7. 参数初始化很重要:将指导分布初始化为低方差¶
优化问题中的初始化可能决定是找到一个好的解决方案还是彻底失败。很难提出一套全面的初始化良好实践,因为好的初始化方案通常高度依赖于具体问题。在随机变分推断中,通常最好将指导分布初始化为低方差。这是因为用于优化 ELBO 的 ELBO 梯度是随机的。如果在 ELBO 优化开始时获得的 ELBO 梯度方差很高,你可能会被引导到参数空间的数值不稳定或不受欢迎的区域。一种防范这种潜在危险的方法是密切关注指导中控制方差的参数。例如,我们通常期望这是一个合理初始化的指导
from pyro.distributions import constraints
def good_guide():
scale = pyro.param("scale", torch.tensor(0.05),
constraint=constraints.positive)
pyro.sample("x", dist.Normal(0.0, scale))
而以下高方差指导很可能导致问题
def bad_guide():
scale = pyro.param("scale", torch.tensor(12345.6),
constraint=constraints.positive)
pyro.sample("x", dist.Normal(0.0, scale))
注意,autoguides 的初始方差可以通过 init_scale
参数控制,例如 AutoNormal
的示例请参阅此处。
8. 探索由 num_particles
、mini-batch size 等控制的权衡¶
如果你的 ELBO 表现出大方差,优化可能会很困难。缓解这个问题的一种方法是增加用于计算每个随机 ELBO 估计的粒子数
elbo = pyro.infer.Trace_ELBO(num_particles=10,
vectorize_particles=True)
(注意,要使用 vectorized_particles=True
,你需要确保你的模型和指导正确地向量化;最佳实践请参阅张量形状教程。)这样做可以降低梯度方差,但会增加计算成本。如果你正在进行数据子采样,mini-batch 大小提供了类似的权衡:更大的 mini-batch 大小可以在增加计算成本的情况下降低方差。虽然最佳选择取决于具体问题,但通常以较少的粒子进行更多梯度步长比以更多粒子进行较少梯度步长更划算。一个重要的注意事项是当你运行在 GPU 上时,在这种情况下(至少对于某些模型),增加 num_particles
或 mini-batch 大小的成本可能是次线性的,在这种情况下增加 num_particles
可能更具吸引力。
9. 如果适用,使用 TraceMeanField_ELBO
¶
Pyro 中基本的 ELBO
实现 Trace_ELBO 使用随机样本来估计 KL 散度项。当解析 KL 散度可用时,可以通过使用解析 KL 散度来降低 ELBO 方差。此功能由 TraceMeanField_ELBO 提供。
10. 考虑对 ELBO 进行归一化¶
默认情况下,Pyro 计算一个未归一化的 ELBO,即计算在完整数据集上(作为条件)计算的对数证据的下界。对于大型数据集,这可能是一个量级很大的数字。由于计算机使用有限精度(例如 32 位浮点数)进行算术运算,大数可能对数值稳定性造成问题,因为它们可能导致精度损失、下溢/上溢等。因此,在许多情况下,将 ELBO 归一化到大致一的量级会很有帮助。这也有助于大致了解你的 ELBO 数值有多好。例如,如果我们有维度为 \(D\) 的 \(N\) 个数据点(例如,\(N\) 个维度为 \(D\) 的实值向量),那么我们通常期望一个合理优化过的 ELBO 的量级为 \(N \times D\)。因此,如果我们用 \(N \times D\) 的因子对 ELBO 进行重新归一化,我们期望 ELBO 的量级为一。虽然这只是一个经验法则,但如果我们使用这种归一化并获得像 \(-123.4\) 或 \(1234.5\) 这样的 ELBO 值,那么很可能出了问题:也许我们的模型严重误设定;也许我们的初始化极其糟糕,等等。有关如何通过归一化常数缩放 ELBO 的详细信息,请参阅本教程。
11. 注意尺度¶
数字的尺度很重要。它们至少有两个重要原因:i) 尺度可以决定某个初始化方案的成败;ii) 如前一节所述,尺度可能影响数值精度和稳定性。
具体来说,假设你正在进行线性回归,即你正在学习一个形如 \(Y = W @ X\) 的线性映射。数据通常带有特定的单位。例如,协变量 \(X\) 的某些分量可能以美元为单位(例如房价),而另一些分量可能以密度为单位(例如每平方英里居民数)。也许第一个协变量的典型值为 \(10^5\),而第二个协变量的典型值为 \(10^2\)。遇到跨越多个数量级的数字时,应该始终注意。在许多情况下,对数据进行归一化使其量级约为一是有意义的。例如,你可以将房价以十万美元为单位进行衡量。
这种数据转换对后续建模和推断有诸多好处。例如,如果已经适当归一化了所有协变量,则可以合理地为权重设置一个简单的各向同性先验分布
pyro.sample("W", dist.Normal(torch.zeros(2), torch.ones(2)))
而不是需要为不同的协变量指定不同的先验协方差
prior_scale = torch.tensor([1.0e-5, 1.0e-2])
pyro.sample("W", dist.Normal(torch.zeros(2), prior_scale))
还有其他好处。现在更容易为你的指导初始化适当的参数。由 pyro.infer.AutoGuide 使用的默认初始化也更有可能适用于你的问题。
12. 保持验证启用¶
默认情况下,Pyro 启用验证逻辑,这有助于调试模型和指导。例如,当分布参数无效时,验证逻辑会通知你。除非你有充分的理由不这样做,否则请保持验证逻辑启用。一旦对模型和推断过程满意,可以使用 pyro.enable_validation 禁用验证。
类似地,在 ELBOs
的上下文中,设置
strict_enumeration_warning=True
当你枚举离散潜在变量时是一个好主意。
13. 张量形状错误¶
如果遇到张量形状错误,请确保已仔细阅读相应的教程。
14. 如果可能,枚举离散潜在变量¶
如果你的模型包含离散潜在变量,对其进行精确枚举可能是有意义的,因为这可以显著降低 ELBO 方差。更多讨论请参阅相应的教程。
15. 一些复杂模型可以从 KL 退火中受益¶
ELBO 的特定形式编码了一种权衡:通过期望对数似然项实现模型拟合,通过 KL 散度实现先验正则化。在某些情况下,KL 散度可能成为难以找到良好最优解的障碍。在这些情况下,在优化过程中对 KL 散度项的相关强度进行退火会有所帮助。更多讨论请参阅深度马尔可夫模型教程。
16. 考虑裁剪梯度或防御性地约束参数¶
模型或指导中的某些参数可能控制对数值问题敏感的分布参数。例如,定义 Gamma 分布的 concentration
和 rate
参数可能表现出这种敏感性。在这些情况下,裁剪梯度或防御性地约束参数可能是有意义的。梯度裁剪的示例请参见此代码片段。关于“防御性”参数约束的一个简单示例是考虑 Gamma
分布的 concentration
参数。此参数必须为正:concentration
> 0。如果我们要确保 concentration
远离零,可以使用带有适当约束的 param
语句
from pyro.distributions import constraints
concentration = pyro.param("concentration", torch.tensor(0.5),
constraints.greater_than(0.001))
这些技巧可以帮助确保你的模型和指导远离参数空间的数值危险区域。