PaperlayerView exact revision on arXiv

Source-linked AI summary

TMB: Automatic Differentiation and Laplace Approximation

Kasper Kristensen, Anders Nielsen, Casper W. Berg, Hans Skaug, Brad Bell

arXiv:1509.00660v1stat.CO

TL;DR

Complex nonlinear random-effects models require efficient integration and differentiation, especially as the numbers of random effects and parameters grow. TMB addresses this by combining C++ template likelihoods, automatic differentiation, and Laplace approximation within R. Compared with ADMB, it delivers similar estimates and uncertainties while becoming one to two orders of magnitude faster for models with random effects.

  • Problem

    Efficient implementation and estimation of complex nonlinear random-effects models require automatic integration of latent variables and higher-order derivatives, with computation remaining practical for large problems.

  • Method

    TMB lets users define the joint negative log-likelihood in C++ and automatically computes derivatives for Laplace approximation and optimization, using sparse derivative subgraphs when appropriate.

  • Results

    One to two orders of magnitude faster: TMB matches ADMB’s estimates and uncertainties while gaining speed for random-effects models, with gains increasing as models grow.

  • Takeaways & Limitations

    TMB provides R users with a fast, general framework for fitting large nonlinear latent-variable models without manually writing second-order derivative code.

  • Takeaways & Limitations

    TMB’s Laplace-based focus excludes models without random effects and models requiring classical numerical quadrature; one reported BLAS-related comparison also involves additional MKL speedups beyond parallelization.

Abstract

from arXiv · show

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.

1. Introduction

TMB brings automatic differentiation, Laplace approximation, and high-performance numerical libraries into an R package for complex random-effects models. It combines these components to simplify implementation and improve computational performance, particularly for large models.

  • Core approach: TMB uses CppAD to compute first-, second-, and possibly third-order derivatives of user-defined C++ functions, while R performs optimization and Laplace approximation.Eigen provides C++ vector and matrix operations, with integration enabled by C++ templates.
  • Motivation: Higher-order automatic differentiation facilitates optimization of Laplace approximations for complex random-effects models, including models with Gaussian Markov random fields.The paper compares TMB and ADMB computation times across random-effects models.
  • Software design: TMB combines CppAD, Matrix, Eigen, and OpenMP rather than implementing all functionality from scratch.The authors state that this design yields better performance, a simpler code base, and easier maintenance.
  • Sparse models: Sparse precision matrices in state-space models and GMRFs enable Laplace approximation with very large numbers of random effects.The paper contrasts TMB and ADMB with INLA, which requires the sparsity structure to be known a priori.
  • Evaluation scope: The paper evaluates performance across models with 1–16 parameters and 40–40,000 random effects.The performance comparison is organized in Section 6, following reviews of Laplace approximation, automatic differentiation, and TMB implementation.

2. The Laplace Approximation

TMB represents random-effects models through a joint negative log-likelihood and approximates the marginal likelihood by integrating random effects around their conditional mode. The resulting objective and derivatives support standard nonlinear optimization, subject to regularity conditions and the scope of Laplace approximation.

  • Model specification: TMB users provide the joint negative log-likelihood f(u, θ) for data, random effects u, and parameters θ as C++ source code.The framework applies to random-effects models for which the Laplace approximation is appropriate.
  • Laplace approximation: The Laplace approximation integrates out random effects by evaluating f at the conditional minimizer û(θ) and using the Hessian H(θ) with respect to u.The marginal likelihood is therefore represented as a function of θ alone.
  • Assumptions: The approximation is widely applicable, but it requires regularity conditions such as uniqueness of the minimizer of f(u, θ).The paper does not discuss these conditions in detail.
  • Scope: Models without random effects and models requiring classical numerical quadrature are outside the paper’s focus.The stated scope boundary concerns the model classes addressed by this Laplace-based framework.
  • Optimization: TMB minimizes the negative log of the Laplace approximation and supplies its derivatives for standard nonlinear optimization algorithms such as BFGS.These calculations produce the maximum-likelihood estimate for θ.
  • Derivative computation: Derivatives of the Laplace objective are difficult because û depends indirectly on θ and the log determinant requires a Cholesky decomposition.Naively applying automatic differentiation without exploiting sparsity would tape O(n^3) floating-point operations.

