31  广义线性模型

Stan 是一个贝叶斯统计建模和计算的概率推理框架,广泛应用于社会、生物、物理、工程、商业等领域,提供统计建模、数据分析和预测能力(Gelman, Lee, 和 Guo 2015),它提供一套成熟的概率编程语言(Carpenter 等 2017)。Stan 还提供自动微分变分推断(ADVI)算法做近似贝叶斯推断获取参数的后验分布,以及拟牛顿法(L-BFGS)优化算法获取参数的惩罚极大似然估计。

经过 10 多年的发展,Stan 已经形成一个相对成熟的生态,在学术界和工业界的影响力都很广泛,下 图 31.1 是 Stan 软件生态的一角。

flowchart TB
  Stan(Stan) --> CmdStan(CmdStan)
  Stan(Stan) --> RStan(RStan)
  RStan --> rstanarm(rstanarm)
  RStan --> brms(brms)
  CmdStan --> CmdStanR(CmdStanR)
  CmdStan --> CmdStanPy(CmdStanPy)
  CmdStan --> MathematicaStan(MathematicaStan)
图 31.1: Stan、CmdStan 和 CmdStanR 的关系

CmdStan 是 Stan 的命令行接口,CmdStanR 是 CmdStan 的 R 语言接口。每次当 Stan 发布新版本时,CmdStan 也会随之发布新版,只需指定新的 CmdStan 安装路径,CmdStanR 就可以使用上,CmdStanR 包与 Stan 是相互独立的更新机制。实际上,CmdStanR 只是负责处理 CmdStan 运行的结果。入门 CmdStanR 后,可以快速转入对 Stan 底层原理的学习,有利于编码符合实际需要的复杂模型,有利于掌握常用的炼丹技巧,提高科研和工作的效率。

rstan 是 Stan 的 R 语言接口,该接口依赖 RcppRcppEigenBHRcppParallelStanHeaders 等 R 包,由于存在众多上游 R 包依赖,RStan 的更新通常滞后于 Stan 的更新,而且滞后很多,不利于及时地使用最新的学术研究成果。 因此,相比于 rstan 包,CmdStanR 更加轻量,可以更快地将 CmdStan 的新功能融入进来,cmdstanr 和 CmdStan 是分离的,方便用户升级和维护。

rstanarmbrmsRStan 的扩展包,各自提供了一套用于表示统计模型的公式语法。它们都支持丰富的统计模型,比如线性模型、广义线性模型、线性混合效应模型、广义线性混合效应模型等。相比于 rstan, 它们使用起来更加方便,因为它内置了大量统计模型的 Stan 实现,即将公式语法翻译成 Stan 编码的模型,然后调用 rstancmdstanr 翻译成 C++,最后编译成动态链接库。除了依赖 rstan 包,rstanarmbrms 还依赖大量其它 R 包,因此,安装、更新都比较麻烦。

顺便一提,类似的用于概率推理和统计分析的框架,还有 Python 社区的 PyMC3TensorFlow Probability

31.1 生成模拟数据

先介绍泊松广义线性模型,包括模拟和计算,并和 Stan 实现的结果比较。

泊松广义线性模型如下:

\[ \begin{aligned} \log(\lambda) &= \beta_0 + \beta_1 x_1 + \beta_2 x_2 \\ Y &\sim \mathrm{Poisson}(u\lambda) \end{aligned} \]

设定参数向量 \(\beta = (\beta_0, \beta_1, \beta_2) = (0.5, 0.3, 0.2)\),观测变量 \(X_1\)\(X_2\) 的均值都为 0,协方差矩阵 \(\Sigma\)

\[ \left[ \begin{matrix} 1.0 & 0.8 \\ 0.8 & 1.0 \end{matrix} \right] \]

模拟观测到的响应变量值和协变量值,添加漂移项

set.seed(2023)
n <- 2500 # 样本量
beta <- c(0.5, 0.3, 0.2)
X <- MASS::mvrnorm(n, mu = rep(0, 2), Sigma = matrix(c(1, 0.8, 0.8, 1), 2))
u <- rep(c(2, 4), each = n / 2)
lambda <- u * exp(cbind(1, X) %*% beta)
y <- rpois(n, lambda = lambda)

