NumPyro中CircularReparam在Predictive使用时的问题分析

问题背景

在使用NumPyro进行贝叶斯建模时，当模型涉及圆形数据（如角度）时，通常会使用VonMises分布作为似然函数。为了提高采样效率，NumPyro提供了CircularReparam重新参数化策略，这是处理圆形变量的推荐方法。

问题现象

当尝试将CircularReparam应用于观测变量（似然函数）并使用Predictive类生成后验预测时，会遇到NotImplementedError错误。这个错误发生在Predictive尝试从重新参数化的分布中采样时。

技术分析

重新参数化的基本原理

重新参数化技巧是概率编程中的重要技术，它通过将随机变量的采样过程分解为确定性变换和基础随机变量的采样，使得梯度计算更加稳定。对于圆形变量，CircularReparam通过将角度变量映射到实数线上进行处理。

问题根源

1. 观测变量的特殊性：观测变量（带有obs参数的sample语句）在模型中扮演特殊角色，它们代表已知数据而非需要采样的随机变量。

2. Predictive的工作机制：Predictive类在生成预测时，会尝试从模型中采样所有未被固定的随机变量。当观测变量被重新参数化后，Predictive试图从基础分布（ImproperUniform）中采样，而该分布未实现采样方法。

3. 设计限制：CircularReparam本意是处理隐变量而非观测变量，对观测变量应用重新参数化会导致不必要的复杂性。

解决方案

正确的做法是仅对模型中的隐变量（非观测变量）应用CircularReparam。对于VonMises似然函数，不需要也不应该进行重新参数化。修改后的模型定义应为：

def model(x, y=None):
    b = numpyro.sample("b", dist.Normal(0, 1))
    sigma = numpyro.sample("sigma", dist.Exponential(1))
    kappa = 1 / sigma**2
    # 直接使用VonMises分布，不进行重新参数化
    numpyro.sample(
        "obs",
        dist.VonMises(loc=b * x, concentration=kappa),
        obs=y,
    )

最佳实践建议

1. 重新参数化的适用对象：仅对模型中的隐变量使用重新参数化策略，特别是那些采样困难的变量。

2. 圆形数据的处理：对于圆形隐变量，优先考虑使用CircularReparam；对于圆形观测数据，直接使用VonMises等圆形分布即可。

3. Predictive的使用：在使用Predictive生成预测时，确保模型中没有对观测变量进行不必要的重新参数化。

4. 错误排查：遇到类似NotImplementedError时，首先检查是否对不应该重新参数化的变量应用了重新参数化策略。

总结

NumPyro中的重新参数化策略是强大的工具，但需要正确使用。理解各种重新参数化策略的适用场景和限制条件，对于构建高效稳定的概率模型至关重要。对于圆形数据，CircularReparam应该仅应用于隐变量，而不应用于观测变量，这样才能保证Predictive等后续分析流程的正常工作。