3. AD and CppAD

CppAD applies operator-overloading automatic differentiation by recording computations as tapes and traversing the resulting computational graph. TMB uses these mechanisms, including optimized derivative tapes and sparse subgraphs, to obtain higher-order derivatives efficiently.

  • Computational graphs: CppAD records a user algorithm as a tape, or computational graph, whose nodes represent variables and the operations computing them.Independent variables form the initial nodes and the function value is the dependent node.
  • AD modes: Forward mode computes values from independent variables toward the output, whereas reverse mode propagates derivative information in the opposite direction.For scalar functions, reverse mode can evaluate derivatives with fewer than four times the operations needed to evaluate the function.
  • Higher-order derivatives: CppAD uses forward and reverse modes of arbitrary order and can tape functions defined in terms of derivatives, enabling TMB to calculate derivatives up to third order.TMB selects derivative schemes based on their performance in this context.
  • Tape optimization: CppAD optimizes tapes by reducing constant expressions and removing code that does not affect dependent variables.These optimizations improve repeated tape evaluation, as occurs in TMB.

4. Software Implementation

TMB combines a C++ user template with R-controlled optimization and sparse linear algebra to evaluate the Laplace approximation and its derivatives. Its implementation uses automatic differentiation tapes, sparse Hessian calculations, and parallelizable computations for high-dimensional random-effects models.

  • Package architecture: The user template defines f(u, θ), while TMB uses the Laplace approximation to estimate parameters and random effects.The template is a C++ function whose floating-point type supports multiple CppAD types.
  • Package architecture: TMB evaluates the objective and derivatives in R, while C++ computes subexpressions such as û(θ) and H(θ).CHOLMOD performs the Laplace approximation through the Matrix package, with optional BLAS linkage.
  • Automatic differentiation: TMB creates tapes T1–T3 for the user function, its derivative, and the nonzero lower-triangular Hessian entries.The sparsity pattern of H(θ) is calculated from the dependency structure before constructing T3.
  • Laplace-gradient computation: The Laplace-gradient computation preserves the cheap gradient principle under the assumption that inner optimization requires at least two evaluations of f′(ξ).The stated assumption is 2 · work (G5) ≤ work (L1).
  • Automatic sparsity detection: TMB targets very high-dimensional problems when the Hessian is sparse, while dense Hessians can require complexity O(n^2).The implementation discusses further reductions toward complexity proportional to the number of Hessian nonzeros.
  • Automatic sparsity detection: TMB computes Hessian columns by restricting reverse sweeps to dependency subgraphs, with complexity O(n_k log n_k) for a subgraph of size n_k.A breadth-first search followed by sorting identifies the subgraph affecting each gradient component.

5. Using the TMB Package

Using TMB involves defining the negative joint log-likelihood in a C++ user template and controlling model setup, optimization, and post-processing from R. The workflow is illustrated with a state-space population model whose latent states are integrated out as random effects.

  • Workflow: The C++ user template defines the negative joint log-likelihood using macros that receive parameters, random effects, and data from R.The template is compiled and linked before evaluation and optimization.
  • Workflow: The R workflow prepares data and initial values, invokes optimization, and post-processes returned results.Typical calls use MakeADFun, nlminb, and sdreport.
  • Example model: The theta logistic example is a state-space model in which all state values u_0, …, u_{n−1} are random effects integrated out of the likelihood.Its parameter vector contains log(r_0), log(ψ), log(K), log(Q), and log(R).
  • R interface: R data and parameter lists must use names matching the DATA_VECTOR and parameter objects declared in the user template.The example reads y, loads TMB, compiles the template, and passes data through MakeADFun.

