TMB: Automatic Differentiation and Laplace Approximation
Total Page：16
File Type：pdf, Size：1020Kb
TMB: Automatic Differentiation and Laplace Approximation Kasper Kristensen Anders Nielsen Casper W. Berg DTU Compute DTU Aqua DTU Aqua Hans Skaug Brad Bell University of Bergen IHME UW 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) (Fournier, Skaug, Ancheta, Ianelli, Magnusson, Maunder, Nielsen, and Sibert 2011). 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 (≈ 106) and parameters (≈ 103). 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. Keywords: automatic differentiation, AD, random effects, latent variables, C++ templates, R . Submitted to Journal of Statistical Software 1. Introduction arXiv:1509.00660v1 [stat.CO] 2 Sep 2015 Calculation of derivatives plays an important role in computational statistics. One classic application is optimizing an objective function; e.g., maximum likelihood. Given a computer algorithm that computes a function, Automatic Differentiation (AD) (Griewank and Walther 2008) is a technique that computes derivatives of the function. This frees the statistician from the task of writing, testing, and maintaining derivative code. This technique is gradually finding its way into statistical software; e.g., the C++ packages ADMB (Fournier et al. 2011), Stan (Stan Development Team 2013) and Ceres Solver (Agarwal and Mierle 2013). These packages implement AD from first principles, rather than using one of the general purpose AD packages that are available for major programming languages such as Fortran, C++, Python, and MATLAB (Bucker¨ and Hovland 2014). The Template Model Builder (TMB) R package 2 TMB: Template Model Builder uses CppAD (Bell 2005) to evaluate first, second, and possibly third order derivatives of a user function written in C++. Maximization of the likelihood, or a Laplace approximation for the marginal likelihood, is performed using conventional R optimization routines. The numerical linear algebra library Eigen (Guennebaud, Jacob et al. 2010) is used for C++ vector and matrix operations. Seamless integration of CppAD and Eigen is made possible through the use of C++ templates. First order derivatives are usually sufficient for maximum likelihood and for hybrid MCMC; e.g., the Stan package which provides both only uses first order derivatives (Stan Development Team 2013). Higher order derivatives calculated using AD greatly facilitate optimization of the Laplace approximation for the marginal likelihood in complex models with random effects; e.g., Skaug and Fournier(2006). This approach, implemented in ADMB and TMB packages, has been used to fit simple random effect models as well as models containing Gaussian Markov random fields (GMRF). In this paper, we compare computation times between these two packages for a range of random effects models. Many statisticians are unfamiliar with AD and for those we recommend reading sections 2.1 and 2.2 of Fournier et al. (2011). The ADMB package is rapidly gaining new users due to its superiority with respect to optimization speed and robustness (Bolker, Gardner, Maunder, Berg, Brooks, Comita, Crone, Cubaynes, Davies, Valpine et al. 2013) compared to e.g. winBUGS (Spiegelhalter, Thomas, Best, and Lunn 2003). The TMB package is built around the same principles, but rather than being coded more or less from scratch, it combines several existing high-performance libraries, to be specific, CppAD for automatic differentiation in C++, Matrix for sparse and dense matrix calculations in R, Eigen for sparse and dense matrix calculations in C++, and OpenMP for parallelization in C++ and Fortran. Using these packages yields better performance and a simpler code-base making TMB easy to maintain. The conditional independence structure in state-space models and GMRFs yields a sparse precision matrix for the joint distribution of the data and the random effects. It is well known that, when this precision matrix is sparse, it is possible to perform the Laplace approximation for models with a very large number of random effects; e.g., INLA (Rue, Martino, and Chopin 2009). The INLA package (http://www.r-inla.org/download) restricts the models to cases where the sparseness structure is known a priori and models can be written in one line of R code. In contrast ADMB requires manual identification of conditional independent likelihood contributions, but is not restricted to any special model class. The TMB package can fit the same models as ADMB, but is better equipped to take maximal advantage of sparseness structure. It uses an algorithm to automatically determine the sparsity structure. Furthermore, in situations where the likelihood can be factored, it enables parallelization using OpenMP (Dagum and Menon 1998). It also allows parallelization through BLAS (Dongarra, Du Croz, Hammarling, and Duff 1990) during Cholesky factorization of large sparse precision matrices. (Note that the BLAS library is written in Fortran.) C++ templates treat variable types as parameters. This obtains the advantages of loose typing because one code base can work on multiple types. It also obtains the advantage of strong typing because types are checked at compile time and the code is optimized for the particular type. CppAD's use of templates enables derivatives calculations to be recorded and define other functions that can then be differentiated. Eigen's use of templates enables matrix calculations where corresponding scalar types can do automatic differentiation. These features are important in the implementation and use of TMB. Kasper Kristensen, Anders Nielsen, Casper W. Berg, Hans Skaug, Brad Bell 3 The rest of this paper is organized as follows: Section2 is a review of the Laplace approx- imation for random effects models. Section3 is a review of automatic differentiation as it applies to this paper. Section4 describes how TMB is implemented. Section5 describes the package from a user's perspective. Section6 compares its performance with that of ADMB for a range of models where the number of parameters is between 1 and 16 and the number of random effects is between 40 and 40; 000. Section7 contains a discussion and conclusion. 2. The Laplace Approximation The statistical framework in this section closely follows that of Skaug and Fournier(2006). Let f(u; θ) denote the negative joint log-likelihood of the data and the random effects. This n m depends on the unknown random effects u 2 R and parameters θ 2 R . The data, be it continuous or discrete, is not made explicit here because it is a known constant for the analysis in this section. The function f(u; θ) is provided by the TMB user in the form of C++ source code. The range of applications is large, encompassing all random effects models for which the Laplace approximation is appropriate. The TMB package implements maximum likelihood estimation and uncertainty calculations for θ and u. It does this in an efficient manner and with minimal effort on the part of the user. The maximum likelihood estimate for θ maximizes Z L(θ) = exp(−f(u; θ)) du Rn w.r.t. θ. Note that the random effects u have been integrated out and the marginal likelihood L(θ) is the likelihood of the data as a function of just the parameters. We useu ^(θ) to denote the minimizer of f(u; θ) w.r.t. u; i.e., u^(θ) = arg min f(u; θ) : (1) u We use H(θ) to denote the Hessian of f(u; θ) w.r.t. u and evaluated atu ^(θ); i.e., 00 H(θ) = fuu(^u(θ); θ) : (2) The Laplace approximation for the marginal likelihood L(θ) is p ∗ n − 1 L (θ) = 2π det(H(θ)) 2 exp(−f(^u; θ)) : (3) This approximation is widely applicable including models ranging from non-linear mixed effects models to complex space-time models. Certain regularity conditions on the joint negative log-likelihood function are required; e.g., the minimizer of f(u; θ) w.r.t. u is unique. These conditions are not discussed in this paper. Models without random effects (n = 0), and models for which the random effects must be integrated out using classical numerical quadratures, are outside the focus of this paper. Our estimate of θ minimizes the negative log of the Laplace approximation; i.e., p 1 − log L∗(θ) = −n log 2π + log det(H(θ)) + f(^u; θ): (4) 2 This objective and its derivatives are required so that we can apply standard nonlinear op- timization algorithms (e.g., BFGS) to optimize the objective and obtain our estimate for θ. 4 TMB: Template Model Builder Uncertainty of the estimate θ^, or any differentiable function of the estimate φ(θ^), is obtained by the δ-method: −1 ^ 0 ^ 2 ∗ ^ 0 ^ > VAR(φ(θ)) = −φθ(θ) r log L (θ) φθ(θ) : (5) A generalized version of this formula is used to include cases where φ also depends on the ran- dom effects, i.e. φ(u; θ), (Skaug and Fournier 2006; Kass and Steffey 1989). These uncertainty calculations also require derivatives of (4).