# An Active Learning Based Robot Kinematic Calibration Framework Using
Gaussian Processes

Ersin Daş1 and Joel W. Burdick1 *“This work was supported by NASA Grant
80NSSC21K1032.”1Ersin Daş and Joel Burdick are with the Department of
Mechanical and Civil Engineering, California Institute of Technology,
Pasadena, CA 91125, USA<EMAIL_ADDRESS>
###### Abstract

Future NASA lander missions to icy moons will require completely automated,
accurate, and data efficient calibration methods for the robot manipulator
arms that sample icy terrains in the lander’s vicinity. To support this need,
this paper presents a Gaussian Process (GP) approach to the classical
manipulator kinematic calibration process. Instead of identifying a corrected
set of Denavit-Hartenberg kinematic parameters, a set of GPs models the
residual kinematic error of the arm over the workspace. More importantly, this
modeling framework allows a Gaussian Process Upper Confident Bound (GP-UCB)
algorithm to efficiently and adaptively select the calibration’s measurement
points so as to minimize the number of experiments, and therefore minimize the
time needed for recalibration. The method is demonstrated in simulation on a
simple 2-DOF arm, a 6 DOF arm whose geometry is a candidate for a future NASA
mission, and a 7 DOF Barrett WAM arm.

##  I Introduction

Kinematic Calibration (KC) of robot manipulators is a classic problem [1, 2,
3]. Calibration was originally motivated by the need for precise positioning
of factory robot tools and end-effectors. With limited real-time sensing,
factory robots depend upon their open-loop positioning precision to reliably
accomplish sophisticated tasks. As further discussed below, there are several
approaches to kinematic calibration, methods to deploy calibration in a
variety of situations, and approaches to better automate the calibration
process.

This paper revisits the calibration problem with new methods because of
proposed NASA missions to the icy moons of Europa and Enceladus, which may
host subsurface oceans. Cryovolcanic activity might transport water/ice from a
subsurface ocean through surface cracks. As part of a Europa lander mission
[4], a robotic arm will gather surface material samples from the vicinity of
these cracks, and analyze the local terrain using tools (scoops, drills,
penetrometers, and bevometers). The icy deposits may contain molecules that
indicate the presence of life in the subsurface oceans.

These missions will likely be battery powered, which severely limits their
duration. Round trip communication delays make remote manipulator operation
impractical. Hence, highly autonomous and efficient science operations are
crucial. While Europa is not a factory environment, kinematic accuracy of the
sampling robot arm is still crucial because proper science interpretation of
the samples requires accurate localization of the their origins. Furthermore,
since deep space missions must rigorously minimize mass, the sampling arm and
its tools will likely be very lightweight, and thus more deformable than a
factory robot.

Clearly, a data-efficient arm self-calibration process is indispensable for
reliable icy moon sample collection and high mission efficiency. At a minimum,
upon landing and first deployment, the sampling arm must be calibrated in-situ
to accommodate any distortions that occur during the multi-year space flight
to the moon. Moreover, these highly automated missions should be resilient to
distortions of the robot and tool geometries that arise from unexpectedly
forceful interactions with the terrain, joint freeze-ups in the $50^{o}K$
surface temperature, and improper loading of the sampling tools. Practically
speaking, icy moon missions require an automated calibration process that (1)
provides sufficient accuracy, especially for the workspace regions where
critical science occurs; (2) requires very few samples of arm data to reach a
high recalibration accuracy, since in-situ recalibration consumes valuable
mission time that should otherwise be devoted to science; and (3) can handle
miscalibrations that are larger than typically experienced by industrial
robots.

Figure 1: Visualization of the proposed active learning-based kinematic
calibration framework on a space robotic arm.

This paper develops an approach to manipulator kinematic calibration based on
Gaussian Processes (GP) [5]. Gaussian Process Regression (GPR) can be used to
learn kinematic residual errors, and allows for arm recalibration after large
arm deformations. More importantly, Gaussian Process Optimization [6]
adaptively and autonomously selects a small set of arm configurations that
provide accurate recalibration. We demonstrate the utility of the proposed
method on a 2-DOF planar manipulator, a 7-DOF Barrett WAM robot arm, and 6-DOF
robot arm used in the OceanWATERS testbed [7].

This paper is organized as follows. After some preliminaries in Sect. II,
Sect. III introduces an active learning-based kinematic calibration framework.
Simulations are presented in Sect. IV, while Sect. V concludes the paper.

Related Work: Kinematic calibration of a robot manipulator has been
extensively studied. Typically, the calibration process requires expert inputs
for experiment design and numerous measurements to optimize an observability
index. After data acquisition, offline optimization-based techniques, such as
iterative least squares regression with linearized kinematic equations [8],
Quadratic-Programming (QP) [9], and the Levenberg-Marquardt algorithm [10] are
used to identify the kinematic parameters. Although these approaches improve
robot positioning accuracy, they are neither data-efficient or capable of
completely autonomous implementation. Both characteristics are needed for a
deep-space application.

Machine learning-based frameworks can reduce positioning errors by considering
issues, such as joint elasticities and backlash, that affect calibration
accuracy. A Product-Of-Exponential kinematic model [11] has been combined with
a GP to compensate for residual deflections [12]. In [13], an Artificial
Neural Network directly learns robot position and orientation without nominal
forward kinematic equations. The post-calibration accuracy of these methods
highly depends upon the quality of the input data, whose selection relies upon
the calibration operator’s knowledge. An unsupervised learning-based
recalibration, which also optimizes the kinematic parameters for manipulator
visual servoing, is proposed in [14].

Many previous works use observability indexes [8] to optimize the choice of
sampling locations. An experimental comparison of the indexes is studied in
[15]. An active calibration algorithm is developed in [16] to find the most
efficient set of calibration poses for the linearization-based techniques. A
convex optimization approach to choosing configurations from a pool of
candidate data points is proposed in [17]. The calibration experiment design
methods in [18, 19], based on partial pose measurement, is performed before
the experiments. Although these methods are efficient for linearization-based
calibration, they do not consider how to choose sampling points in the
presence of larger inaccuracies. In [20], an online calibration approach
combines an unscented Kalman filter and iterative particle filter with a
probabilistic sampling strategy. A Kalman filter-based online calibration
method is presented in [21] for legged robots.

##  II Preliminaries

We use the standard notation: $\mathbb{R}$ and $\mathbb{R}^{+}$ represent the
set of real and positive real numbers, respectively. The Euclidean norm of a
matrix is denoted by $\|\cdot\|$. $f\sim GP(\mu,k)$ denotes that the function
$f(\cdot)$ is sampled from a Gaussian Process (which is reviewed in section
II-D).

###  II-A Forward Kinematics

A forward kinematic function maps a robot’s joint space to its workspace, and
it is a crucial part of closed-loop kinematic motion analysis. We adopt the
widely used D-H convention that defines the relative position and orientation
of the robot links and end effector using a four-parameter
representation.111The methods of this paper can be readily adapted to a
product-of-exponentials representation of manipulator kinematics. Given an $n$
jointed robot and the D-H parameters for the $i^{th}$ link: joint angle
$\theta_{i}$, twist angle $\alpha_{i}$, link length $a_{i}$, link offset
$d_{i}$, we define the transformation matrix ${{}^{i-1}T_{i}\in SE(3)}$, which
consists of a rotation matrix ${{}^{i-1}R_{i}\in SO(3)}$ and a translation
vector ${{}^{i-1}P_{i}\in\mathbb{R}^{3}}$, between neighboring robot links,
link ${i-1}$ and link ${i}$, as

${}^{i-1}T_{i}=\left[\begin{array}[]{c|c}^{i-1}R_{i}&{}^{i-1}P_{i}\\\
\cline{1-2}\cr 0~{}~{}0~{}~{}0&1\end{array}\right]=$ (3)
$\displaystyle\begin{bmatrix}\cos{\theta_{i}}&-\sin{\theta_{i}}\cos{\alpha_{i}}&\sin{\theta_{i}}\sin{\alpha_{i}}&a_{i}\cos{\theta_{i}}\\\
\sin{\theta_{i}}&\cos{\theta_{i}}\cos{\alpha_{i}}&-\cos{\theta_{i}}\sin{\alpha_{i}}&a_{i}\sin{\theta_{i}}\\\
0&\sin{\alpha_{i}}&\cos{\alpha_{i}}&d_{i}\\\ 0&0&0&1\end{bmatrix}$ (4)

for $i=1,~{}2\dots,~{}n$, where $\theta_{i}$ is variable if the joint is
revolute, while $d_{i}$ is variable if the joint is prismatic. Then, the
forward kinematics function is given by

$^{0}T_{\mathcal{T}}=~{}^{0}T_{1}(\theta_{1},\phi_{1})^{1}T_{2}(\theta_{2},\phi_{2})\ldots^{n-1}T_{n}(\theta_{n},\phi_{n})^{n}T_{\mathcal{T}},$
(5)

where
${\theta=\left[\theta_{1}~{}\ldots~{}\theta_{n}\right]^{T}\in\mathbb{R}^{n}}$,
${\phi=\left[\alpha_{i}~{}d_{i}~{}a_{i}\right]^{T}\in\mathbb{R}^{3n}}$, and
$\mathcal{T}$ represents the tool frame. This function gives the orientation
and position of the tool frame with respect to the base frame, i.e.,
${}^{0}R_{\mathcal{T}}\in SO(3),~{}^{0}P_{\mathcal{T}}\in\mathbb{R}^{3}$, and
can be written more compactly as

$\underbrace{\begin{bmatrix}{}^{0}Q_{\mathcal{T}}&{}^{0}P_{\mathcal{T}}\end{bmatrix}^{T}}_{F}=f(\theta,\phi),$
(6)

where ${}^{0}Q_{\mathcal{T}}$ is a four-parameter quaternion representation of
the rotation matrix [22]. Note that the quaternion conversion from a rotation
matrix may be discontinuous, and therefore complex to model by Gaussian
Processes[23]. To circumvent this problem, we utilize the continuous
quaternion conversion framework proposed in [24].

###  II-B Kinematic Calibration Problems

We now formalize two different calibration problems with respect to the
kinematic uncertainty structure. In the following definitions, let
$\tilde{F}=~{}[^{0}\tilde{Q}{}_{\mathcal{T}}~{}^{0}\tilde{P}{}_{\mathcal{T}}]^{T}$
represent the actual (uncertain, after distortion) kinematic function.

The Parametric calibration problem: Consider a perturbed D-H convention
parameter set

$[\theta~{}\phi]^{T}=\left[\theta_{i}+\Delta\theta_{i}~{}\alpha_{i}+\Delta\alpha_{i}d_{i}+\Delta
d_{i}~{}a_{i}+\Delta a_{i}\right]^{T},$ (7)
$\Delta=\left[\Delta\theta_{i}~{}\Delta\alpha_{i}~{}\Delta d_{i}~{}\Delta
a_{i}\right]^{T}\in\mathbb{R}^{4n},$ (8)