6. Case Studies

Across case studies with varied distributions, sparsity structures, and inner optimization problems, TMB generally matched ADMB's results while often running substantially faster. The largest gains appeared in larger problems and through parallelization of BLAS and automatic-differentiation computations.

  • Case-study design: Case studies covered varied distribution families, sparsity structures, and convex or non-convex inner optimization problems.Convex inner problems used a Newton optimizer, whereas non-convex problems generally required more iterations and specially adapted optimizers.
  • Accuracy: TMB and ADMB produced differences small enough to attribute to optimization termination criteria and floating-point roundoff.Both packages were stable with respect to initial values, validating their maximum-likelihood, Laplace-approximation, and uncertainty computations.
  • Speed: In six of nine cases, TMB's speedup over ADMB exceeded 20, while the spatial case achieved 1.5 and longlinreg achieved 0.9.Speedup is defined as ADMB execution time divided by TMB execution time; the longlinreg case had no random effects.
  • Scaling: For larger multivariate random-walk problems, time steps increased from 100 to 1600 and state-vector dimension from 3 to 24 to assess scaling.The authors report that even greater speedups than those in Table 4 should be expected for larger problems.
  • Performance analysis: Profiling separated cases dominated by sparse Cholesky and inverse-subset routines from cases dominated by automatic-differentiation calculations.Intel MKL parallel BLAS produced a factor 10 speedup for ar1 4D, while four cores produced a factor 2 speedup for sdv multi and ten cores produced a sixfold speedup in a larger test.
  • Gradient cost: Across all case studies, evaluating the Laplace approximation and its gradient took less than 2.8 times the function-only evaluation, within the theoretical upper bound of 4.The factor was as low as 2.1 for sam, which had the highest number of parameters.

7. Discussion

TMB provides a fast, flexible R framework for models with and without random effects by combining established high-performance libraries with automatic Laplace approximation and parallel computation. Compared with ADMB, it produces nearly identical estimates and uncertainties while offering much greater speed for larger random-effects models.

  • Discussion: TMB combines R, CppAD, Eigen, BLAS, and CHOLMOD with new code for the Laplace approximation and its derivatives.Users avoid writing the second-order derivative code needed for the Hessian matrix.
  • Performance comparison: One to two orders of magnitude faster performance was observed for models with random effects, with gains increasing as models grew larger.For models without random effects, TMB and ADMB had similar performance on a single core.
  • Performance comparison: TMB produced practically identical estimates and uncertainties to ADMB, providing validation of both independently coded implementations.The comparison covered maximum likelihood, Laplace approximation, and uncertainty computations.
  • Parallel computation: Multiple cores provide further speedups through parallel BLAS computations and OpenMP-enabled user template functions.TMB supports parallel computation both in matrix operations and in the user-specified template function.
  • Additional uses: TMB can also evaluate arbitrary C++ template functions in R and automatically provide their derivatives.This use case supports fast function and derivative evaluation beyond negative log-likelihoods.
  • Comparison with INLA: TMB offers broader model coverage than INLA for the case studies discussed, including non-convex examples that would be difficult to implement in INLA.INLA efficiently targets a general class of Gaussian Markov random-field models but would handle only some of the listed cases.
  • Conclusion: TMB’s modular design allows advances in its component libraries to be adopted quickly while distributing testing and maintenance among independent developers.The conclusion presents this modularity as part of TMB’s framework for estimation in complex statistical models.

8. Supplementary Material

The supplementary material illustrates how TMB parallelizes automatic differentiation by splitting tape accumulation across threads and then processing later derivative calculations independently.

  • 8. Supplementary Material: TMB splits the computational tape so separate threads accumulate even and odd terms before combining their results.The sum of the two thread-specific nodes matches the corresponding result in the original tape.
  • 8. Supplementary Material: After accumulation, threads independently perform sparsity detection and Hessian calculations.The figure describes these operations as part of TMB’s automatic parallelization workflow.
Loading 1509.00660v1…