# JAXFit: Trust Region Method for Nonlinear Least-Squares Curve Fitting on the
GPU

Lucas R. Hofer<EMAIL_ADDRESS>Milan Krstajić  Robert P. Smith
Clarendon Laboratory, University of Oxford, Parks Road, Oxford OX1 3PU, United
Kingdom

###### Abstract

We implement a trust region method on the GPU for nonlinear least squares
curve fitting problems using a new deep learning Python library called JAX.
Our open source package, JAXFit, works for both unconstrained and constrained
curve fitting problems and allows the fit functions to be defined in Python
alone—without any specialized knowledge of either the GPU or CUDA programming.
Since JAXFit runs on the GPU, it is much faster than CPU based libraries and
even other GPU based libraries, despite being very easy to use. Additionally,
due to JAX’s deep learning foundations, the Jacobian in JAXFit’s trust region
algorithm is calculated with automatic differentiation, rather than than using
derivative approximations or requiring the user to define the fit function’s
partial derivatives.

##  I Introduction

Nonlinear least squares regression (NLSQ), colloquially known as curve
fitting, is an essential tool in many fields and is used to determine the set
of parameters which cause a model function to best match observational data.
If the model function represents the data well, the extracted parameters often
allow useful information to be recovered from the observed data. For example,
images of laser beams are often fit with 2D Gaussians to measure the beams’
size [1, 2] and in biomedicine NLSQ regression is used to determine the
location of cells in microscopy images [3].

Unlike linear least squares regression, NLSQ regression cannot be solved
analytically and numerical methods must be used instead. Algorithms for
performing NLSQ regression have existed for centuries with Gauss improving on
Newton’s method to give the Gauss-Newton method [4]. Levenberg [5]—and later
Marquardt [6]—added a damping term to the Gauss-Newton method to form the
Levenberg-Marquardt method. By reformulating the Levenberg-Marquardt algorithm
as a trust region problem [7], Moré [8] developed a robust NLSQ algorithm
which is still used today. Although the aforementioned methods address NLSQ
regression problems without parameter constraints, Coleman and Li [9]
developed a constrained NLSQ regression method, which they further expanded
for large data [10].

The advent of modern computing has greatly enhanced the utility of these NLSQ
regression algorithms and there are several well established software packages
such as MATLAB’s curvefit toolbox [11], MINPACK [12] and SciPy optimize [13]
which have have been robustly tested. These packages rely on the central
processing unit (CPU) for all of the computation and when regressing a
relatively small number of data points (e.g. $<10^{4}$) these algorithms run
quickly; however, regressing a large number of data points, which is quite
common when fitting image data, can drastically increase the fit time. The
advent of graphical processing units (GPUs)—which can parallelize operations
on their many cores—means these NLSQ regression algorithms can be
significantly sped up as most of the computation time is spent on
parallelizable operations such as matrix multiplication.

In the past, programming the GPU required specialized knowledge, usually of
CUDA, which is likely why only a few frameworks implement NLSQ regression
algorithms on the GPU [14, 15]. Even Gpufit [16], which allows its Levenberg-
Marquardt algorithm to be called from MATLAB or Python, requires the user to
define both custom fit functions as well as the function’s partial derivatives
with CUDA before rebuilding the entire project with a C++ compiler.

In this paper we introduce our open source package JAXFit [17], which
implements both unconstrained and constrained NLSQ regression algorithms on
the GPU and solely in Python by using a relatively new deep learning package
called JAX [18]. JAX executes Python defined operations on the GPU and
additionally has in-built automatic differentiation of functions [19], which
removes the need to explicitly define the fit function’s partial derivatives.
JAXFit is easy to use—working as a drop in replacement for SciPy’s curve fit
function—and achieves significant speed-ups due to the GPU.

The rest of the paper is laid out as follows: in Section II we detail the
trust region method algorithm. For readers interested in the implementation
alone, we recommend skipping to Section III in which the speed-up of the trust
region method using JAX is described. Finally, in Section IV we evaluate the
fitting speed of JAXFit in comparison with both SciPy and Gpufit.

##  II Trust Region Method Algorithm

We use SciPy’s trust region method algorithm—which is an amalgamation of both
Moré’s [8, 7] and others’ [20, 9] work—as the basis for our own trust region
method (TRM) implementation. SciPy is the default library for curve fitting in
Python and since, to our knowledge, the SciPy TRM algorithm has been only
partially explained in a blog post [21] by the author, we describe the
unconstrained TRM algorithm below. JAXFit also implements SciPy’s adapted
version of the Coleman-Li [9, 10] method for constrained NLSQ regression, but
we neglect to detail it.

We first define a function $h(y;\mathbf{x})$ which relies on an independent
variable $y$ and $n$ fixed parameters $\left\\{x_{1},x_{2},...,x_{n}\right\\}$
which comprises a parameter vector $\mathbf{x}$. If we have $v$ number of
observations $\left\\{y_{i},z_{i}\right\\}$ which we want to model with $h$ we
first define a residual function $r_{i}$

$r_{i}(x)=h(y_{i},\mathbf{x})-z_{i}\text{.}$ (1)

Now, to find the parameters $\mathbf{x}$ which best approximate the observed
data, the least squares of the residual functions are minimized as follows
xf(x) = 12∑_i^v r_i^2(x)=12r^T(x)r(x). where $\mathbf{r}(\mathbf{x})$ is a
vector containing the $v$ residual functions. For linear functions Eq. 1 can
be analytically solved; however, for nonlinear functions numerical methods
must be used and these generally involve iterative steps in which the
parameter vector $\mathbf{x}$ is successively updated with a parameter update
vector $\mathbf{w}$.

The problem then is to find the parameter update vector $\mathbf{w}$ which
minimizes $f(\mathbf{x}+\mathbf{w})$. To do this, a quadratic model of
$f(\mathbf{x}+\mathbf{w})$ is formed using Taylor’s theorem [20],

$f(\mathbf{x}+\mathbf{w})\approx f(\mathbf{x})+\nabla
f(\mathbf{x})^{T}\mathbf{w}+\frac{1}{2}\mathbf{w}^{T}H(\mathbf{x})\mathbf{w}\text{,}$
(2)

where the expansion has been truncated after the quadratic term. The gradient
$\nabla f(\mathbf{x})$ is given by

$\nabla f(\mathbf{x})=\sum_{i=1}^{v}\nabla
r_{i}(\mathbf{x})r_{i}(\mathbf{x})=J(\mathbf{x})^{T}\mathbf{r}(\mathbf{x})$
(3)

where the Jacobian $J(\mathbf{x})$ is

$J(\mathbf{x})=\begin{bmatrix}\frac{\partial r_{1}}{\partial
x_{1}}&\frac{\partial r_{1}}{\partial x_{2}}&\dots&\frac{\partial
r_{1}}{\partial x_{n}}\\\\[4.30554pt] \frac{\partial r_{2}}{\partial
x_{1}}&\frac{\partial r_{2}}{\partial x_{2}}&\dots&\frac{\partial
r_{2}}{\partial x_{n}}\\\\[4.30554pt]
\vdots&\vdots&\ddots&\vdots\\\\[4.30554pt] \frac{\partial r_{v}}{\partial
x_{1}}&\frac{\partial r_{v}}{\partial x_{2}}&\dots&\frac{\partial
r_{v}}{\partial x_{n}}\end{bmatrix}\text{.}$

The Hessian, $H(\mathbf{x})$, in Eq. 2 is given by

$H(\mathbf{x})=\sum_{i=1}^{v}\nabla r_{i}(\mathbf{x})\nabla
r_{i}(\mathbf{x})^{T}+\sum_{i=1}^{v}r_{i}(\mathbf{x})\nabla^{2}r_{i}(\mathbf{x})\text{;}$
(4)

however, since either the Laplacians or residuals in the second term of Eq. 4
are often quite small, it is common in NLSQ algorithms to ignore it and
approximate the Hessian with only the first term [20]. Additionally, the first
term can be rewritten using the Jacobian such that the approximate Hessian is
$H(\mathbf{x})\approx J^{T}(\mathbf{x})J(\mathbf{x})$. Finally Eq. 2 is
rewritten such that $\mathbf{x}$ is implicit in the various terms

$m(\mathbf{w})=f+\mathbf{g}^{T}\mathbf{w}+\frac{1}{2}\mathbf{w}^{T}B\mathbf{w}\text{,}$
(5)

where $m(\mathbf{w})\approx f(\mathbf{x}+\mathbf{w})$, $f=f(\mathbf{x})$ is
the function, $\mathbf{g}=\nabla f(\mathbf{x})$ is the gradient and
$B=J^{T}(\mathbf{x})J(\mathbf{x})$ is the approximate Hessian.

Figure 1: Visualization of the trust region method algorithm. The parameter
vector $\mathbf{x}$ is iteratively updated to find the minimum of the function
$f(\mathbf{x})$ (solid contour lines, see Eq. 1). (a) A single update step $k$
in which an approximate quadratic model $\hat{m}(\mathbf{p_{k}})$ (dashed
contour lines) of the function is minimized within a trust region of radius
$\Delta$ (dotted line) to determine the parameter update vector
$\mathbf{p_{k}}$ (see Algorithm 1). (b) Multiple update steps—with their
respective trust regions—showing the convergence of the algorithm from the
seed parameter vector $\mathbf{x_{0}}$ to the final parameter vector
$\mathbf{x_{f}}$ at the minimum of $f(\mathbf{x})$. Although these trust
region sizes are all the same, they will often differ in practice.

