[论文解读] TMB: Automatic Differentiation and Laplace Approximation

Kasper Kristensen, Anders Nielsen+3 more|RePEc: Research Papers in Economics|Sep 2, 2015

Drilling and Well Engineering被引用 35

一句话总结

TMB 是一个 R 包，通过使用自动微分（最高至三阶）和 Laplace 近似来积分随机效应，实现对复杂非线性混合模型的快速、自动拟合。通过利用 C++ 模板、稀疏矩阵计算以及通过 OpenMP 和 BLAS 实现的并行化，TMB 在大型模型中（尤其是具有大量随机效应的模型）相比 ADMB 实现了 1.5–100 倍的速度提升。

ABSTRACT

TMB is an open source R package that enables quick implementation of complex nonlinear random effect (latent variable) models in a manner similar to the established AD Model Builder package (ADMB, admb-project.org). In addition, it offers easy access to parallel computations. The user defines the joint likelihood for the data and the random effects as a C++ template function, while all the other operations are done in R; e.g., reading in the data. The package evaluates and maximizes the Laplace approximation of the marginal likelihood where the random effects are automatically integrated out. This approximation, and its derivatives, are obtained using automatic differentiation (up to order three) of the joint likelihood. The computations are designed to be fast for problems with many random effects (~10^6) and parameters (~10^3). Computation times using ADMB and TMB are compared on a suite of examples ranging from simple models to large spatial models where the random effects are a Gaussian random field. Speedups ranging from 1.5 to about 100 are obtained with increasing gains for large problems. The package and examples are available at http://tmb-project.org.

研究动机与目标

• 为 R 用户提供一个高性能、用户友好的框架，用于拟合具有随机效应的复杂非线性混合模型。
• 克服有限差分近似在梯度计算中速度慢且不准确的局限性。
• 通过使用自动微分和稀疏矩阵技术，实现对最多包含 10^6 个随机效应和 10^3 个参数的大规模模型的高效计算。
• 通过在模型评估和矩阵运算中支持并行化，进一步加速多核系统上的计算。
• 在保持 C++ 级性能的同时，通过使用 CppAD、Eigen 和 CHOLMOD 等高性能库，实现与 R 的无缝集成。

提出的方法

• 用户使用模板在 C++ 中定义联合对数似然函数，从而实现类型灵活性和编译时优化。
• 通过 CppAD 实现的自动微分可精确计算对数似然函数的一阶、二阶和三阶导数（最高至三阶）。
• 通过积分随机效应，对边际似然应用 Laplace 近似，并自动计算该近似的导数。
• 通过 Eigen 和 CHOLMOD 库处理稀疏矩阵运算，利用状态空间模型和 GMRF 模型中常见的条件独立结构。
• 通过 OpenMP 实现用户定义函数评估的并行化，通过 BLAS 实现大规模稀疏精度矩阵的 Cholesky 分解的并行化。
• 框架与 R 集成以实现数据输入和后续处理，而计算密集型部分则卸载到编译后的 C++ 代码中执行。

实验结果

研究问题

• RQ1 能否在 R 包中高效结合自动微分与 Laplace 近似，以加速复杂非线性混合模型的基于似然的推断？
• RQ2 在随机效应数量不断增加的模型中，TMB 与 ADMB 在速度和准确性方面的表现如何比较？
• RQ3 在潜在变量模型的精度矩阵中，多大程度上可以利用其稀疏性，以实现大规模问题（如 10^6 个随机效应）的可扩展计算？
• RQ4 通过 OpenMP 和 BLAS 实现的并行化能否显著减少大规模模型的计算时间？其性能随核心数量的增加如何变化？
• RQ5 评估 Laplace 近似及其梯度的计算成本与理论上限相比如何？是否满足‘廉价梯度’原则？

主要发现

• TMB 的计算速度比 ADMB 快 1.5 至 100 倍，且在大型模型中性能提升尤为显著。
• 对于包含最多 10^6 个随机效应和 10^3 个参数的模型，TMB 凭借对稀疏性和自动微分的高效利用，保持了高速度和数值稳定性。
• Laplace 近似及其梯度的计算时间不足近似本身时间的 2.8 倍，证实了‘廉价梯度’原则，并保持在理论上限 4 以内。
• 使用 10 个核心相比 1 个核心，使一个大型模型的运行速度提升了 6 倍，证明了通过 OpenMP 实现的并行化效率很高。
• 在所有测试案例中，TMB 和 ADMB 的估计值和标准误几乎完全一致，验证了实现的正确性。
• TMB 使用维护良好的库（如 CppAD、Eigen、BLAS）的模块化设计，确保了长期可维护性，并能快速采纳新的性能改进。