31.2 拟合泊松模型

拟合泊松回归模型

fit_poisson_glm <- glm(y ~ X, family = poisson(link = "log"), offset = log(u))
summary(fit_poisson_glm)

Call:
glm(formula = y ~ X, family = poisson(link = "log"), offset = log(u))

Coefficients:
            Estimate Std. Error z value Pr(>|z|)    
(Intercept) 0.488932   0.009427   51.86   <2e-16 ***
X1          0.289984   0.014298   20.28   <2e-16 ***
X2          0.214846   0.014420   14.90   <2e-16 ***
---
Signif. codes:  0 '***' 0.001 '**' 0.01 '*' 0.05 '.' 0.1 ' ' 1

(Dispersion parameter for poisson family taken to be 1)

    Null deviance: 6052.9  on 2499  degrees of freedom
Residual deviance: 2675.5  on 2497  degrees of freedom
AIC: 10773

Number of Fisher Scoring iterations: 4
# 对数似然函数值
log_poisson_lik <- logLik(fit_poisson_glm)
# 计算 AIC AIC(fit_poisson_glm)
-2 * c(log_poisson_lik) + 2 * attr(log_poisson_lik, "df")
[1] 10772.79

下面用 Stan 编码泊松回归模型,模型代码如下:

data {
  int<lower=1> k;
  int<lower=0> n;
  matrix[n, k] X;
  array[n] int<lower=0> y;
  vector[n] log_offset;
}
parameters {
  vector[k] beta;
  real alpha;
}
model {
  target += std_normal_lpdf(beta);
  target += std_normal_lpdf(alpha);
  target += poisson_log_glm_lpmf(y | X, alpha + log_offset, beta);
}
generated quantities {
  vector[n] log_lik; // pointwise log-likelihood for LOO
  vector[n] y_rep;   // replications from posterior predictive dist
  for (i in 1 : n) {
    real y_hat_i = alpha + X[i] * beta + log_offset[i];
    log_lik[i] = poisson_log_lpmf(y[i] | y_hat_i);
    y_rep[i] = poisson_log_rng(y_hat_i);
  }
}

Stan 代码主要分三部分:

  1. 数据部分 data:声明模型的输入数据,数据类型、大小、约束。

  2. 参数部分 parameters:类似数据部分,声明模型的参数,参数类型、大小。

  3. 模型部分 model:指定模型参数的先验分布。

  4. 生成量 generated quantities:拟合模型获得参数估计值后,计算一些统计量。

下面准备数据

nchains <- 4 # 4 条迭代链
# 给每条链设置不同的参数初始值
inits_data <- lapply(1:nchains, function(i) {
  list(
    alpha = runif(1, 0, 1),
    beta = runif(2, 1, 10)
  )
})

# 准备数据
poisson_d <- list(
  n = 2500, # 观测记录的条数
  k = 2, # 协变量个数
  X = X, # N x 2 矩阵
  y = y, # N 向量
  log_offset = log(u)
)

编译模型,抽样获取参数的后验分布

# 加载 cmdstanr 包
library(cmdstanr)
# 编译模型
mod_poisson <- cmdstan_model(
  stan_file = "code/poisson_log_glm.stan",
  compile = TRUE,
  cpp_options = list(stan_threads = TRUE)
)
# 采样拟合模型
fit_poisson_stan <- mod_poisson$sample(
  data = poisson_d, # 观测数据
  init = inits_data, # 迭代初值
  iter_warmup = 1000, # 每条链预处理迭代次数
  iter_sampling = 2000, # 每条链总迭代次数
  chains = nchains, # 马尔科夫链的数目
  parallel_chains = 1, # 指定 CPU 核心数,可以给每条链分配一个
  threads_per_chain = 1, # 每条链设置一个线程
  show_messages = FALSE, # 不显示迭代的中间过程
  refresh = 0, # 不显示采样的进度
  seed = 20222022 # 设置随机数种子,不要使用 set.seed() 函数
)
# 迭代诊断
fit_poisson_stan$diagnostic_summary()
$num_divergent
[1] 0 0 0 0

$num_max_treedepth
[1] 0 0 0 0