where $\Delta$ is the difference between the nominal (idealized) and actual
D-H parameters. The parametric calibration problem is —s— [ϕ]^T ∈R^4n∥ ~F-F
∥^2 [ϕ^*]^T =

The Non-parametric calibration problem: Consider the nominal forward kinematic
function $f_{n}(\theta,\phi)$ with an unstructured additive uncertainty term
$\Delta f(\theta,\phi)$ such that $f(\theta,\phi)=f_{n}(\theta,\phi)+\Delta
f(\theta,\phi)$. Then, the non-parametric calibration problem is defined as
Δf(θ)∥ ~F-F ∥^2 Δf^* (θ) = That is, we want to find the residual function
$\Delta f^{*}(\theta)$ that minimizes the calibration error with respect to
the forward kinematic function that uses idealized D-H parameters.

Kinematic calibration methods are driven by error between the measured end-
effector data and the computed forward kinematics data. Therefore, an expert
proposes enough experiments to provide sets of input joint angles $\theta$ and
measurements $\tilde{F}$ to solve the given optimal error reduction problem.
Then, the objective function of the optimization problems in (8), (II-B) is
redefined on a discrete set of data as

$\hat{\Delta}F=\begin{bmatrix}\tilde{F}_{1}-F_{1}&\ldots&\tilde{F}_{m}-F_{m}\end{bmatrix}^{T}~{}\in\mathbb{R}^{7m},$
(9)

where $m$ is the number of measurements.

###  II-C Classical Linearized Recalibration

The calibration cost function given in (8) is generally nonlinear and non-
convex depending on the robot kinematic structure. Thus, most classical
recalibration techniques are based on a linearized version of the parametric
calibration problem. By assuming relatively small parametric deflections, one
can linearize the calibration function (6) as

$\underbrace{\tilde{F}-F}_{\Delta
F\in\mathbb{R}^{7}}=\underbrace{\begin{bmatrix}\dfrac{\partial
f}{\partial\alpha_{i}}&\dfrac{\partial f}{\partial d_{i}}&\dfrac{\partial
f}{\partial a_{i}}\end{bmatrix}}_{J_{\phi}\in\mathbb{R}^{7\times
3n}}\underbrace{\begin{bmatrix}\Delta\alpha_{i}&\Delta d_{i}&\Delta
a_{i}\end{bmatrix}^{T}}_{\Delta{\phi}\in\mathbb{R}^{3n}},$ (10)

where $J_{\phi}$ is the Jacobian matrix. Therefore, when the nominal D-H
parameters and their possible lower bound
${\Delta{\phi}_{lb}\in\mathbb{R}^{3n}}$ and upper bound
${\Delta{\phi}_{ub}\in\mathbb{R}^{3n}}$ are given, the linearization-based KC
problem can be solved via the following constrained QP: —s— Δϕ ∈R^3n∥ ~F - F -
J_ϕ Δϕ ∥^2 Δϕ^* = Δϕ_lb ≤Δϕ ≤Δϕ_ub Then, we obtain
$\phi^{*}=\phi+\Delta{\phi}^{*}$. This QP is iteratively solved until an
acceptable error convergence is achieved [9].

###  II-D Gaussian Processes (GPs)

Gaussian Processes provide a stochastic, data-driven, supervised machine
learning approach to specify the relations between input and output data sets
of smooth functions through Bayesian inference [5]. A GP function,
${f(x):X\xrightarrow[]{}\mathbb{R}}$, defined on a data domain $X$, is fully
specified by its mean function ${\mu(x):X\xrightarrow[]{}\mathbb{R}}$ and a
covariance/kernel function ${k(x,x^{\prime}):X\times
X\xrightarrow[]{}\mathbb{R}}$:

$\displaystyle\mu(x)$
$\displaystyle\triangleq\mathop{{}\mathbb{E}}\left[f(x)\right],$ (11)
$\displaystyle k(x,x^{\prime})$
$\displaystyle\triangleq\mathop{{}\mathbb{E}}\left[(f(x)-\mu(x))(f(x^{\prime})-\mu(x^{\prime}))\right].$
(12)

We say that $f(\cdot)$ is distributed according to:

$f(x)\sim GP(\mu(x),k(x,x^{\prime})).$ (13)

Given a collection of training data $\mathcal{D}\triangleq\left\\{X_{N}\in
X,~{}Y_{N}\in\mathbb{R}^{N}\right\\}$ with inputs
$X_{N}=\left[x_{1},\ldots,x_{N}\right]$, outputs
$Y_{N}=\left[f(x_{1})+\epsilon_{1},\ldots,f(x_{N})+\epsilon_{N}\right]^{T}$
and i.i.d zero-mean Gaussian noises
${\epsilon_{i}\sim\mathcal{N}(0,\sigma^{2}_{\epsilon})}$, a GP predicts a set
of outputs based on an input test data set
${\tilde{X}=\left[\tilde{x}_{1},\ldots,\tilde{x}_{M}\right]\in X}$. The
posterior of the GP over the observations,
$\tilde{f}\big{|}\tilde{X},\mathcal{D}$, is also a GP distribution with a mean
$\tilde{\mu}$ and variance $\tilde{\sigma}^{2}$ given by

$\displaystyle\tilde{f}$
$\displaystyle\big{|}\tilde{X},\mathcal{D}\sim\mathcal{N}\big{(}\tilde{\mu},\tilde{\sigma}^{2}\big{)},$
(14) $\displaystyle{\tilde{\mu}}$
$\displaystyle=K_{\tilde{X}X_{N}}(K_{X_{N}X_{N}}+\sigma^{2}_{\epsilon}I)^{-1}~{}Y_{N},$
(15) $\displaystyle{\tilde{\sigma}^{2}}$
$\displaystyle=K_{\tilde{X}\tilde{X}}-K_{\tilde{X}{X}_{N}}(K_{X_{N}X_{N}}+\sigma^{2}_{\epsilon}I)^{-1}K_{{X}_{N}\tilde{X}},$
(16)

where $I$ is the identity matrix, ${K_{X_{N}X_{N}}\in\mathbb{R}^{N\times N}}$,
${K_{\tilde{X}\tilde{X}}\in\mathbb{R}^{M\times M}}$,
${K_{\tilde{X}X_{N}}\in\mathbb{R}^{M\times N}}$,
${K_{X_{N}\tilde{X}}\in\mathbb{R}^{N\times M}}$ are the covariance matrices
for the set of points, which measure the correlations between the inputs with
${[K_{X_{N}X_{N}}]_{i,j}=k(x_{i},x_{j})}$,
${[K_{\tilde{X}\tilde{X}}]_{i,j}=k(\tilde{x}_{i},\tilde{x}_{j})}$,
${[K_{\tilde{X}X_{N}}]_{i,j}=k(\tilde{x}_{i},{x}_{j})}$,
${[K_{X_{N}\tilde{X}}]_{i,j}=k({x}_{i},\tilde{x}_{j})}$.

Briefly, the kernel function reflect the features of the functions to be
learned, such as smoothness and input correlations. Our experiments use the
commonly used squared-exponential (SE) kernel given by

$k(x_{i},x_{j})=\sigma^{2}_{f}~{}exp\bigg{(}\dfrac{-\|x_{i}-x_{j}\|^{2}}{2l^{2}}\bigg{)}+\sigma^{2}_{n}\delta(x_{i},x_{j}),$
(17)

where ${\theta_{h}=[l~{}\sigma_{f}~{}\sigma_{n}]^{T}}$ are the tunable
hyperparameters and $\delta(x_{i},x_{j})$ is the Kronecker delta function. The
optimal hyperparameters can be found by minimizing the negative log marginal
likelihood function [5]. For data set $\mathcal{D}$, this function is defined
as

$\displaystyle\log
p(Y_{N}\big{|}X_{N},\theta_{h})=-\dfrac{1}{2}Y_{N}^{T}(K_{X_{N}X_{N}}(\theta_{h})+\sigma^{2}_{\epsilon}I)^{-1}Y_{N}$
$\displaystyle-\dfrac{1}{2}\log\big{|}(K_{X_{N}X_{N}}(\theta_{h})+\sigma^{2}_{\epsilon}I)\big{|}-\dfrac{N}{2}\log
2\pi.$ (18)

Finally, the optimal hyperparameters $\theta_{h}^{*}$, and
$\sigma_{\epsilon}^{*}$ can be computed via the gradient-descent based
solution of the following typically non-convex optimization problem: —s—
θ_h,σ_ϵ-logp (Y_N — X_N , θ_h). [ θ_h^* σ_ϵ^* ]^T =

##  III Active Learning Based Kinematic Calibration Approach

###  III-A GPs for Non-parametric Calibration

In this section, the residuals representing the inaccuracies in the forward
kinematics are modelled via GPs. Consequently, we will be able to achieve
effective non-parametric robot kinematic learning and calibration using noisy
sensor measurements in a stochastic framework.

The kinematic residual errors can be captured with seven GPs along each
orientation and translation direction. In particular, each component of matrix
$F$ in (6) is modelled as a single-output GP with a mean and kernel function
as

$\underbrace{\begin{bmatrix}{}^{0}\tilde{Q}_{\mathcal{T}}\\\
{}^{0}\tilde{P}_{\mathcal{T}}\end{bmatrix}}_{\tilde{F}}=f(\theta,\phi)+\begin{bmatrix}GP_{1}(\mu(\theta),k(\theta,\theta^{\prime}))\\\
\vdots\\\ GP_{7}(\mu(\theta),k(\theta,\theta^{\prime}))\end{bmatrix},$ (19)

where $\theta$ is the vector of joint angles, and $\tilde{F}$ is the measured
tool-frame poses. The regressed GP functions can predict the residual error
and its variance in each direction of motion, at a new query input
$\theta^{*}$, by using (15) and (16).

Our GP-based modeling of the residual errors in each motion direction is
simplistic. One should expect coupling in the errors across these directions.
Multi-output Gaussian Processes (MOGPs) [25] and vector-valued GPs [26] are
suitable solutions to this problem. However, due to the computational scaling
of these methods, we assume single-output GPs. As shown in Section IV, the
Gaussian distribution property holds in our applications. Thus, the decoupled
GP assumption is reasonable and computationally efficient.

###  III-B GPs for Parametric Calibration

We briefly note that GP regression can be used for parametric calibration, all
input parameters of the kinematic function in (6) can be modelled as a single-
output GP. To this end, we modify (6) as

$\begin{bmatrix}\theta&\phi\end{bmatrix}^{T}=f^{-1}(^{0}Q_{\mathcal{T}},^{0}P_{\mathcal{T}}),$
(20)

where the D-H parameters, including joint angles, are outputs, and the tool
frame poses are inputs. By assuming that each component in
${[\theta~{}\phi]^{T}\in\mathbb{R}^{4n}}$ can be modeled by a GP, the
parametric calibration problem can be written as