A parameter update $\mathbf{w}$ is now needed such that the quadratic model in
Eq. 5 is minimized. Since, the model is only good as a local approximation, a
trust region radius $\Delta$ is defined within the parameter space such that
wm(w) = f + g^Tw +12w^TBw ∥Dw∥≤Δ where $D$ is a—usually diagonal—scaling
matrix used to make the trust region an ellipsoid rather than a sphere as the
different parameters $\left\\{x_{1},x_{2},...,x_{n}\right\\}$ comprising
$\mathbf{x}$ can have vastly different scalings. However, according to Moré
[7], this is equivalent to solving the problem in a scaled parameter space
where the trust region is a sphere (see Fig. 1a) and the scaled update vector
is $\mathbf{p}=D\mathbf{w}$. In this case Eq. 1 is equivalent to p^m(p)=f +
^g^Tp +12p^T^Bp ∥p∥≤Δ . where $\mathbf{\hat{g}}=\hat{J}^{T}\mathbf{r}$,
$\hat{B}=\hat{J}^{T}(\mathbf{x})\hat{J}(\mathbf{x})$ and
$\hat{J}(\mathbf{x})=J(\mathbf{x})D^{-1}$ with the $\hat{}$ symbol denoting
the scaled space.

Moré and Sorensen [7] showed that the parameter update $\mathbf{p}$ which is a
solution to Eq. 1 is also the solution to the following

$\displaystyle\left(\hat{B}+\alpha I\right)\mathbf{p}$
$\displaystyle=-\mathbf{\hat{g}}$ (6)
$\displaystyle\alpha\left(\Delta-\left\lVert\mathbf{p}\right\rVert\right)$
$\displaystyle=0$ (7) $\displaystyle\left(\hat{B}+\alpha I\right)$ is positive
semidefinite

where $I$ is the identity matrix and the Levenberg-Marquardt parameter is a
scalar $\alpha\geq 0$. This is the problem we solve.

A basic outline of the TRM algorithm (see Algorithm 1) can now be given. After
being seeded with an initial parameter vector $\mathbf{x_{0}}$ and initial
trust radius $\Delta_{0}$, the parameter vector is iteratively updated for
each step $k$ using Eq. 6-7 until a given accuracy is reached (see Fig. 1b). A
trial update step with $\alpha=0$ is first attempted for which
$\mathbf{p_{t}}=-\hat{B}^{-1}\mathbf{\hat{g}}$ and if the trial update vector
is within the trust region ($\left\lVert\mathbf{p_{t}}\right\rVert\leq\Delta$)
then it is accepted as the parameter update vector. However, if it lies
outside the trust region then the trial parameter update vector must be
discarded as in actuality $\alpha>0$. Furthermore, in this case, the
Levenberg-Marquardt parameter $\alpha$ must then be explicitly solved for as a
sub-problem.

To solve for $\alpha$, we follow Moré’s [8] numerical method in which an
initial guess $\alpha_{0}=0$ is iteratively updated with $q$ steps (see
Algorithm 2). Since, according to Eq. 7, the parameter update vector lies on
the trust region radius ($\left\lVert\mathbf{p}\right\rVert-\Delta=0$) when
$\alpha>0$, we iteratively update $\alpha_{q}$ until $\mathbf{p}$ lies within
a distance $\sigma\Delta$ of the trust radius—where $\sigma$ is the numerical
accuracy. Formally this is $\left|\phi(\alpha)\right|\leq\sigma\Delta$ where
$\phi(\alpha_{q})=\left\lVert\mathbf{p}\right\rVert-\Delta$. Solving Eq. 6 for
$\mathbf{p}$ allows $\phi(\alpha_{q})$ to be rewritten as

$\phi(\alpha_{q})=\left\lVert\left(\hat{B}+\alpha_{q}I\right)^{-1}\mathbf{\hat{g}}\right\rVert-\Delta\text{.}$
(8)

Using Eq. 8 and $\phi^{\prime}\left(\alpha_{q}\right)$, the derivative of
$\phi\left(\alpha_{q}\right)$ with respect to $\alpha_{q}$, the update to the
Levenberg-Marquardt parameter is given by

$\alpha_{q+1}=\alpha_{q}-\left(\frac{\phi\left(\alpha_{q}\right)+\Delta}{\Delta}\right)\left(\frac{\phi\left(\alpha_{q}\right)}{\phi^{\prime}\left(\alpha_{q}\right)}\right)$
(9)

$\mathbf{x_{k}}\leftarrow\mathbf{x_{0}}$;

$\Delta_{k}\leftarrow\Delta_{0}$;

while  _not accuracy condition_ do

$\mathbf{p_{t}}\leftarrow-\hat{B}^{-1}\mathbf{\hat{g}}$;

if  _$\left\lVert\mathbf{p_{t}}\right\rVert\leq\Delta_{k}$_ then

$\mathbf{p_{k}}\leftarrow\mathbf{p_{t}}$;

else

$\alpha_{k}\leftarrow$ LM parameter, Algorithm 2;

$\mathbf{p_{k}}\leftarrow\left(\hat{B}+\alpha_{k}I\right)^{-1}\mathbf{\hat{g}}$;

end if

$\mathbf{x_{k+1}}$, $\Delta_{k+1}\leftarrow$ Update, Algorithm 3
$(\Delta_{k},\mathbf{x_{k}},\mathbf{p_{k}})$;

end while

Algorithm 1 Overall trust region method algorithm. The Levenberg-Marquardt
(LM) parameter and update sub-problems are detailed in Algorithm 2 and
Algorithm 3 respectively.

$\alpha_{q}\leftarrow\alpha_{0}$;

$u_{q}\leftarrow\frac{\left\lVert\mathbf{\hat{g}}\right\rVert}{\Delta}$;

$l_{q}\leftarrow-\frac{\phi\left(\alpha_{0}\right)}{\phi^{\prime}\left(\alpha_{0}\right)}$;

while  _$\left|\phi(\alpha_{q})\right|\leq\sigma\Delta$_ do

$\alpha_{t}\leftarrow\alpha_{q}-\left(\frac{\phi\left(\alpha_{q}\right)+\Delta}{\Delta}\right)\left(\frac{\phi\left(\alpha_{q}\right)}{\phi^{\prime}\left(\alpha_{q}\right)}\right)$;

if  _$l_{q}\leq\alpha_{t}\leq u_{q}$_ then

$\alpha_{q+1}\leftarrow\alpha_{t}$;

else

$\alpha_{q+1}\leftarrow\text{max}\left\\{0.001u_{q},\left(l_{q}u_{q}\right)^{\frac{1}{2}}\right\\}$;

end if

$l_{q+1}\leftarrow\text{max}\left\\{l_{q},\alpha_{q}-\frac{\phi\left(\alpha_{q}\right)}{\phi^{\prime}\left(\alpha_{q}\right)}\right\\}$;

if  _$\phi\left(\alpha_{q}\right) <0$_ then

$u_{q+1}\leftarrow\alpha_{q}$;

else

$u_{q+1}\leftarrow u_{q}$;

end if

end while

Algorithm 2 Solving for the Levenberg-Marquardt parameter $\alpha$ when the
update vector $\mathbf{p}$ lies on the trust region radius.

$\mathbf{w_{k}}\leftarrow D^{-1}\mathbf{p_{k}}$;

if  _$\gamma(\mathbf{x_{k}},\mathbf{w_{k}}) >0.75$_ then

$\mathbf{x_{k+1}}\leftarrow\mathbf{x_{k}}+\mathbf{w_{k}}$;

$\Delta_{k+1}\leftarrow 2\Delta_{k}$;

else if  _$0.25\leq\gamma(\mathbf{x_{k}},\mathbf{w_{k}})\leq 0.75$_ then

$\mathbf{x_{k+1}}\leftarrow\mathbf{x_{k}}+\mathbf{w_{k}}$ ;

$\Delta_{k+1}\leftarrow\Delta_{k}$;

else

$\mathbf{x_{k+1}}\leftarrow\mathbf{x_{k}}$;

$\Delta_{k+1}\leftarrow 0.25\left\lVert\mathbf{p_{k}}\right\rVert$;

end if

Algorithm 3 Determining the updated parameter vector $\mathbf{x_{k+1}}$ and
trust region radius $\Delta_{k+1}$.

subect to the condition that $l_{q}\leq\alpha_{q+1}\leq u_{q}$, where $l_{q}$
and $u_{q}$ are the lower and upper bounds [8] respectively (see Algorithm 2).
If the update lies outside these bounds, it is instead given by

$\alpha_{q+1}=\text{max}\left\\{0.001u_{q},\left(l_{q}u_{q}\right)^{\frac{1}{2}}\right\\}\text{.}$
(10)

When given a reasonable initial value for $\alpha_{0}$ this method for
determining the Levenberg-Marquardt parameter generally converges to a
solution within a few iterations. This solution, in conjunction with Eq. 6, is
then used to calculate the scaled update vector $\mathbf{p}$ and thereafter
the unscaled parameter update vector $\mathbf{w}=D^{-1}\mathbf{p}$.

Once the parameter update vector has been found, it is important to determine
how much it improves the accuracy and whether $\mathbf{w}$ should be accepted
as an update to the parameter vector $\mathbf{x}$. To do this, the following
equation is used [20]

$\gamma(\mathbf{x},\mathbf{w})=\frac{f(\mathbf{x})-f(\mathbf{x}+\mathbf{w})}{m(0)-m(\mathbf{w})}$
(11)

which gives the ratio of the actual reduction over the predicted reduction. If
$\gamma(\mathbf{x},\mathbf{w})$ is larger than a set threshold value, then the
update $\mathbf{w}$ is accepted and the trust region radius is either kept the
same or increased (see Algorithm 3). However, if
$\gamma(\mathbf{x},\mathbf{w})$ is lower than the threshold value then the
parameter vector is kept the same and the size of the trust region radius is
decreased.

##  III Implementing the Trust Region Algorithm with JAX

Our goal when designing JAXFit was to have a fast, yet easy to use NLSQ
regression algorithm. We therefore decided to use Python as it is a popular
data analysis language featuring several libraries which implement array
operations on the GPU. Python’s popularity as a data analysis language is due
in large part to open source packages such as NumPy [22], SciPy [13] and
Pandas [23]. NumPy provides efficient routines for array operations on the
CPU, including an extensive set of linear algebra functions, whereas SciPy
complements NumPy with a large set of scientific functions and algorithms.

