Fitting Matérn smoothness parameters using automatic differentiation

Gaussian process modeling is a ubiquitous tool in a variety of disciplines. One attractive feature of a Gaussian process is that it is specified by its first two moments: for a Gaussian random field $Z$ with (often multivariate) index $\bm{x}$, one needs only to select a mean function $\mathbb{E}Z(\bm{x})=:m(\bm{x})$¹¹1For convenience we consider only mean-zero processes, however extensions to parametric mean functions are straightforward. and a positive-definite covariance function $K$ such that $\text{Cov}(Z(\bm{x}),Z(\bm{x}^{\prime}))=:K(\bm{x},\bm{x}^{\prime})$ to completely specify the multivariate law. A classical application of Gaussian process modeling in the mean-zero setting with data $\bm{z}$ is to select a parametric family of positive-definite functions $\left\{K_{\bm{\theta}}\right\}_{\bm{\theta}\in\bm{\Theta}}$ and then optimize the negative log-likelihood over the parametric family to obtain the estimator

It is common practice to then treat this MLE as a known true value and compute subsequent estimators like interpolants or forecasts using the Gaussian law specified by the MLE.

One of the most popular isotropic covariance functions, the Matérn covariance function [21], is given by

where $\Gamma$ is the gamma function, $\mathcal{K}_{\nu}$ is the second-kind modified Bessel function [8], and $\sigma$, $\rho$, and $\nu$ are scale, range, and smoothness parameters respectively. The Matérn class is distinguished from other standard covariance functions by the flexibility that the parameter $\nu$ gives, as it controls the mean-square differentiability of the process [30]. While the above functional form is complicated, the Matérn covariance is most naturally motivated by its Fourier transform pair (also called the spectral density), which for $d$-dimensional processes has the form $S(\bm{f})=\tau^{2}(\zeta^{2}+\left\|\bm{f}\right\|^{2})^{-\nu-d/2}\ ,$ where $\tau$ and $\zeta$ are functions of $(\sigma,\rho,\nu)$ and $d$. As it turns out, $\nu$ has significant effects on the behavior of derived quantities. Interpolants and forecasts, for example, are fundamentally different when $\nu<1$, corresponding to a process that is not mean-square differentiable, and when $\nu\geq 1$, corresponding to a process with at least one mean-square derivative. Special concern about whether processes are (mean-square) differentiable dates back to at least Matheron [22] in the geostatistical community. For a complete discussion on the importance of this parameter and the value in estimating it, see [30].

A widely acknowledged challenge in GP modeling is that the log-likelihood surface created by $\ell(\bm{\theta}\;|\;\bm{z})$ is non-convex and difficult to optimize over. Figure 1 provides an example of the log-likelihood surface (profiled in $\sigma$, see Appendix B) for a Matérn process that illustrates the non-quadratic shape of level surfaces and the difficulty of estimating both smoothness and range parameters. A similar but arguably worse challenge occurs for estimating the range and scale even when the smoothness parameter is known [36].

An effective way to combat the log-likelihood’s lack of convexity is to optimize using derivative information. Until recently the use of derivatives to aid Gaussian process optimization was not particularly common (see [31, 23, 11, 16] for examples of separate strategies for optimization at scale using approximated derivative information). The gradient of $\ell(\bm{\theta}\;|\;\bm{z})$ is computable in closed form and is given by

where we use the shorthand $\bm{\Sigma}_{j}(\bm{\theta}):=\frac{\partial}{\partial\theta_{j}}\bm{\Sigma}(\bm{\theta})$. Setting aside the computational challenges this presents for large matrices, computing this quantity requires derivatives of the covariance function with respect to all parameters in order to assemble the $\bm{\Sigma}_{j}(\bm{\theta})$ matrices. This is the step in the process where attempting to estimate $\nu$ from data becomes problematic. Until recently, no closed form expressions for the derivative $\partial_{\nu}\mathcal{K}_{\nu}$ were available in the literature, and by extension accurate computation of the derivatives $\partial_{\nu}\mathcal{M}_{\nu}$ was not possible. While these derivatives have been studied [6], including very recent work providing completely closed-form derivatives [12], the expressions are cumbersome, for example involving higher-order generalized hypergeometric functions, and as such are difficult to implement reliably. Considering that these functions are significantly more general than $\mathcal{K}_{\nu}$, existing software implementations for them are understandably complex and in our experience come at significant speed costs on at least some parts of the domain, even if they are reliably accurate on the entire domain of $\mathcal{K}_{\nu}$.