$\begin{bmatrix}\tilde{\theta}\\\
\tilde{\phi}\end{bmatrix}=\begin{bmatrix}{\theta_{0}}\\\
{\phi_{0}}\end{bmatrix}+\begin{bmatrix}GP_{1}(\mu(F),k(F,F^{\prime}))\\\
\vdots\\\ GP_{4n}(\mu(F),k(F,F^{\prime}))\end{bmatrix},$ (21)

where ${\tilde{\theta},~{}\tilde{\phi},~{}F}$ are the measurements,
${{\theta_{0}},~{}{\phi_{0}}}$ are the nominal values of the D-H parameters.

As compared to the non-parametric approach in (19), the joint angle vector is
also assumed to be uncertain. Therefore, Eq. (21) can model the possible joint
angle offsets. But, the number of GPs needed in this approach increases with
the number of robot joints. For example, for a 7 DOF robot, 28 different GPs
must be defined. Because of these drawbacks, we leave the parametric approach
to our future studies.

###  III-C GP-UCB for Non-parametric Experiment Design

We are concerned not only with fitting a GP model to calibration data, but
also with designing a calibration experiment that minimizes the number of
samples. such optimization is a classical experimental design problem that
aims to find the best set of measurement points [27, 28]. When using GP
regression-based calibration alone, a user must select the training and test
data points, which are essential for the method’s accuracy. This section
adopts a Gaussian Process Upper Confidence Bound (GP-UCB) [6] active learning
algorithm to adaptively select a sequence of informative measurements to
minimize the residual error.

Modeling the residual kinematic error with a GP, as in Section III-A, brings a
significant advantage for experiment design, since it encodes residual
function uncertainty in the confidence bounds [29]. With this feature, the GP-
UCB algorithm uses currently available measurements and the residual function
confidence bounds to decide the next test point that optimizes the utility
function $V_{t}(x):X\to\mathbb{R}$

$\displaystyle{V_{t}(x)=\ }$
$\displaystyle{\mu_{t-1}(x)}+{\beta_{t}}^{1/2}\sigma_{t-1}(x),$ (22)
$\displaystyle{x_{t}^{*}=\ }$ $\displaystyle\operatorname*{arg\,max}_{x\in
X}~{}V_{t}(x),$ (23)

where $\beta_{t}\in\mathbb{R}^{+}$ is an iteration-varying parameter [6] that
weights residual function uncertainty in the selection of the next sampling
point, $\sigma_{t-1}(x)$ is the square root of the variance, and
$\mu_{t-1}(x)$ is the predicted value of the mean function at input $x$.
Hence, GP-UCB chooses sampling points that trade off exploitation of joint
space regions with high calibration performance against exploration of regions
with high uncertainty about the residual function.

The GP-UCB algorithm has an appealing no regret property [6]. I.e., the
algorithm is guaranteed to find the globally optimal residual functions given
sufficient examples. Our experiments below show that some other calibration
experiment design approaches become stuck in local minima. Moreover, its
instantaneous regret is proportional to $T^{-1/2}$, where $T$ is the number of
experiments. This property implies that the calibration error steadily
decreases with each sample. While this paper focuses on sequential sampling
selection, a variation of this approach, entitled GP-BUCB [30], proposes
batches of experiments, while still retaining a no-regret property and similar
instantaneous regret performance.

Since there are seven different GPs with the same input vector, we modify the
utility function as the sum of UCBs:

$V_{t}^{T}(x)=\sum_{i=1}^{7}{\mu_{t-1}^{i}(x)}+{\beta_{t}}^{1/2}\sigma_{t-1}^{i}(x),$
(24)

where ${\mu_{t-1}^{i}(x)},~{}\sigma_{t-1}^{i}(x)$ are the mean and variance of
the $i^{th}$ output given in (19). Hence, the resulting new putatively optimal
data sampling point is obtained by considering all GPs simultaneously. The
authors in [31] show that this approximation does not violate the no-regret
property of the UCB algorithm under mild conditions.

The basic pseudo-code of the proposed GP-UCB-based calibration approach is
presented in Algorithm 1, which starts with any available GP prior, and the
next measurement is chosen to maximize the utility function (24).

Algorithm 1 GP-UCB for non-parametric arm calibration

0: Input set $\theta\in X$ GP prior $\mu_{0},~{}\sigma_{0},~{}k$ Kernel
function $k(\cdot,\cdot)$ Parameter $\beta_{t}$

0: $\theta^{*}$

for $t=1,2,\ldots$, do

Choose ${\theta_{t}^{*}=\ }\operatorname*{arg\,max}_{\theta\in
X}~{}V_{t}^{T}(\theta)$

Compute $\sigma_{t}$ via Equation (16)

Obtain ${\Delta F}_{t}=(\tilde{F}(\theta_{t})-F(\theta_{t}))+\epsilon_{t}$

Compute $\mu_{t}$ via Equation (15)

end for

The GP residual is updated from the new measurement (using Eqs. (15) and
(16)). This iteration continues until a stopping criterion is satisfied. For
the algorithm parameter ${\beta_{t}}$, which is defined in Theorem 2 of [6],
we use a strategy proposed in [6] as

$\displaystyle\beta_{t}=$ $\displaystyle
2log(t^{2}2\pi^{2}/(3\delta_{\beta}))~{}+$ $\displaystyle
2d_{\beta}log\big{(}t^{2}d_{\beta}a_{\beta}b_{\beta}r_{\beta}\sqrt{log(4d_{\beta}/\delta_{\beta})}\big{)}.$
(25)

Figure 2: Simulation examples. (a) 2-DOF Manipulator. (b) 7-DOF Barrett WAM
manipulator arm in OWLAT. (c) 6-DOF OceanWATERS Lander Robot Arm. (d), (e),
(f) The evaluations of the calibration errors with respect to the number of
iterations for 2-DOF, 7-DOF, and 6-DOF robots, respectively. (g), (h), (i) The
errors in position $x$ obtained from the test samples for 2-DOF, 7-DOF, and
6-DOF robots, respectively.

##  IV Experiments

We evaluate the performance of the GP-based residual regression and GP-UCB-
based sampling selection in three simulated robotic arms with perturbed
kinematic parameters. We compare our results with three other techniques: 1) A
griding-based D-optimal design algorithm with a convex approximation proposed
in [17]. This algorithm utilizes a pool of possible joint angle sets. Then,
the D-optimality criteria is adopted using the Jacobian matrix $J_{\phi}$
given in (10) and this pool. In [17], it is shown that D-optimal experimental
design leads to better performance than random sampling and the DEMAX
algorithm. 2) Expected improvement (EI) based sampling algorithm [32]. Like
the UCB utility function in (24), we modify the EI algorithm’s utility to sum
the expected improvements in each direction. 3) Random sampling.

###  IV-A 2-DOF Manipulator

We start with an easy-to-understand simulated calibration of a 2R planar robot
(Fig. 2-(a)). We seek to learn the forward kinematic residuals while choosing
a small set of sampling points. The uncertain D-H parameters of the arm are
$d_{1}=d_{2}=0\pm 0.1~{}m$, $a_{1}=a_{2}=1\pm 0.2~{}m$,
$\alpha_{1}=\alpha_{2}=0\pm 0.1~{}rad$,
$\theta_{1},\theta_{2}\in[-3,~{}3]~{}rad$. Even in this simple example,
linearization-based methods may not yield an accurate calibration in the
presence of additive uncertainties, since these approaches are formulated to
mainly address constant inaccuracies, i.e., ${\phi=\phi_{0}+e_{\phi}}$.

We assume precise measurements of the joint angles ($\theta_{1},\theta_{2}$)
and the orientation and position of tool frame ($F\in\mathbb{R}^{7}$) (e.g.,
by a calibrated high-resolution camera). We run Algorithm 1 for
$t=1,2,\ldots,30$ iterations, and choose the $\beta_{t}$ in (III-C) and an SE
kernel (17) with the hyperparameters optimized using (II-D) for all simulation
examples.

Fig. 2-(d) plots the calibration error, $\|\tilde{F}-F\|$, versus the number
of samples selected by GP-UCB. As more samples are collected, the calibration
error norm converges to zero, confirming the no-regret property of this
approach. EI-based sampling initially concentrates on the exploration of the
calibration error functions. But after 10 steps, it encounters a local
minimum, and does not continue to explore. For the D-optimal sampling, we
start with 10 randomly selected data points, since this method requires as
many data points as the number of unknowns. Although this method initially
offers better performance, the convergence rate is slow. While random sampling
also provides high-precision kinematic error modelling, it requires
significantly more sampling points than the active learning approach, as shown
in Fig. 2-(d). This confirms that GP-UCB is intelligently selecting its
experiments. The post-calibrated performance of GP-UCB is shown in Fig. 2-(g)
for the $x$ position error.

To demonstrate our methods ability to handle larger kinematic errors, and
noisy measurements, the planar robot’s D-H parameters are perturbed at a mean
level starting from $10\%$ up to to $200\%$. Secondly, a normally distributed
noise with a variance of 0.01 has been added to each output. For comparison,
we used the Kalibrot toolbox [9], which relies on classical linearized
recalibration. Fig. 3-(a) shows that our method learns and tolerates
increasing kinematic uncertainties and is sufficiently robust to sensor noise.
Although Kalibrot is robust against sensor noise, its calibration error
increases roughly linearly with parameter uncertainty.

Figure 3: (a) Kalibrot vs GP-UCB. (b) Histograms of calibration error
distribution along x-axis before the calibration process.

###  IV-B 7-DOF Barrett WAM Robot Arm

The OWLAT testbed was developed to evaluate autonomous techniques for Ocean
Worlds icy moon lander missions [33]. A 7-DOF Barrett WAM manipulator
simulates robotic sampling operations (Fig. 2-(b)) that might take place on an
icy moon mission. The arm’s kinematic parameters are given in Table I with
additive uncertainty bounds.

Fig. 2-(e) shows the simulated convergence of the forward kinematic error
estimate versus GP-UCB sample number. The results in this figure suggest that
the calibration performance and data efficiency of our framework are better
than other sampling methods for this 7-DOF arm, especially for small sample
number. Although the EI algorithm gave the second best results for this robot
arm, this method’s lack of design parameters is a weakness. The gridding-based
D-optimal algorithm fails in this simulation. Fig. 2-(h) shows the post-
calibrated kinematic accuracy on 50 test points linearly spaced along the $x$
position.