Although SciPy features NLSQ fitting, its original—and currently
default—implementation simply wraps MINPACK’s [12] Fortran Levenberg-Marquardt
algorithm for unconstrained NLSQ problems. Nikolay Mayorov [21] significantly
improved SciPy’s curve fitting code with a more modern approach by writing a
trust region algorithm for both unconstrained (based on More [8, 7], see
Section II) and constrained NLSQ problems [9, 9] solely in Python.

Unlike NumPy for CPU array operations, a similar unifying Python library for
array operations on the GPU has not been established. We decided to use a
relative newcomer, JAX, for GPU array operations—which since its release by
Google in 2018 has seen rapid adoption by both the machine learning and
scientific computing communities. JAX describes itself as “NumPy on the GPU”
and it implements almost all of NumPy’s functionality and a significant
portion of SciPy’s on the GPU.

In addition to running on the GPU, JAX also speeds up operations through use
of Just in Time (JIT) compiliation in which JIT wrapped Python functions are
compiled to optimized machine code at run time. When coupled with a GPU this
massively speeds-up execution times as JAX converts the JIT wrapped functions
to highly optimized Accelerated Linear Algebra (XLA) code [18] which runs
efficiently on the GPU hardware. Thus, JAX allows code to be written in a
high-level language (Python), but with the performance of a hardware specific
language (XLA).

Although for large array operations JAX is significantly faster than NumPy, it
is often slower for small array operations due to the overhead of transferring
data to the GPU. Therefore, we only move large array operations (array sizes
$\geq v$, where $v$ is the number of data points being fit) to JAX and the
GPU, while leaving the small array operations to NumPy and the CPU. The
operations we run with JAX include calculating the function $f(\mathbf{x})$,
the gradient $\nabla f(\mathbf{x})$, the quadratic model $m(\mathbf{w})$ and
the loss function.

We also to rely on JAX to calculate the Jacobian $J(\mathbf{x})$ and to
perform singular value decomposition (SVD, see Appendix B) of the Jacobian
matrix—which is a robust method for solving the linear least squares problem
in Eq. 6 111Even though SciPy uses a LSMR algorithm [31] for solving the least
squares problem when the Jacobian matrix is both large and sparse, JAX does
not currently support most SciPy sparse matrix operations and we therefore
implement only SVD.. When calculating the Jacobian, we use JAX’s in-built
automatic differentiation [25] rather than either requiring the user to
provide partial derivatives (as in Gpufit) or using approximations to the
partial derivatives (the SciPy default).

Even though JAX is both fast and easy to use, its dynamism comes at a price.
When creating optimized machine code, JAX traces out a computational graph of
the function using abstract tracer arrays for the function inputs—which can be
computationally expensive. However, the computational overhead pays off when
the code is repeatedly called.

The functional style of programming used in SciPy’s TRM code caused
significant retracing to occur and to avoid this we incorporated the TRM
algorithm into a Python object. The object’s JAX functions are traced only
when first called and successive calls to the object’s curve fit function run
without tracing. Additionally, since changing the size of a JIT function’s
input array also causes retracing to occur, we developed an array zero-padding
method which fixes the size of the JIT functions’ input arrays (see Appendix
A).

##  IV Fit Performance Comparison

We compared the fit speed of JAXFit against both SciPy and Gpufit by creating
a simulated dataset. The dataset is comprised of 2D rotated, elliptical
Gaussians (see Figure 2) which feature seven fitting parameters [26]. We chose
fifteen data lengths (number of residuals $v$) evenly distributed between
$v=10^{4}$ and $v=8\times 10^{6}$ on a logarithmic scale and simulated fifty-
one 2D Gaussian images for each data length. The seven fitting parameters for
each 2D Gaussian image were drawn from a uniform, random distribution.
Gaussian distributed noise was also added to each image with a mean of zero
and a fixed standard deviation for all the images in the dataset.

Figure 2:  Dataset sample images. (a) Image with the smallest data length,
$v=10^{4}$. (b) Image with the largest data length, $v=8\times 10^{6}$.
Although the simulated Gaussian noise in (b) has the same standard deviation
as the noise in (a), it is visually smoothed out due to the large number of
data points.

JAXFit and SciPy use the same TRM algorithm which enables us to compare the
speed of the large array operations—which run on the GPU and CPU for JAXFit
and SciPy respectively. The dataset’s simulated Gaussian images were fit for
each data length in turn using both JAXFit and SciPy. We then averaged the fit
speed for each of the fifteen data lengths (see Fig. 3) to show the operation
speed vs. the data length; however, for each data length the first fit
(tracing) was not included in the time average as tracing can generally be
avoided (see Appendix A).

Figure 3:  Comparison of SciPy vs. JAXFit (JAX) for individual operations. (a)
Speed of the more computationally expensive operations—including function
evaluation $f(\mathbf{x})$, Jacobian calculation $J(\mathbf{x})$ and singular
value decomposition (SVD) of the Jacobian—as a function of the data length.
Additionally, the time it takes to transfer fit data from the CPU to the GPU
is shown. (b) Other operations, including the gradient calculation $\nabla
f(\mathbf{x})$ and evaluation of the quadratic model $m(\mathbf{w})$ as a
function of the data length. The summative CPU to GPU array transfer time for
one iteration step within the trust region method algorithm is also shown.

As expected, the JAXFit operations—–including the function evaluations
$f(\mathbf{x})$, calculation of Jacobians $J(\mathbf{x})$, SVD of the Jacobian
matrix, calculation of the gradient $\nabla f(\mathbf{x})$ and evaluation of
the model function $m(\mathbf{w})$—were significantly faster than SciPy’s.
Although the JAXFit operations generally scaled linearly with the data length,
a few of the operations including SVD and the gradient calculation showed a
flat region for small data lengths. This is likely due to not all the GPU
resources being utilized. However, since the CPU resources are generally fully
utilized, almost all the SciPy operations scale linearly—irrespective of the
data length.

Figure 4: Overall fit speed of JAXFit, SciPy and Gpufit as a function of the
data length. JAXFit (solid-circle) and Gpufit (dotted-triangle) are much
faster than SciPy (dash-square) as they utilize the GPU rather than CPU.
Although Gpufit is slightly faster for a small data lengths ($v<3\times
10^{4}$), JAXFit is significantly faster for large data lengths.

In addition to the computation time, we measured the time required to move the
fit data from the CPU to the GPU (see Fig. 3a), as well as the time to
transfer data back from the GPU to CPU (see Fig. 3b). As expected, the CPU to
GPU time scales linearly with the size of the data being transferred, whereas
the GPU to CPU transfer time is roughly constant—since the size of the arrays
transferred back are solely dependent on the number of fitting parameters $n$.

We also perform an overall fit speed comparison of JAXFit, SciPy and Gpufit
for unconstrained NLSQ regression (see Fig. 4). Using each method we measured
the overall fit time for all the simulated Gaussian images in the dataset. We
again ignored the first image for each data length and averaged the remaining
fifty fit times. As expected, JAXFit and Gpufit are much faster than SciPy due
to the GPU vs. CPU hardware advantage.

When comparing the GPU methods, Gpufit is slightly faster than JAXFit if the
data length is small (${v<3\times 10^{4}}$), as its Levenberg-Marquardt
algorithm runs entirely in CUDA on the GPU. However, as the data length
increases, this advantage quickly disappears and JAXFit significantly
outperforms Gpufit when the data length is large—almost an order of magnitude
faster for $v>10^{6}$. Since the number of iterations it takes the algorithms
in JAXFit and Gpufit to converge is roughly equal when using the same fit
data, JAXFit’s better performance is likely due to the highly optimized XLA
library underpinning JAX, which uses the GPU hardware resources more
efficiently.

Finally, we note that these experiments were all run in a Google Colab
notebook featuring a single 2 GHz Intel Xeon CPU core and a NVIDIA Tesla
V100-SXM2-16GB GPU—which is already several years old and is roughly
equivalent to NVIDIA’s latest series of gaming GPUs (RTX3080/RTX3090).
Although the speed of JAXFit will depend primarily on the GPU specifications,
even when using a lower spec’d GPU, JAXFit is much faster than CPU based NLSQ
regression methods.

##  V Conclusion

We have implemented a trust region method for both unconstrained and
constrained nonlinear least squares (NLSQ) curve fitting on the GPU using a
deep learning Python library, JAX. Our open source package, JAXFit [17],
allows the user to dynamically define fit functions in Python and, unlike
other GPU NLSQ libraries, requires no specialized knowledge of CUDA or the
GPU. Furthermore, the partial derivatives which comprise the Jacobian in the
trust region algorithm are calculated with automatic differentiation—rather
than either requiring the user to define the partial derivatives or using
approximations.

JAXFit significantly outperforms SciPy’s CPU based NSLQ curve fitting function
due to its use of the GPU. Furthermore, JAXFit is easier to use and generally
faster than another GPU fitting library (Gpufit) since it makes more efficient
use of the GPU resources. However, JAXFit’s simplicity comes with the
computational overhead of JAX’s function tracing. Although the initial trace
is unavoidable, JAXFit is generally able to avoid retracing even when the size
of the fit data changes (see Appendix A) and thus remains very fast.

For ease of use, JAXFit is designed as a drop-in replacement for SciPy’s curve
fit function. Additionally, JAXFit incorporates features not previously
described such as handling fit data uncertainty (see Appendix C), scaling of
the Jacobian and implementation of various loss functions—see the Github
repository for all the available features [17].

##  Appendix A Avoiding Retracing of Dynamic Array Sizes

JAX needs to retrace a JIT function whenever the size of the input array
changes. Although this isn’t a problem for neural networks—which it was
primarily constructed for—it can be an issue for curve fitting. For example in
our laboratory, we fit images of cold atom clouds [27] with 2D Gaussians or 2D
parabolas and since the size of the clouds can be very different, we often
need to fit different sized data arrays.

