贝叶斯回归并行化是否适用？R环境模型提速方案问询

问题背景 #

在R环境中运行包含1000条链的贝叶斯回归时速度过慢，尝试用并行化优化性能。原本认为马尔可夫链属于可独立执行的任务，适合并行化，但实际使用后速度反而比不使用并行化更差。现寻求除并行化外的其他模型提速方法。

环境与可复现代码 #

依赖库与数据集 #

library(rstan)
library(rstanarm)
library(tidyverse)
library(stringr)
library(openxlsx)
library(tictoc)
library(parallel)

tasas_ori <- data.frame(
  gedad = c(rep(c("A", "B", "C", "D", "E", "F"), 6)), 
  employment = c(rep(c(20:25), 6),
                 rep(c(20:25) + 1, 6),
                 rep(c(20:25) + 2, 6),
                 rep(c(20:25) + 3, 6),
                 rep(c(20:25) + 4, 6),
                 rep(c(20:25) + 5, 6)),
  anio = c(rep(2006, 6), rep(2012, 6), rep(2014, 6),
           rep(2018, 6), rep("2018 - aj.", 6), rep(2023, 6))
)


tasas <- tasas_ori %>% 
  filter(!(anio %in% c("2018-aj.", "2018"))) %>% 
  mutate(anio = as.numeric(anio) - 2006)

tasas2 <- tasas_ori %>% 
  filter(!(anio %in% c("2018-aj."))) %>% 
  mutate(anio = as.numeric(anio) - 2006)

并行实现代码 #

set.seed(1984)

tic()

num_cores <- parallel::detectCores()
cl <- makeCluster(num_cores - 2)
registerDoParallel(cl)


for (i in 1:3) {
  
  # Parallel loop to fit stan_glm models
  bayes_reg <- foreach(i = 1:1, .combine = "c") %dopar% {
    library(rstan)  # Load rstan within the parallel worker
    library(rstanarm)  # Load rstanarm within the parallel worker
    library(dplyr)
    
    stan_glm(as.formula(paste("employment", "~ anio")),
             data = tasas %>% filter(gedad == unique(tasas$gedad)[1]),
             chain = 50,
             # prior = normal(normal_reg$coefficients[2], summary(normal_reg)$coefficients[4]),
             refresh = 0
    )
  }
  
  # Stop the parallel processing
  
  summary(bayes_reg, digits = 5) %>% print()
  
  coefficients_by[i] <- as.data.frame(bayes_reg$coefficients)
  predictions_by[i] <- as.data.frame(predict(bayes_reg, newdata = tasas2 %>%
                                               filter(gedad == unique(tasas$gedad)[i])))
  
}

stopCluster(cl)
toc()

非并行实现代码 #

# 按年龄组回归

coefficients_by <- vector("list")
predictions_by <- vector("list")

bayes <- vector("list")

set.seed(1984)

tic()


for (i in 1:3) {
  
  cat(paste("
", "**** 贝叶斯回归", "-", "年龄组: ", unique(dci_tasas$gedad)[i], "****", "
"))
  
  bayes_reg <- stan_glm(as.formula(paste("employment", "~ anio")), 
                        data  = dci_tasas %>% 
                          filter(gedad == unique(dci_tasas$gedad)[i]), 
                        chain = 50, 
                        # prior_intercept = normal(normal_reg$coefficients[1], summary(normal_reg)$coefficients[3]),
                        # prior = normal(normal_reg$coefficients[2], summary(normal_reg)$coefficients[4]),
                        refresh = 0
  )
  
  
  summary(bayes_reg, digits = 5) %>% print()
  
  coefficients_by[i] <- as.data.frame(bayes_reg$coefficients)
  predictions_by[i] <- as.data.frame(predict(bayes_reg, newdata = dci_tasas2 %>% 
                                               filter(gedad == unique(dci_tasas$gedad)[i])))
  
}
toc()
  

stopCluster(cl)
toc()

问题分析与优化方案 #

为什么你的并行化反而变慢？ #

你的并行实现存在几个关键问题：

• 并行逻辑错误：foreach(i = 1:1)只执行了一次模型拟合，完全没有利用并行计算的优势，反而增加了集群启动、库加载的额外开销。
• 重复加载依赖：每个并行worker都重新加载rstan、rstanarm等库，这会带来大量冗余的初始化时间。
• 错误的并行层级：rstanarm本身支持通过cores参数直接并行运行多条链（这才是链并行的正确方式），不需要手动用foreach搭建集群，手动集群会和Stan内部的并行机制冲突。

正确的链并行写法应该是在stan_glm中指定cores参数：

bayes_reg <- stan_glm(
  employment ~ anio,
  data = your_data,
  chains = 50,
  cores = parallel::detectCores() - 2,  # 直接指定并行核数
  refresh = 0
)

其他提速方法 #

1. 优化Stan模型的采样设置

• 减少不必要的迭代：如果模型收敛速度快，可以降低warmup（预热迭代数），比如总迭代数1000的话，warmup设为300（前提是确认收敛）。
• 调整采样算法：对于简单模型，使用algorithm = "optimizing"（MAP估计）代替MCMC采样，速度会快很多；或者用algorithm = "meanfield"或"fullrank"的变分推断，比MCMC快几个数量级，适合快速探索。
• 关闭不必要的诊断：保持save_warmup = FALSE（默认设置），减少磁盘IO和内存占用。

2. 简化模型结构

• 减少参数数量：先用经典回归筛选显著变量，再构建贝叶斯模型，避免冗余变量增加计算量。
• 使用更简单的先验：优先选择共轭先验或简单无信息先验，复杂的分层/自定义先验会增加计算开销。

3. 预编译模型

rstanarm默认会缓存编译后的模型，但可以手动预编译并保存，避免重复编译：

# 预编译模型（仅编译不采样）
compiled_model <- stan_glm(
  employment ~ anio,
  data = tasas,
  chains = 0
)
# 保存编译后的模型
saveRDS(compiled_model, "compiled_model.rds")
# 后续使用时加载
compiled_model <- readRDS("compiled_model.rds")
# 基于预编译模型采样
bayes_reg <- update(compiled_model, data = your_subset_data, chains = 50)

4. 硬件与环境优化

• 使用多线程BLAS：将R的BLAS后端换成OpenBLAS或MKL，利用多核加速矩阵运算。比如Windows下安装Microsoft R Open，Linux下编译配置OpenBLAS。
• 增加内存：如果模型占用内存大，提升RAM可以减少磁盘交换，加快运算速度。
• GPU加速：对于大规模模型，使用cmdstanr结合CUDA环境实现GPU加速（需额外配置）。

5. 数据优化

• 减少数据量：用子集数据验证模型，或对数据按组别/时间聚合，降低计算规模。
• 转换数据类型：将字符型变量转为因子型，删除冗余列，减少内存占用和计算开销。

6. 改用更高效的工具

• 使用cmdstanr代替rstanarm：cmdstanr是Stan官方R接口，更轻量、性能更好，支持更多优化选项。
• 使用brms：语法类似rstanarm，但在部分场景下优化更充分，支持自动并行和更多模型类型。

内容的提问来源于stack exchange，提问作者Santiago Valdivieso