7-DOF Barrett WAM Robot Arm
---
Joint | $\boldsymbol{1}$ | $\boldsymbol{2}$ | $\boldsymbol{3}$ | $\boldsymbol{4}$ | $\boldsymbol{5}$ | $\boldsymbol{6}$ | $\boldsymbol{7}$
$\boldsymbol{d}$ | $0\pm 0.01$ | $0\pm 0.02$ | $0.55\pm 0.2$ | $0\pm 0.03$ | $0.3\pm 0.2$ | $0\pm 0.04$ | $0.06\pm 0.06$
$\boldsymbol{a}$ | $0\pm 0.01$ | $0\pm 0.03$ | $0.045\pm 0.01$ | $-0.045\pm 0.01$ | $0\pm 0.07$ | $0\pm 0.1$ | $0\pm 0.01$
$\boldsymbol{\alpha}$ | $-\pi/2\pm 0.2$ | $\pi/2\pm 0.2$ | $-\pi/2\pm 0.3$ | $\pi/2\pm 0.2$ | $-\pi/2\pm 0.1$ | $\pi/2\pm 0.1$ | $0\pm 0.3$
$\boldsymbol{\theta}$ | $\theta_{1}$ | $\theta_{2}$ | $\theta_{3}$ | $\theta_{4}$ | $\theta_{5}$ | $\theta_{6}$ | $\theta_{7}$
6-DOF OceanWATERS Lander Robot Arm
$\boldsymbol{d}$ | $0\pm 0.01$ | $0\pm 0.02$ | $0\pm 0.06$ | $-0.15\pm 0.01$ | $0\pm 0.04$ | $0\pm 0.06$ | -
$\boldsymbol{a}$ | $0.16\pm 0.1$ | $0.37\pm 0.1$ | $0.05\pm 0.02$ | $0.463\pm 0.1$ | $-0.238\pm 0.1$ | $0.225\pm 0.02$ | -
$\boldsymbol{\alpha}$ | $\pi/2\pm 0.2$ | $0\pm 0.2$ | $\pi\pm 0.3$ | $0\pm 0.2$ | $0\pm 0.1$ | $\pi/2\pm 0.1$ | -
$\boldsymbol{\theta}$ | $\theta_{1}$ | $\theta_{2}$ | $\theta_{3}$ | $\theta_{4}$ | $\theta_{5}$ | $\theta_{6}$ | -
TABLE I: DH parameters for the 7-DOF Barrett WAM Robot and 6-DOF OceanWATERS
Lander Robot Arm.

###  IV-C 6-DOF OceanWATERS Lander Robot Arm

We additionally demonstrate the proposed algorithm on the open-source
OceanWATERS simulation testbed[7] (see Fig. 2-(c) and Fig. 1), which is a
likely design for a future lander mission. The arm’s nominal kinematic
parameters and our tested uncertainty bounds are given in Table I.

Since our goal is to fit a GP to the residuals, we present the distribution of
forward kinematic errors for 1500 sample points, chosen by the GP-UCB
algorithm, in Fig. 3-(b). The figure demonstrates that the kinematic error has
a nearly normal distribution with zero mean. This result confirms that the use
of a Gauss Process framework is warranted.

Fig. 2-(f) compares the calibration performance of GP-UCB and the other
approaches over 600 sampling iterations. The figure shows that GP-UCB is an
efficient sampling method for this robot, since fewer data points are needed
for accurate recalibration. The EI and D-optimal algorithms reduce the norm of
the calibration error; but they become ”stuck” on one sample point after
$\sim$300 samples. The post-calibration accuracy of the GP model is
demonstrated on 50 linearly spaced $x$-axis test sample points in Fig. 2-(i).

##  V Conclusions and Future Work

Motivated by the proposed use of manipulators in future icy moon sampling
missions, this study aimed to accurately learn the manipulator forward
kinematic error residual in an online, non-parametric, and data-efficient way.
GP regression models the kinematic inaccuracies and their confidences, while
an online sequential GP-UCB algorithm efficiently selects the optimal
measurement in a way that adapts to the specific errors experienced by the
robot. Our numerical studies on three different arm geometries found that this
framework generally outperformed other methods in the limit of a small number
of calibration experiments, which is vital for deep space missions. In future
work, we will investigate specific GP kernels that better capture the
uncertainty properties of manipulator kinematics, and study the effects of
measurement noise on calibration accuracy under GP and classical calibration
frameworks. We will also evaluate if coupled Multivariate Gaussian Process
models [34, 26, 35] may further improve the calibration accuracy and sample
efficiency.

## References

  * [1] J. M. Hollerbach, “A survey of kinematic calibration,” in _The robotics review 1_ , 1989, pp. 207–242.
  * [2] W. Khalil and E. Dombre, _Modeling identification and control of robots_. CRC Press, 2002.
  * [3] A. Nubiola and I. A. Bonev, “Absolute calibration of an ABB IRB 1600 robot using a laser tracker,” _Robotics and Computer-Integrated Manufacturing_ , vol. 29, no. 1, pp. 236–245, 2013.
  * [4] K. P. Hand, C. B. Phillips, A. Murray, J. Garvin, E. Maize, R. Gibbs, G. Reeves, A. San Martin, G. Tan-Wang, J. Krajewski, _et al._ , “Science goals and mission architecture of the Europa lander mission concept,” _The Planetary Science Journal_ , vol. 3, no. 1, p. 22, 2022.
  * [5] C. R. Rasmussen and C. K. Williams, _Gaussian processes for machine learning_. The MIT Press, 2006.
  * [6] N. Srinivas, A. Krause, S. M. Kakade, and M. Seeger, “Gaussian process optimization in the bandit setting: No regret and experimental design,” _arXiv preprint arXiv:0912.3995_ , 2009.
  * [7] D. Catanoso, A. Chakrabarty, J. Fugate, U. Naal, T. M. Welsh, and L. J. Edwards, “OceanWATERS lander robotic arm operation,” in _2021 IEEE Aerospace Conference (50100)_. IEEE, 2021, pp. 1–11.
  * [8] J. M. Hollerbach and C. W. Wampler, “The calibration index and taxonomy for robot kinematic calibration methods,” _The international journal of robotics research_ , vol. 15, no. 6, pp. 573–591, 1996.
  * [9] F. Cursi, W. Bai, and P. Kormushev, “Kalibrot: A simple-to-use MATLAB package for robot kinematic calibration,” in _2021 IEEE/RSJ International Conference on Intelligent Robots and Systems (IROS)_. IEEE, 2021, pp. 8852–8859.
  * [10] T. S. Lembono, F. Suárez-Ruiz, and Q.-C. Pham, “SCALAR: Simultaneous calibration of 2-D laser and robot kinematic parameters using planarity and distance constraints,” _IEEE Transactions on Automation Science and Engineering_ , vol. 16, no. 4, pp. 1971–1979, 2019.
  * [11] F. C. Park, “Computational aspects of the product-of-exponentials formula for robot kinematics,” _IEEE Transactions on Automatic Control_ , vol. 39, no. 3, pp. 643–647, 1994.
  * [12] W. Jing, P. Y. Tao, G. Yang, and K. Shimada, “Calibration of industry robots with consideration of loading effects using Product-Of-Exponential (POE) and Gaussian Process (GP),” in _2016 IEEE International Conference on Robotics and Automation (ICRA)_. IEEE, 2016, pp. 4380–4385.
  * [13] F. Cursi, W. Bai, W. Li, E. M. Yeatman, and P. Kormushev, “Augmented neural network for full robot kinematic modelling in SE(3),” _IEEE Robotics and Automation Letters_ , vol. 7, no. 3, pp. 7140–7147, 2022.
  * [14] N. Gothoskar, M. Lázaro-Gredilla, Y. Bekiroglu, A. Agarwal, J. B. Tenenbaum, V. K. Mansinghka, and D. George, “DURableVS: Data-efficient Unsupervised Recalibrating Visual Servoing via online learning in a structured generative model,” _arXiv preprint arXiv:2202.03697_ , 2022.
  * [15] A. Joubair and I. A. Bonev, “Comparison of the efficiency of five observability indices for robot calibration,” _Mechanism and Machine Theory_ , vol. 70, pp. 254–265, 2013.
  * [16] Y. Sun and J. M. Hollerbach, “Observability index selection for robot calibration,” in _2008 IEEE international conference on robotics and automation_. IEEE, 2008, pp. 831–836.
  * [17] K. Kamali and I. A. Bonev, “Optimal experiment design for elasto-geometrical calibration of industrial robots,” _IEEE/ASME Transactions on Mechatronics_ , vol. 24, no. 6, pp. 2733–2744, 2019.
  * [18] Y. Wu, A. Klimchik, S. Caro, B. Furet, and A. Pashkevich, “Geometric calibration of industrial robots using enhanced partial pose measurements and design of experiments,” _Robotics and Computer-Integrated Manufacturing_ , vol. 35, pp. 151–168, 2015.
  * [19] H. Wang, T. Gao, J. Kinugawa, and K. Kosuge, “Finding measurement configurations for accurate robot calibration: Validation with a cable-driven robot,” _IEEE Transactions on Robotics_ , vol. 33, no. 5, pp. 1156–1169, 2017.
  * [20] G. Du, Y. Liang, C. Li, P. X. Liu, and D. Li, “Online robot kinematic calibration using hybrid filter with multiple sensors,” _IEEE Transactions on Instrumentation and Measurement_ , vol. 69, no. 9, pp. 7092–7107, 2020.
  * [21] S. Yang, H. Choset, and Z. Manchester, “Online kinematic calibration for legged robots,” _IEEE Robotics and Automation Letters_ , vol. 7, no. 3, pp. 8178–8185, 2022.
  * [22] R. M. Murray, Z. Li, and S. S. Sastry, _A mathematical introduction to robotic manipulation_. CRC press, 2017\.
  * [23] Y. Zhou, C. Barnes, J. Lu, J. Yang, and H. Li, “On the continuity of rotation representations in neural networks,” in _Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition_ , 2019, pp. 5745–5753.
  * [24] J. Wu, “Optimal continuous unit quaternions from rotation matrices,” _Journal of Guidance, Control, and Dynamics_ , vol. 42, no. 4, pp. 919–922, 2019.
  * [25] W. Bruinsma, E. Perim, W. Tebbutt, S. Hosking, A. Solin, and R. Turner, “Scalable exact inference in multi-output Gaussian processes,” in _International Conference on Machine Learning_. PMLR, 2020, pp. 1190–1201.
  * [26] M. Hutchinson, A. Terenin, V. Borovitskiy, S. Takao, Y. Teh, and M. Deisenroth, “Vector-valued Gaussian processes on Riemannian manifolds via gauge independent projected kernels,” _Advances in Neural Information Processing Systems_ , vol. 34, pp. 17 160–17 169, 2021.
  * [27] Y. Sun and J. M. Hollerbach, “Active robot calibration algorithm,” in _2008 IEEE international conference on robotics and automation_. IEEE, 2008, pp. 1276–1281.
  * [28] J. Hollerbach, W. Khalil, and M. Gautier, “Model identification,” in _Springer handbook of robotics_. Springer, 2016, pp. 113–138.
  * [29] E. Schulz, M. Speekenbrink, and A. Krause, “A tutorial on Gaussian process regression: Modelling, exploring, and exploiting functions,” _Journal of Mathematical Psychology_ , vol. 85, pp. 1–16, 2018.
  * [30] T. Desautels, A. Krause, and J. W. Burdick, “Parallelizing exploration-exploitation tradeoffs in Gaussian process bandit optimization,” _Journal of Machine Learning Research_ , vol. 15, pp. 3873–3923, 2014.
  * [31] B. Shen, R. Gnanasambandam, R. Wang, and Z. J. Kong, “Multi-task Gaussian process upper confidence bound for hyperparameter tuning and its application for simulation studies of additive manufacturing,” _IISE Transactions_ , pp. 1–13, 2022.
  * [32] B. Shahriari, K. Swersky, Z. Wang, R. P. Adams, and N. De Freitas, “Taking the human out of the loop: A review of Bayesian optimization,” _Proceedings of the IEEE_ , vol. 104, no. 1, pp. 148–175, 2015.
  * [33] H. Nayar, A. Goel, A. Boettcher, M. Hans, J. Sawoniewicz, A. Gaut, S. Higa, H. Ono, A. Jain, C. Lim, _et al._ , “Development of a lander autonomy testbed for ocean worlds missions,” in _Earth and Space 2021_ , 2021, pp. 531–540.
  * [34] T. de Wolff, A. Cuevas, and F. Tobar, “MOGPTK: The multi-output Gaussian process toolkit,” _Neurocomputing_ , vol. 424, pp. 49–53, 2021.
  * [35] R. Moriconi, M. P. Deisenroth, and K. Sesh Kumar, “High-dimensional Bayesian optimization using low-dimensional feature spaces,” _Machine Learning_ , vol. 109, no. 9, pp. 1925–1943, 2020.


