如何在R的optim()中优化同时返回函数值与梯度的自定义函数？

背景与函数验证 #

我实现了一个用于计算复杂混合模型拉普拉斯近似边际似然的函数，底层采用Eigen/RcppEigen稀疏矩阵与C++自动微分库实现。可通过以下代码安装galamm包的开发版本获取该函数：

remotes::install_github("LCBC-UiO/galamm", 
ref = "7b0b521bb6a9aeb2c93da39dc7a07fef61d81211")

包中的marginal_likelihood函数接受参数值后，返回包含两个元素的列表：

• 模型的deviance（偏差）
• 参数的梯度

尽管该函数为心理测量学中的潜变量模型设计，这里以线性混合模型为例，验证其返回的deviance与lme4::lmer计算结果一致（注：marginal_likelihood是底层函数，不直接面向用户）：

library(lme4)
#> Loading required package: Matrix
library(galamm)

# 使用lme4模块化拟合函数构建线性混合模型
lmod <- lFormula(Reaction ~ Days + (Days|Subject), sleepstudy, REML = FALSE)
devfun <- do.call(mkLmerDevfun, lmod)
opt <- optimizeLmer(devfun)
mod1 <- mkMerMod(environment(devfun), opt, lmod$reTrms, fr = lmod$fr)

# 验证galamm::marginal_likelihood在极大似然估计处返回正确值
mod2 <- marginal_likelihood(
  y = sleepstudy$Reaction,
  trials = 0, # 此处未使用
  X = lmod$X,
  Zt = lmod$reTrms$Zt,
  Lambdat = lmod$reTrms$Lambdat,
  beta = fixef(mod1),
  theta = getME(mod1, "theta"),
  theta_mapping = lmod$reTrms$Lind - 1L,
  lambda = numeric(), # 此处未使用因子载荷
  lambda_mapping_X = integer(),
  lambda_mapping_Zt = integer(),
  family = "gaussian"
)


-2 * logLik(mod1)
#> 'log Lik.' 1751.939 (df=6)
mod2$deviance
#> [1] 1751.939

# marginal_likelihood同时返回theta、beta和lambda的梯度
mod2
#> $deviance
#> [1] 1751.939
#> 
#> $gradient
#> [1] -1.392026e-03 -1.483929e-03 -1.407926e-03  1.327937e-14 -1.724646e-13

Created on 2022-08-18 with reprex v2.0.2

当前优化方案的权宜之计与问题 #

我希望利用marginal_likelihood返回的$deviance和$gradient，通过optim()寻找边际极大似然估计。由于梯度由C++自动微分计算，重复计算会造成资源浪费，因此我实现了以下临时方案，虽能得到正确结果，但存在诸多问题：

library(lme4)
#> Loading required package: Matrix
library(galamm)
lmod <- lFormula(Reaction ~ Days + (Days|Subject), sleepstudy, REML = FALSE)

theta_inds <- seq(from = 1, to = length(lmod$reTrms$theta))
beta_inds <- seq(from = max(theta_inds) + 1, length.out = ncol(lmod$X))
bounds <- c(lmod$reTrms$lower, rep(-Inf, length(beta_inds)))

optfun <- function(par){
  ml <- marginal_likelihood(
    y = sleepstudy$Reaction,
    trials = 0, # 此处未使用
    X = lmod$X,
    Zt = lmod$reTrms$Zt,
    Lambdat = lmod$reTrms$Lambdat,
    beta = par[beta_inds],
    theta = par[theta_inds],
    theta_mapping = lmod$reTrms$Lind - 1L,
    lambda = numeric(), # 此处未使用因子载荷
    lambda_mapping_X = integer(),
    lambda_mapping_Zt = integer(),
    family = "gaussian"
  )
  # 全局更新梯度
  grad <<- ml$gradient
  ml$deviance
}

# 直接从全局环境获取梯度
gradfun <- function(par){
  grad
}

# 初始化梯度为NULL
grad <- NULL
par_init <- c(lmod$reTrms$theta, rep(1, length(beta_inds)))
opt <- optim(par_init, fn = optfun, gr = gradfun, lower = bounds,
             method = "L-BFGS-B")
opt
#> $par
#> [1]   0.92992479   0.01811702   0.22274704 251.40024419  10.46528199
#> 
#> $value
#> [1] 1751.939
#> 
#> $counts
#> function gradient 
#>      116      116 
#> 
#> $convergence
#> [1] 0
#> 
#> $message
#> [1] "CONVERGENCE: REL_REDUCTION_OF_F <= FACTR*EPSMCH"

# 验证估计结果与lme4一致：
mod1 <- lmer(Reaction ~ Days + (Days|Subject), sleepstudy, REML = FALSE)

plot(getME(mod1, "theta"), opt$par[theta_inds]); abline(0, 1)

plot(fixef(mod1), opt$par[beta_inds]); abline(0, 1)

Created on 2022-08-18 with reprex v2.0.2

该方法虽能收敛到正确解，但存在以下问题：

• 使用全局赋值更新梯度，不符合良好编程规范，易引发意外变量冲突
• 依赖L-BFGS-B算法在评估目标函数后立即计算梯度的特性，若在optim()中设置hessian = TRUE，会得到零矩阵——因为迭代结束后计算Hessian时，函数与梯度的调用顺序不再符合预期

提问 #

当待优化函数可同时返回函数值和参数梯度时，有没有更优的方法使用optim()或optimx()、optimr()等工具进行最小化？

内容的提问来源于stack exchange，提问作者Øystein S