Stochastic Taylor Derivative Estimator: Efficient amortization for arbitrary differential operators

随机泰勒导数估计器：任意微分算子的高效摊销

Zekun Shi Affiliation:  National University of Singapore Affiliation:  Sea AI Lab Email:  shizk@sea.com ,    Zheyuan Hu Affiliation:  National University of Singapore Email:  e0792494@u.nus.edu ,    Min Lin Affiliation:  Sea AI Lab Email:  linmin@sea.com ,    Kenji Kawaguchi Affiliation:  National University of Singapore Email:  kenji@nus.edu.sg

Abstract

摘要

Optimizing neural networks with loss that contain high-dimensional and high-order differential operators is expensive to evaluate with back-propagation due to $\mathcal{O}(d^{k})$ scaling of the derivative tensor size and the $\mathcal{O}(2^{k-1}L)$ scaling in the computation graph, where $d$ is the dimension of the domain, $L$ is the number of ops in the forward computation graph, and $k$ is the derivative order. In previous works, the polynomial scaling in $d$ was addressed by amortizing the computation over the optimization process via randomization. Separately, the exponential scaling in $k$ for univariate functions ($d=1$) was addressed with high-order auto-differentiation (AD). In this work, we show how to efficiently perform arbitrary contraction of the derivative tensor of arbitrary order for multivariate functions, by properly constructing the input tangents to univariate high-order AD, which can be used to efficiently randomize any differential operator. When applied to Physics-Informed Neural Networks (PINNs), our method provides >1000$\times$ speed-up and >30$\times$ memory reduction over randomization with first-order AD, and we can now solve 1-million-dimensional PDEs in 8 minutes on a single NVIDIA A100 GPU[^1]. This work opens the possibility of using high-order differential operators in large-scale problems.

若损失函数含高维、高阶微分算子，用反向传播求值代价很高：导数张量规模按 $\mathcal{O}(d^{k})$ 增长，计算图规模按 $\mathcal{O}(2^{k-1}L)$ 增长，其中 $d$ 是定义域维数，$L$ 是前向计算图中的运算数，$k$ 是导数阶数。先前工作通过随机化把计算摊销到优化过程中，缓解了 $d$ 上的多项式增长；对单变量函数（$d=1$）则用高阶自动微分（AD）应对 $k$ 上的指数增长。本文说明：只要恰当构造一元高阶 AD 的输入切向量，就能对多元函数任意阶导数张量做任意收缩，从而高效随机化任意微分算子。用于物理信息神经网络（PINN）时，相对一阶 AD 的随机化，本方法加速超过 1000 倍、显存降低超过 30 倍，并可在单块 NVIDIA A100 上于 8 分钟内求解百万维 PDE[^1]。这为在大规模问题中使用高阶微分算子打开了可能。

1 Introduction

1 引言

In many problems, especially in Physics-informed machine learning [19, 33], one needs to solve optimization problems where the loss contains differential operators:

在许多问题中，尤其是物理信息机器学习 [19, 33]，需要求解损失里含微分算子的优化问题：

$$\argmin_{\theta}f(\mathbf{x},u_{\theta}(\mathbf{x}),\mathcal{D}^{\alpha^{(1)}}u_{\theta}(\mathbf{x}),\dots,\mathcal{D}^{\alpha^{(n)}}u_{\theta}(\mathbf{x})),\quad u_{\theta}:\mathbb{R}^{d}\to\mathbb{R}^{d^{\prime}}. \tag{1}$$

In this above, $\mathcal{D}^{\alpha}=\frac{\partial^{\absolutevalue{\alpha}}}{\partial x_{1}^{\alpha_{1}},\dots,\partial x_{d}^{\alpha_{d}}}$, $\alpha=(\alpha_{1},\alpha_{2},\dots,\alpha_{d})$ is a multi-index, $u_{\theta}$ is some neural network parameterized by $\theta$, and $f$ is some cost function. When either the differentiation order $k$ or the dimensionality $d$ is high, the objective function above is expensive to evaluate with back-propagation (backward mode AD) in both memory and computation: the size of the derivative tensor has scaling $\order{d^{k}}$, and the size of the computation graph has scaling $\order{2^{k-1}L}$ where $L$ is the number of ops in the forward computation graph.

上式中 $\mathcal{D}^{\alpha}=\frac{\partial^{\absolutevalue{\alpha}}}{\partial x_{1}^{\alpha_{1}},\dots,\partial x_{d}^{\alpha_{d}}}$，$\alpha=(\alpha_{1},\alpha_{2},\dots,\alpha_{d})$ 是多重指标，$u_{\theta}$ 是由 $\theta$ 参数化的神经网络，$f$ 是代价函数。当微分阶数 $k$ 或维数 $d$ 很高时，用反向传播（反向模式 AD）求目标函数在显存与计算上都很贵：导数张量规模为 $\order{d^{k}}$，计算图规模为 $\order{2^{k-1}L}$，$L$ 是前向图中的运算数。

There have been several efforts to tackle this curse of dimensionality. One line of work uses randomization to amortize the cost of computing differential operators with AD over the optimization process so that the $d$ in the above scaling becomes a constant for the case of $k=2$. Stochastic Dimension Gradient Descent (SDGD) [13] randomizes over the input dimensions where in each iteration, the partial derivatives are only calculated for a minibatch of sampled dimensions with back-propagation. In [12, 21, 15], the classical technique of Hutchinson Trace Estimator (HTE) [16] is used to estimate the trace of Hessian or Jacobian to inputs. Others choose to bypass AD completely to reduce the complexity of computation. In [31], the finite difference method is used for estimating the Hessian trace. Randomized smoothing [11, 14] uses the expectation over Gaussian random variable as ansatz, so that its derivatives can be expressed as another expectation Gaussian random variable via Stein’s identity [39]. However, compared to AD, the accuracy of these methods is highly dependent on the choice of discretization.

已有若干工作试图打破维数灾难。一类做法用随机化，把用 AD 算微分算子的代价摊到优化过程中，使 $k=2$ 时上述缩放里的 $d$ 变成常数。随机维梯度下降（SDGD）[13] 对输入维做随机采样，每轮只用反向传播算一小批维上的偏导。[12, 21, 15] 用经典 Hutchinson 迹估计（HTE）[16] 估计 Hessian 或对输入的 Jacobian 的迹。另一些工作干脆绕过 AD 以降复杂度：[31] 用有限差分估 Hessian 迹；随机平滑 [11, 14] 以高斯期望作 ansatz，再经 Stein 恒等式 [39] 把导数写成另一高斯期望。但相对 AD，这些方法的精度高度依赖离散化选择。

