第3章随机变量与分布

第3章：随机变量与分布：NumPyro基础、分布对象、采样与对数概率、自定义分布

好，咱们进入正题。这一章聊的是概率编程的基石——随机变量和分布。说白了，贝叶斯建模就是跟分布打交道。你定义模型，本质上就是在组合各种分布。NumPyro把这件事做得特别优雅，咱们今天就来拆解一下。

3.1 NumPyro的分布模块：从零开始

NumPyro的分布都在 numpyro.distributions 里。我个人习惯把它简写成 dist，这样代码清爽很多。

import numpyro
import numpyro.distributions as dist
import jax.numpy as jnp

你可能会问：为什么不用Pyro？嗯，Pyro基于PyTorch，而NumPyro基于JAX。JAX有个杀手锏——自动微分和JIT编译。这意味着你的贝叶斯模型可以跑得飞快，还能轻松上GPU。我在项目中遇到过几次，同样的模型，NumPyro比Pyro快了好几倍，尤其是数据量大的时候。

3.2 分布对象：不只是概率

NumPyro里的每个分布都是一个对象。它不只是存个概率值，而是封装了采样、对数概率、参数检查等一系列功能。

举个例子，正态分布：

normal_dist = dist.Normal(loc=0.0, scale=1.0)

这个 normal_dist 就是一个分布对象。你可以对它做三件事：

采样：从分布里抽数据
计算对数概率：给定一个值，算它在这个分布下的对数概率密度
检查参数：看看分布设得合不合理

我刚开始用的时候，总觉得「分布对象」这个概念有点抽象。后来想明白了——它就像一个概率生成器，你给它参数，它就能告诉你所有概率相关的信息。

3.3 采样：从分布中抽取样本

采样用 .sample() 方法。注意，NumPyro的采样需要传入一个随机数生成器（PRNG key）。这是JAX的风格，确保可复现性。

import jax.random as random

key = random.PRNGKey(42)
samples = normal_dist.sample(key, sample_shape=(5,))
print(samples)
# 输出类似：[-0.184,  0.031,  1.023, -0.456,  0.789]

这里 sample_shape=(5,) 表示抽5个样本。你可以抽任意形状，比如 (100, 3) 就是100行3列的矩阵。

小技巧： 如果你需要多次采样，记得每次用不同的key。可以用 random.split(key) 来生成子key。我踩过这个坑——用同一个key采样两次，结果一模一样，还以为模型出bug了。

3.4 对数概率：模型的核心

贝叶斯推断的本质，就是计算后验概率。而计算后验，离不开对数概率。NumPyro用 .log_prob() 方法搞定。

value = jnp.array([0.5, -0.3, 1.2])
log_probs = normal_dist.log_prob(value)
print(log_probs)
# 输出：每个值对应的对数概率密度

为什么会用到对数概率？直接算概率不好吗？嗯，这里有个数值问题。概率值通常很小，比如0.0001，乘起来容易下溢。取对数之后，乘法变加法，数值稳定多了。我在做高维模型时，经常遇到概率值小到浮点数都表示不了的情况，用对数概率就没事。

  核心要点： 在NumPyro里，模型的对数概率是自动累加的。你只需要定义每个变量的分布，框架会自动计算联合对数概率。这就是概率编程的魔力。

3.5 常用分布一览

NumPyro支持几十种分布。我挑几个最常用的列出来：

分布名称	参数	适用场景
Normal	loc, scale	连续数据，误差项
Bernoulli	probs	二分类，0/1数据
Categorical	probs	多分类，离散选择
Poisson	rate	计数数据，事件发生次数
Gamma	concentration, rate	正数数据，方差先验
Beta	concentration1, concentration0	概率的先验，0-1数据

每个分布都有对应的 .sample() 和 .log_prob() 方法。用法完全一致，只是参数不同。

3.6 自定义分布：当内置分布不够用

现实问题里，内置分布不一定够用。比如你遇到一个特殊的概率密度函数，或者需要组合多个分布。这时候就得自己写自定义分布。

NumPyro里自定义分布很简单。你只需要继承 dist.Distribution，实现两个方法：.sample() 和 .log_prob()。

举个例子，假设我们需要一个「截断正态分布」——只在 [a, b] 区间内有值：

class TruncatedNormal(dist.Distribution):
    def __init__(self, loc, scale, low, high):
        self.loc = loc
        self.scale = scale
        self.low = low
        self.high = high
        # 计算归一化常数
        self.norm_const = dist.Normal(loc, scale).cdf(high) - dist.Normal(loc, scale).cdf(low)
        super().__init__(batch_shape=(), event_shape=())

    def sample(self, key, sample_shape=()):
        # 用拒绝采样实现
        samples = dist.Normal(self.loc, self.scale).sample(key, sample_shape)
        # 截断到 [low, high]
        return jnp.clip(samples, self.low, self.high)

    def log_prob(self, value):
        # 计算对数概率，考虑归一化
        raw_log_prob = dist.Normal(self.loc, self.scale).log_prob(value)
        # 超出范围的概率为 -inf
        out_of_bounds = (value < self.low) | (value > self.high)
        raw_log_prob = jnp.where(out_of_bounds, -jnp.inf, raw_log_prob)
        return raw_log_prob - jnp.log(self.norm_const)

注意几个关键点：

batch_shape 和 event_shape 要设对。简单分布一般用空元组。
log_prob 必须返回正确的对数概率，包括归一化。
采样方法可以简单，但效率别太差。拒绝采样在小样本时没问题。

避坑指南： 我曾经在自定义分布里忘了处理边界情况，结果模型跑出来一堆NaN。排查了半天才发现是 log_prob 在边界处返回了错误的值。记住：超出支持域的值，对数概率必须是 -inf，不是0。

3.7 知识体系总览

这一章的内容其实就围绕一个核心——分布对象。我画了张图帮你理清思路：

你看，分布对象就三个核心功能。采样用来生成数据，对数概率用来计算模型似然，参数检查用来确保设置合理。内置分布直接拿来用，不够就自己写，或者组合现有的分布。

3.8 实战小贴士

最后分享几个实战经验：

批量操作：NumPyro的分布支持向量化。你可以传一个向量作为参数，它会自动广播。比如 dist.Normal(jnp.array([0, 1]), 1.0) 就创建了两个正态分布。
可复现性：每次运行前固定PRNG key。我习惯在脚本开头设一个全局key，然后每次采样前用 random.split 生成子key。
调试技巧：如果模型跑出来结果奇怪，先检查 log_prob 是不是返回了合理值。我经常写个小测试，手动算一下对数概率，跟NumPyro的结果对比。

嗯，这一章就到这。分布是概率编程的砖瓦，把砖瓦搞明白了，后面盖房子就顺手了。

📚 课程目录