To circumvent the JAX retracing process for these different sized arrays,
JAXFit allows the user to define a fixed input data size $s$. When the fit
data passed to JAXFit is smaller than this fixed size then the fit data is
padded with dummy data to account for the difference between the actual data
length $v$ and the fixed data length $s$.

import jax.numpy as jnpimport jax.jit as jit@jitdef masked_jac(coords, args,
data_mask) J = jac_func(coords, args) J_masked = jnp.where(data_mask, J, 0)
return J_masked

Figure 5: Example of masking dummy data within a JIT wrapped Python function.

This does require the user to pick a sensible value for $s$. If $s\gg v$, then
significant time is wasted processing the dummy data. On the other hand, if
the actual data size is larger than the fixed input data size ($v>s$) then the
JAXFit functions must be retraced. Thus, when all the data being analyzed is
roughly the same size, the optimal fixed size is $s=v_{\text{max}}$ where
$v_{\text{max}}$ is the largest data size being fit. However, if the fit data
sizes differ significantly then it can be more efficient to instantiate
multiple curve fit objects—each with a different fixed $s$ to handle the range
of data sizes.

Lastly, simply padding the input fit data with dummy data would lead to
incorrect fits. Thus, we’ve added masking operations (see Listing 5)
throughout the TRM algorithm so that the dummy data doesn’t affect the fit,
but merely operates as a placeholder to keep JAX from retracing the functions.

##  Appendix B Solving Linear Least Squares Equation with Singular Value
Decomposition

Singular value decomposition [28] factors any matrix $M$ which has size
$v\times n$ into the product of three matrices

$M=U\Sigma V^{T}$ (12)

where $U$ is a $v\times v$ orthonormal matrix, $\Sigma$ is a $v\times n$
diagonal matrix and $V$ is a $n\times n$ orthonormal matrix. Since both $U$
and $V$ are orthonormal $U^{T}U=UU^{T}=I$ (with the same expression for $V$)
where $I$ is the identity matrix.

The special properties of SVD allow the least squares problem (see Eq. 6) in
the trust region method algorithm to be solved. We start by rewriting Eq. 6
explicitly in terms of the scaled Jacobian $J$ as

$\left(J^{T}J+\alpha I\right)\mathbf{p}=-J^{T}\mathbf{r}\text{,}$ (13)

where the scaled space’s hat notation $\hat{}$ has been dropped. Now using
SVD, the Jacobian is $J=U\Sigma V^{T}$ and its transpose is
$J^{T}=V\Sigma^{T}U^{T}$. Multiplying these two together and using the
orthonormality of the matrix $U$ gives $J^{T}J=V\Sigma^{T}\Sigma V^{T}$.
Plugging all this into Eq. 13

$\left(V\Sigma^{T}\Sigma V^{T}+\alpha
I\right)\mathbf{p}=V\Sigma^{T}U^{T}\mathbf{r}$ (14)

and multiplying from the left on both sides of this equation by $V^{T}$ yields

$\left(\Sigma^{T}\Sigma+\alpha
I\right)V^{T}\mathbf{p}=\Sigma^{T}U^{T}\mathbf{r}\text{.}$ (15)

Finally, solving Eq. 15 for $\mathbf{p}$ gives

$\mathbf{p}=V\left(\Sigma^{T}\Sigma+\alpha
I\right)^{-1}\Sigma^{T}U^{T}\mathbf{r}$ (16)

which is significantly simpler to work with as $\Sigma^{T}\Sigma+\alpha I$ is
diagonal and therefore easily invertible.

The SVD of the scaled Jacobian is done once per parameter update step $k$ (see
Algorithm 1), as is calculating both $\Sigma^{T}\Sigma$ and
$\Sigma^{T}U^{T}\mathbf{r}$. Even though Eq. 16 may be called multiple times
within a single update step $k$ when solving for the Levenberg-Marquardt
parameter (see Algorithm 2), the matrices involved in this repeated
calculation are all of size $n\times n$ and the operation is therefore
computationally inexpensive.

##  Appendix C Dealing with Uncertainty

Sometimes the data points $\left\\{y_{i},z_{i}\right\\}$ which are being fit
with the function $h(y;\mathbf{x})$ have different weightings/errors
$\sigma_{i}$ which need to be taken into consideration. These errors modifiy
Eq. 1 to the following form [29]

$f(\mathbf{x})=\frac{1}{2}\sum_{i}^{v}\left[\frac{h(y_{i},\mathbf{x})-z_{i}}{\sigma_{i}}\right]^{2}=\frac{1}{2}\mathbf{r}^{T}(\mathbf{x})W\mathbf{r}(\mathbf{x})$
(17)

where the residual vector is explicitly defined using Eq. 1 and where $W$ is
the weighting matrix given by

$W=\begin{bmatrix}\frac{1}{\sigma_{11}^{2}}&0&\dots&0\\\\[4.30554pt]
0&\frac{1}{\sigma_{22}^{2}}&\dots&0\\\\[4.30554pt]
\vdots&\vdots&\ddots&\vdots\\\\[4.30554pt]
0&0&\dots&\frac{1}{\sigma_{vv}^{2}}\end{bmatrix}\text{.}$ (18)

However, this diagonal weighting matrix is a special case and the more general
form is $W=C^{-1}$ [29] where $C$ is the symmetric covariance matrix given by

$C=\begin{bmatrix}\sigma_{11}^{2}&\sigma_{12}^{2}&\dots&\sigma_{1v}^{2}\\\\[4.30554pt]
\sigma_{21}^{2}&\sigma_{22}^{2}&\dots&\sigma_{2v}^{2}\\\\[4.30554pt]
\vdots&\vdots&\ddots&\vdots\\\\[4.30554pt]
\sigma_{v1}^{2}&\sigma_{v2}^{2}&\dots&\sigma_{vv}^{2}\end{bmatrix}\text{.}$
(19)

By using Eq. 17 instead of Eq. 1 and working through the math in Section II
the gradient (Eq. 3) and approximated Hessian are modified to be

$\mathbf{g}=\nabla f(\mathbf{x})=J(\mathbf{x})^{T}W\mathbf{r}(\mathbf{x})$
(20)

and

$H(\mathbf{x})\approx B=J^{T}(\mathbf{x})WJ(\mathbf{x})\text{.}$ (21)

respectively. This causes Eq. 6 to have the following form

$\left[J^{T}WJ+\alpha I\right]\mathbf{p}=-J^{T}W\mathbf{r},$ (22)

where, again, the scaled space’s hat notation $\hat{}$ has been removed for
readability. The weighting matrix in Eq. 22 seemingly complicates the
algorithmic implementation in Section II; however, we show below that Eq. 22
can be rewritten in the same form as Eq. 6 thereby keeping the algorithm the
same.

Since by definition the covariance matrix $C$ is real, Hermitian and positive-
definite it can be factorized via Cholesky decomposition as $C=LL^{T}$ where
$L$ is a real, lower triangular matrix. The weighting matrix can then be
defined in terms of this factorization
$W=C^{-1}=\left(L^{T}\right)^{-1}L^{-1}=\left(L^{-1}\right)^{T}L^{-1}$ and
substituted into Eq. 22

$\left[J^{T}\left(L^{-1}\right)^{T}L^{-1}J+\alpha
I\right]\mathbf{p}=-J^{T}\left(L^{-1}\right)^{T}L^{-1}\mathbf{r}\text{.}$ (23)

By using matrix transpose and multiplication rules Eq. 23 is rewritten as

$\left[\tilde{J}^{T}\tilde{J}+\alpha
I\right]\mathbf{p}=-\tilde{J}^{T}\tilde{\mathbf{r}}\text{.}$ (24)

where $\tilde{J}=L^{-1}J$ and $\tilde{\mathbf{r}}=L^{-1}\mathbf{r}$—which
successfully gives us the same form as Eq. 6.

Thus, when the user gives a covariance matrix, the weighted residual vector
$\tilde{\mathbf{r}}$ and weighted Jacobian $\tilde{J}$ are determined before
following the algorithm in Section II. To do this, the Cholesky decomposition
is first performed to find $L$. Since calculating the inverse of $L$ can be
computationally expensive the following two equations are solved to find
$\tilde{J}$ and $\tilde{\mathbf{r}}$

$\displaystyle L\tilde{J}$ $\displaystyle=J$ (25) $\displaystyle
L\tilde{\mathbf{r}}$ $\displaystyle=\mathbf{r}\text{.}$ (26)

However, it is rare to use the full covariance matrix and often only the
diagonal elements are known—which results in the weighting matrix given by Eq.
18. In this case the inverted form of $L$ is easily found to be

$L^{-1}=\begin{bmatrix}\frac{1}{\sigma_{11}}&0&\dots&0\\\\[4.30554pt]
0&\frac{1}{\sigma_{22}}&\dots&0\\\\[4.30554pt]
\vdots&\vdots&\ddots&\vdots\\\\[4.30554pt]
0&0&\dots&\frac{1}{\sigma_{vv}}\end{bmatrix}\text{.}$ (27)

which means $\tilde{J}$ and $\tilde{\mathbf{r}}$ can be calculated without
having to solve Eqs. 25-26. Additionally, this allows vector rather than
matrix multiplication to be used

$\displaystyle\tilde{J}$
$\displaystyle=L^{-1}J=\left(\mathbf{S}^{-1}\right)^{T}J$ (28)
$\displaystyle\tilde{\mathbf{r}}$
$\displaystyle=L^{-1}\mathbf{r}=\left(\mathbf{S}^{-1}\right)^{T}\mathbf{r}$
(29)

which is a less computationally expensive operation. $\mathbf{S}$ is a vector
containing the diagonal elements of $L$