The non-convex optimization problem posed by maximum likelihood for Matérn covariance functions may not be practically solvable using only gradient information. In an attempt to introduce higher-order information to further aide in optimization, several recent works have used the expected Fisher information matrix as a proxy for Hessians [11, 16]. True Hessians (also referred to as observed information matrices), however, while coming with an additional computational burden, are more effective than expected information matrices in our experience as well as arguably being preferable as proxies for second-order information of estimated parameters [7]. They can be computed as

where $\bm{\Sigma}_{j,k}$ is the second derivative $\frac{\partial^{2}}{\partial\theta_{j}\partial\theta_{k}}\bm{\Sigma}(\bm{\theta})$. The derivative in the last quadratic form can be expanded using matrix calculus but is left in this form for clarity (see the appendix of [11] for the full expression). As can be seen, the Hessian of $\ell(\bm{\theta})$ requires second derivatives of $\mathcal{K}_{\nu}$, making the second-order optimization for the Matérn class of covariance functions more challenging still. To our knowledge, no existing literature gives explicit forms for these second derivatives. These considerations indicate that a primary bottleneck in performing maximum likelihood estimation for Matérn models is the robust evaluation of the second-kind modified Bessel function, $\mathcal{K}_{\nu}$ and its first two derivatives. While we have not directly verified that the inaccuracies of finite-difference derivatives pose a substantial problem in gradient-based Bayesian methods, such as Hamiltonian Monte Carlo [25], it is possible that the improvements we discuss here could also provide meaningful gains in those methods.

$\mathcal{K}_{\nu}$ is a special function defined as one of the two linearly independent solutions of a second order differential equation [8]. Alternative expressions for $\mathcal{K}_{\nu}$, such as series expansions, integrals on the real line, path integrals, as well as confluent and generalized hypergeometric functions have been used in existing special function libraries [8, 2, 32]. The primary difficulty of practical implementations is that no single strategy for accurate and efficient evaluation is valid on the entire real axis and a domain partitioning is required to obtain an optimal method to evaluate $\mathcal{K}_{\nu}$. Current libraries focus entirely on evaluations of Bessel functions, neglecting their derivatives, and even if they provide access to the source implementation, they are not easily extended to the problem of computing derivatives. Due to this difficulty in computing the derivatives of $\mathcal{K}_{\nu}$ with respect to the order $\nu$, it is common practice to fix $\nu$ prior to performing the optimization instead of treating it as a parameter to be estimated (see, for example, the popular GPyTorch software library [10], or the RobustGaSP softare [15, 14], both of which only offer a few special half-integer values of $\nu$). This often constitutes a significant sacrifice, as $\nu$ can have important practical implications for many tasks performed after maximum likelihood estimation, such as interpolation [30].

In an attempt to offer the functionality of estimating $\nu$ from data, several software libraries for fitting Matérn models provide simple finite difference approximations for the derivative approximation as $\partial_{\nu}\mathcal{K}_{\nu}(x)\approx h^{-1}(\mathcal{K}_{\nu+h}(x)-\mathcal{K}_{\nu}(x))$, where a typical choice is $h=10^{-6}$. For special function libraries that provide only function evaluations this is one of the few possible choices. And for a suitable choice of $h$, these approximations may be reasonably accurate when $x$ is not near the origin and $\nu$ is not overly large, although finite difference approximations are known to be prone to round-off errors. This is particularly problematic for second derivatives, as those calculations involve a divisor of $h^{2}$. Combined with the finite precision representation of the numerator, this division can lead to severe floating point errors. To our knowledge no software library currently attempts to offer second derivative information for $\mathcal{K}_{\nu}$ with respect to $\nu$, or, by extension, for $\mathcal{M}_{\nu}$. Considering that the evaluation of $\mathcal{K}_{\nu}$ is the dominant cost in evaluating the entire Matérn covariance function, finite differences, which require at least twice as many evaluations as direct computations, also come with a burdensome computational cost. A more efficient approach, essentially free of round-off errors for first derivatives, is the complex step method [29, 9]. This was studied for modified Bessel function order derivatives in [20], however not for edge cases such as $\nu\in\mathbb{Z}$, since they are less problematic for complex step method evaluations but are highly significant for the current work. Adaptive finite difference methods, uncommon in the statistical community, are more reliable since they select $h$ adaptively based on function values in the vicinity of differentiation points. But while they provide higher accuracy, it comes at a significant performance cost due to the many function evaluations required. And so while a sufficiently high-order adaptive finite difference method can be used to solve the problem we discuss here, the performance cost is so significant that we do not consider them further.