In this work, we address the scaling issue in both $d$ and $k$ for the optimization problem in Eq. 1 at the same time, by proposing an amortization scheme that can be efficiently evaluated via high-order AD, which we call Stochastic Taylor Derivative Estimator (STDE). Our main contributions are:

本文同时应对式 (1) 中 $d$ 与 $k$ 两边的缩放问题：提出可经高阶 AD 高效求值的摊销方案，称为随机泰勒导数估计器（STDE）。主要贡献如下：

• • We demonstrate how Taylor mode AD [6], a high-order AD method, can be used to amortize the optimization problem in Eq. 1. Specifically, we show that, with properly constructed input tangents, the univariate Taylor mode can be used to contract multivariate functions’ derivative tensor of arbitrary order;
• • We provide a comprehensive procedure for randomizing arbitrary differential operators with STDE, while previous works mainly focus on the Laplacian operator, and we provide abundant examples of STDE constructed for operators in common PDEs;
• • STDE encompass and generalizes previous methods like SDGD [13] and HTE [16, 12]. We also prove that HTE-type estimator cannot be generalized beyond fourth order differential operator;
• • We determine the efficacy of STDE experimentally. When applied to PINN, our method provides a significant speed-up compared to the baseline method SDGD [13] and the backward-free method like random smoothing [11]. Due to STDE’s low memory requirements and reduced computation complexity, PINNs with STDE can solve 1-million-dimensional PDEs on a single NVIDIA A100 40GB GPU within 8 minutes, which shows that PINNs have the potential to solve complex real-world problems that can be modeled as high-dimensional PDEs. We also provide a detailed ablation study on the source of performance gain of our method.

• • 说明如何用高阶 AD 方法——Taylor 模式 AD [6]——摊销式 (1) 的优化；具体地，在恰当构造输入切向量后，一元 Taylor 模式可用来收缩多元函数任意阶导数张量；
• • 给出用 STDE 随机化任意微分算子的完整流程（先前工作多集中在 Laplace 算子），并给出常见 PDE 算子上的大量 STDE 构造例子；
• • STDE 涵盖并推广 SDGD [13] 与 HTE [16, 12]；我们还证明 HTE 型估计器无法推广到四阶以上微分算子；
• • 用实验验证 STDE 的效力。用于 PINN 时，相对基线 SDGD [13] 以及无反向的随机平滑 [11] 有显著加速；因显存需求低、计算复杂度降，带 STDE 的 PINN 可在单块 NVIDIA A100 40GB 上 8 分钟内求解百万维 PDE，表明 PINN 有望处理可建模为高维 PDE 的复杂现实问题。我们也对性能增益来源做了细消融。

2 Related works

2 相关工作

High-order and forward mode AD

The idea of generalizing forward mode AD to high-order derivatives has existed in the AD community for a long time [5, 18, 40, 22]. However, accessible implementation for machine learning was not available until the recent implementation in JAX [6, 7], which implemented the Taylor mode AD for accelerating ODE solver. There are also efforts in creating the forward rule for a specific operator like the Laplacian [24, 23]. Randomization over the linearized part of the AD computation graph was considered in [30]. Forward mode AD can also be used to compute neural network parameter gradient as shown in [2].

高阶与前向模式 AD

把前向模式 AD 推广到高阶导数，在 AD 社区早已有之 [5, 18, 40, 22]。但面向机器学习、可直接用的实现直到近期 JAX [6, 7] 才出现——其为加速 ODE 求解实现了 Taylor 模式 AD。也有工作为 Laplace 等特定算子写前向规则 [24, 23]。[30] 考虑对 AD 计算图线性化部分做随机化。前向模式 AD 也可用来算神经网络参数梯度，见 [2]。

Randomized Gradient Estimation

Randomization [28, 29, 8] is a common technique for tackling the curse of dimensionality for numerical linear algebra computation, which can be applied naturally in amortized optimization [1]. Hutchinson trace estimator [16] is a well-known technique, which has been applied to diffusion model [37] and PINNs [12]. Another case that requires gradient estimation is when the analytical form of the target function is not available (black box), which means AD cannot be applied. The method of zeroth-order optimization [25] can be used in this case, as it only requires evaluating the function at arbitrary input. It is also useful when the function is very complicated like in the case of a large language model [27].

随机梯度估计

随机化 [28, 29, 8] 是数值线性代数里对付维数灾难的常用手段，自然可用于摊销优化 [1]。Hutchinson 迹估计 [16] 广为人知，已用于扩散模型 [37] 与 PINN [12]。另一类需要梯度估计的情形是目标函数解析式不可得（黑盒），因而无法用 AD；此时可用零阶优化 [25]，只要求在任意输入上求函数值。对大语言模型这类极复杂函数也有用 [27]。

3 Preliminaries and discussions

3 预备知识与讨论

3.1 First-order auto-differentiation (AD)

3.1 一阶自动微分（AD）

AD is a technique for evaluating the gradient of composition of known analytical functions commonly called primitives. In an AD framework, a neural network $F_{\theta}:\mathbb{R}^{d}\to\mathbb{R}^{d^{\prime}}$ is constructed as the composition of primitives $F_{i}$ that are parameterized by some parameters $\theta_{i}$. In this section, we will consider the neural networks with linear computation graphs like $F=F_{L}\circ F_{L-1}\circ\dots\circ F_{1}$, but the results generalize to arbitrary directed acyclic graphs (DAGs). We will assume that all hidden dimensions are $h$. See Appendix B for more details on first-order AD.

AD 用于求由已知解析函数（常称原语）复合而成的函数的梯度。在 AD 框架中，神经网络 $F_{\theta}:\mathbb{R}^{d}\to\mathbb{R}^{d^{\prime}}$ 由带参数 $\theta_{i}$ 的原语 $F_{i}$ 复合而成。本节先考虑线性计算图 $F=F_{L}\circ F_{L-1}\circ\dots\circ F_{1}$，结论可推广到任意有向无环图（DAG）。假设所有隐层维数为 $h$。一阶 AD 更多细节见附录 B。

Forward mode AD

Each primitives $F_{i}$ is linearized as the Fréchet (directional) derivative $\partial F_{i}:\mathbb{R}^{h}\to\mathrm{L}(\mathbb{R}^{h},\mathbb{R}^{h})$, which computes the Jacobian-vector-product (JVP): $\partial F_{i}(\mathbf{a})(\mathbf{v})=\evaluated{\pdv{F}{\vb{x}}}_{\vb{a}}\mathbf{v}$, where $\mathbf{a}$ is referred to as the primal and $\mathbf{v}$ the tangent. $\partial F_{i}$ form a linearized computation graph (third row in Fig. 3), that computes the JVP of the composition $\partialderivative{F}{\vb{x}}\mathbf{v}$:

前向模式 AD

每个原语 $F_{i}$ 被线性化为 Fréchet（方向）导数 $\partial F_{i}:\mathbb{R}^{h}\to\mathrm{L}(\mathbb{R}^{h},\mathbb{R}^{h})$，计算 Jacobian–向量积（JVP）：$\partial F_{i}(\mathbf{a})(\mathbf{v})=\evaluated{\pdv{F}{\vb{x}}}_{\vb{a}}\mathbf{v}$，其中 $\mathbf{a}$ 称主值，$\mathbf{v}$ 称切向量。$\partial F_{i}$ 构成线性化计算图（图 3 第三行），用于计算复合的 JVP $\partialderivative{F}{\vb{x}}\mathbf{v}$：

$$\partialderivative{F}{\vb{x}}\mathbf{v}=\partial F(\mathbf{x})(\mathbf{v})=[\partial F_{L}\circ\partial F_{L-1}\circ\dots\circ\partial F_{1}](\mathbf{x})(\mathbf{v}). \tag{2}$$

By setting the tangent to $\mathbf{v}$ one of the standard basis of $\mathbb{R}^{d}$, JVP computes one column of the Jacobian $D_{F}$, so the full Jacobian can be computed with $d$ JVPs. Each JVP call requires $\order{\max(d,h)}$ memory as only the current activation $\mathbf{y}_{i}$ and tangent $\mathbf{v}_{i}$ are needed to carry out the computation, and the computation complexity is usually in the same order as the forward computation graph. In the case of MLP, both the forward and the linearized graph have a complexity of $\order{dh+(L-1)h^{2}}$.

把切向量 $\mathbf{v}$ 设为 $\mathbb{R}^{d}$ 的某个标准基向量时，JVP 给出 Jacobian $D_{F}$ 的一列，故完整 Jacobian 需 $d$ 次 JVP。每次 JVP 只需 $\order{\max(d,h)}$ 显存——只需当前激活 $\mathbf{y}_{i}$ 与切向量 $\mathbf{v}_{i}$；计算复杂度通常与前向图同阶。对 MLP，前向图与线性化图复杂度均为 $\order{dh+(L-1)h^{2}}$。

Backward mode AD

Each primitives $F_{i}$ is linearized as the adjoint of the Fréchet derivative $\partial^{\top}F_{i}$ instead, which computes the vector-Jacobian-product (VJP): $\partial^{\top}F_{i}(\mathbf{a})(\mathbf{v}^{\top})=\mathbf{v}^{\top}\evaluated{\pdv{F}{\vb{x}}}_{\vb{a}}$ where $\mathbf{v}^{\top}$ is the cotangent. The linearized computation graph now runs in the reverse order:

反向模式 AD

此时每个原语 $F_{i}$ 线性化为 Fréchet 导数的伴随 $\partial^{\top}F_{i}$，计算向量–Jacobian 积（VJP）：$\partial^{\top}F_{i}(\mathbf{a})(\mathbf{v}^{\top})=\mathbf{v}^{\top}\evaluated{\pdv{F}{\vb{x}}}_{\vb{a}}$，$\mathbf{v}^{\top}$ 为余切。线性化计算图按相反顺序运行：

$$\mathbf{v}^{\top}\partialderivative{F}{\vb{x}}=\partial^{\top}F(\mathbf{x})(\mathbf{v}^{\top})=[\partial^{\top}F_{1}(\mathbf{x})\circ\dots\circ\partial^{\top}F_{L-1}(\mathbf{y}_{L-2})\circ\partial^{\top}F_{L}(\mathbf{y}_{L-1})](\mathbf{v}^{\top}), \tag{3}$$

which is also clear from Fig. 3. Furthermore, due to this reversion, we first need to do a forward pass to obtain the evaluation trace $\{\mathbf{y}_{i}\}_{i=1}^{L}$ before we can invoke the VJPs $\partial^{\top}F_{i}$, which apparent as shown in Eq. 3. Hence the number of sequential computations is twice as much compared to forward mode. The memory requirement becomes $\order{d+(L-1)h}$ as we need to store the entire evaluation trace. Similar to JVP, VJP computes one row of $J_{F}$ at a time, so the full Jacobian $\partialderivative{F}{\vb{x}}$ can be computed using $d^{\prime}$ VJPs. When optimizing scalar cost functions $\ell(\theta):\mathbb{R}^{n}\to\mathbb{R}$ of the network parameters $\theta$, backward mode efficiently trades off memory with computation complexity as $d^{\prime}=1$ and only $1$ VJP is needed to get the full Jacobian. Furthermore, all parameter $\theta_{i}$ can use the same cotangent $\mathbf{v}^{\top}$, whereas with forward mode, separate tangent for each parameter $\theta_{i}$ is needed.

图 3 也清楚显示这一点。由于顺序反转，须先前向得到求值轨迹 $\{\mathbf{y}_{i}\}_{i=1}^{L}$，才能调用 VJP $\partial^{\top}F_{i}$，见式 (3)。因此串行计算次数约为前向模式的两倍。显存变为 $\order{d+(L-1)h}$，因需存整条求值轨迹。与 JVP 类似，VJP 每次给出 $J_{F}$ 的一行，完整 Jacobian $\partialderivative{F}{\vb{x}}$ 需 $d^{\prime}$ 次 VJP。优化对参数 $\theta$ 的标量代价 $\ell(\theta):\mathbb{R}^{n}\to\mathbb{R}$ 时，反向模式在显存与计算间更划算：$d^{\prime}=1$，一次 VJP 即得完整 Jacobian；且所有 $\theta_{i}$ 可共用同一余切 $\mathbf{v}^{\top}$，而前向模式则需为每个 $\theta_{i}$ 单独准备切向量。

3.2 Inefficiency of the first-order AD for high-order derivative on inputs

3.2 一阶 AD 对输入高阶导数的低效

Figure 1: The computation graph of computing second order gradient by repeated application of backward mode AD, for a function $F(\cdot)$ with $4$ primitives ($L=4$), which computes the Hessian-vector-product. Red nodes represent the cotangent nodes in the second backward pass. With each repeated application of VJP the length of sequential computation doubles.

图 1：对含 4 个原语（$L=4$）的函数 $F(\cdot)$，反复用反向模式 AD 求二阶梯度（即 Hessian–向量积）的计算图。红节点表示第二次反向中的余切。每多做一次 VJP，串行计算长度翻倍。

High-order input derivatives $\partialderivative[k]{u_{\theta}}{\vb{x}}$ for scalar $u_{\theta}$ can be implemented as repeated applications of first-order AD, but this approach will exhibit fundamental inefficiency that cannot be remedied by randomization.