# DOF: Accelerating High-order Differential Operators with Forward Propagation

Ruichen Li1, Chuwei Wang2,∗ Haotian Ye3,∗ Di He1, Liwei Wang1,4,†
1National Key Laboratory of General Artificial Intelligence, Peking University
2California Institute of Technology 3Stanford University
4Center for Machine Learning Research, Peking University
Equal contributions.Corresponding to<EMAIL_ADDRESS>
###### Abstract

Solving partial differential equations (PDEs) efficiently is essential for
analyzing complex physical systems. Recent advancements in leveraging deep
learning for solving PDE have shown significant promise. However, machine
learning methods, such as Physics-Informed Neural Networks (PINN), face
challenges in handling high-order derivatives of neural network-parameterized
functions. Inspired by Forward Laplacian, a recent method on accelerating
Laplacian computation, we propose an efficient computational framework,
Differential Operator with Forward-propagation (DOF), for calculating general
second-order differential operators without losing any precision. We provide
rigorous proof of the advantages of our method over existing methods,
demonstrating two times improvement in efficiency and reduced memory
consumption on any architectures. Empirical results illustrate that our method
surpasses traditional automatic differentiation (AutoDiff) techniques,
achieving 2x improvement on the MLP structure and nearly 20x improvement on
the MLP with Jacobian sparsity.

##  1 Introduction

Partial differential equations (PDEs) play a pivotal role in understanding and
predicting the behavior of physical systems. While classical numerical methods
have proven effective in some cases, they can be prohibitively challenging
when dealing with complicated problems, e.g., turbulence in fluid dynamics
[20] and high dimensional equations [7]. Recently, the advent of deep
learning[13] has spurred a wave of innovations in leveraging neural networks
(NN) for numerical solutions of PDEs. These NN-based approaches have been
applied to various problems, including fluid dynamics [3; 12], high-
dimensional optimal control problems [22; 8], and quantum many-body
problems[18; 10]. Notable works, such as the Physics-Informed Neural Network
(PINN)[19], showcase the potential of neural networks in capturing the
underlying physics of systems governed by PDEs.

Among these works, the central idea is to parameterize the solution as a
neural network, and optimize this expressive network with the guidance from
PDE. The ever-growing AutoDiff packages[2; 17; 1] enables convenient
calculation of associated quantities such as residual losses and derivatives,
avoiding discretization errors in classical methods. However, unlike most
computer vision or natural language processing tasks in which first-order
derivatives are sufficient for optimizations, in PDE-relevant problems one has
to deal with high-order derivatives. This raises a significant challenge as
common AutoDiff packages are computationally intensive under this
circumstance[16; 15]. There have been attempts to address this issue [9; 11;
21; 4]. However, they resort to either randomized methods[9; 11] or numerical
differentiation[21; 4], which introduce unsatisfying statistical errors and
are limited in problems where high precision is not demanded.

Recently, Li et al. [14] proposed a computational framework, Forward Laplacian
(FL). This framework is designed specifically for accelerating Laplacian
operator computation and thus can significantly boost the computation of
Laplacian-relevant PDE like Schrödinger equation in quantum chemistry[18].
Remarkably, Forward Laplacian is proven to be precision-preserved as it
introduces no statistical errors at all. Consequently, it is natural to ask
whether the idea of FL can be leveraged for other PDEs associated with high-
order derivatives computation.

To fill this gap, we develop a new computational framework named Differential
Operator with Forward-propagation (DOF). DOF shares a similar computational
procedure as FL, yet can be applied to compute general second-order
differential operators. We demonstrate that DOF outperforms conventional
AutoDiff methods in memory and computational cost by a large margin, both
theoretically and practically. In practice, DOF can accelerate NN-based
solvers across a wide range of PDEs such as the non-homogeneous heat equation
and Klein-Gordon equation.

The contribution of this paper is summarized as follows:

  * •

We generalize the Forward Laplacian method and propose DOF to precisely
compute arbitrary second-order differential operators of neural networks
(DOF).

  * •

We demonstrate that DOF improves computation efficiency and memory consumption
simultaneously, regardless of the architecture of neural networks, both
theoretically and empirically. The improvement can be significant in common
architectures like MLP.

##  2 Method

In PINNs and many other NN-based PDE solvers, the solution of a PDE,
$\phi({\mathbf{x}})$, is parameterized as neural networks,
$\phi({\mathbf{x}}):=\phi({\mathbf{x}};\theta)$, where $\theta$ represent the
NN parameters. These methods necessitate computing the high-order derivatives
of a neural network. It has been shown that standard AutoDiff methods are not
efficient for the high-order derivatives calculation[16; 15].

In this work, we focus on the calculation of the second-order differential
operators, which have the form

$\mathcal{L}:\phi({\mathbf{x}})\to\sum_{1\leq i,j\leq
N}a_{ij}({\mathbf{x}})\partial^{2}_{ij}\phi({\mathbf{x}})+\sum_{i\leq
N}b_{i}({\mathbf{x}})\partial_{i}\phi({\mathbf{x}})+c({\mathbf{x}})\phi({\mathbf{x}}).$
(1)

Here, $a_{ij},b_{i},c:\mathbb{R}^{N}\rightarrow\mathbb{R}$ are coefficients in
the second order operator, and $N$ is the input dimension (the time variable
is comprised in ${\mathbf{x}}$ for evolution equations). In practice, the
first term dominates the computation cost. Thus, for brevity and clarity, we
will only focus on the case when $b_{i}\equiv c\equiv 0$ in the following
discussions. We always denote the symmetric matrix
$\big{(}a_{ij}({\mathbf{x}})\big{)}_{i,j}$ as $A({\mathbf{x}})$.

###  2.1 Forward Laplacian

In standard AutoDiff packages, the second-order operator calculation is based
on the Hessian matrix $H=(\partial_{ij}\phi(x))_{i,j}$. We call those methods
Hessian-based methods. They use multiple Jacobian calculations to derive the
Hessian matrix, resulting in a huge computation cost. Recently, Li et al. [14]
proposed a new computational framework, Forward Laplacian (FL), which
primarily focuses on accelerating the calculation of Laplacian, i.e., $A\equiv
I_{N}$. Below we briefly review this method.

FL computes the Laplacian with one efficient forward pass, avoiding redundant
calculation in the Hessian-based approach. Following the notation in Li et al.
[14], we describe FL in a computation graph $\mathcal{G}$. The node set
$V=\\{v^{i}|i=0,1,...,M\\}$ represents the operations or variables used in a
neural network. We use the abbreviation $i\to j$ if there is a directed edge
from $v^{i}$ to $v^{j}$ in $\mathcal{G}$, and we denote operations as $F$,
e.g., $v^{j}=F_{j}(\\{v^{i}:i\to j\\})$ for all $j\geq 0$. Notice that the
node indices are arranged according to the topological order, i.e., for all
$i\to j$, we have $i<j$. The output of $\phi$ is denoted by $v^{M}$. In
addition, $\nabla$ and $\Delta$ represent the gradient operator and Laplacian
operator with respect to the input, respectively. Detailed notations can be
found in the appendix A.

Specifically, according to node dependency, FL sequentially computes the
Laplacian tuple $(v^{i},\nabla v^{i},\Delta v^{i})$ associated with each node.
In a simplified case where $v^{i}$ depends only on $v^{i-1}$, i.e.,
$v^{i}=F_{i}(v^{i-1})$, the graph can be represented as a chain. We can
compute the output tuple in a forward pass:

$({\mathbf{x}},\nabla{\mathbf{x}},\Delta{\mathbf{x}})\rightarrow(v^{0},\nabla
v^{0},\Delta v^{0})\rightarrow\cdots\rightarrow(v^{M},\nabla v^{M},\Delta
v^{M})$ (2)

The propagation rule of Laplacian tuple is derived through the chain rule:

$v^{i}=F_{i}(v^{i-1}),\leavevmode\nobreak\ \nabla
v^{i}=\partial_{v^{i-1}}F_{i}\nabla v^{i-1},\leavevmode\nobreak\ \Delta
v^{i}=\partial_{v^{i-1}}^{2}F_{i}|\nabla
v^{i-1}|^{2}+\partial_{v^{i-1}}F_{i}\Delta v^{i-1}$ (3)

For the general computation graph, we can generalize eq. 3 to the following
formula:

$\displaystyle v^{j}$ $\displaystyle=F_{j}(\\{v^{i}:i\to j\\})$ (4)
$\displaystyle\nabla v^{j}$ $\displaystyle=\sum\limits_{i:i\to
j}\frac{\partial F_{j}}{\partial v^{i}}\nabla v^{i}$ (5) $\displaystyle\Delta
v^{j}$ $\displaystyle=\sum\limits_{\begin{subarray}{c}i,l\\\ i\to j\ l\to
j\end{subarray}}\frac{\partial^{2}F_{j}}{\partial v^{i}\partial v^{l}}\nabla
v^{i}\cdot\nabla v^{l}+\sum\limits_{i:i\to j}\frac{\partial F_{j}}{\partial
v^{i}}\Delta v^{i}.$ (6)

Then we can sequentially compute the Laplacian tuple for each node in the
topological order according to eqs. 4, 5 and 6. As shown in Li et al. [14],
this approach outperforms the Hessian-based approach in numerous types of
computational graphs. It has been successfully applied to solving the
Schrodinger equation, resulting in over a magnitude of acceleration.