Arguably the most powerful differentiation approach when derivatives are not computable by hand or easily programmed manually is automatic differentiation (AD), which avoids issues such as round-off errors and allows highly efficient implementations. In this work, we develop a new implementation of $\mathcal{K}_{\nu}$ that is purposefully designed to be automatically differentiable with high accuracy. The negligible accumulation of errors in the first-order derivatives using AD computations further enables highly accurate computations of second-order derivatives of $\mathcal{K}_{\nu}$ with respect to $\nu$. We take this approach here, making use of the Julia [5] programming language, which we consider to be an ideal tool due to its type system and powerful AD tools such as ForwardDiff.jl [27].

In the next section, we provide an overview of AD, focusing on the challenges associated with an automatically differentiable implementation of $\mathcal{K}_{\nu}$. In Section 3, we present our choices for direct evaluations and their derivatives on the entire domain of definition of $\mathcal{K}_{\nu}$. Next, in Section 4, we verify the accuracy and efficiency of the derivatives as well as the Matérn covariance. Finally, we provide a demonstration of second-order optimization in Section 5, where we compare the results of optimizing using derivatives of the log-likelihood built using our AD implementation to optimization results building the same objects with finite difference derivatives. We conclude that optimization using the software we provide here yields correct MLEs, and that the AD-generated derivatives are fast and sufficiently accurate that even very complicated derived functions such as $H\ell(\bm{\theta}\;|\;\bm{z})$ are computed to high accuracy in settings where building the same objects using finite difference derivatives results in a complete loss of precision.

Automatic differentiation (AD) refers to the algorithmic transformation of function implementations $f$ into code that computes the derivative $f^{\prime}$. In the forward-mode paradigm, which is the only one that will be discussed in this work, an implementation of $f$, representing a multivariate vector function $\mathbf{y}=f({\mathbf{x}}),\mathbb{R}^{n}\mapsto\mathbb{R}^{m}$ with $n$ inputs and $m$ outputs is transformed by an AD tool into an implementation of the product ${\dot{{\mathbf{y}}}}=\nabla f({\mathbf{x}})\cdot{\dot{{\mathbf{x}}}}=J_{v}\left({\mathbf{x}},{\dot{{\mathbf{x}}}}\right)$, where $J_{v}$ is the Jacobian vector product of the Jacobian $J({\mathbf{x}})$ and the vector ${\dot{{\mathbf{x}}}}$. The tangent vectors ${\dot{{\mathbf{x}}}}$ and ${\dot{{\mathbf{y}}}}$ are the input directional derivatives and the output sensitivities, respectively. This technique can be employed recursively by applying AD to the Jacobian vector product implementation ${\dot{{\mathbf{y}}}}=J_{v}({\mathbf{x}},{\dot{{\mathbf{x}}}})$ with inputs ${\mathbf{x}}$ and ${\dot{{\mathbf{x}}}}$. We refer to [13] for an extensive introduction to AD.

An AD tool applies differentiation at the programming language intrinsic function level, having hard-coded differentiation rules covering most algebraic operations, some special functions, and optionally higher-level constructs like linear solvers. In Julia, the programming language of choice in the current work, these rules are implemented by ChainRules.jl [34]. The tool then generates code that computes the function value and derivative value at each program statement. As an example, we illustrate as pseudocode a forward-mode AD tool which generates $f^{\prime}(x)$ from the implementation of $y=f(x_{1},x_{2})=e^{(x_{1}\cdot x_{2})}$:

function(x1, x2)
    tmp = x1*x2
    return exp(tmp)
end

function(x1, dx1, x2, dx2)
    tmp = x1*x2
    dtmp = dx1 * x2 + x1 * dx2
    return exp(tmp), exp(tmp)*dtmp
end