$ebfmi
[1] 1.172526 1.114983 1.010774 1.125754
# 输出结果
fit_poisson_stan$summary(c("alpha", "beta", "lp__"))
# A tibble: 4 × 10
  variable      mean    median      sd     mad        q5      q95  rhat ess_bulk
  <chr>        <num>     <num>   <num>   <num>     <num>    <num> <num>    <num>
1 alpha        0.489     0.489 0.00948 0.00953     0.473  5.04e-1  1.00    4448.
2 beta[1]      0.290     0.290 0.0143  0.0144      0.267  3.14e-1  1.00    3144.
3 beta[2]      0.214     0.214 0.0145  0.0147      0.191  2.38e-1  1.00    3133.
4 lp__     -5388.    -5388.    1.20    1.02    -5390.    -5.39e+3  1.00    3232.
# ℹ 1 more variable: ess_tail <num>

31.3 参数后验分布

加载 bayesplot 包,bayesplot 包提供一系列描述数据分布的绘图函数,比如绘制散点图 mcmc_scatter()\(\beta_1\)\(\beta_2\) 的联合分布

library(ggplot2)
library(bayesplot)
mcmc_scatter(fit_poisson_stan$draws(c("beta[1]", "beta[2]"))) +
  theme_classic() +
  labs(x = expression(beta[1]), y = expression(beta[2]))

图 31.2: \(\beta_1\)\(\beta_2\) 的联合分布

如果提取采样的数据,也可使用 ggplot2 包绘图,不局限于 bayesplot 设定的风格。

beta_df <- fit_poisson_stan$draws(c("beta[1]", "beta[2]"), format = "draws_df")
ggplot(data = beta_df, aes(x = `beta[1]`, y = `beta[2]`)) +
  geom_density_2d_filled() +
  facet_wrap(~.chain, ncol = 2) +
  theme_classic() +
  labs(x = expression(beta[1]), y = expression(beta[2]))

图 31.3: \(\beta_1\)\(\beta_2\) 的联合分布

\(\beta_1\)\(\beta_2\) 的热力图

mcmc_hex(fit_poisson_stan$draws(c("beta[1]", "beta[2]"))) +
  theme_classic() +
  labs(x = expression(beta[1]), y = expression(beta[2]))

图 31.4: \(\beta_1\)\(\beta_2\) 的热力图

各个参数的轨迹图

mcmc_trace(fit_poisson_stan$draws(c("beta[1]", "beta[2]")),
  facet_args = list(
    labeller = ggplot2::label_parsed, strip.position = "top", ncol = 1
  )
) +
  theme_classic()

图 31.5: 各个参数的轨迹图

可以将模型参数的后验分布图展示出来

mcmc_dens(fit_poisson_stan$draws(c("beta[1]", "beta[2]")),
  facet_args = list(
    labeller = ggplot2::label_parsed, strip.position = "top", ncol = 1
  )
) +
  theme_classic()

图 31.6: 各个参数的分布图(密度图)

岭线图就是将各个参数的后验分布图放在一起。

mcmc_areas_ridges(x = fit_poisson_stan$draws(), pars = c("beta[1]", "beta[2]")) +
  scale_y_discrete(labels = scales::parse_format()) +
  theme_classic()

图 31.7: 各个参数的分布图(岭线图)

参数的 \(\hat{R}\) 潜在尺度收缩因子

bayesplot::rhat(fit_poisson_stan, pars = "alpha")
   alpha 
1.000897 

后验预测诊断的想法是检查根据拟合模型生成的随机数 \(y^{rep}\) 与真实观测数据 \(y\) 的接近程度。为直观起见,可以用一系列描述数据分布的图来可视化检验。

y 是真实数据,yrep 是根据贝叶斯拟合模型生成的数据。下图是真实数据的密度图和50组生成数据的密度图。

# 抽取 yrep 数据
yrep <- fit_poisson_stan$draws(variables = "y_rep", format = "draws_matrix")
pp_check(y, yrep = yrep[1:50, ], fun = ppc_dens_overlay) +
  theme_classic()

图 31.8: 后验预测诊断图(密度图)

31.4 模型评估指标

