PyMC中var_names参数在pm.sample()中的异常行为分析

在PyMC 5.13版本中，pm.sample()函数新增了var_names参数，用于指定需要存储在跟踪(trace)中的变量名称。根据文档描述，该参数默认会包含所有自由变量和确定性变量。然而，在实际使用中发现该参数存在非预期的行为，可能导致采样结果出现偏差。

问题现象

通过一个具体案例可以清晰地观察到这个问题。考虑一个包含批次效应和温度效应的贝塔回归模型，其中包含以下变量：

• 随机效应：b_batch（批次效应）
• 固定效应：b_temp（温度效应）
• 精度参数：kappa
• 确定性变量：mu（线性预测的逆logit变换）

当不使用var_names参数时，采样结果正常；但当指定var_names=["b_batch", "b_temp", "kappa"]时，采样结果似乎忽略了似然函数，变成了从先验分布中采样。

技术分析

深入研究发现，问题的根源在于BaseTrace类的初始化过程。当指定var_names参数时，PyMC尝试创建一个仅包含指定变量的计算函数，但这一过程存在两个关键问题：

1. 编译函数时使用了on_unused_input="ignore"选项，导致部分输入被忽略
2. 指定的变量名直接用于选择随机变量，而没有正确处理它们对应的值变量

具体来说，model.compile_fn在编译计算图时，如果只传入部分变量而没有正确处理依赖关系，会导致计算图不完整。当启用on_unused_input="ignore"时，系统会忽略缺失的依赖，从而产生随机结果；若禁用该选项，则会直接报错提示输入变量不在计算图中。

解决方案思路

正确的实现方式应该是：

1. 首先获取所有需要的未观测值变量（包括指定变量及其依赖）
2. 将这些变量转换为对应的值变量（类似model.unobserved_value_vars的处理方式）
3. 然后根据var_names参数过滤需要存储在跟踪中的变量

这种处理方式确保了计算图的完整性，同时又能灵活控制最终存储的变量子集。

验证方法

为了确保修复的有效性，建议采用以下验证方案：

1. 使用相同随机种子分别运行带和不带var_names参数的采样
2. 比较两种情况下公共变量的后验分布
3. 确保两者的统计特性一致（均值、方差等）

这种验证方法能够有效确认var_names参数仅影响存储的变量，而不影响实际的采样过程。

总结

var_names参数的设计初衷是提供更灵活的后验采样存储控制，但在实现时需要特别注意计算图的完整性。正确的实现应确保采样过程不受存储变量选择的影响，同时又能精确控制最终保存的变量子集。这一问题提醒我们在修改采样器核心功能时，需要全面考虑变量间的依赖关系和计算图的构建机制。