标量 $u_{\theta}$ 对输入的高阶导数 $\partialderivative[k]{u_{\theta}}{\vb{x}}$ 可通过反复套用一阶 AD 实现，但这种方式有根本性低效，随机化也无法消除。

Repeating backward mode AD

With each repeated application of backward mode AD, the new evaluation trace will include the cotangents from the previous application of backward AD, so the length of sequential computation doubles. Furthermore, the size of the cotangent also grows by $d$ times. Therefore applying backward mode AD has additional memory cost of $\order{d+(L-1)h}$ and additional computation cost of $\order{2dh+2(L-1)h^{2}}$, which is clear from Fig. 1. In general, with $k$ repeated applications of backward mode AD will incur $\order{2^{k-1}(d+(L-1)h)}$ memory cost and $\order{2^{k}(dh+(L-1)h^{2})}$ computation cost. And $\order{d^{k-1}}$ calls are needed to evaluate the entire derivative tensor. So both memory and compute scale exponentially in derivative order $k$

反复反向模式 AD

每多做一次反向模式 AD，新的求值轨迹会并入上一次反向的余切，串行计算长度翻倍；余切规模也按 $d$ 倍增长。因而每次额外反向带来 $\order{d+(L-1)h}$ 显存与 $\order{2dh+2(L-1)h^{2}}$ 计算，见图 1。一般地，$k$ 次反复反向的显存为 $\order{2^{k-1}(d+(L-1)h)}$，计算为 $\order{2^{k}(dh+(L-1)h^{2})}$；完整导数张量还需 $\order{d^{k-1}}$ 次调用。故显存与计算都对导数阶数 $k$ 指数增长。

Repeating forward mode AD

Consider $u_{\theta}:\mathbb{R}^{d}\to\mathbb{R}$. The input tangent dimension is $d$ on the first application of forward mode AD, but on the second application, it will become $d\times d$ since we are now computing the forward mode AD for $\nabla u_{\theta}:\mathbb{R}^{d}\to\mathbb{R}^{d}$. So the size of the input tangent with $k$ repeated application is $\order{d^{k}}$, so it grows exponentially. This is also inefficient.

反复前向模式 AD

考虑 $u_{\theta}:\mathbb{R}^{d}\to\mathbb{R}$。第一次前向时输入切向量维数为 $d$；第二次则变成 $d\times d$，因为此时是对 $\nabla u_{\theta}:\mathbb{R}^{d}\to\mathbb{R}^{d}$ 做前向。$k$ 次反复后输入切向量规模为 $\order{d^{k}}$，同样指数增长，同样低效。

Mixed mode AD schemes are also likely inefficient

See more detail in Appendix C.

混合模式 AD 方案多半也不高效

详见附录 C。

3.3 Stochastic Dimension Gradient Descent

3.3 随机维梯度下降

SDGD [13] amortizes high-dimensional differential operators by computing only a minibatch of derivatives in each iteration. It replaces a differential operator $\mathcal{D}$ with a randomly sampled subset of additive terms, where each term only depends on a few input dimensions

SDGD [13] 通过每轮只算一小批导数，摊销高维微分算子。它把微分算子 $\mathcal{D}$ 换成随机采样的若干加性项，每项只依赖少数输入维

$$\mathcal{D}:=\sum_{j=1}^{N_{\mathcal{D}}}\mathcal{D}_{j}\approx\frac{N_{\mathcal{D}}}{\absolutevalue{J}}\sum_{j\in J}\mathcal{D}_{j}:=\tilde{\mathcal{D}_{J}}, \tag{4}$$

where $\tilde{\mathcal{D}_{J}}$ denotes the SDGD operator that approximates the true operator $\mathcal{D}$, $J$ is the sampled index set, and $\absolutevalue{J}$ is the batch size. For example, in $d$-dimensional Poisson equation, $N_{\mathcal{D}}=d$, $\mathcal{D}=\sum_{j=1}^{d}\partialderivative[2]{x_{j}}$, and the additive terms are $\mathcal{D}_{j}=\partialderivative[2]{x_{j}}$.

其中 $\tilde{\mathcal{D}_{J}}$ 是逼近真算子 $\mathcal{D}$ 的 SDGD 算子，$J$ 是采样指标集，$\absolutevalue{J}$ 是批大小。例如 $d$ 维 Poisson 方程中 $N_{\mathcal{D}}=d$，$\mathcal{D}=\sum_{j=1}^{d}\partialderivative[2]{x_{j}}$，加性项为 $\mathcal{D}_{j}=\partialderivative[2]{x_{j}}$。

$\tilde{\mathcal{D}_{J}}$ are cheaper to compute than $\mathcal{D}$ due to reduced dimensionality: for each sampled index, by treating all other input as constant we get a function with scalar input and output. For a given index set $J$, the memory requirements are reduced from $\order{2^{k-1}(d+(L-1)h)}$ to $\order{\abs{J} (2^{k-1}(1+(L-1)h))}$, and the computation complexity reduces to $\order{\abs{J}2^{k}(h+(L-1)h^{2})}$. This reduction is significant when $d\gg h$ as in the experimental setting of SDGD [13], but the exponential scaling in $k$ persists.

因维数降低，$\tilde{\mathcal{D}_{J}}$ 比 $\mathcal{D}$ 更便宜：对每个采样指标，把其余输入当常数，得到标量入、标量出函数。给定 $J$，显存从 $\order{2^{k-1}(d+(L-1)h)}$ 降到 $\order{\abs{J} (2^{k-1}(1+(L-1)h))}$，计算降到 $\order{\abs{J}2^{k}(h+(L-1)h^{2})}$。在 SDGD [13] 实验设定 $d\gg h$ 时这一降幅显著，但 $k$ 上的指数缩放仍在。

3.4 Univariate Taylor mode AD

3.4 一元 Taylor 模式 AD

One way to define high-order AD is by determining how the high-order Taylor expansion of a univariate function changes when mapped by primitives. Firstly, the Fréchet derivative $\partial F$ can be rewritten to operate on a space curve $g:\mathbb{R}\to\mathbb{R}^{d}$ that passes through the primal $\mathbf{a}$, i.e. $g(t)=\mathbf{a}$, and has tangent $g^{\prime}(t)=\mathbf{v}$:

定义高阶 AD 的一种方式是：看一元函数的高阶 Taylor 展开在经原语映射后如何变化。首先，Fréchet 导数 $\partial F$ 可改写为作用在穿过主值 $\mathbf{a}$ 的空间曲线 $g:\mathbb{R}\to\mathbb{R}^{d}$ 上，即 $g(t)=\mathbf{a}$，且切向量 $g^{\prime}(t)=\mathbf{v}$：