###  2.2 DOF for the Second-order Differential Operators

We now formally propose DOF to efficiently compute all kinds of second-order
operators. For brevity, we denote $A({\mathbf{x}})$ as $A$ and
$a_{ij}({\mathbf{x}})$ as $a_{ij}$ in the following discussion. To compute
$\sum_{i,j}a_{ij}\partial_{i}\partial_{j}\phi({\mathbf{x}})$, we first
decompose the coefficient matrix $A$ into $L^{\top}DL$ such that $D$ is a
diagonal matrix whose diagonal elements are all $\pm 1$ and $0$. As $A$ is a
symmetric matrix, this decomposition can be done easily. For instance, we can
eigen-decompose $A=S^{\top}\Sigma S$, where $S$ is an orthogonal matrix and
$\Sigma$ is the diagonal eigenvalue matrix, and choose $L=|\Sigma|^{1/2}S$ and
$D=\text{sgn}(\Sigma)$.

During the computation, for each node $v^{k}$, we compute tuple
$(v^{j},{\mathbf{g}}^{j},s^{j}):=(v^{j},L\nabla v^{j},\mathcal{L}v^{j})$. We
can derive the following formula by chain rule:

$\displaystyle v^{j}$ $\displaystyle=F_{j}(\\{v^{i}:i\to j\\})$ (7)
$\displaystyle{\mathbf{g}}^{j}$ $\displaystyle=\sum\limits_{i:i\to
j}\frac{\partial F_{j}}{\partial v^{i}}{\mathbf{g}}^{i}$ (8) $\displaystyle
s^{j}$ $\displaystyle=\sum\limits_{\begin{subarray}{c}i,l\\\ i\to j\ l\to
j\end{subarray}}\frac{\partial^{2}F_{j}}{\partial v^{i}\partial
v^{l}}{\mathbf{g}^{i}}^{\top}D\mathbf{g}^{l}+\sum\limits_{i:i\to
j}\frac{\partial F_{j}}{\partial v^{i}}s^{i}$ (9)

We can derive $\mathcal{L}\phi$ by sequentially applying this propagation rule
to each node in the topological order. We will prove that this method
outperform the Hessian-based methods in both the memory usage and computation
cost for any neural network architecture:

######  Theorem 2.1.

The computation cost (counted in FLOPs) of DOF is at most half that of
Hessian-based methods for any neural network architecture.

######  Theorem 2.2.

The memory consumption of DOF ($\mathcal{M}_{1}$) is smaller than that of
Hessian-based methods ($\mathcal{M}_{2}$) for any neural network architecture.

Specifically, $\mathcal{M}_{1}\lesssim\frac{2}{L}\mathcal{M}_{2}$ for an
$L$-layer MLP.

The proof for theorems 2.1 and 2.2 could be found in appendices B and D. We
remark that the memory and computation consumption of DOF can be further
reduced in some specific network architectures. See section 3.2 for details.

To better understand the DOF method, we discuss the implementation of DOF on
two special classes of second-order operators.

#### Elliptic Operator.

For elliptic operator, the coefficient matrix $A$ is positive so we have
$D=I_{N}$. As a result, eqs. 7, 8 and 9 are reduced to eqs. 4, 5 and 6. The
only difference between DOF and Forward Laplacian here is the initial value of
the tuple (i.e. the tuple at node $v^{j}$ for $j=0$). Thus, we utilize some
existing Forward Laplacian package[23; 6] to compute the elliptic operator.

#### Low-rank Coefficient Matrix.

It is well-known that computing the Hessian-vector product can be faster than
computing the entire Hessian matrix in the standard AutoDiff package. Namely,
computing Hessian-vector product takes $\mathcal{O}(1/N)$ cost of computing
the entire Hessian matrix. Thus, in the standard AutoDiff package, if the
coefficient matrix is a low-rank matrix, the calculation of the second-order
operator can be accelerated through multiple Hessian-vector products.
Similarly, the DOF method also accomplishes this acceleration when dealing
with a low-rank coefficient matrix.

If the coefficient matrix $A$ is a rank-$r$ matrix, we can eliminate the
columns and rows associated with the zero eigenvalues in $L$ and $D$. Then we
have $L^{\prime}\in\mathbb{R}^{r\times N}$ and
$D^{\prime}\in\mathbb{R}^{r\times r}$ that still satisfy
$A=L^{\prime\top}D^{\prime}L^{\prime}$. Thus, while the propagation rule of
DOF remains the same, the dimension of ${\mathbf{g}}^{k}$ is reduced from $N$
to $r$. According to the analysis in the appendices B and D, this reduction in
the dimension will lead to an $\mathcal{O}(r/N)$ reduction in both the memory
usage and the computation cost.

##  3 Results

Table 1: Comparison between DOF and Hessian-based method on the MLP Operator | GPU Memory Usage (MB) | Time (ms)
---|---|---
Hessian |  DOF |  ratio |  Hessian |  DOF |  ratio
Elliptic |  10421 |  3165 |  3.3 |  196.7 |  106.6 |  1.8
Low-rank |  10427 |  2141 |  4.9 |  196.2 |  55.8 |  3.5
General |  10429 |  3181 |  3.3 |  197.4 |  122.2 |  1.6
Table 2: Comparison between DOF and Hessian-based method on the MLP with Jacobian sparsity Operator | GPU Memory Usage (MB) | Time (ms)
---|---|---
Hessian |  DOF |  ratio |  Hessian |  DOF |  ratio
Elliptic |  21401 |  997 |  21.5 |  366.4 |  18.9 |  19.4
Low-rank |  21401 |  869 |  24.6 |  366.2 |  12.7 |  28.9
General |  21401 |  997 |  21.5 |  366.6 |  18.9 |  19.4

In this section, we compare DOF with the standard Hessian-based approach on
different operators and network architectures. We study three kinds of second-
order operators: elliptic operator, elliptic operator with low-rank
coefficient matrix, and general operator. As for the architecture, we choose
the standard MLP structure and the MLP with Jacobian sparsity. Details could
be found in appendix E.

###  3.1 Comparison on the MLP

The benchmark results for the MLP are shown in table 1. For both the elliptic
and general operators, DOF demonstrates significant efficiency improvements,
halving the computation cost and reducing memory usage by a third compared to
the Hessian-based method, aligning with our theoretical predictions. In the
case of the low-rank operator, DOF achieves even greater acceleration,
attributed to the dimensionality reduction of $\mathbf{g}$.

###  3.2 Comparison on the MLP with Jacobian Sparsity

As discussed in Li et al. [14], the cost of Forward Laplacian method is
significantly reduced when the Jacobian of intermediate component is sparse.
We notice that this sparse Jacobian property also exists in some advanced PINN
network architecture [5]. As a forward AutoDiff method, DOF can also leverage
this property to significantly accelerate the calculation. The memory and time
comparison are listed in table 2 and the architecture details can be found in
appendix E. Compared with the Hessian-based method, DOF can significantly
reduce both the memory and computation consumption in the MLP with Jacobian
sparsity, showing a great potential in applying our method to the advanced
machine learning-based PDE solver.

## References

  * Abadi et al. [2016] Martín Abadi, Paul Barham, Jianmin Chen, Zhifeng Chen, Andy Davis, Jeffrey Dean, Matthieu Devin, Sanjay Ghemawat, Geoffrey Irving, Michael Isard, et al.  Tensorflow: a system for large-scale machine learning.  In _Osdi_ , volume 16, pp. 265–283. Savannah, GA, USA, 2016.
  * Bradbury et al. [2018] James Bradbury et al.  JAX: composable transformations of Python+NumPy programs, 2018.  URL http://github.com/google/jax.
  * Cai et al. [2021] Shengze Cai, Zhiping Mao, Zhicheng Wang, Minglang Yin, and George Em Karniadakis.  Physics-informed neural networks (pinns) for fluid mechanics: A review.  _Acta Mechanica Sinica_ , 37(12):1727–1738, 2021.
  * Chiu et al. [2022] Pao-Hsiung Chiu, Jian Cheng Wong, Chinchun Ooi, My Ha Dao, and Yew-Soon Ong.  Can-pinn: A fast physics-informed neural network based on coupled-automatic–numerical differentiation method.  _Computer Methods in Applied Mechanics and Engineering_ , 395:114909, 2022.
  * Cho et al. [2023] Junwoo Cho, Seungtae Nam, Hyunmo Yang, Seok-Bae Yun, Youngjoon Hong, and Eunbyung Park.  Separable physics-informed neural networks.  _Advances in Neural Information Processing Systems_ , 2023.
  * Gao et al. [2023] Nichoals Gao, Jonas Köhler, and Adam Foster.  folx - forward laplacian for jax, 2023.  URL http://github.com/microsoft/folx.
  * Han et al. [2018] Jiequn Han, Arnulf Jentzen, and Weinan E.  Solving high-dimensional partial differential equations using deep learning.  _Proceedings of the National Academy of Sciences_ , 115(34):8505–8510, 2018.
  * Han et al. [2019] Jiequn Han, Linfeng Zhang, and Weinan E.  Solving many-electron schrödinger equation using deep neural networks.  _Journal of Computational Physics_ , 399:108929, December 2019.  ISSN 00219991.  doi: 10.1016/j.jcp.2019.108929.  URL http://arxiv.org/abs/1807.07014.
  * He et al. [2023] Di He, Shanda Li, Wenlei Shi, Xiaotian Gao, Jia Zhang, Jiang Bian, Liwei Wang, and Tie-Yan Liu.  Learning physics-informed neural networks without stacked back-propagation.  In _International Conference on Artificial Intelligence and Statistics_ , pp. 3034–3047. PMLR, 2023.
  * Hermann et al. [2020] Jan Hermann, Zeno Schätzle, and Frank Noé.  Deep-neural-network solution of the electronic Schrödinger equation.  _Nature Chemistry_ , 12(10):891–897, 2020.
  * Hu et al. [2023] Zheyuan Hu, Zekun Shi, George Em Karniadakis, and Kenji Kawaguchi.  Hutchinson trace estimation for high-dimensional and high-order physics-informed neural networks.  _arXiv preprint arXiv:2312.14499_ , 2023.
  * Kovachki et al. [2021] Nikola Kovachki, Zongyi Li, Burigede Liu, Kamyar Azizzadenesheli, Kaushik Bhattacharya, Andrew Stuart, and Anima Anandkumar.  Neural operator: Learning maps between function spaces.  _arXiv preprint arXiv:2108.08481_ , 2021.
  * LeCun et al. [2015] Yann LeCun, Yoshua Bengio, and Geoffrey Hinton.  Deep learning.  _nature_ , 521(7553):436–444, 2015.
  * [14] Ruichen Li, Haotian Ye, Du Jiang, Xuelan Wen, Chuwei Wang, Zhe Li, Xiang Li, Di He, Ji Chen, Weiluo Ren, et al.  A computational framework for neural network-based variational monte carlo with forward laplacian.  _Nature Machine Intelligence_.  URL https://www.nature.com/articles/s42256-024-00794-x.
  * Meng et al. [2021] Chenlin Meng, Yang Song, Wenzhe Li, and Stefano Ermon.  Estimating high order gradients of the data distribution by denoising.  _Advances in Neural Information Processing Systems_ , 34:25359–25369, 2021.
  * Pang et al. [2020] Tianyu Pang, Kun Xu, Chongxuan Li, Yang Song, Stefano Ermon, and Jun Zhu.  Efficient learning of generative models via finite-difference score matching.  _Advances in Neural Information Processing Systems_ , 33:19175–19188, 2020.
  * Paszke et al. [2017] Adam Paszke, Sam Gross, Soumith Chintala, Gregory Chanan, Edward Yang, Zachary DeVito, Zeming Lin, Alban Desmaison, Luca Antiga, and Adam Lerer.  Automatic differentiation in pytorch.  2017\.
  * Pfau et al. [2020] D. Pfau, J.S. Spencer, A.G. de G. Matthews, and W.M.C. Foulkes.  Ab-initio solution of the many-electron schrödinger equation with deep neural networks.  _Physical Review Research_ , 2:033429, 2020.  doi: 10.1103/PhysRevResearch.2.033429.  URL https://link.aps.org/doi/10.1103/PhysRevResearch.2.033429.
  * Raissi et al. [2019] Maziar Raissi, Paris Perdikaris, and George E Karniadakis.  Physics-informed neural networks: A deep learning framework for solving forward and inverse problems involving nonlinear partial differential equations.  _Journal of Computational Physics_ , 378:686–707, 2019.
  * Rogallo & Moin [1984] Robert S Rogallo and Parviz Moin.  Numerical simulation of turbulent flows.  _Annual review of fluid mechanics_ , 16(1):99–137, 1984.
  * Sirignano & Spiliopoulos [2018] Justin Sirignano and Konstantinos Spiliopoulos.  Dgm: A deep learning algorithm for solving partial differential equations.  _Journal of computational physics_ , 375:1339–1364, 2018.
  * Wang et al. [2022] Chuwei Wang, Shanda Li, Di He, and Liwei Wang.  Is ${L}^{2}$ physics informed loss always suitable for training physics informed neural network?  In _Advances in Neural Information Processing Systems_ , 2022.  URL https://openreview.net/forum?id=cy1TKLRAEML.
  * Ye et al. [2023] Haotian Ye, Ruichen Li, and Du Jiang.  Lapjax, 2023.  URL https://github.com/YWolfeee/lapjax.