$\mathbf{S}=\begin{bmatrix}\sigma_{11}\\\ \sigma_{22}\\\ \vdots\\\
\sigma_{vv}\end{bmatrix}\text{.}$ (30)

Similar to SciPy’s curve fit, JAXFit allows the user to give either
$\mathbf{S}$ or $C$ for the fit data error when calling the curve fit
function. At the beginning of the algorithm $\left(\mathbf{S}^{-1}\right)^{T}$
is calculated if given $\mathbf{S}$, whereas if $C$ is given it is factorized
via Cholesky decomposition to yield $L$. During the iterative part of the
algorithm (see Algorithm 1), the weighted residual vector $\tilde{\mathbf{r}}$
must be determined for each parameter update step $k$ using either Eq. 26 and
$\left(\mathbf{S}^{-1}\right)^{T}$ or Eq. 29 and $L$—depending on whether
$\mathbf{S}$ or $C$ has been given.

The weighted Jacobian $\tilde{J}$ also needs to be calculated for each
parameter update step $k$. By default JAXFit calculates the Jacobian with
automatic differentiation and $\tilde{J}$ is therefore calculated by simply
passing the weighted residual function $\tilde{\mathbf{r}}$ to JAX’s autodiff
function. However, JAXFit also allows users to explicitly define the Jacobian
matrix—in terms of $h(y;\mathbf{x})$—in which case $\tilde{J}$ must be solved
for explicitly using either Eq. 25 or Eq. 28.

## Funding

This work was supported by EPSRC Grant Nos. EP/P009565/1 and EP/TO19913/1, the
John Fell Oxford University Press (OUP) Research Fund and the Royal Society.

## Acknowledgements

L.R. thanks Oliver Karnbach and Mark Ijspeert for helpful discussions.

## Disclosures

The authors declare no conflicts of interest. After finishing this paper, we
became aware of a recently released Levenberg-Marquardt algorithm using JAX
[30].

## Data Availability

We make code available in [17].

## References

  * Siegman [1998] A. E. Siegman, How to (maybe) measure laser beam quality, in  _DPSS (Diode Pumped Solid State) Lasers: Applications and Issues_ (Optica Publishing Group, 1998) p. MQ1.
  * Hofer  _et al._ [2022a] L. R. Hofer, M. Krstajić, and R. P. Smith, Measuring laser beams with a neural network, Applied Optics 61, 1924 (2022a).
  * Bates  _et al._ [2008] M. Bates, B. Huang, and X. Zhuang, Super-resolution microscopy by nanoscale localization of photo-switchable fluorescent probes, Current Opinion in Chemical Biology 12, 505 (2008).
  * Gauss [1877] C. F. Gauss,  _Theoria motus corporum coelestium in sectionibus conicis solem ambientium_, Vol. 7 (FA Perthes, 1877).
  * Levenberg [1944] K. Levenberg, A method for the solution of certain non-linear problems in least squares, Quarterly of Applied Mathematics 2, 164 (1944).
  * Marquardt [1963] D. W. Marquardt, An algorithm for least-squares estimation of nonlinear parameters, Journal of the Society for Industrial and Applied Mathematics 11, 431 (1963).
  * Moré and Sorensen [1983] J. J. Moré and D. C. Sorensen, Computing a trust region step, SIAM Journal on Scientific and Statistical Computing 4, 553 (1983).
  * Moré [1978] J. J. Moré, The Levenberg-Marquardt algorithm: implementation and theory, in  _Numerical analysis_ (Springer, 1978) pp. 105–116.
  * Coleman and Li [1996] T. F. Coleman and Y. Li, An interior trust region approach for nonlinear minimization subject to bounds, SIAM Journal on Optimization 6, 418 (1996).
  * Branch  _et al._ [1999] M. A. Branch, T. F. Coleman, and Y. Li, A subspace, interior, and conjugate gradient method for large-scale bound-constrained minimization problems, SIAM Journal on Scientific Computing 21, 1 (1999).
  * MATLAB [2010] MATLAB,  _version 7.10.0 (R2010a)_ (The MathWorks Inc., Natick, Massachusetts, 2010).
  * Moré  _et al._ [1980] J. J. Moré, B. S. Garbow, and K. E. Hillstrom,  _User guide for MINPACK-1_, Tech. Rep. (CM-P00068642, 1980).
  * Virtanen  _et al._ [2020] P. Virtanen, R. Gommers, T. E. Oliphant, M. Haberland, T. Reddy, D. Cournapeau, E. Burovski, P. Peterson, W. Weckesser, J. Bright, _et al._ , SciPy 1.0: fundamental algorithms for scientific computing in Python, Nature Methods 17, 261 (2020).
  * Gaunt [2015] A. L. Gaunt,  _Degenerate Bose gases: tuning interactions & geometry_, Ph.D. thesis, University of Cambridge (2015).
  * Zhu and Zhang [2013] X. Zhu and D. Zhang, Efficient parallel Levenberg-Marquardt model fitting towards real-time automated parametric imaging microscopy, PloS One 8, e76665 (2013).
  * Przybylski  _et al._ [2017] A. Przybylski, B. Thiel, J. Keller-Findeisen, B. Stock, and M. Bates, Gpufit: An open-source toolkit for GPU-accelerated curve fitting, Scientific Reports 7, 1 (2017).
  * Hofer  _et al._ [2022b] L. R. Hofer, M. Krstajić, and R. P. Smith, JAXFit, GitHub (2022b).
  * Frostig  _et al._ [2018] R. Frostig, M. J. Johnson, and C. Leary, Compiling machine learning programs via high-level tracing, Systems for Machine Learning , 23 (2018).
  * Bradbury  _et al._ [2018] J. Bradbury, R. Frostig, P. Hawkins, M. J. Johnson, C. Leary, D. Maclaurin, G. Necula, A. Paszke, J. VanderPlas, S. Wanderman-Milne, and Q. Zhang, JAX: composable transformations of Python+NumPy programs (2018).
  * Nocedal and Wright [1999] J. Nocedal and S. J. Wright,  _Numerical optimization_ (Springer, 1999).
  * Mayorov [2015] N. Mayorov, Trust region reflective algorithm (2015).
  * Harris  _et al._ [2020] C. R. Harris, K. J. Millman, S. J. Van Der Walt, R. Gommers, P. Virtanen, D. Cournapeau, E. Wieser, J. Taylor, S. Berg, N. J. Smith, _et al._ , Array programming with NumPy, Nature 585, 357 (2020).
  * McKinney [2011] W. McKinney, Pandas: a foundational Python library for data analysis and statistics (2011).
  * Note [1] Even though SciPy uses a LSMR algorithm [31] for solving the least squares problem when the Jacobian matrix is both large and sparse, JAX does not currently support most SciPy sparse matrix operations and we therefore implement only SVD.
  * Baydin  _et al._ [2018] A. G. Baydin, B. A. Pearlmutter, A. A. Radul, and J. M. Siskind, Automatic differentiation in machine learning: a survey, Journal of Machine Learning Research 18, 1 (2018).
  * Hofer  _et al._ [2017] L. R. Hofer, R. V. Dragone, and A. D. MacGregor, Scale factor correction for Gaussian beam truncation in second moment beam radius measurements, Optical Engineering 56, 043110 (2017).
  * Hofer  _et al._ [2021] L. R. Hofer, M. Krstajić, P. Juhász, A. L. Marchant, and R. P. Smith, Atom cloud detection and segmentation using a deep neural network, Machine Learning: Science and Technology 2, 045008 (2021).
  * Martin and Porter [2012] C. D. Martin and M. A. Porter, The extraordinary SVD, The American Mathematical Monthly 119, 838 (2012).
  * Gavin [2019] H. P. Gavin, The Levenberg-Marquardt algorithm for nonlinear least squares curve-fitting problems, Department of Civil and Environmental Engineering, Duke University 19 (2019).
  * Blondel  _et al._ [2021] M. Blondel, Q. Berthet, M. Cuturi, R. Frostig, S. Hoyer, F. Llinares-López, F. Pedregosa, and J.-P. Vert, Efficient and modular implicit differentiation, arXiv preprint arXiv:2105.15183 (2021).
  * Fong and Saunders [2011] D. C.-L. Fong and M. Saunders, LSMR: An iterative algorithm for sparse least-squares problems, SIAM Journal on Scientific Computing 33, 2950 (2011).


footnote

# A Method for Controlling Extrapolation when Visualizing and Optimizing the
Prediction Profiles of Statistical and Machine Learning Models

Jeremy Ash,
Laura Lancaster,
and
Chris Gotwalt
JMP Division, SAS Institute

###### Abstract

We present a novel method for controlling extrapolation in the prediction
profiler in the JMP software. The prediction profiler is a graphical tool for
exploring high dimensional prediction surfaces for statistical and machine
learning models. The profiler contains interactive cross-sectional views, or
profile traces, of the prediction surface of a model. Our method helps users
avoid exploring predictions that should be considered extrapolation. It also
performs optimization over a constrained factor region that avoids
extrapolation using a genetic algorithm. In simulations and real world
examples, we demonstrate how optimal factor settings without constraint in the
profiler are frequently extrapolated, and how extrapolation control helps
avoid these solutions with invalid factor settings that may not be useful to
the user.

Keywords: Prediction Profiler, Extrapolation, Applicability Domain, Outlier
Analysis, Hotelling’s $T^{2}$, Genetic Algorithm, Response Surface

##  1 Introduction

When predictive models contain many factors, the prediction surface of the
model is high dimensional and can be challenging to visualize. This makes it
difficult to explore how predictions will change in response to changes in
factor settings. A graphical tool that addresses this problem is the
prediction profiler (Jones, 1991, 2021), hereafter referred to as the
profiler, which allows the user to visualize predictions over a high
dimensional factor space. Figure 1A shows an example of a profiler in the JMP
software. The diabetes data and the least squares model shown in this figure
will be discussed in more detail later in the results section. For each model
factor, a cross sectional view, or profiler trace, of the prediction surface
is shown. A profiler trace shows how predictions change as a function of a
model factor, with all other factors set at fixed values. Sliders for each
factor allow users to explore how the cross-sectional views change as the
factor settings change. Users can also set desirability functions, or
optimization goals, for each response in the profiler (e.g., maximize,
minimize, or match target) and perform optimization. Models with multiple
responses can be accommodated in the profiler, where a row of plots
corresponds to each response.