$$\partial F(g(t))(g^{\prime}(t))=\evaluated{\pdv{F}{\vb{x}}}_{\vb{x}=g(t)}g^{\prime}(t)=\derivative{t}[F\circ g](t). \tag{5}$$

This shows that the $\partial$ (JVP) is the same as the univariate chain rule. The tuple $J_{g}(t):=(g(t),g^{\prime}(t))$ can be thought of as the first-order expansion of $g$ which lives in the tangent bundle of $F$. Treating $F$ as the smooth map between manifolds, we can define the pushforward $\differential F$ which pushes the first order expansion of $g$ (i.e. $J_{g}(t)$) forward to the first order expansion of $F\circ g$ (i.e. $J_{F\circ g}(t)$):

这说明 $\partial$（JVP）与一元链式法则相同。元组 $J_{g}(t):=(g(t),g^{\prime}(t))$ 可视为 $g$ 的一阶展开，落在 $F$ 的切丛上。把 $F$ 当作流形间光滑映射，可定义前推 $\differential F$，把 $g$ 的一阶展开 $J_{g}(t)$ 前推到 $F\circ g$ 的一阶展开 $J_{F\circ g}(t)$：

$$\differential F(J_{g}(t))=J_{F\circ g}(t)=\left([F\circ g](t),\derivative{t}[F\circ g](t)\right)=(F(\mathbf{a}),\partial F(\mathbf{a})(\mathbf{v})). \tag{6}$$

Naturally, to extend this to higher orders, one can consider the $k$th order expansion of the input curve $g$, which is equivalent to the tuple $J_{g}^{k}(t):=(g(t),g^{\prime}(t),g^{\prime\prime}(t),\dots,g^{(k)}(t))=(\mathbf{a},\mathbf{v}^{(1)},\mathbf{v}^{(2)},\dots,\mathbf{v}^{k})$ known as the $k$-jet of $g$ where $\mathbf{v}^{j}$ is called the $j$th order tangent of $g$. $J_{g}^{k}$ lives in the $k$th order tangent bundle of $F$, and we can define the $k$th-order pushforward $\differential^{k}F$ :

自然地，推广到更高阶可考虑输入曲线 $g$ 的 $k$ 阶展开，即元组 $J_{g}^{k}(t):=(g(t),g^{\prime}(t),g^{\prime\prime}(t),\dots,g^{(k)}(t))=(\mathbf{a},\mathbf{v}^{(1)},\mathbf{v}^{(2)},\dots,\mathbf{v}^{k})$，称为 $g$ 的 $k$-jet，其中 $\mathbf{v}^{j}$ 是 $g$ 的 $j$ 阶切向量。$J_{g}^{k}$ 落在 $F$ 的 $k$ 阶切丛上，可定义 $k$ 阶前推 $\differential^{k}F$：

$$\displaystyle\differential^{k}F(J_{g}^{k}(t))= \\ \displaystyle J_{F\circ g}^{k}(t)=\left([F\circ g](t),\partialderivative{t}[F\circ g](t),\partialderivative[2]{t}[F\circ g](t),\dots,\partialderivative[k]{t}[F\circ g](t)\right) \\ \displaystyle= \\ \displaystyle(F(\mathbf{a}),\partial F(\mathbf{a})(\mathbf{v}^{(1)}),\partial^{2}F(\mathbf{a})(\mathbf{v}^{(1)},\mathbf{v}^{(2)}),\dots,\partial^{k}F(\mathbf{a})(\mathbf{v}^{(1)},\dots,\mathbf{v}^{(k)})),$$

which pushes the $k$th order expansion of $g$ (i.e. $J_{g}^{k}$) forward to the $k$th order expansion of $F\circ g$ (i.e. $J_{F\circ g}^{k}$). $\partial^{k}F=\partialderivative[k]{t}[F\circ g](t)$ is the $k$-th order Fréchet derivative, whose analytical formula is given by the high-order univariate chain rule known as the Faa di Bruno’s formula (Eq. 43).

它把 $g$ 的 $k$ 阶展开 $J_{g}^{k}$ 前推到 $F\circ g$ 的 $k$ 阶展开 $J_{F\circ g}^{k}$。$\partial^{k}F=\partialderivative[k]{t}[F\circ g](t)$ 是 $k$ 阶 Fréchet 导数，其解析式由一元高阶链式法则——Faa di Bruno 公式（式 43）——给出。

Since $J_{g}^{k}$ contains all information needed to evaluate $\partialderivative[j]{t}[F\circ g](t)$ for any $j\leq k$, the map $\differential^{k}F$ is well-defined. $\differential^{k}$ defines a high-order AD: we can compute $\differential^{k}F$ of arbitrary composition $F$ from the $k$th-order pushforward of the primitives $\differential^{k}F_{i}$, since $\differential^{k}$ is an homomorphism of the group $(\{F_{i}\},\circ)$:

因 $J_{g}^{k}$ 含计算任意 $j\leq k$ 的 $\partialderivative[j]{t}[F\circ g](t)$ 所需全部信息，$\differential^{k}F$ 良定义。$\differential^{k}$ 定义了高阶 AD：任意复合 $F$ 的 $\differential^{k}F$ 可由原语的 $k$ 阶前推 $\differential^{k}F_{i}$ 复合得到，因为 $\differential^{k}$ 是群 $(\{F_{i}\},\circ)$ 的同态：

$$\differential^{k}[F_{2}\circ F_{1}](J_{g}^{k}(t))=J_{F_{2}\circ F_{1}\circ g}^{k}(t)=\differential^{k}F_{2}(J_{F_{1}\circ g}^{k}(t))=[\differential^{k}F_{2}\circ\differential^{k}F_{1}](J_{g}^{k}(t)). \tag{8}$$

This approach of composing $\differential^{k}$ of primitives is also known as the Taylor mode AD. For more details on Taylor mode AD, see Appendix D.

这种把原语的 $\differential^{k}$ 复合起来的做法即 Taylor 模式 AD。更多细节见附录 D。

4 Method

4 方法

From the previous discussion, it is clear that the exponential scaling in $k$ for the problem described in Eq. 1 cannot be mitigated by amortization alone. Although high-order AD methods like Taylor mode AD [6] can address this scaling issue, it is only defined for univariate functions. In this section, we describe a method that addresses the scaling issue in $k$ and $d$ simultaneously when amortizing Eq. 1 by seeing univariate Taylor mode AD as contractions of multivariate derivative tensor.

