NumPyro中AutoNormal等自动引导对变长子样本的支持问题分析

问题背景

在Pyro的概率编程框架NumPyro中，自动变分推断(AutoVI)提供了一系列自动引导(guide)实现，如AutoNormal、AutoDelta和AutoGuideList等。这些自动引导在实现随机变分推断(SVI)时存在一个限制：当子样本(subsample)大小在不同log_density评估过程中变化时，会抛出异常。

问题现象

当使用变长子样本进行训练时，例如在mini-batch训练场景下，以下自动引导会出现问题：

• AutoNormal
• AutoDelta
• AutoGuideList

而AutoDiagonalNormal等其他自动引导则能正常工作。

技术分析

问题的根源在于这些自动引导使用了_create_plates方法，该方法依赖于原型追踪(prototype trace)来确定子样本大小。具体来说：

1. 这些引导在第一次调用时会创建原型追踪，记录当时的子样本大小
2. 当后续使用不同大小的子样本时，引导期望的子样本大小与实际输入不匹配
3. 导致形状广播(broadcast)失败，抛出"Model and guide shapes disagree"错误

关键代码路径在autoguide.py中，这些引导类通过_create_plates方法创建plate上下文，而plate的subsample_size在初始化后就被固定。

解决方案

对于需要处理变长子样本的场景，正确的使用方式是显式提供create_plates参数：

create_plates = lambda n, x, subsample_size=None: numpyro.plate("n", n, subsample_size=subsample_size)
AutoNormal(..., create_plates=create_plates)

这种方式允许plate根据每次调用时实际的subsample_size动态调整，而不是固定在初始化的值。

最佳实践建议

1. 当模型涉及子采样时，总是显式提供create_plates参数
2. 确保模型和引导中的子采样行为一致
3. 对于简单的变分近似，考虑使用AutoDiagonalNormal等不受此限制的引导
4. 在开发过程中，注意检查子采样维度是否如预期变化

总结

NumPyro的自动引导系统为变分推断提供了强大支持，但在处理变长子样本时需要特别注意。理解不同引导类的行为差异和限制条件，可以帮助开发者避免这类形状不匹配的问题，构建更灵活的变分推断流程。