The model factor space in the profiler is set by default to have rectangular
boundary constraints for continuous factors that are given by the range of the
data. Users may also add mixture constraints or other linear constraints. The
model factor space for categorical factors is the discrete grid of factor
levels. Defining the factor space in this way does not prevent the user from
exploring prediction points that should be considered extrapolation. The
particular class of extrapolations that we focus on in this paper are
violations of the correlation structure of the data. Figure 1B-C provides an
example were the predicted response has been maximized and there are strong
violations of the factor correlation structure. Extrapolation can be dangerous
for a number of reasons:

  * •

Predictions rely on model assumptions that may become more invalid the further
the prediction is from the original data.

  * •

Prediction uncertainty increases the further the prediction is from the
original data.

  * •

Extrapolations may involve factor combinations that cannot be realized
physically. For example, in Figure 1C, total cholesterol is a function of the
sum of HDL and LDL. It is not possible to observe high LDL and low total
cholesterol as observed in the prediction point obtained in the profiler.

In earlier iterations of the profiler, it was easy to achieve extrapolated
predictions and difficult to know that the predictions were extrapolated. One
could plot the predictions on accompanying plots, such as a scatterplot matrix
of the factor variables, but this process was often tedious, especially when
the number of factors was large. The goal of this paper is to provide a method
for determining if predictions are extrapolated directly within the profiler.

There are a wide variety of methods used to assess whether a prediction point
should be considered extrapolation. The fields that have developed the
methodology often have both high dimensional and highly correlated factors.
For example, in cheminformatics (Eriksson et al., 2003; Netzeva et al., 2005)
and chemometrics (Rousseeuw et al., 2006), the chemical features are often
highly correlated. In these areas it is common for new data to be collected
from subpopulations that are not observed in the training data, so it is often
necessary to screen out any data that should be considered extrapolation prior
to performing prediction. Similar methodology can also be found in other
fields, such as ecology (Mesgaran et al., 2014; Bartley et al., 2019).

See Mathea et al. (2016) for a review of the popular metrics used for defining
extrapolation in predictive models. Some examples are leverage, Hotelling’s
$T^{2}$, convex hull based methods, nearest neighbor based estimation of local
density, kernel density estimation and 1-class SVM. We utilize two metrics in
this study: the well known leverage metric, which is specific to least squares
linear models, and a novel regularized Hotelling’s $T^{2}$ metric, which is
designed for machine learning models in general. The basic workflow for our
method is:

  * •

Compute an extrapolation measure for all observations in the training data.

  * •

Use the extrapolation measures to establish a threshold, beyond which
prediction points are classified as extrapolation.

  * •

Use the classifications to either restrict profiler traces so that
extrapolations are inaccessible to the user, or warn the user when they move
into an extrapolated region.

  * •

If a user requests optimization of a desirability function, perform
optimization subject to the extrapolation constraint.

The structure of the paper is as follows. Section 2 presents details on the
extrapolation control workflow: the distance metrics used, the procedure for
obtaining thresholds for the determination of extrapolation, and the genetic
algorithm used for optimization. Section 3 evaluates the workflow using
simulation studies. Section 4.1 demonstrates the extrapolation control
workflow for least squares linear models on a real world example. Section 4.2
demonstrates the extrapolation control workflow for general machine learning
models on a real world example.

Figure 1: Example of a prediction profiler for a least squares model fit to
the diabetes data discussed in Section 4.1. Predictions are shown for the
response variable on the y-axis. The model factors are shown on the x-axis.
Only a subset of the factors are shown for simplicity. 95% onfidence intervals
are shown in grey. The desirability function has been selected to increase
linearly with the response. A) Continuous factors are initialized to their
mean values in the training data and categorical factors are initialized to
the highest frequency level. B) Factor settings when the response variable (Y)
is maximized. C) Scatterplot for three factor variables. The training data
(grey) and the prediction point where the response is maximized (red). Since
total cholesterol is a function of the sum of HDL and LDL, the high LDL and
low total cholesterol observed in the prediction point is not physically
realizable.

##  2 Methods

###  2.1 Least squares models

For least squares linear models we chose leverage, a well known metric used to
identify outliers in linear models. Let
${\bm{\mathbf{{H}}}}={\bm{\mathbf{{X}}}}({\bm{\mathbf{{X}}}}^{T}{\bm{\mathbf{{X}}}})^{-1}{\bm{\mathbf{{X}}}}^{T}$
be the hat matrix for a linear model with $n\times p$ design matrix
${\bm{\mathbf{{X}}}}$. $h_{ii}$, which is the $i$th diagonal element of
${\bm{\mathbf{{H}}}}$, is the leverage for observation $i$ in the training
data. A typical point will have average leverage, which is
$\frac{1}{n}\sum_{i=1}^{n}h_{ii}=\frac{p}{n}$. The leverage for a new
prediction point is computed as
$h_{pred}={\bm{\mathbf{{x}}}}_{pred}^{T}({\bm{\mathbf{{X}}}}^{T}{\bm{\mathbf{{X}}}})^{-1}{\bm{\mathbf{{x}}}}_{pred}$,
where $x_{pred}$ is a $p$-dimensional new prediction observation.

One interpretation of $h_{pred}$ is that is the multivariate distance from the
center of the training data in the factor space. Another interpretation is
that it is the scaled prediction variance. That is, as the prediction point
moves further away from the center of the data (in a multivariate sense) the
uncertainty of the prediction increases.

There are two criteria commonly used in the statistical literature (Cook,
1977; Bartley et al., 2019) for determining if a prediction point with
leverage $h_{pred}$ should be considered extrapolation:

  * •

$h_{pred}>k\cdot max(h_{ii})$, where $k$ is a customizable multiplier. A
typical value of $k$ is 1. The $max(h_{ii})$ is the leverage of the furthest
point on the convex hull of the training data in factor space. When $k=1$,
prediction points beyond the threshold are outside of the convex hull.

  * •

$h_{pred}>l\cdot\frac{p}{n}$, where $l$ is a customizable multiplier. Typical
values of $l$ are 2 and 3. Recall that $\frac{p}{n}$ is the average leverage
for the training data.

###  2.2 General machine learning models

When selecting a generalized extrapolation control metric for models other
than linear models, we had a number of criteria. First, we needed a metric
that was computationally efficient for a large number of factors and
observations. This was necessary to maintain the interactivity of the profiler
traces and to preform optimization efficiently. Next, we wanted to be able to
support continuous, categorical and ordinal factors. We also wanted to utilize
observations with missing cells, because some modeling methods include these
observations in the fit. We wanted a method that was robust to linear
dependencies in the data, these occur when the number of variables is larger
than the number of observations, for example. We also wanted something that
was easy to automate, without much user input. And finally, we wanted
something that was easy to explain to non-statisticians.

We restricted our focus to extrapolation control methods that were
unsupervised. That is, we only flag a prediction point as extrapolation if it
is far outside the distribution of the training data in factor space. We do
this so that our extrapolation control metric will generalize to wide variety
of machine learning models. This was also necessary to be consistent across
model profilers, so that profilers can be linked and models can be ensembled.

We also restrict our consideration to extrapolation control methods that
protect against violations of correlation structure in the training data. This
type of extrapolation is the major concern in the use cases we focused on,
where popular methodology is based on $T^{2}$ and square prediction error
(SPE) metrics for principal components analysis and partial least squares
models (Eriksson et al., 2003; Rousseeuw et al., 2006). We discuss future
generalizations of this approach to protect against extrapolations between
clusters of data in the discussion.

The multivariate distance interpretation of leverage suggested Hotelling’s
$T^{2}$ as a distance metric for general extrapolation control. In fact, there
is a monotonic relationship between Hotelling’s $T^{2}$ and leverage. Since we
are no longer in linear models, this metric does not have the same connection
to prediction variance. So instead of relying on the extrapolation thresholds
used for linear models, we make distributional assumptions about $T^{2}$ to
determine an upper control limit to be used as the extrapolation threshold.

The formula for Hotelling’s $T^{2}$ is
$T^{2}=(x-\bar{x})^{T}\hat{\Sigma}^{-1}(x-\bar{x})$. The mean and covariance
matrix for the factors are estimated using the training data for the model. If
$p<n$ and the factors are multivariate normal (MVN), then $T^{2}_{pred}$ for a
prediction point has an F distribution:
$T_{pred}^{2}\sim\frac{(n+1)(n-1)p}{n(n-p)}F(p,n-p)$. However, we wanted the
method to generalize to data with complex data types such as a mix of
continuous and categorical variables, data sets where $p>n$, and data sets
where there are a large fraction of missing values. Instead of determining the
distribution of $T^{2}_{pred}$ in each of these scenarios, we use a simpler
and more conservative control limit that we found works well in practice: a
3-sigma control limit that uses the empirical distribution of $T^{2}$ based on
the training data. The control limit, that we use as our extrapolation
threshold, can be calculated by the following equation:
$UCL=\bar{T}^{2}+3\hat{\sigma}_{T^{2}}$ where $\bar{T}^{2}$ \+ is the mean of
training data $T^{2}$ and $\hat{\sigma}_{T^{2}}$ is the standard deviation of
the training data $T^{2}$.

Next, we describe a novel approach to computing the Hotelling’s $T^{2}$, which
we refer to as Regularized $T^{2}$. One complication with the standard
Hotelling’s $T^{2}$ is that it is undefined when $p>n$, because there are too
many parameters in the covariance matrix to estimate with the available data.
To address this we use a regularized covariance matrix estimator developed by
Schäfer and Strimmer (2005). This estimator was originally developed for the
estimation of covariance matrices for high-dimensional genomics data. The
estimator has been shown to produce a more accurate covariance matrix estimate
when $p$ is much larger than $n$. The estimator is:

