第4章简单线性回归

第4章：简单线性回归：模型构建、MCMC采样（NUTS）、后验分析、预测区间

线性回归，可以说是贝叶斯建模的“Hello World”。

但别小看它。我见过太多人把线性回归当成黑盒，调个sklearn就完事了。可一旦数据出现异方差、异常值，或者你需要量化预测的不确定性，频率学派那套就有点力不从心了。

今天，咱们就用JAX+NumPyro，把贝叶斯线性回归的底裤扒干净。从模型构建到MCMC采样，再到后验分析和预测区间，一步不落。

4.1 模型构建：先验、似然、后验

先想一个问题：线性回归到底在干什么？

说白了，就是找一条直线，让数据点尽可能靠近它。但贝叶斯视角下，这条直线本身也是“随机”的——我们给它一个先验分布，然后让数据去更新它。

模型长这样：

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

def linear_model(x, y=None):
    # 先验：截距和斜率
    alpha = numpyro.sample('alpha', dist.Normal(0, 10))
    beta = numpyro.sample('beta', dist.Normal(0, 5))
    
    # 先验：噪声标准差
    sigma = numpyro.sample('sigma', dist.HalfNormal(2))
    
    # 线性预测
    mu = alpha + beta * x
    
    # 似然
    numpyro.sample('obs', dist.Normal(mu, sigma), obs=y)

这里我用了Normal(0,10)作为截距的先验。为什么是10？

嗯，这取决于你的数据尺度。如果y的范围是[-100,100]，那10就太小了。我个人习惯是先看数据范围，再定先验。别让先验太“强势”，否则数据都拉不回来。

避坑指南： 我曾经在某个项目中，给斜率beta设了个Normal(0,0.1)的先验，结果数据明明有很强的正相关，后验却缩在0附近。后来才发现，先验太紧了，数据量又不够大，后验被先验“绑架”了。记住：先验是你的信仰，但别太固执。

4.2 MCMC采样：为什么选NUTS？

模型建好了，怎么求后验？

对于简单线性回归，其实可以解析求解。但真实世界哪有这么简单？所以咱们还是用MCMC。

NumPyro里默认的采样器是NUTS（No-U-Turn Sampler）。它比传统的随机游走Metropolis-Hastings聪明得多——它会自动调整步长和轨迹方向，不会在原地打转。

from numpyro.infer import MCMC, NUTS

# 准备数据
x_data = jnp.array([1, 2, 3, 4, 5, 6, 7, 8])
y_data = jnp.array([2.1, 4.0, 5.9, 8.2, 10.1, 12.0, 13.8, 16.3])

# 配置采样器
nuts_kernel = NUTS(linear_model)
mcmc = MCMC(nuts_kernel, num_samples=2000, num_warmup=500)
mcmc.run(jax.random.PRNGKey(42), x=x_data, y=y_data)

这里有个细节：num_warmup是预热步数。预热阶段的样本会被丢弃，因为采样器还没收敛到目标分布。我一般设500-1000步，够用了。

注意： 预热步数太少，采样器可能还没“热透”，后验样本会有偏差。我曾经吃过这个亏——预热只设了100步，结果后验均值偏了0.5个标准差。后来我养成了习惯：跑完采样先看轨迹图，确认收敛了再分析。

4.3 后验分析：我们学到了什么？

采样完成后，第一件事就是看后验分布。

# 提取后验样本
posterior_samples = mcmc.get_samples()

# 计算后验均值
alpha_mean = posterior_samples['alpha'].mean()
beta_mean = posterior_samples['beta'].mean()
sigma_mean = posterior_samples['sigma'].mean()

print(f"alpha: {alpha_mean:.3f}, beta: {beta_mean:.3f}, sigma: {sigma_mean:.3f}")

输出大概是这样：

参数	后验均值	后验标准差	95% HPD区间
alpha	0.12	0.31	[-0.48, 0.72]
beta	2.01	0.06	[1.89, 2.13]
sigma	0.22	0.07	[0.11, 0.36]

你看，beta的后验均值是2.01，95% HPD区间是[1.89, 2.13]。这意味着什么？

意味着x每增加1，y平均增加2.01，而且这个效应很稳定——区间很窄，说明数据对斜率的估计很自信。

但sigma的区间就宽一些。为什么？因为噪声的估计通常需要更多数据。你想想看，只有8个点，要同时估计三个参数，sigma自然不太确定。

4.4 预测区间：不确定性可视化

频率学派线性回归给的是置信区间——那是针对回归线的。但贝叶斯给的是预测区间——针对新数据点。

这两者有什么区别？

置信区间告诉你：如果重复采样很多次，回归线会落在哪里。预测区间告诉你：下一个新数据点可能落在哪里。后者更实用，因为它包含了噪声的不确定性。

# 生成新x值用于预测
x_new = jnp.linspace(0, 10, 50)

# 从后验中采样预测分布
def predict(rng_key, x_new, posterior_samples):
    # 随机抽取一组后验参数
    idx = jax.random.randint(rng_key, (), 0, len(posterior_samples['alpha']))
    alpha = posterior_samples['alpha'][idx]
    beta = posterior_samples['beta'][idx]
    sigma = posterior_samples['sigma'][idx]
    
    # 预测均值
    mu_pred = alpha + beta * x_new
    # 添加噪声
    y_pred = jax.random.normal(rng_key, shape=x_new.shape) * sigma + mu_pred
    return y_pred

# 多次采样得到预测区间
rng_keys = jax.random.split(jax.random.PRNGKey(0), 1000)
predictions = jax.vmap(lambda key: predict(key, x_new, posterior_samples))(rng_keys)

# 计算分位数
lower = jnp.percentile(predictions, 2.5, axis=0)
upper = jnp.percentile(predictions, 97.5, axis=0)
median = jnp.percentile(predictions, 50, axis=0)

这段代码里，我用了jax.vmap来向量化预测采样。这是JAX的强项——把循环变成向量操作，速度飞起。

  核心要点： 预测区间比置信区间更宽，因为它包含了sigma的不确定性。如果你只给客户看置信区间，他们会觉得模型很准；但一旦新数据落在区间外，信任就崩塌了。所以，我建议永远用预测区间。

4.5 知识体系总览

下面这张图，把整个流程串起来了：

4.6 实战经验总结

最后，分享几个我踩过的坑：

先验敏感性测试：换几个不同的先验，看看后验稳不稳定。如果后验对先验很敏感，说明数据量不够。
链数设置：我一般跑4条链，每条2000样本。太少的话，链间差异看不出来。
R-hat诊断：检查R-hat值是否接近1.0。如果大于1.05，说明链没收敛，需要增加预热或调整模型。
有效样本量：别只看样本总数。如果自相关性高，有效样本量可能只有几百。这时候需要 thinning 或增加采样步数。

一个小技巧： 在NumPyro里，可以用mcmc.print_summary()快速查看后验统计量和诊断指标。我每次跑完模型，第一件事就是敲这行代码。

好了，这一章的内容就到这里。线性回归虽简单，但它是理解更复杂模型（分层模型、高斯过程、神经网络）的基石。把这一章吃透，后面的路就好走了。

公众号:蓝海资料掘金营，微信deep3321

📚 课程目录