To devise a robust implementation for $\mathcal{K}_{\nu}$ and derivatives $\partial_{\nu}^{k}\mathcal{K}_{\nu}$, with $k=1,2$, we partition the domain of definition, $\mathbb{R}_{+}^{2}$ (this is sufficient since $\mathcal{K}_{\nu}(x)=\mathcal{K}_{-\nu}(x)$), into several intervals. This strategy of breaking up the domain and applying different strategies for different regions is ubiquitious in the implementation of all special functions, even ones as basic as trigonometric functions, and certainly has been applied in every commonly used implementation of Bessel functions (see, for example, the domain partitioning of the AMOS library [2]). A particular goal in our domain partitioning and method selection, however, is to obtain accurate and efficient automatic derivatives as well as accurate evaluations. In Figure 2, we summarize the domain partitioning based on intervals for pairs $(x,\nu)$, which splits real axis into different sub-intervals, i.e. $\mathbb{R}=(0,a_{1})\cup[a_{1},a_{2})\cup[a_{2},a_{3})\cup[a_{3},\infty)$, with recommended selections for the parameters $a_{i},\ i=1,2,3$. Even in the most rigorous settings of error analysis, however, there is an element of personal judgment in the exact choices of domain partitioning parameters. While in broad strokes they are guided by when different representations that can be exactly evaluated by a computer are accurate, such as large-argument expansions being picked when the argument is “large enough”, they are also guided by practical metrics like numerical tests. This is the case in this work, and as such the boundary values we recommend here can be altered slightly without major impacts on accuracy.

Each interval is studied in its associated subsection, accompanied by discussions on implementation details for the selected method in that interval (typically a suitable series-type representation) and adjustments necessary to expedite and preserve the accuracy of automatic derivatives. Each separate expansion is truncated at different levels $\{t_{j}\}$, and these parameters have a meaningful impact on both efficiency and accuracy. In the present work we chose the $\{t_{j}\}$ levels that provide close to computer precision. A user willing to sacrifice a few digits of accuracy to increase the efficiency may prefer to truncate the series at lower levels than the suggested values, but a comprehensive study of accuracy and efficiency impacts at such levels may be lengthy and is not considered in this work. Special values of the order $\nu$, such as integers and half-integers must be be considered separately. Since $\mathcal{K}_{\nu}$ for these values of $\nu$ is commonly expressed in a limiting form, substantial work is required to develop an implementation that avoids the limit branch problem. As the rescaled function $x^{\nu}\mathcal{K}_{\nu}(x)$ appears in the Matérn covariance function, we also provide discussion of special optimizations for this function in Section A of the appendix.

This section will focus on the verification and efficiency of our implementation of $\mathcal{K}_{\nu}$ and its derivatives, which we will denote either $\partial_{\nu}^{1}\mathcal{K}_{\nu}$ or $\partial_{\nu}^{2}\mathcal{K}_{\nu}$. We compare our implementation, denoted $\tilde{\mathcal{K}_{\nu}}$, with the one provided by the AMOS library [2], $\mathcal{K}_{\nu}^{A}$, using a reference solution $\mathcal{K}_{\nu}^{R}$, considered to be the ground truth and computed using the arbitrary-precision Arb library [18]. Derivatives of up to second order computed in two ways will be compared, namely AD-generated derivatives of our implementation, $\partial_{\nu}^{k\text{AD}}\tilde{\mathcal{K}_{\nu}}$, and non-adaptive finite difference-based derivatives of the AMOS library, $\partial_{\nu}^{k\text{FD}}\mathcal{K}_{\nu}$. These two values will be compared against a reference derivative computed using tenth-order adaptive finite difference methods of $\mathcal{K}_{\nu}^{A}$, denoted as $\partial_{\nu}^{k\text{R}}{\color[rgb]{0,0,0}\definecolor[named]{pgfstrokecolor}{rgb}{0,0,0}\pgfsys@color@gray@stroke{0}\pgfsys@color@gray@fill{0}\mathcal{K}_{\nu}}$, computed using FiniteDifferences.jl [35].

To offer access to a wider audience, we have packaged the code in the Julia language as the library BesselK.jl²²2https://github.com/cgeoga/BesselK.jl. The source repository additionally contains example R language bindings as well as the scripts used to generate the results and figures in this work.