##  Appendix A Preliminary

In this section, we briefly review computation framework in deep learning and
auto differentiation, and introduce the notations we will use in the following
sections.

###  A.1 Computation Graph

The computation graph serves as a descriptive language of deep learning models
across various deep learning toolkits, including PyTorch [17], TensorFlow [1],
and Jax [2].

In the computation graph $\mathcal{G}$ associated with a function
$\phi({\mathbf{x}}),\ {\mathbf{x}}=(x_{1},...x_{N})\in\mathbb{R}^{N}$(we will
always use $N$ for input dimension), the edges represent function arguments,
and nodes represent operations or variables. Note that $\mathcal{G}$ is always
a directed acyclic graph. We use $\\{{v}^{-1},...{v}^{-N}\\}$ to represent the
external nodes of $\mathcal{G}$ (i.e., the input ${\mathbf{x}}$ of the neural
network-parameterized function $\phi$), and use $\\{v^{0},...v^{M}\\}$ to
represent the internal nodes, sorted in topological orders. A node can be
referred to as a neuron in the network. Specifically, we use $v^{M}$ to serve
as the network output $\phi$. We use the abbreviation $i\to j$ if there is a
directed edge from $v^{i}$ to $v^{j}$ in $\mathcal{G}$. Furthermore, we denote
operations as $F$, e.g., $v^{j}=F_{j}(\\{v^{i}:i\to j\\})$ for all $j\geq 0$.

######  Example A.1.

Take Multi-Layer Perceptron (MLP) function $\phi({\mathbf{x}})$ as an example.

$\phi$ has the form $\phi=F_{L}\circ F_{L-1}\circ...\circ F_{0}$, where
$F_{l}=(F_{l,1},...F_{l,N_{l+1}}):\mathbb{R}^{N_{l}}\to\mathbb{R}^{N_{l+1}}$
are the mapping in each layer with $N_{0}=N$ and $N_{L+1}=1$. Let
${\mathbf{u}}^{l}=(u_{1}^{l},u_{2}^{l},...u^{l}_{N_{l}})\in\mathbb{R}^{N_{l}}$
be the vector consist of the neurons in the $l$-th layer, we have
${\mathbf{u}}^{0}={\mathbf{x}}$ and
${\mathbf{u}}_{l+1}=F_{l}({\mathbf{u}}^{l})=\sigma(W^{l}{\mathbf{u}}^{l}+b^{l})$
where $W^{l}\in\mathbb{R}^{N_{l+1}\times N_{l}},\
b^{l}\in\mathbb{R}^{N_{l+1}}$ are network parameters and $\sigma$ is a
nonlinear function operated element-wise.

In this setting, the nodes $\\{v^{i}\\}_{i}$ in the computation graph are

$\displaystyle x_{1},...x_{N},\ u^{1}_{1},...u^{1}_{N_{1}},\
u^{1.5}_{1},...u^{1.5}_{N_{1}},\
u^{2}_{1},...,u^{2}_{N_{2}},...,u^{{L+1}}_{1}(=\phi({\mathbf{x}})),$ (10)

and the operations generating the node representing $u_{i}^{l}$ and
$u_{i}^{l-0.5}$ are $u_{i}^{l}=\sigma(u_{i}^{l-0.5})$ and
$u^{l-0.5}_{i}=\sum_{j=1}^{N_{l-1}}W^{l-1}_{ij}u_{j}^{l-1}+b^{l-1}_{i}$,
respectively.

###  A.2 Auto Differentiation

In most machine learning toolkits, auto differentiation(AutoDiff) implemented
with back propagation algorithm is applied to compute the derivatives of
neural network functions.

For a function $\phi({\mathbf{x}})$, this method first performs forward
propagation to obtain the value of each variable $v^{j}$, and construct the
computation graph $\mathcal{G}$. Next, a backward process is employed by
creating a new computation graph $\hat{\mathcal{G}}$ where node
$\hat{v}^{i}\in\hat{\mathcal{G}}$ represents the operation to calculate
$\frac{\partial\phi}{\partial v^{i}}$, $i=M,...,-N$. The associated
computations are

$\frac{\partial\phi}{\partial v^{i}}=\sum\limits_{j:i\to j}\frac{\partial
F_{j}}{\partial v^{i}}\frac{\partial\phi}{\partial v^{j}},\ i=M-1,...,-N.$
(11)

##  Appendix B Proof of theorem 2.1

Our proof follows similar idea to what is discussed in section 4.2 in Li et
al. [14], except that we give the proof for general second order operator
here.

Recall that we are going to compute eq. 1 and we only need to analyze the
computation cost of its first term. In the following, we will use the notation
described in appendix A for computation graphs.The previous auto
differentiation method obtains $\mathcal{L}\phi({\mathbf{x}})$ through
computing the Hessian matrix and its inner product with $A$. It first performs
forward propagation to obtain the value of each variable $v^{j}$. Next, a
standard backward process is employed by creating a new computation graph
$\hat{\mathcal{G}}$ where node $\hat{v}^{i}\in\hat{\mathcal{G}}$ represents
the operation to calculate $\frac{\partial\phi}{\partial v^{i}}$,
$i=M,...,-N$. The associated computations are

$\frac{\partial\phi}{\partial v^{i}}=\sum\limits_{j:i\to j}\frac{\partial
F_{j}}{\partial v^{i}}\frac{\partial\phi}{\partial v^{j}},\ i=M-1,...,-N.$
(12)

The Hessian matrix is then obtained through the following forward-mode
Jacobian calculation along $\mathcal{G}$ and $\hat{\mathcal{G}}$,
respectively:

$\displaystyle\nabla v^{i}$ $\displaystyle=\sum\limits_{j:j\to
i}\frac{\partial F_{i}}{\partial v^{j}}\nabla v^{j},\ i=-N,...M\quad$ (13)
$\displaystyle\nabla\hat{v}^{i}=\nabla\frac{\partial\phi}{\partial v^{i}}=$
$\displaystyle\sum\limits_{\begin{subarray}{c}j,l\\\ i\to j\\\ l\to
j\end{subarray}}\frac{\partial^{2}F_{j}}{\partial v^{l}\partial
v^{i}}\frac{\partial\phi}{\partial v^{j}}\nabla v^{l}+\sum_{j:i\to
j}\frac{\partial F_{j}}{\partial v^{i}}\nabla\frac{\partial\phi}{\partial
v^{j}},\ i=M-1,...-N$ (14)

where we always use $\nabla$ to denote $\nabla_{\mathbf{x}}$ for simplicity.

The bottleneck of Hessian computation comes from eq. 13 and eq. 14.

For every connected nodes $(j\to i)$, it takes $N$ float multiplication to
yield $\frac{\partial F_{i}}{\partial v^{j}}\nabla v^{j}$. Therefore eq. 13
takes $N|E|$ floating point operations (FLOPs) if we only count
multiplications, where $E$ denotes the set of edges in $\mathcal{G}$. For eq.
14, the second term also takes $N|E|$ FLOPs. To calculate the computational
cost of the first term in eq. 14, we first introduce two notations, $T$ and
$R$, which are both sets of ordered tuples:

$\displaystyle T$ $\displaystyle=\\{(i,l,j)|i\to j,\ l\to j,\
\frac{\partial^{2}F_{j}}{\partial v^{i}\partial v^{l}}\neq 0\\},$ (15)
$\displaystyle R$ $\displaystyle=\\{(i,l)|\exists j\ s.t.\ (i,l,j)\in T\\}.$

The AutoDiff method sums over $j$ first to obtain $\sum\limits_{j:i\to j,l\to
j}\frac{\partial^{2}F_{j}}{\partial v^{l}\partial
v^{i}}\frac{\partial\phi}{\partial v^{j}}$ for all $(i,l)\in R$, and then sums
over $l$. For the first step, by leveraging the symmetry of Hessian matrix, it
spends 1 FLOPS for a pair of $(i,l,j)$ and $(l,i,j)$ in $T$. Thus it spends
$0.5|T|$ FLOPs in total. For the second step, for any $(i,l)\in R$, it takes
$N$ FLOPs to multiply a vector $\nabla v^{l}$ with a scalar. Thus this steps
takes $N|R|$ FLOPs in total.Consequently, the total FLOPs for the previous
method is about $N(|R|+2|E|)+0.5|T|$.

Next we analyze the computation cost of DOF. For readers’ convenience, we
repeat the propagation scheme here:

During the computation, for each node $v^{k}$, we maintain a tuple
$(v^{k},{\mathbf{g}}^{k},s^{k}):=(v^{k},L\nabla v^{k},\mathcal{L}v^{k})$. The
propagation rule of this tuple is:

$\displaystyle v^{j}$ $\displaystyle=F_{j}(\\{v^{i}:i\to j\\})$ (16)
$\displaystyle{\mathbf{g}}^{j}$ $\displaystyle=\sum\limits_{i:i\to
j}\frac{\partial F_{j}}{\partial v^{i}}{\mathbf{g}}^{i}$ (17) $\displaystyle
s^{j}$ $\displaystyle=\sum\limits_{\begin{subarray}{c}i,l\\\ i\to j\ l\to
j\end{subarray}}\frac{\partial^{2}F_{j}}{\partial v^{i}\partial
v^{l}}{\mathbf{g}^{i}}^{\top}D\mathbf{g}^{l}+\sum\limits_{i:i\to
j}\frac{\partial F_{j}}{\partial v^{i}}s^{i}$ (18)

For the proposed DOF method, we perform forward propagation along
$\mathcal{G}$ to obtain $\phi({\mathbf{x}})$, $L\nabla\phi({\mathbf{x}})$
(recall that $L$ is a matrix that comes from the decomposition of $A$) and
$\mathcal{L}\phi({\mathbf{x}})$. The second term $L\nabla\phi({\mathbf{x}})$
is calculated along $\mathcal{G}$ according to eq. 17. The third term, i.e.,
the target differential operator, is calculated along $\mathcal{G}$ according
to eq. 18.

The computational cost of the DOF method is dominated by eq. 17 and eq. 18. As
previously discussed, eq. 17 takes $N|E|$ FLOPs. For the first term in eq. 18,
we decompose its calculation into two steps. First, we compute
$\\{{\mathbf{g}}^{i\top}D{\mathbf{g}}^{l}\\}_{i\leq l,\ (i,l)\in R}$. Since
$D$ is a diagonal matrix only with diagonal element in $\\{0,\pm 1\\}$, this
computation takes $0.5\mathrm{rank}(D)|R|\leq 0.5N|R|$ FLOPs in total. Next,
following the topological order of $\mathcal{G}$, we sum over $i\leq l$ for
each $j$, deriving the first term in eq. 18. By leveraging the symmetry of the
Hessian matrix and Gram matrix, we reduce this computation by a factor of 2,
which is $0.5|T|$ FLOPs. The computational cost of the second term is
negligible compared with the first term since $s^{i}$ is a scalar.

Summing the computational cost of all the terms, we have that the DOF method
uses at most $0.5N(|R|+2|E|)+0.5|T|$ FLOPs.

In practice, a large percentage of operations are linear transformations, and
for any linear operation $F_{j}$, $\frac{\partial^{2}F_{j}}{\partial
v^{i}\partial v^{l}}=0$ for any $i\to j,\ l\to j$. This means the value $|T|$
is much smaller than $N|R|$ and $N|E|$. Thus, our method is about two times
faster than the previous Hessian-based methods for computing second order
differential operators of general neural network functions.

##  Appendix C Case Study for MLP

In this section, we follow the notation in appendix A for an MLP $\phi$ and
show that there is further speedup comparing with the general 2x result stated
in theorem 2.1.

For MLP, We could explicitly compute

$|E|=\sum_{l=0}^{L}N_{l}N_{l+1},\ |T|\leq\sum_{l=0}^{L}N_{l+1}N_{l}(N_{l}-1),\
|R|=\sum_{l=0}^{L}N_{l}(N_{l}-1).$ (19)

To compute the first term in eq. 18, note that

$\displaystyle\frac{\partial F_{k,i}}{\partial u^{k}_{j}}$
$\displaystyle=\sigma^{\prime}(W^{k}{\mathbf{u}}^{k}+b^{k})_{i}W^{k}_{ij}$
(20) $\displaystyle\frac{\partial^{2}F_{k,i}}{\partial u^{k}_{j}\partial
u^{k}_{l}}$
$\displaystyle=\sigma^{\prime\prime}(W^{k}{\mathbf{u}}^{k}+b^{k})_{i}W^{k}_{ij}W^{k}_{i,l},$
(21)

so we actually have

$\displaystyle\sum_{j,l}\frac{\partial^{2}F_{k,i}}{\partial u^{k}_{j}\partial
u^{k}_{l}}(L\nabla u^{k}_{j})^{\top}D(L\nabla u^{k}_{l})$ (22)
$\displaystyle=$ $\displaystyle(\frac{\sigma^{\prime\prime}}{\sigma^{\prime
2}}(W^{k}{\mathbf{u}}^{k}+b^{k}))_{i}(L\nabla u^{k+1}_{i})^{\top}D(L\nabla
u^{k+1}_{i}).$ (23)

This gives an alternative way to compute the first term in (18), whose total
computation cost is reduced from $\mathrm{r}(D)|R|$ to
$\mathrm{r}(D)\sum_{l=0}^{L}N_{l+1}$ FLOPs, where $\mathrm{r}$ stands for
$rank$.

Recall that the first term of (18) is one of the dominant calculations, this
manner certainly boost the efficiency.

##  Appendix D Proof of theorem 2.2

We focus on the peak memory usage caused by the forward-mode Jacobian
calculation, i.e., the storage of $\nabla v^{i}$(or ${\mathbf{g}}^{i}$), which
usually dominates the memory cost. In a typical forward-mode Jacobian
calculation, we write $\nabla v^{i}$ into memory when the algorithm is applied
to $v_{i}$ and remove it from memory when all the direct subnodes of $v^{i}$
have been computed.

We denote

$\tau(i):=\max\\{j:i\to j\\}.$ (24)

Then at the moment our DOF forward propagation comes to the node $v^{j}$, the
memory consumption is

$C(j):=N\sum_{i:i\leq j\leq\tau(i)}1.$ (25)

If we further denote $\mathcal{M}_{1}$ to be the peak memory usage when
executing the forward-mode DOF, we have

$\mathcal{M}_{1}=\max\limits_{j}C(j).$ (26)

From eq. 25 we clearly see that $\mathcal{M}_{1}\leq N|V|$, here $V$ is the
set of nodes in the computation graph $\mathcal{G}$. Furthermore, for any
$j<M,\ C(j)\leq N(|V|-1)$. As a result, $\mathcal{M}_{1}=N|V|$ if and only if
every nodes is pointing to the end node $v^{M}$, which corresponds to either
one-layer linear model or a multivariate elementary function taking in input
${\mathbf{x}}$ and directly giving the output. These extreme neural network
functions rarely occur in deep learning literature.

By contrast, in the Hessian calculation in AutoDiff, the computation graph
becomes a combination of $\mathcal{G}$ and $\hat{\mathcal{G}}$. As suggested
by eq. 14, the node $\hat{v}^{i}$ is a direct subnode of $v^{i}$.
Consequently, at the moment when $\nabla v^{M}$ is written into memory, every
$\nabla v^{i}$ for $v^{i}\in\mathcal{G}$ have been written into memory and
could not be released since $\hat{v}^{i}$ have not been computed yet. As a
result, the peak memory usage(denoted as $\mathcal{M}_{2}$) is strictly larger
than $N|V|$, and therefore larger than the peak memory of Forward Laplacian.

In the specific case when $\phi$ is MLP. For any node $u_{i}^{k},\
k\in\frac{1}{2}\mathbb{N}$,

$\tau(\\#u_{i}^{k})=\begin{cases}\\#u^{k+0.5}_{N_{k+1}},\ k=l\\\
\\#u^{k+0.5}_{i},\ k=l+0.5,\ \ l\in\mathbb{N},\end{cases}$ (27)

where $\\#$ denotes the order index of the node representing certain neurons
in $\mathcal{G}$. Thus,

$\frac{1}{N}C(\\#u_{i}^{k})=\begin{cases}N_{l}+1,\ k=l\\\ N_{l-1}+i,\
k=l+0.5,\ \ l\in\mathbb{N}.\end{cases}$ (28)

Thus we conclude that $\mathcal{M}_{1}\leq
N\max\limits_{l}N_{l}+N_{l-1}\lesssim
N\frac{2}{L}\sum_{l}N_{l}=N\frac{2}{L}|V|\leq\frac{2}{L}\mathcal{M}_{2}.$ ∎

##  Appendix E Experiments Settings

#### Hardware.

All the results are evaluated on a single NVIDIA Tesla V100 GPU.

#### Network structure.

We use the MLP and MLP with Jacobian sparsity to benchmark different methods.
The MLP with Jacobian sparsity means we split each data into some small blocks
${\mathbf{x}}=({\mathbf{x}}_{1},{\mathbf{x}}_{2},...{\mathbf{x}}_{k})$ and
independently operate each block with an MLP. The output is a sum of product
of each MLP output, i.e.

$\text{output}=\sum_{d}\prod_{i=1}^{k}[\text{MLP}^{i}(\mathbf{x}_{i})]_{d}.$

Here $i$ refers to the index of block and $d$ refers to the index of the
output in each MLP. Detailed hyperparameters can be found in table 3.

Table 3: Hyperparameters |  MLP |  MLP with Jacobian sparsity
---|---|---
hidden dimension |  256 |  256
input dimension |  64 |  64
#layer |  8 |  8
#blocks |  - |  16
output dimension for each MLP |  - |  8

#### Coefficient matrix used for each experiments.

For the MLP structure, the coefficient matrix $A$ is a $64\times 64$ matrix.
For the MLP with Jacbobian sparsity, the coefficient matrix $A$ is a $64\times
64$ block diagonal matrix. The coefficient matrices are listed in table 4.
Here, $(\alpha_{ij})_{i,j}$ is a $64\times 64$ matrix and
$(\sigma_{ij})_{i,j}$ is a $4\times 4$ matrix. Both $\alpha_{ij}$ and
$\sigma_{ij}$ are drawn from a standard normal distribution.
$(\delta_{ij})_{i,j}$ is the identity matrix such that $\delta_{ij}=0\text{ if
}i\neq j\text{ else }1$. $s_{i}=-1\text{ if }i=0\text{ else }1$. .

Table 4: Coefficient matrix Structure |  Elliptic |  Low-rank |  general
---|---|---|---
MLP |  $a_{ij}=\sum_{k=1}^{64}\alpha_{ik}\alpha_{jk}$ |  $a_{ij}=\sum_{k=1}^{32}\alpha_{ik}\alpha_{jk}$ |  $a_{ij}=\delta_{ij}s_{i}$
MLP with sparsity |  $a_{il,jm}=\delta_{lm}\sum_{k=1}^{4}\sigma_{ik}\sigma_{jk}$ |  $a_{il,jm}=\delta_{lm}\sum_{k=1}^{2}\sigma_{ik}\sigma_{jk}$ |  $a_{il,jm}=\delta_{lm}\delta_{ij}s_{i}$