由前文可见，仅靠摊销无法缓解式 (1) 中 $k$ 上的指数缩放。Taylor 模式 AD [6] 等高阶 AD 可应对该缩放，但只对一元函数定义。本节把一元 Taylor 模式 AD 视为多元导数张量的收缩，从而在摊销式 (1) 时同时解决 $k$ 与 $d$ 的缩放。

4.1 Univariate Taylor mode AD as contractions of multivariate derivative tensor

4.1 一元 Taylor 模式 AD 作为多元导数张量的收缩

$\differential F$ projects the Jacobian of $F$ to $\mathbb{R}^{d^{\prime}}$ with a 1-jet $J_{g}(t)$. Similarly, $\differential^{k}F$ contracts a set of derivative tensors to $\mathbb{R}^{d^{\prime}}$ with a $k$-jet $J_{g}^{k}$. We can expand $\partialderivative[k]{t}F\circ g$ with Eq. 43 to see the form of the contractions. For example, $\partial F$ is JVP, and $\partial^{2}F$ contains a quadratic form of the Hessian $D_{F}^{2}$:

$\differential F$ 用 1-jet $J_{g}(t)$ 把 $F$ 的 Jacobian 投影到 $\mathbb{R}^{d^{\prime}}$；同理，$\differential^{k}F$ 用 $k$-jet $J_{g}^{k}$ 把一组导数张量收缩到 $\mathbb{R}^{d^{\prime}}$。用式 43 展开 $\partialderivative[k]{t}F\circ g$ 可看出收缩形式。例如 $\partial F$ 即 JVP，$\partial^{2}F$ 含 Hessian $D_{F}^{2}$ 的二次型：

$$\partial^{2}F(\mathbf{a})(\mathbf{v}^{(1)},\mathbf{v}^{(2)})=\partialderivative[2]{t}[F\circ g](t)=D_{F}(\mathbf{a})\mathbf{v}^{(2)}+D^{2}_{F}(\mathbf{a})_{d^{\prime},d_{1},d_{2}}v^{(1)}_{d_{1}}v^{(1)}_{d_{2}}. \tag{9}$$

From Eq. 43, one can always find a $J^{l}_{g}$ with large enough $l\geq k$ such that there exists $k\leq l^{\prime}\leq l$ with $\partial^{l^{\prime}}F(J_{g}^{l^{\prime}})=D^{k}_{F}(\mathbf{a})\cdot\otimes_{i=1}^{k}\mathbf{v}^{(v_{i})}$ where $v_{i}\in[1,k]$, by setting some tangents $\mathbf{v}^{(v_{i})}$ to the zero vector. That is, arbitrary derivative tensor contraction is contained within a Fréchet derivative of high-order, which can be efficiently evaluated through Taylor mode AD.

由式 43，总能找到足够大的 $l\geq k$ 的 $J^{l}_{g}$，使存在 $k\leq l^{\prime}\leq l$ 满足$\partial^{l^{\prime}}F(J_{g}^{l^{\prime}})=D^{k}_{F}(\mathbf{a})\cdot\otimes_{i=1}^{k}\mathbf{v}^{(v_{i})}$（$v_{i}\in[1,k]$），只需把部分切向量 $\mathbf{v}^{(v_{i})}$ 置零。亦即：任意导数张量收缩都含于某个高阶 Fréchet 导数中，可用 Taylor 模式 AD 高效求值。

How large $l$ should be depends on how off-diagonal the operator is. If the operator is diagonal (i.e. contains no mixed partial derivatives), $l=k$ is enough. If the operator is maximally non-diagonal, i.e. it is a partial derivative where all dimensions to be differentiated are distinct, then the minimum $l$ needed is $(1+k)k/2$. For more details, please refer to Appendix F where a general procedure for determining the jet structure is discussed.

$l$ 需多大取决于算子有多「非对角」。若算子对角（无混合偏导），$l=k$ 即可；若最大程度非对角——对两两不同的维求偏导——则最小 $l$ 为 $(1+k)k/2$。细节与确定 jet 结构的一般流程见附录 F。

Figure 2: The computation graph of $\differential^{2}F$ for $F$ with $4$ primitives. Parameters $\theta_{i}$ are omitted. The first column from the left represents the input 2-jet $J_{g}^{2}(t)=(\mathbf{x},\mathbf{v}^{(1)},\mathbf{v}^{(2)})$, and $\differential^{2}F_{1}$ pushes it forward to the 2-jet $J_{F_{1}\circ g}^{2}(t)=(\mathbf{y}_{1},\mathbf{v}_{1}^{(1)},\mathbf{v}_{1}^{(2)})$ which is the subsequent column. Each row can be computed in parallel, and no evaluate trace needs to be cached.

图 2：含 4 个原语的 $F$ 上 $\differential^{2}F$ 的计算图。参数 $\theta_{i}$ 省略。最左列是输入 2-jet $J_{g}^{2}(t)=(\mathbf{x},\mathbf{v}^{(1)},\mathbf{v}^{(2)})$，$\differential^{2}F_{1}$ 将其前推为下一列的 2-jet $J_{F_{1}\circ g}^{2}(t)=(\mathbf{y}_{1},\mathbf{v}_{1}^{(1)},\mathbf{v}_{1}^{(2)})$。各行可并行计算，且无需缓存求值轨迹。

4.2 Estimating arbitrary differential operator by pushing forward random jets

4.2 通过前推随机 jet 估计任意微分算子

Next, we show how to use the above facts to construct a stochastic estimator derivative operator. Differential operators can be evaluated through derivative tensor contraction. The action of the derivative $\mathcal{D}^{\alpha}=\frac{\partial^{\absolutevalue{\alpha}}}{\partial x_{1}^{\alpha_{1}},\dots,\partial x_{d}^{\alpha_{d}}}$ on function $u$ can be identified with the derivative tensor slice $D_{u}^{\absolutevalue{\alpha}}(\mathbf{a})_{\alpha}$. Differential operator $\mathcal{L}$ can be written as a linear combination of derivatives: $\mathcal{L}=\sum_{\alpha\in\mathcal{I}(\mathcal{L})}C_{\alpha}\mathcal{D}^{\alpha}$, where $\mathcal{I}(\mathcal{L})$ is the set of tensor indices representing terms included in the operator $\mathcal{L}$. For simplicity we only consider $k$th order differential operator, i.e. $\absolutevalue{\alpha}=k\in\mathbb{N}$ for all $\alpha$. For scalar $u:\mathbb{R}^{d}\to\mathbb{R}$, we can identify a $k$th order differential operator $\mathcal{L}$ with the following tensor dot product