$\displaystyle\hat{\Sigma}=(1-\hat{\lambda})\hat{{\bm{\mathbf{{U}}}}}+\hat{\lambda}\hat{{\bm{\mathbf{{D}}}}}.$

This is simply a weighted combination of the full sample covariance matrix
($\hat{{\bm{\mathbf{{U}}}}}$) and a constrained target covariance matrix
($\hat{{\bm{\mathbf{{D}}}}}$). The target matrix we used is a diagonal matrix,
with the sample variances of the factor variables on the diagonal and zeroes
on the off-diagonal. This shrinkage estimator is guaranteed to be positive
definite when $p>n$. This is necessary to compute Hotelling’s $T^{2}$ when
$p>n$. For the $\hat{\lambda}$ weight parameter, Schafer and Strimmer derived
an analytical expression that minimizes the mean squared error of the
estimator asymptotically.

Schaffer and Strimmer proposed several possible target matrices. The diagonal
target matrix $\hat{{\bm{\mathbf{{D}}}}}$ assumes that variables are
uncorrelated as the prior. This works well for a general extrapolation control
method, as we do not assume any correlation structure between the variables
without prior knowledge of the data. Also, when there is little data to
estimate the covariance matrix, the elliptical extrapolation control
constraints are expanded by the shrinkage estimator with the diagonal target
matrix. This results in a more conservative test statistic for extrapolation
control. That is, when there is little data available to estimate the
covariances, tests based on Regularized $T^{2}$ are less likely to label
prediction points as extrapolation. This is sensible, as covariances may be
observed by chance when training data is limited.

Another complication that needed to be addressed was how to compute
Hotelling’s $T^{2}$ when there is missing data. Many predictive modeling
methods can utilize observations with missing values for prediction, such as
random forests and boosted trees. When the sample size of the training data is
small and there is missing data, it would be ideal to do pair-wise deletion
instead of row-wise deletion to estimate the covariance matrix. This increases
the amount of data that is available to estimate the entries. Let
$\hat{{\bm{\mathbf{{U}}}}}=((u^{kl}))$ be the covariance matrix estimator. The
following formulas show the pairwise deletion method:

$\begin{gathered}\bar{x}^{k}=\frac{\sum_{i=1}^{n}x_{i}^{k}\mathbbm{1}(x^{k}_{i}\neq
NA)}{\mathbbm{1}(x^{k}_{i}\neq NA)}\\\
u^{kl}=\frac{\sum_{i=1}^{n}(x_{i}^{k}-\bar{x}^{k})(x_{i}^{l}-\bar{x}^{l})\mathbbm{1}(x^{k}_{i}\neq
NA,x^{l}_{i}\neq NA)}{\mathbbm{1}(x^{k}_{i}\neq NA,x^{l}_{i}\neq
NA)}\end{gathered}$

Previously, the profiler allowed constrained optimization with linear
constraints. Since extrapolation control is a non-linear constraint, the
optimization problem is more challenging. A genetic algorithm has been
implemented to perform the optimization.

##  3 Simulation Studies

To evaluate the extrapolation control performance of Regularized $T^{2}$, we
performed a simulation study. First, we simulate a factor matrix with a low
rank approximation. We do this to evaluate our ability to detect violations of
the correlation structure in the data:

${\bm{\mathbf{{X}}}}_{n\times p}={\bm{\mathbf{{U}}}}_{n\times
r}{\bm{\mathbf{{D}}}}_{r\times p}+{\bm{\mathbf{{e}}}}_{n\times 1}$

where $r$ is the desired the rank, and each element of ${\bm{\mathbf{{U}}}}$,
${\bm{\mathbf{{D}}}}$, and ${\bm{\mathbf{{e}}}}$ is i.i.d. standard normal
(Figure 2).

Figure 2: Example grid of prediction points classified as extrapolation (red)
or not extrapolation (green).

For each simulated data set, we chose the pair of most correlated variables.
We then extend a grid of equally spaced points from the center of the data to
the corner of the box constraint in the profiler, which is the range of the
factors according to the training data. Since the factor variables are
multivariate normal (MVN) with known covariance matrix, $T^{2}$ is
$\chi^{2}$-distributed. We then use the probability that $T^{2}_{pred}$ for a
grid point is from the same distribution as the data to classify points as
extrapolation or not extrapolation:

$\left\\{\begin{array}[]{@{}ll@{}}\text{``extrapolation"},&\text{if}\
P(T^{2}\geq T^{2}_{pred})<.05\\\ \text{``not extrapolation"},&\text{if}\
P(T^{2}\geq T^{2}_{pred})\geq.05\end{array}\right.$

To evaluate Regularized $T^{2}$ extrapolation control in terms of both false
positive rate (FPR) and true positive rate (TPR), we simulated data with
increasing sample sizes using the true MVN distribution. We then used the
simulated data to compute a 3-sigma limit, and we determined how well we were
able to classify the new prediction points as extrapolated or not. For all of
the scenarios shown in Figure 3 the FPR is less than .05, so only the TPR for
extrapolated points are shown. On the x-axis, the new prediction points along
the grid are ranked by the extent of extrapolation. Along the panels, the
number of variables and the rank of the factor matrix is varied. Figure 3A
shows a case when there are 20 factor variables with rank 10. Because the
threshold is conservative, notice that the TPR is low when the extrapolation
is mild. When there is limited training data, the Regularized $T^{2}$ is less
likely to label predictions as extrapolation. This is desirable because
covariances may be observed by chance when training data is limited. However,
the TPR does improve the larger the sample size.

Figure 3: True positive rates of classification of extrapolated points in
simulated data consisting of entirely continuous factors. Various sample
sizes, number of factors and rank of factor matrix are shown. The new
prediction points are ranked by the extent of extrapolation. 100 simulation
replicates were performed. Error bars show 95% confidence intervals.

Also note that the TPRs are high in all cases when $p=20$ and $n=20$. This is
when $p=n$, and when the rank of the covariance matrix, $r_{cov}$, is equal to
$n-1$. In these cases, Hotelling’s $T^{2}$ is undefined, since the covariance
is not full rank. If a generalized inverse, such as the Moore-Penrose inverse,
is used to invert the covariance matrix, then the Hotelling’s $T^{2}$ for all
training set observations will be a constant, $(2r_{cov})^{1/2}$. Our 3-sigma
control limit for Hotelling’s $T^{2}$ will be too large are too small
depending on $r_{cov}$. Figure S.1 shows a case when the Hotelling’s $T^{2}$
control limit is too small, leading to extrapoltion control with high FPRs,
and a case when the limit is too large, leading to extrapolation control with
low TPRs. Regularized $T^{2}$ achieves improved extrapolation control in both
cases.

When $p=100$ and $r=25$ (Figure 3B), the TPRs are low when $p$ is much larger
than $n$. This is a hard problem for distance based methods due to the curse
of dimensionality, but the TPRs do improve as the sample size increases.
Figure 3C-D shows how TPR is affected by varying rank with fixed p. The TPR is
better when the intrinsic dimensionality of the data is low.

We also evaluated Regularized $T^{2}$ extrapolation control on simulated data
with a mix of continuous and categorical variables (Figure 4). We simulated an
entirely continuous factor matrix using the same procedure as previously
described. We then selected a number of variables ($p_{cat}$) to transform
into categorical variables by randomly selecting the number of categories from
2 to 4 for each variable and using equally spaced quantiles to discretize each
variable. Figure 4 shows two scenarios where the number of factor variables
are 20 and 50. Half of the variables are categorical, and the rank of the
factor matrix is half of the number of variables. When $p=20$, our method is
able to obtain high TPRs for points that are at least moderately extrapolated,
but the problem becomes more challenging when $p>>n$.

Figure 4: True positive rates of classification of extrapolated points in
simulated data consisting of a mix of categorical and continuous factors. The
new prediction points are ranked by the extent of extrapolation. 100
simulation replicates were performed. Error bars show 95% confidence
intervals.

##  4 Applications

###  4.1 Diabetes data: Least square model example

We demonstrate extrapolation control for least squares models using a diabetes
data set originally analyzed in Efron et al. (2004). In the data, ten factor
variables have been measured on 422 diabetes patients. The variables include
age, gender, body mass index, average blood pressure, and six blood serum
measurements. Note that this data has mix of continuous and categorical
factors, as gender is categorical. The response (Y) is a continuous measure of
diabetes progression. 133 observations were randomly held out as a validation
set. A least squares model was fit to the remaining observations (Training
$R^{2}$ = .54, Validation $R^{2}$ = .44).

The extrapolation control feature in the prediction profiler has two settings:
warn and constrain. The warn setting warns users when factors are set to
values that are beyond the extrapolation threshold. Figure 5A shows profiler
traces for six out of the ten factors for simplicity. The factors are set to
the values found to maximize the response in Figure 1. The leverage for this
prediction point (8.62) is far above the maximum leverage for the training
data (.18). This has triggered the extrapolation warning.

The warn setting is a permissive mode that only warns the user that prediction
points may be suspect. A stricter option is constrain. This restricts the
profiler traces such that extrapolated regions of the factor space are
inaccessible (Figure 5B). This is similar to how the profile traces are
restricted in the presence of linear constraints, except the leverage
constraint is a non-linear, elliptical constraint.

Extrapolation control is particularly important when optimizing desirability
or using a model to match a target value. To optimize with extrapolation
control, we simply set extrapolation control to warn or constrain prior to
optimization. Figure 5B shows the factor settings when desirability has been
optimized with extrapolation control. The unconstrained solution shows a large
improvement in shrinkage, but the predicted value is far outside of the range
observed in the training data. Figure 5C shows how the solution has been
constrained to remain within the correlation structure of HDL and Total
Cholesterol. Figure 6 shows that extrapolation has been controlled for many
other pairs of variables.

Figure 5: Extrapolation control in a profiler for a least squares linear model
fit to the diabetes data. (A) Profiler traces with extrapolation control set
to “warn” and factors at the settings found to maximize the response in Figure
1. Leverage is far above the extrapolation threshold, which triggers the
extrapolation warning. (B) Profiler traces with extrapolation control set to
“constrain”. Factors are set to the prediction point maximizing the response,
subject to the extrapolation control constraint. (C) The prediction point
without extrapolation control (red) violates the correlation structure of LDL
and Total Cholesterol, while the extrapolation control prediction (green) does
not. Figure 6: Scatterplot matrix showing the optimal prediction point without
extrapolation control (red) and with extrapolation control (green) for the
diabetes data.

###  4.2 Metallurgy data: Machine learning model example

We demonstrate generalized extrapolation control using a manufacturing data
set. The data are from a company that manufactures steel drive shafts for the
automotive industry using powder metallurgy. The surface of the parts must be
free of pitting, or porosity, and if any pitting is observed the entire
production run must be scrapped. The company wishes to find the conditions
that minimize the chance of failure.

Key process variables and features of the parts have been recorded for 6,253
production runs. The manufacturing process begins with a metal powder.
Pressure is applied to compact the powder into a desired part shape, and one
key process variable is the amount of compaction pressure. Next, a sintering
furnace heats the metal particles so that they bind together, which
strengthens the metal part. Another process variable is sintering temperature,
which is the average temperature in the sintering furnace. Two response
variables have been collected. One is surface condition, a binary response
indicating whether the product was a failure or not, and another is shrinkage,
a continuous response that is associated with surface condition, where large
values of shrinkage are associated with a high failure rate.

We trained a boosted neural network to predict both response variables as a
function of the process variables. Weak learners were desirable for the
boosting iterations, so at each iteration we trained a neural network with one
hidden layer and three neurons. Each neuron had a TanH activation function. 20
boosting iterations were used. To validate our model, we held out
approximately 25% (1,563 observations) as a validation set. We performed
stratified sampling with both responses as stratification variables.

Next, we assess whether the model had reasonable prediction performance for
the response variables. Figures S.2 shows performance for the binary surface
condition response. Because the classes were highly imbalanced (the frequency
of failures was .045 in the entire dataset), precision and recall were used as
the primary metrics for model evaluation. On the training set, a threshold of
.05 on the predicted probability of failure was found to achieve a balance
between recall (.75) and precision (.14). With this threshold, the model also
achieved reasonable performance on the validation set (recall = .69, precision
= .13). Figures S.3 shows the model obtained good prediction performance for
continuous shrinkage variable, as well (Training $R^{2}$ = .84, Validation
$R^{2}$ = .83).

Having established that the model obtains acceptable prediction performance,
we proceed to finding the factor settings that minimize the chance of failure.
According to surface condition, we maximize the probability of pass and
minimize the probability of fail. We also minimize shrinkage. The overall
desirability that is maximized for all responses is defined as the geometric
mean of the desirability functions for the individual responses. Figure 7
shows the optimized factor settings with extrapolation control, compared to
the unconstrained solution. At the top of Figure 7 are the Regularized $T^{2}$
value at the optimized prediction point with extrapolation control turned on.
Note that it is below the 3-sigma control limit. Figure 8 shows how the
unconstrained solution violates the correlation structure of many of the
process variables, and how extrapolation control option corrects this. Note
that only a subset of the process variables with the largest extrapolations
for the unconstrained solution are shown (Figure S.4 shows the full set of
process variables). At the bottom of the Figure 7 is the difference in the
desirability between the extrapolation control solution and the unconstrained
solution. The desirabilities are practically equivalent, and the prediction
point with extrapolation control obtains much more reasonable factor settings.

Analyzing the same data set, we demonstrate how Regularized $T^{2}$ can be
used when there are missing values. We introduce at random missing values to
50 percent of the entries in the factor matrix. The boosted neural network
with the same parameters was trained with the informative missing values
option turned on. This option mean imputes the missing values and augments the
model factors with additional indicator variables, indicating the presence and
absence of missing values for each factor. This was necessary to train a
neural model that utilized observations with missing values. By default,
neural network models in JMP will remove all observations with missing values
prior to training. Figures S.5 and S.6 demonstrate that the model obtained
reasonable prediction performance, although performance did detoriate due to
the large amount of missing values. Figures S.7 shows that optimization with
extrapolation control obtains a less extrapolated prediction point in
comparison to the unconstrained solution.

Figure 7: The optimal prediction point in the boosted neural network profiler
with extrapolation control. Since factor settings that minimize failure rate
are desired, we maximize the probability of passing, minimize the probability
of failure, and minimize shrinkage. There is little difference between the
desirabilities for the extrapolation control optimal prediction point and the
unconstrained optimal prediction point. Figure 8: Scatterplot matrix showing
the optimal prediction point without extrapolation control (red) and with
extrapolation control (green) according to the boosted neural network trained
on the metallurgy data without missing values. Optimization with extrapolation
control obtains a less extrapolated prediction point in comparison to the
unconstrained solution. A subset of factors with large violation of
correlation structure are shown. See Figure S.4 for the full set of variables.

## References

  * Bartley et al. (2019) Bartley, M. L., E. M. Hanks, E. M. Schliep, P. A. Soranno, and T. Wagner (2019).  Identifying and characterizing extrapolation in multivariate response data.  PloS one 14(12), e0225715.
  * Cook (1977) Cook, R. D. (1977).  Detection of influential observation in linear regression.  Technometrics 19(1), 15–18.
  * Efron et al. (2004) Efron, B., T. Hastie, I. Johnstone, R. Tibshirani, et al. (2004).  Least angle regression.  Annals of statistics 32(2), 407–499.
  * Eriksson et al. (2003) Eriksson, L., J. Jaworska, A. P. Worth, M. T. Cronin, R. M. McDowell, and P. Gramatica (2003).  Methods for reliability and uncertainty assessment and for applicability evaluations of classification-and regression-based qsars.  Environmental health perspectives 111(10), 1361–1375.
  * Jones (1991) Jones, B. (1991).  An interactive graph for exploring multidimensional response surfaces.  In 1991 Joint Statistical Meetings, Atlanta, Georgia.
  * Jones (2021) Jones, B. (2021).  The prediction profiler at 30.  Quality Engineering, 1–8.
  * Mathea et al. (2016) Mathea, M., W. Klingspohn, and K. Baumann (2016, may).  Chemoinformatic Classification Methods and their Applicability Domain.  Molecular Informatics 35(5), 160–180.
  * Mesgaran et al. (2014) Mesgaran, M. B., R. D. Cousens, and B. L. Webber (2014).  Here be dragons: a tool for quantifying novelty due to covariate range and correlation change when projecting species distribution models.  Diversity and Distributions 20(10), 1147–1159.
  * Netzeva et al. (2005) Netzeva, T. I., A. P. Worth, T. Aldenberg, R. Benigni, M. T. Cronin, P. Gramatica, J. S. Jaworska, S. Kahn, G. Klopman, C. A. Marchant, et al. (2005).  Current status of methods for defining the applicability domain of (quantitative) structure-activity relationships: The report and recommendations of ecvam workshop 52.  Alternatives to Laboratory Animals 33(2), 155–173.
  * Rousseeuw et al. (2006) Rousseeuw, P. J., M. Debruyne, S. Engelen, and M. Hubert (2006).  Robustness and outlier detection in chemometrics.  Critical reviews in analytical chemistry 36(3-4), 221–242.
  * Schäfer and Strimmer (2005) Schäfer, J. and K. Strimmer (2005).  A shrinkage approach to large-scale covariance matrix estimation and implications for functional genomics.  Statistical Applications in Genetics and Molecular Biology 4(1), 1–30.

SUPPLEMENTAL MATERIALS

Figure S.1: False positive rates (A, C) and true positive rates (B, D) for
classification of extrapolation status of new prediction points using
simulated data consisting of entirely continuous factors. The prediction
points are ranked by the extent of extrapolation. 100 simulation replicates
were performed. Error bars show 95% confidence intervals. (A-B) the
Hotelling’s $T^{2}$ extrapolation threshold is too small, leading to
extrapolation control with high FPRs. (C-D) the Hotelling’s $T^{2}$
extrapolation threshold is too large, leading to extrapolation control with
low TPRs. Regularized $T^{2}$ achieves improved extrapolation control in both
cases. Figure S.2: Classification performance for surface condition according
to the boosted neural network trained on the metallurgy data without missing
values. The optimal probability threshold for failure status was selected to
be .05 using the training data. This threshold obtained a balance between
precision and recall (sensitivity) in both the training and validation set.
Figure S.3: Prediction performance for shrinkage according to the boosted
neural network trained on the metallurgy data without missing values. Figure
S.4: Scatterplot matrix showing the optimal prediction point without
extrapolation control (red) and with extrapolation control (green) according
to the boosted neural network trained on the metallurgy data without missing
values. Optimization with extrapolation control obtains a less extrapolated
prediction point in comparison to the unconstrained solution. Figure S.5:
Classification performance for surface condition according to the boosted
neural network trained on the metallurgy data with 50% missing values. The
optimal probability threshold for failure status was selected to be .05 using
the training data. This threshold obtained a balance between precision and
recall (sensitivity) in both the training and validation set. Figure S.6:
Prediction performance for shrinkage according to the boosted neural network
trained on the metallurgy data with 50% missing values. Figure S.7:
Scatterplot matrix showing the optimal prediction point without extrapolation
control (red) and with extrapolation control (green) according to the boosted
neural network trained on the metallurgy data with 50% missing values.
Optimization with extrapolation control obtains a less extrapolated prediction
point in comparison to the unconstrained solution.