loo 包可以计算 WAIC

fit_poisson_waic <- loo::waic(fit_poisson_stan$draws(variables = "log_lik"))
print(fit_poisson_waic)

Computed from 8000 by 2500 log-likelihood matrix

          Estimate   SE
elpd_waic  -5386.4 37.7
p_waic         2.9  0.1
waic       10772.8 75.5

loo 包推荐使用 LOO-CV ,它还提供诊断信息、有效样本量和蒙特卡罗估计。

fit_poisson_loo <- fit_poisson_stan$loo(variables = "log_lik", cores = 2)
print(fit_poisson_loo)

Computed from 8000 by 2500 log-likelihood matrix

         Estimate   SE
elpd_loo  -5386.4 37.7
p_loo         2.9  0.1
looic     10772.8 75.5
------
Monte Carlo SE of elpd_loo is 0.0.

All Pareto k estimates are good (k < 0.5).
See help('pareto-k-diagnostic') for details.

31.5 brms

对于常见的统计模型,brms 包都内置了预编译的 Stan 程序,下面用 brms 包的函数 brm() 拟合带上述漂移项的泊松广义线性模型,参数估计结果和 Base R 函数 glm() 的几乎一致,因编译和抽样的过程比较花费时间,速度不及 Base R。

# brms
dat <- data.frame(y = y, X = X, u = u)
colnames(dat) <- c("y", "x1", "x2", "u")
fit_poisson_brm <- brms::brm(y ~ x1 + x2 + offset(log(u)),
  data = dat, family = poisson(link = "log")
)
fit_poisson_brm
 Family: poisson 
  Links: mu = log 
Formula: y ~ x1 + x2 + offset(log(u)) 
   Data: dat (Number of observations: 2500) 
  Draws: 4 chains, each with iter = 2000; warmup = 1000; thin = 1;
         total post-warmup draws = 4000

Population-Level Effects: 
          Estimate Est.Error l-95% CI u-95% CI Rhat Bulk_ESS Tail_ESS
Intercept     0.49      0.01     0.47     0.51 1.00     2509     2171
x1            0.29      0.01     0.26     0.32 1.00     1771     1645
x2            0.21      0.01     0.19     0.24 1.00     1727     1847

Draws were sampled using sampling(NUTS). For each parameter, Bulk_ESS
and Tail_ESS are effective sample size measures, and Rhat is the potential
scale reduction factor on split chains (at convergence, Rhat = 1).

计算 LOO-CV

loo::loo(fit_poisson_brm)
Computed from 4000 by 2500 log-likelihood matrix

         Estimate   SE
elpd_loo  -5386.3 37.8
p_loo         2.9  0.1
looic     10772.6 75.5
------
Monte Carlo SE of elpd_loo is 0.0.

All Pareto k estimates are good (k < 0.5).
See help('pareto-k-diagnostic') for details.

looic 指标的作用类似 AIC 指标,所以也几乎相同的。

brms::pp_check(fit_poisson_brm)

31.6 习题

  1. 分析挑战者号航天飞机 O 型环数据。DAAG 包的 orings 数据集记录美国挑战者号航天飞机 O 型环在不同温度下发生 Erosion 腐蚀和 Blowby 串气的失效数量。 图 31.9 展示航天飞机 O 型环在不同温度下失效的分布图(条件密度图):随着温度升高,O 型环越来越不容易失效。请分别用 Base R 函数 glm()cmdstanr 包建模分析 O 型环数据。

    代码
    # data(orings, package = "DAAG")
    orings <- readRDS(file = "data/orings.rds")
    ggplot(orings, aes(x = Temperature, y = after_stat(count))) +
      geom_density(aes(fill = Total > 0), position = "fill", bw = 2) +
      scale_y_continuous(labels = scales::label_percent()) +
      scale_fill_grey(labels = c("TRUE" = "是", "FALSE" = "否")) +
      theme_classic() +
      labs(x = "温度", y = "比例", fill = "失效")

    图 31.9: 航天飞机 O 型环在不同温度下失效的条件密度图
  2. 根据 章节 27 的数据,建立贝叶斯空间广义线性混合模型,用 Stan 预测核辐射强度的分布。