下面用上述事实构造随机导数算子估计器。微分算子可通过导数张量收缩求值。导数 $\mathcal{D}^{\alpha}=\frac{\partial^{\absolutevalue{\alpha}}}{\partial x_{1}^{\alpha_{1}},\dots,\partial x_{d}^{\alpha_{d}}}$ 作用在 $u$ 上，等同于导数张量切片 $D_{u}^{\absolutevalue{\alpha}}(\mathbf{a})_{\alpha}$。微分算子 $\mathcal{L}$ 可写成导数线性组合：$\mathcal{L}=\sum_{\alpha\in\mathcal{I}(\mathcal{L})}C_{\alpha}\mathcal{D}^{\alpha}$，$\mathcal{I}(\mathcal{L})$ 是算子所含项的张量指标集。为简单起见只考虑 $k$ 阶微分算子，即对所有 $\alpha$ 有 $\absolutevalue{\alpha}=k\in\mathbb{N}$。对标量 $u:\mathbb{R}^{d}\to\mathbb{R}$，$k$ 阶微分算子 $\mathcal{L}$ 可等同于如下张量点积

$$\mathcal{L}u(\mathbf{a})=\sum_{\alpha\in\mathcal{I}(\mathcal{L})}C_{\alpha}\mathcal{D}^{\alpha}u(\mathbf{a})=\sum_{d_{1},\dots,d_{k}}D^{k}_{u}(\mathbf{a})_{d_{1},\dots,d_{k}}C_{d_{1},\dots,dk}(\mathcal{L})=D^{k}_{u}(\mathbf{a})\cdot\mathbf{C}(\mathcal{L}), \tag{10}$$

where $d_{i}\in[1,d],i\in[1,k]$ is the tensor index on the $i$th axis, , and $\mathbf{C}(\mathcal{L})$ is a tensor of the same shape as $D^{k}_{u}(\mathbf{a})$ that equals $C_{\alpha}$ when $d_{1},\dots,d_{k}$ matches the multi-index $\alpha\in\mathcal{I}(\mathcal{L})$ and $0$ otherwise. We call $\mathbf{C}(\mathcal{L})$ the coefficient tensor of $\mathcal{L}$. For example, the coefficient tensor of the Laplacian $\laplacian$ is the $d$-dimensional identity matrix $\mathbf{I}$. More complicated operators can be built as $f(\mathbf{x},u,\mathcal{D}_{k_{1}}u,\dots,\mathcal{D}_{k_{n}}u)$ where $f$ is arbitrary function.

其中 $d_{i}\in[1,d],i\in[1,k]$ 是第 $i$ 轴上的张量指标，$\mathbf{C}(\mathcal{L})$ 与 $D^{k}_{u}(\mathbf{a})$ 同形：当 $d_{1},\dots,d_{k}$ 匹配多重指标 $\alpha\in\mathcal{I}(\mathcal{L})$ 时等于 $C_{\alpha}$，否则为 0。称 $\mathbf{C}(\mathcal{L})$ 为 $\mathcal{L}$ 的系数张量。例如 Laplace 算子 $\laplacian$ 的系数张量是 $d$ 维单位阵 $\mathbf{I}$。更复杂的算子可写成 $f(\mathbf{x},u,\mathcal{D}_{k_{1}}u,\dots,\mathcal{D}_{k_{n}}u)$，$f$ 任意。

Any derivative tensor contractions $D^{k}_{u}(\mathbf{a})\cdot\mathbf{C}(\mathcal{L})$ can be estimated through random contraction, which can be implemented efficiently as pushing forward random jets from an appropriate distribution. With random $(\mathbf{v}^{(1)},\dots,\mathbf{v}^{(k)})$, we have

任意导数张量收缩 $D^{k}_{u}(\mathbf{a})\cdot\mathbf{C}(\mathcal{L})$ 都可通过随机收缩估计，并高效实现为：从适当分布前推随机 jet。对随机 $(\mathbf{v}^{(1)},\dots,\mathbf{v}^{(k)})$，有

$$\mathbb{E}[D^{k}_{u}(\mathbf{a})_{d_{1},\dots,d_{k}}v_{d_{1}}^{(v_{1})}\dots v_{d_{k}}^{(v_{k})}]=D^{k}_{u}(\mathbf{a})_{d_{1},\dots,d_{k}}\mathbb{E}[v_{d_{1}}^{(v_{1})}\dots v_{d_{k}}^{(v_{k})}]=D^{k}_{u}(\mathbf{a})\cdot\mathbb{E}\left[\otimes_{i=1}^{k}\mathbf{v}^{(v_{i})}\right] \tag{11}$$

where $\otimes$ denotes Kronecker product, $v_{d_{i}}^{(v_{i})}\in[1,k]$ is the $d_{i}$ dimension of the $v_{i}$th order tangent in the input $k$-jet. Eq. 11 is an unbiased estimator of the $k$th order operator $\mathcal{L}u=D^{k}_{u}(\mathbf{a})\cdot\mathbf{C}(\mathcal{L})$ when

其中 $\otimes$ 表示 Kronecker 积，$v_{d_{i}}^{(v_{i})}\in[1,k]$ 是输入 $k$-jet 中第 $v_{i}$ 阶切向量的第 $d_{i}$ 维分量。当满足下式时，式 11 是 $k$ 阶算子 $\mathcal{L}u=D^{k}_{u}(\mathbf{a})\cdot\mathbf{C}(\mathcal{L})$ 的无偏估计：

$$\mathbb{E}[v_{d_{1}}^{(v_{1})}\dots v_{d_{k}}^{(v_{k})}]=C_{d_{1},\dots,dk}(\mathcal{L}). \tag{12}$$

For example, the condition for unbiasedness for the Laplacian $\laplacian$ is $\mathbb{E}[\mathbf{v}^{(a)}\mathbf{v}^{(b)\top}]=\mathbf{I}$. As discussed, one can always find a $J^{l}_{g}$ with large enough $l\geq k$ such that $\partial^{l}F(J_{g}^{l})=D^{k}_{F}(\mathbf{a})\cdot\otimes_{i=1}^{k}\mathbf{v}^{(v_{i})}$, so with a distribution $p$ over the input $l$-jet $J_{g}^{l}$ that satisfies the unbiasedness condition (Eq. 12), we have

例如 Laplace 算子 $\laplacian$ 的无偏条件是 $\mathbb{E}[\mathbf{v}^{(a)}\mathbf{v}^{(b)\top}]=\mathbf{I}$。如前所述，总能找到足够大的 $l\geq k$ 使 $\partial^{l}F(J_{g}^{l})=D^{k}_{F}(\mathbf{a})\cdot\otimes_{i=1}^{k}\mathbf{v}^{(v_{i})}$，因而对满足无偏条件（式 12）的输入 $l$-jet 分布 $p$，有

$$\mathbb{E}_{J_{g}^{l}\sim p}[\partial^{l}u(J_{g}^{l})]=\mathbb{E}[v_{d_{1}}^{(v_{1})}\dots v_{d_{k}}^{(v_{k})}]=D^{k}_{u}(\mathbf{a})\cdot\mathbf{C}(\mathcal{L})=\mathcal{L}u(\mathbf{a}), \tag{13}$$

which means $\mathcal{L}u(\mathbf{a})$ can be approximated by the sample mean of the pushforwards of random $l$-jet drawn from $p$, which can be computed efficiently via Taylor mode AD. We call this method Stochastic Taylor Derivative Estimator (STDE). The advantages of STDE are:

亦即 $\mathcal{L}u(\mathbf{a})$ 可用从 $p$ 抽取的随机 $l$-jet 前推的样本均值逼近，可用 Taylor 模式 AD 高效计算。称此方法为随机泰勒导数估计器（STDE）。STDE 的优点如下：

1. General: STDE can be applied to differential operators of arbitrary order and dimensionality.
2. Scalable: The scaling issue in the dimensionality $d$ and the derivative order $k$ are addressed at the same time. From the example computation graph (Fig. 2) we see that, for one call to $\differential^{k}F$, the memory requirement has scaling of $\order{kd}$ and the computation complexity has scaling $\order{k^{2}dL}$. Like first-order forward mode AD, the derivative tensor $D_{u}^{k}$ is never fully computed and stored. Combined with randomization, the polynomial scaling in $d$ will be removed.
3. Parallelizable: The number of sequential computations does not grow with the order as can be seen in Fig. 2, and the computation can be trivially vectorized and parallelized since the pushforward of sample jets can be computed independently, and it uses the same computation graph ($\differential^{k}u$);

1. 通用：STDE 可用于任意阶、任意维的微分算子。
2. 可扩展：同时解决维数 $d$ 与导数阶数 $k$ 的缩放。由图 2 可见，单次调用 $\differential^{k}F$ 的显存为 $\order{kd}$、计算为 $\order{k^{2}dL}$；与一阶前向 AD 一样，从不完整计算并存储导数张量 $D_{u}^{k}$。再结合随机化，可去掉 $d$ 上的多项式缩放。
3. 可并行：串行计算次数不随阶数增长（见图 2）；样本 jet 的前推彼此独立且共用同一计算图 $(\differential^{k}u)$，可平凡向量化与并行。

4.3 Constructing STDE for high-order differential operators with sparse random jets

4.3 用稀疏随机 jet 为高阶微分算子构造 STDE

Note that all coefficient tensor has the following additive decomposition:

注意所有系数张量都有如下加性分解：

$$\mathbf{C}(\mathcal{L})=\sum_{d_{1},\dots,d_{k}\in\mathcal{I}(\mathcal{D})}C_{d_{1},\dots,d_{k}}\mathbf{e}_{d_{1}}\otimes\dots\otimes\mathbf{e}_{d_{k}} \tag{14}$$

where $\mathbf{e}_{i}$ is the $i$th standard basis. For example, if the input dimension $d$ is $3$, then $\mathbf{e}_{2}=\matrixquantity[\lx@physics@matrix 0,1,0\endlx@physics@matrix]^{\top}$. As discussed before, there exists a $J_{g}^{k}$ whose pushforward under $\partial^{l}u$ is equivalent to contracting $D_{u}^{k}$ with $\otimes_{i=1}^{k}\mathbf{e}_{d_{i}}$. We call $k$-jet consisting of only standard basis and the zero vector $\mathbf{0}$ sparse. Therefore the discrete distribution $p$ over the sparse $k$-jets in Eq. 14 satisfies the unbiasedness condition 12

其中 $\mathbf{e}_{i}$ 是第 $i$ 个标准基。例如输入维 $d=3$ 时，$\mathbf{e}_{2}=\matrixquantity[\lx@physics@matrix 0,1,0\endlx@physics@matrix]^{\top}$。如前所述，存在 $J_{g}^{k}$，其在 $\partial^{l}u$ 下的前推等价于用 $\otimes_{i=1}^{k}\mathbf{e}_{d_{i}}$ 收缩 $D_{u}^{k}$。只含标准基与零向量 $\mathbf{0}$ 的 $k$-jet 称为稀疏。因而式 14 中稀疏 $k$-jet 上的离散分布 $p$ 满足无偏条件 12

$$p(\otimes_{i=1}^{k}\mathbf{e}_{d_{i}})=C_{d_{1},\dots,d_{k}}/Z,\quad d_{1},\dots,d_{k}\in\mathcal{I}(\mathcal{L}), \tag{15}$$

where $Z$ is the normalization factor and we identify $\otimes_{i=1}^{k}\mathbf{e}_{d_{i}}$ with the corresponding $k$-jet $J_{u}^{k}$.

其中 $Z$ 是归一化因子，并把 $\otimes_{i=1}^{k}\mathbf{e}_{d_{i}}$ 与对应的 $k$-jet $J_{u}^{k}$ 等同。

4.3.1 Differential operator with easy to remove mixed partial derivatives

Next, we show some concrete examples for constructing STDE with sparse random jets.

4.3.1 混合偏导易于消去的微分算子

下面给出用稀疏随机 jet 构造 STDE 的若干具体例子。

Laplacian

From Eq. 9 we know that the quadratic form of Hessian can be computed through $\partial^{2}$ by setting $\mathbf{v}^{(2)}=\mathbf{0}$ and $\mathbf{v}^{(1)}=\mathbf{e}_{j}$. Therefore, the STDE of the Laplacian operator is given by

Laplace 算子

由式 9，令 $\mathbf{v}^{(2)}=\mathbf{0}$、$\mathbf{v}^{(1)}=\mathbf{e}_{j}$，即可经 $\partial^{2}$ 得到 Hessian 二次型。因此 Laplace 算子的 STDE 为

$$\tilde{\nabla^{2}}_{J}u_{\theta}(\mathbf{a})=\frac{d}{\absolutevalue{J}}\sum_{j\in J}\partialderivative[2]{x_{j}}u_{\theta}(\mathbf{a})=\frac{d}{\absolutevalue{J}}\sum_{j\in J}\partial^{2}u_{\theta}(\mathbf{a})(\mathbf{e}_{j},\mathbf{0})=\frac{d}{\absolutevalue{J}}\sum_{j\in J}\differential^{2}u_{\theta}(\mathbf{a},\mathbf{e}_{j},\mathbf{0})_{[2]} \tag{16}$$