Collapsing Taylor Mode Automatic Differentiation
Authors: Felix Dangel
Tim Siebert
Marius Zeinhofer
Andrea Walther
Venue: NeurIPS 2025 poster ()
Keywords: Automatic differentiation
Taylor mode
Differential operators
Higher-order derivatives
PINNs
VMC
Laplacian
Primary area: machine_learning_for_sciences
TL;DR: We accelerate Taylor mode for practically relevant differential operators by collapsing Taylor coefficients; this can be done automatically with compute graph simplifications

Abstract: Computing partial differential equation (PDE) operators via nested backpropagation is expensive, yet popular, and severely restricts their utility for scientific machine learning.
  Recent advances, like the forward Laplacian and randomizing Taylor mode automatic differentiation (AD), propose forward schemes to address this.
  We introduce an optimization technique for Taylor mode that 'collapses'  derivatives by rewriting the computational graph, and demonstrate how to apply it to general linear PDE operators, and randomized Taylor mode.
  The modifications simply require propagating a sum up the computational graph, which could---or should---be done by a machine learning compiler, without exposing complexity to users.
  We implement our collapsing procedure and evaluate it on popular PDE operators, confirming it accelerates Taylor mode and outperforms nested backpropagation.

Collapsing Taylor Mode Automatic Differentiation
Felix Dangel"
Tim Siebert"
Vector Institute Toronto, Canada
Coanbblld--Univensee zu Berlin and
Zuse Institute Berlin
fdangal@@vectoriimstttuuu...
Berlin, Germany
tim. eebeer@@hu-berinn
Marius Zeinhofer
Andrea Walther
ETH Zurich
Humboldt- Universstiiii zu Berlin and
Zurich, Switzerland,
Zuse Institute Berlin
marius.. ze www.eeeeeeeaaaaeeeeoo.oeeeeeeeee math. ethz. ch
Berlin, Germany
andrea wal ther@matrh htw.:/aathhhhhhbblaeeeee hwwberlinnn...
Abstract
Computing partial differential equation (PDE) operators via nested backpropagation is expensive, yet popular, and severely restricts their utility for scientific machine learning. Recent advances, like the forward Laplacian and randomizing Taylor mode automatic differentiation (AD), propose forward schemes to address this. We introduce optimization technique for Taylor mode that "collapses" derivatives by rewriting the computational graph, and demonstrate how to apply it to general linear PDE operators, and randomized Taylor mode. The modifications simply require propagating a sum up the computational graph, which could-or or should- be done by machine learning compiler, without exposing complexity users. We implement our collapsing procedure and evaluate it on popular PDE operators, confirming it accelerates Taylor mode and outperforms nested backpropagation.
Introduction
Using neural networks to learn functions constrained by physical laws a popular trend in scientific machine learning [4, 16, 17, 19, 24, 25, 29]. Typically, the Physics is encoded through partial differential equations (PDEs) that the neural net must satisfy. The associated loss functions require evaluating differential operators w.r.t. the net's input, rather than weights. Evaluating these differential operators remains a computational challenge, especially if they contain high-order derivatives.
Computing PDE operators. Two important fields that build PDE operators variational Monte-Carlo (VMC) simulations and Physics-informed neural networks (PINNs). VMC employs neural networks as ansatz for the Schrödinger equation [4, 16, 24] and demands computing the net's Laplacian (the Hessian trace) for the Hamiltonian's kinetic term. PINNs represent PDE solutions as neural net and train it by minimizing the residuals of the governing equations [19, 25]. For instance, Kolmogorov equations like the Fokker-Plan Planck and Sacck-Scholes equation require weighted second-order derivatives on high-dimensioonl spatial domains [17, 28]. Other PINNs for elasticity problems use the biharmonic operator [7, 17, 27, 31], which contains fourth-order derivatives.
Is back propagaton all we need? Although nesting first-order automatic differentiation (AD) to compute high-order derivatives scales exponentially w.r.t. the degree in time and memory [27, $3.2].
39th Conference Neural Information Processing Systems (NeurIPS 2025).
"Equal contribution

---

Nested 1 order AD
Derivative degree
Taylor forward propagation
1.5
Time [s]
Vanilla Tay-
1.0
lor mode
ho
Forward
0.5
Laplacian
0.0
1500
(
(
Batch size
Figure 1:
A Vanilla Taylor mode not
enough to beat nested -rdorder AD. Il- lustrated for computing the Laplacian of tanh-activated 50 768 768 + 512 512 MLP with JAX (+ jit) on GPU (details in SG). We show how automati- cally obtain the specialized forward Lapla- cian through simple graph transformations that "collapse" vanilla Taylor mode.
{
(
9K 1
K
Figure 2: Collapsed Taylor mode di- rectly propagates the sum of highest de- gree coefficients. Visualized for pushing 4 K -ejets through R R function (K' yields the forward Laplacian).
{
(
(g
K (ours)
Σn uxx
Σ hxg
E,
this approach is common practice: it is easy to implement in ML libraries, and their backpropagation highly optimized. A promising alternative is Taylor mode AD [or simply Taylor mode, 13, $13), introduced to the ML community in 2019, which scales polynomially w.r.t. the degree in time and memory [14]. However, we observe empirically that vanilla Taylor mode may not be enough to beat nesting (fig. I): evaluating the Laplacian of 5-layer MLP, using JAX's Taylor mode is 50% slower than nested buckpropagation that computes, then traces, the Hessian via Hessian- veetor products [23]. This calls into question the relevance of Taylor mode for computing common PDE operators.
The advent of forward schemes. Recent works have successfully demonstrated the potential of
modified forward propagation schemes, though. For the Laplacian, Li et al. [20, 21] developed special forward propagation framework called the forward Laplacian, whose JAX implementation [12] is roughly twice fast as nested first-order AD (fig. 1). While the forward Laplacian does not rely on Taylor mode, recent work pointed out connection [6]; it remains unclear, though, efficient forward schemes exist for other differential operators, and how they relate to Taylor mode. Concurrently, Shi et al. [27] derived stochastic approximations differential operators in high dimensions by evaluating Taylor mode along suitably sampled random directions.
Irrespective stochastic or exact computation, their core, these popular PDE operators are linear: we must evaluate derivatives along multiple directions, then sum them. Based on this lincarity, we identify an optimization technique to rewrite the computational graph of standard Taylor mode that applicable to general linear PDE operators and randomized Taylor mode:
We propose optimizing standard Taylor mode by collapsing the highest Taylor coefficients, directly propagating their sum, rather than propagating then summing (fig. 2). Our approach contains the forward Laplacian as special case, is applicable to randomized Taylor mode, and also general linear PDE operators, which show using the techniques from Griewank et al. [14].
2. We show how to collapse standard Taylor mode by simple graph rewrites based on linearity. This leads to clean separation of concepts: Users can build their computational graph using standard Taylor mode, then rewrite it to collapse it. Due to the simple nature of our proposed rewrites, they could easily absorbed into the just-in-time (JTT) compilation of MI. frameworks without introducing new interface or exposing complexity to users.

---

3. We empirically demonstrate that collapsing Taylor mode accelerates standard Taylor mode. We implement a Taylor mode library² for PyTorch [22] that realizes the graph simplifications with torch. fx [26]. On popular PDE operators, we empirically find that, compared to standard Taylor mode, collapsed Taylor mode achieves superior performance that well-aligned with the theoretical expectation, while consistently outperforming nested first-order AD.
Our work takes important step towards the broader adoption of Taylor mode as viable alternative to nested first-order AD for computing PDE operators, while being as easy to use.
Background: Introduction to Taylor Mode AD
Taylor mode AD (or, simply, Taylor mode) computes higher-order derivatives- needed, e.g., for PDE pperators-throu propagation of Taylor coefficients according to the chain rule.
Scalar case. To illustrate Taylor mode, consider the scalar function f IR R and extend the input variable r to a path x(t) with r(0) = zo. whose form is a univariate Taylor polynomial degree K, xxt) Σ k=u kT Jh with 2y the k-th Taylor coefficient, If f. is smooth enough, we evaluate
d' Taylor coefficients of the transformed path f(x(t)) = =Ex= kt fk with fx == dek f(x(t)) The chain rule provides the coefficients' propagation rules. E.g., for degree K 3 get
fo f(0o), f1 = (x)))21
fz f(xo) 0f(z))
(1)
fs = 0==00 + 30 (10) + (1o))
Faà Di Bruno [9] provided the general formula for fk. and Fraenkel [11] extended to the multivariate
case [see also 15]. It serves as foundation for Taylor mode to compute higher-order derivatives
[e.g.. 13, $13]: setting I = 12 23 = 0 yields fi O (f()) /2 = 0² f(zo), fa = 0 fxxo). We
call the univariate Taylor polynomial of function xx) of degree K, represented by the coefficients (zo...., Tx), the K jet of x, following the terminology of JAX's Taylor mode [2].
Notation for multivariate case. consider the general case of computing higher-order derivatives, e.g., PDE operators, of a vector-to-vector function f IR + RC, This requires additional notation generalize eq. (1). Given K vectors vI...., UK E IR we write their tensor product
k= 2k
en e & uK E (IRP) with entries [00 1K
di... )
[wi]d ...... [uK]dx
for d1..... . dx E {1,.... D), and compactly write y @k _
U. We define the inner product of two tensors A, B E (IR (RR as the Euclidean inner product of their flattened versions
(A. B) = ΣΣ d1 d2
IiAla,da...d, dx
[B] ddddda.....dd dx E R.
dx
We allow broadcasting in eq. (2): if one tensor has more dimensions but matching trailing dimensions, we take the inner product for each component of the leading dimensions. This allows to express contractions with derivative tensors of vector-valued functions, e.g., contracting the k-th derivative tensor 0 f(a) E C (RK) such that (A, 8* f(xo)) E R
Multivariate case & composition. Evaluating the K-jet of f at E RD starts with the extension of 2o to a smooth path x RD with (0) = Ny. Formally, the K-jet of f is defined as
k
th
d with fk = dt f(x(t))leo
RK ((K =
Σ fx
kl
k=0
and requires the K-jet of æ. (J xt)
# 0k As we are interested in the coefficients, we
will slightly abuse the K --jet as mapping (za.... [k] + (fo..... fk) (see fig. 3 for an illustration).
As is common for AD, propagating the coefficients is broken down into composing f of atomic functions with known derivatives and the chain rule. In the simplest case, let g o hh R >
Available at https: pttttp::/giiie h/ghthub.//t/.////////pg http://gtthbbhom.o.tee wwwm////aam////ooccoeeeeeeeee
3

---

Extent input to smooth path xt))
f
Path in output space f(x(t)
Figure 3: Taylor mode propagates Taylor coefficients of a path in input space. This results in the function-transformer path's 's Taylor coefficients. The Taylor expansion of degree K called a K-jet; hence Taylor mode propagates the input K jet to the output K K-jet.
J
Jk
K-jet Ek=o bs ak as (Ii0. ØK
Taylor mode
K-jet
Σ === E! (fo....: :
e fk
R. + RC for two elemental functions g, h. Given the input K-jet for æ, the coefficients he de h(x(()))e= follow from the generalized Faà di Bruno formula (spelled out for some ks in $A)
dt
h4 Σ
v(o)
1 h, @ E. SEO
with v(o) = (II.eo nz!) (TI.Eo s!)
(3)
@Epart(k)
Here, part(k) is the integer partitioning of k (a of sets), is multiplicity function, and n, counts occurrences of in a set a (e.g., (2((1, 1,3)) +3)==2 =2 and n3 = 1). Propagating the h&s through results in the K-jet for In summary, the propagation scheme is (with ø R he R f& R
ho h(x0))
h1 @h(aa)), σ1)
h2 8 h(
at (8h(z)))
22
q
hg Σ v(0) (golh(z0)) 2 x @Epart(K)
sEa
(4)
90 g (ha)
f(o)
(gg ho)) h1)
fi = (0)) xt
92 8 gho).
h1 202 (8g())))
fz (02 f( (2)) x1 82 - x22)
9K Σ v(a) agoll @ h. sEo
fK v(o)) a f(zo), (Epart(K)
& x SEO
(Kpart(K)
which describes the forward propagation of a single K-jet. However, computing popular PDE operators requires propagating multiple K-jets in parallel, then summing their results. We propose pull this accumulation inside Taylor mode's propagation scheme, thereby collapsing it.
Collapsing Taylor Mode AD
We now deseribe how collapse the Taylor mode AD computation of popular linear PDE operators and their stochastic approximations proposed [27], and provide a general recipe for computing and collapsing general linear differential operators by interpolation, using earlier work from Griewank et al. [14]. At its core, our procedure uses the linearity of the highest Taylor coefficient's propagation rule. allows collapse coefficients along multiple directions and directly propagate their sum, rather than propagating then summing, yielding substantial reductions in computational cost.
3.1 Exploiting Linearity to Collapse Taylor Mode AD
To derive proposed method, we start with sum of Kthhorder directional derivatives of the function of along R directions ( R which is a common building block for all our PDE operators:
R
Σ
(ak f(x)) o o E R
(5)
dim( = D, and the unit vectors v,- = e, E RD as (e.g., the exact Laplacian uses = R= directions; see $3.2 below). Instead of nesting K calls sstorreer AD, we can use K K-jets to calculate

---

each summand of eq. (5) with Taylor mode. In total, we need R K K-jets, and have to set the r-th jet's coefficients Forr zp. 1, V, and I2,r =.. KK (eq. (D13) applies this to eq. (4)).
Standard Taylor mode propagates 1 + vectors through every node of the computational graph (the Oth component is shared across all jets, fig. 2). This gives the output jets { ((fkr) ,=1 from which we only select the highest-degree coefficients (fk.)>= then sum them to obtain eq. (5).
The approach we propose here exploits that the K-th derivative of h is åg times the K-th derivative h plus other lower-order terms in h. Therefore, g o is linear in the K-th derivative of h. Mathematically speaking, there is special element in the set of integer partitions part(K), namely the trivial partition = { , which contributes the term v()) (Qg( hx.r) to eq. (3). This is the only term that uses the input jet's highest coefficient K.r, and its dependency is linear. Separating it in the highest coefficient's forward propagation, we get (using v(ō)= 1)
E JK.
Σ 9K
E v(oo (0% g(o)), sEa
(lo)ggth)) @ har
Σ (Og(ha); hKg)
(εpan(())()))
and since the propagation rule is linear w.r.t. hgrr, can pull the summation inside:
2 9K. = = E v(o) ( go)))(()))))))))))))))))) @ herr ( sEo
( EnK. (6)
1=1 (Epar((()M))
This is the key insight of our work: The summed highest-degree outpul coefficients depend on the summed highest-degree input coefficients (as well as all lower-degree coefficients). The reason is linearity in Faà di Bruno's formula. Hence, to compute the sum Σ, 9K,r we can directly propagate the sum Σ, h,, collapsing coefficients over all directions. We call this collapsed Taylor mode AD.
Collapsed Taylor mode propagates only 1 + (K 1)R + vectors through every node in the computational graph (see fig. and eq. (D14) which applies this to eq. (3)). These savings of R 1 coefficients are significant improvements over standard Taylor mode, as we show below. In the following, we discuss how to collapse the Taylor mode computation of various PDE operators.
3.2 Linear Second-order Operators
Laplacian. The Laplace operator plays a central role in Physics and engineering, including electro- staties, fluid dynamics, heat conduction, and quantum mechanies [10, 24]. It contains the Hessian trace of each element of a function, i.e., for f R + RC, it
0 f( (za). 2
(exact)
Af(x)) = f()))
(7a)
80 (xo), v 22) (stochastic)
with the d-th standard basis vector ea used for exact computation, and random vectors v, drawn
ii.dd from distribution with unit variance (e.g. Rademacher or standard Gaussian) for stochastic
estimation. By pattern-aacciing eq. (7a) with eq. (5) we conclude that K 2, and the following choices for computing the Laplacian with standard Taylor mode:
{Wood 2o.
F1.d Edi
F2,d 0))2z
1+ D vectors
(exact)
(7b)
(( Zg.
11, Uor
22 = 0))52
1+5+ vectors
(stochastic)
Collapsing standard Taylor mode yields 1 D + (exact) and 1 S (stochastic) vectors. In
fact, the collapsed Taylor mode for the exact Laplacian is the forward Laplacian from Li et al. [21] (see eq. (D16) for detailed presentation of the forward propagation). Note how we can seamlessly also collapse the stochastie approximation
the sampled directions, which is currently not done.
Weighted Laplacian. A natural generalization of the Laplacian involves contracting with a positive semi-definite matrix aa E RDD rather than the identity. D may represent the diffusion tensor in Kolmogorov-type PDEs like the Bokker-Planck equation [17], and can depend wo
5

---

[8]. The weighted Laplacian contains the weighted Hessian's 's trace Tr(oo [fle) for each output [1......pp)
element of f. ra)k(D) R. and therefore
R Rx x
R
8 f sT
(exact)
ADf(a = ( D) ERG
(8a)
02 f f(x)), (ovo) )02)
(stochastic)
Computing requires evaluating the following 2-jets with standard Taylor mode:
( 30
11r 8r:
22,r -)] 1 R vectors (exact)
(8b)
((ro,s 2o,
21,s 0
22,5 = 0)] 1+5+3 S vectors
(stochastic)
Our collapsed Taylor mode 1 R (exact) and + S+++1 (stochastic) vectors. This yields the modified forward Laplacian from Li et al. [20]; collapsing the stochastic variant speeds up the Hutchinson trace estimator from Hu et al. [17]. For indefinite D, can simply apply this scheme the positive and negative cigen-spaces (however, such weightings are not used in practise).
3.3 Collapsed Taylor Mode for Arbitrary Mixed Partial Derivatives
So far, discussed operators that result from contracting the second-order derivative tensor with a coefficient matrix (I or D) that can conveniently be written as sum of vector outer products. For orders higher than two, the coefficient tensor in general nof easily be decomposed such. Hence, we extend our framework to also cover differential operators containing mixed-partial derivatives by evaluating suitable family of jets using the interpolation result of Griewank al. [14). As illustrative example, we will use the biharmonic operator with a 4-dimensional coefficient tensor:
D
Σ
602 (0t f(zo), Cd,
& ex
(exact)
A f ERC
=1
27
00 f 1841)
(stochastic)
3S
We can directly collapse the stochastic version: draw S standard normal vectors Ui..... Us and propagate the coefficients {(x0.. = wy 11, e = Us 22 = 33 = 14 = 0)} With standard Taylor mode, this uses 1 + 4$ vectors; collapsed Taylor mode uses + 35 + vectors. For the exact biharmonic operator, however, we need to develop an approach to compute mixed partials.
General approach. Assume want to compute linear differential operator of degree K. We can do so by contracting the K-th order derivative tensor O f(xo) with a coefficient tensor (IR P We can always express this tensor in tensor product basis, such that
D2 Σ
(Ø f(x)) C)
8 f(x0), Ua 2 ...
0i1 E C
(10)
where the multi-index entries i == (f,, , i1) sum to and Dj < D. For the exact biharmonic operator (eq. (9)), we identify K 4, 2, (2, Dj D2 D, Cdj and Vd, Cdg. From the Faá di Bruno formula, we know that we can only compute terms of the form (a f(x)), o with a K-jet. The challenge in eq. (10) that includes terms where not all directions coincide (e.g., for the biharmonic we have / = 2 different directions).
Fortunately, Griewank al. [14] derived approach to reconstruct such mixed-diree direction terms by linearly combining a family of K K-jets that is determined by all vectors j E N' whose entries sum to K, see fig. 4 for illustration for the biharmonic (5 members). The K-jets along these directions are then combined with coefficients Yi.j E I, whose definition we provide in $E. In summary, we get
1
@K
(a f(xo)) 241 60.. ...88 8021 10
Σ
Vh.j K!
O f(xo);
Σ Wa. Lile
(11)
JEN |

---

This construction allows us to rewrite eq. (10) as
@K
Σ d,=1
Σ
Σ
Tig K!
OK f(xo);
Σ vid [i];
d==1
JEN'
sind-since the coefficients Ne. only depend on the problem structure (K, I and i) and nor on the function f and the directions Va;; [14]-we we can pull out the inner sum to obtain the final expression
D1
Dr
@K
Y
jk f(xo),
Σ vd [i]
(12)
K!
JENI. =
d1=1
d1=1
We can evaluate eq. (12) with standard Taylor mode: For each j. compute D, K K-jets with coefficients 2o. 21 = E, Vd. [i]..22 = =... æk 0. The sums from the tensor basis expansion can be collapsed with our proposed optimization, removing [ D, vectors from the propagation for each j. After repeating for each member j of the interpolation family, we form the linear combination using the "i.js, which yields the desired differential operator. We can often exploit symmetries in the Ji and basis vectors to further reduce the number of K K-jets (see $E. for complete example).
Applied to the biharmonic operator. Let us now illustrate the key steps of applying eq. (12) to the exact biharmonic operator eq. (9) (full procedure in $E. Figure illustrates the 5 multi-indices, 22
characterizing the 4-jets we need to interpolate (0+ f(x)), ed @02) and their coefficients "a.- Their definition,
eq. (E17), shows the equality of Y for (4, 0) and (0, 4), as well as
j = (3, 1) and j = (1.3). 3). Exploiting those symmetries reduces the number of interpolation terms from 5 to 3 (eq. (E19)), corresponding to D + D2 + D 4-jets. Removing houbly-computt
terms
brings down the number of 4-jets to D + 1)))((DDD)+++(((())00
(eq. (E21)). Translated to vectors,
standard Taylor mode propagates 1 4D 4D -) 4/2 1) 6D -2D 1 vectors.
After collapsing, we get 1 1+3 + +00+++++: 1--100D---11-+ 9/222 3/2D + vectors.
This demonstrates the relevance of collapsing: it achieves a 25 % reduction in the quadratic coefficient.
lil2
i == (2 2)
(4.00)) = 13/192
Summary & relation other approaches for com- puting mixed partials. The scheme we propose based on Griewank al. [14]'s interpolation result allows to calculate general linear differential operators beyond Laplacians, and is amenable to collapsing. Admitiedly, eq. (12) seems daunting at first glance. However, it (i) of- fers a one-fits mit-fits-lll recipe to construct schemes for general linear PDE operators, and (ii) does not jets of order K' > K to compute K- order derivatives. It is pos- sible to derive more "pedagogical" approaches, which however require hand-crafted interpolation rules case by case, and propagation of higher-order jets which is costly (see $E.2 for a pedagogical example using less efficient 6-jets to compute the biharmonic operator, or [27, $F) for other operators).
74,(1,) = 1/3
(4.(2.2) = -5/8
(4.
(1.0))
1
2
3
[i
Figure 4: Illustration of eq. (12) for the biharmonic operator, i.e., the values of j with ||أ|| 4 and their coefficients 7 to interpolate the desired mixed partials.
4 Implementation & Experiments
Here, we describe our implementation of the Taylor mode collapsing process and empirically validate its performance improvements on the previously discussed operators.
Design decisions & limitations. JAX [3] already offers an albeit experimental-TBee implementation [2]. However, we found challenging to capture the computation graph and modify it using JAX's public interface. In contrast, PyTorch [22] provides torch. fx [26], which offers user-friendly interface to capture and transform computational graphs purely in Python. Hence, we re-implemented Taylor mode in PoTccc,, taking heavy inspiration from the JAX implementation.
mode
7

---

Laplacian
Weighted Laplacian (D 50)
Biharmonic
(D 50)
(D =
Time [s]
1.0
1.0
1.0
0.5
0.0
0.0
0.0
Exact
Mem. [GiB]
5
0
0
1000
2000
0
1000
2000
0
500
1000
Nested -order
Standard Taylor
Batch size
Batch size
Batch size
Collapsed (ours)
(N 2048)
(N 2048)
(N = 256)
Differentiable
Non-diff.
Time [s]
1.0
0.5
0.5
Mem. [GiB]
10
10
Stochastic
20
0
40
20
Samples
Samples
Samples
Figure 5: Collapsed Taylor mode accelerates standard Taylor mode and outperforms nested 1 AD. Exact computation varies the batch size, stochastic computation fixes a batch size and varies the samples such that < D (Laplacians), and 2 35 < 9/2D2 2 3/2D + (biharmonic operator); we could compute exactly otherwise. Opaque markers are non-differentiable computations.
This deliberate choice imposes certain limitations. First, of now, our Taylor mode in PyTorch supports only a small number of primitives, because the Taylor arithmetic in eq. (3) needs to be implemented case by case (this of course also applies to JAX's Taylor mode, which has broader operator coverage). Second, while our Taylor mode implementation is competitive with JAX's, s, we did not fully optimize it (e.g., we do not use in-place operations, and we do not implement the efficient schemes from Griewank & Walther [13, $13], but stick to Faà di Bruno (eq. (3))). Given our implementation's superiority compared to nested first-order AD that we demonstrate below, these are promising future efforts that will further improve performance, and we believe that making Taylor mode available to the PyTorch community is also an important step towards establishing its use.
Usage (overview in SB). Our implementation takes a Py Torch function (c.g., a neural net) and first captures its computational graph using torch. fx's symbolic tracing mechanism. Then, it replaces each operation with its Taylor arithmetic, which yields the computational graph of the function's K-jet. Users can then write function compute their differential operator with this vanilla Taylor mode. Collapsing is achieved using a function simplify, which traces the computation again, rewrites the graph, and propagates the summation of highest coefficients up to its leafs. This requires one backward traversal through the graph (&C presents a detailed example). The simplified graph produces the same result, but propagates summed coefficients, i.e., uses collapsed Taylor mode.
Experimental setup. We empirically validate our proposed collapsing approach in Py Porch. We compare standard Taylor mode with collapsed Taylor mode and nested 1d-order AD on an Nvidia RTX 6000 GPU with 24 GiB memory. To implement the (weighted) Laplacian and its stochastic counterpart, we use Nictor-Hessian-veedd recommended [5, 13]. For the biharmonic operator, we simply nest two VHVPs. For the weighted Laplacian's 's coefficient matrix, we choose full-rank diagonal matrix ($F.2 shows results for rank-
products (VHVPs) in forward-over-revent
order, as
8

---

Table 1: Benchmark from fig. 5 in numbers. We fit linear functions and report their slopes, i.e.. how much runtime and memory increase when incrementing the batch size or random samples. We show two significant digits and bold values are best according to parenthesized values.
Per-datum
Weighted Laplacian
Mode
Implementation
Laplacian
Biharmonic
-sample cost
Nested 1"-order
0.61 (1.0x)
0.60 (1.0x)
(1.0x)
Time [ms]
Standard Taylor
0.56 (0.93x)
0.57 (0.94x)
0.90 (0.72x)
Collapsed (ours)
0.29 (0.48x)
0.29 (0.48x)
0.69 (0.55x)
Nested --order
4.4 (1.0x)
4.4 (1.0x)
7.9 (1.0x)
Exact
Mem. [MiB]
Standard Taylor
4.6 (1.0x)
4.6 (1.0x)
7.7 (0.98x)
(differentiable)
Collapsed (ours)
2.1 (0.47x)
2.1 (0.47x)
4.8 (0.61x)
Nested lo-order
2.2 (L.0x)
2.2 (1.0x)
(1.00)
Mem. [MiB]
Standard Taylor
1.2 (0.54x)
1.2 (0.4l)
1.2 (0.36x)
(non-diff.)
Collapsed (ours)
0.90 (0.41x)
0.90 (0.41x)
1.1 (0.33x)
Nested 1d -orree
(1.0x)
(1.0x)
(1.0x)
Time [ms]
Standard Taylor
23 (0.97x)
23 (0.97x)
6,6 (0.15x)
Collapsed (ours)
12 (0.49x)
12 (0.49x)
4.9 (0.11x)
Nested 1-oorder
180 (1.0x)
180 (1.0x)
210 (1.0x)
Stochastic
Mem. [MiB]
Standard Taylor
200 (1.2x)
200 (1.2x)
64 (0.30x)
(differentiable)
Collapsed (ours)
(0.50x)
89 (0.50x)
38 (0.18x)
Nested 1"-order
89 (1.0x)
90 (L.0x)
86 (1.0x)
Mem. [MiB]
Standard Taylor
48 (0.54x)
48 (0.53x)
15 (0.17x)
(non-diff.)
Collapsed (ours)
36 (0.40x)
36 (0.40x)
13 (0.16x)
deficient weightings). To avoid confounding factors, all implementations executed without compilation (our JAX experiments with the Laplacian in &G confirm that jit does not affect the relative performance). As common for PINNs [e.g., 6, 27], we 768 + 768 + 512 + 512 + with tanh activations and trainable parameters O, and compute the PDE operators on batches of size N. We measure three performance metries: (1) runtime reports the smallest execution time of 50 repetitions.
5-layer MLP Se D
(2) Peak memory mion-differentiable) measures the maximum allocated GPU memory when computing the PDE operator's ' value (e.g., used in VMC [24]) inside a torch. no orad context. (3) Peak memory (differentiable) is the maximum memory usage when computing the PDE operator inside torch. enable_ e_adddd context, which allows backpropagation to 0 (required for training PINNs, or alternative VMC works [30, 32)). This demands saving intermediates, which uses more memory but does not affect runtime. As memory allocation does not fluctuate much, we measure it in a single run.
Results. Figure 5 visualizes the growth in computational resources w.r.t. the batch size (exact) and random samples (stochastic) for fixed dimensions D. Runtime and memory increase linearly in both, as expected. We quantify the results by fitting linear functions and reporting their slopes (i.c., time and memory added per datum/sample) in table We make the following observations:
Collapsed Taylor mode accelerates standard Taylor mode. The measured performance dif- ferences correspond well with the theoretical estimate from counting the number of forward- propagated vectors. E.g.. for the exact Laplacian, adding one datum introduces 1+2D new vectors. For D = 50, their ratio is (2+2)/(1+2D) z 0.51. Empirically, we measure that adding one datum adds 0.56 ms standard, and 0.29 ms to collapsed, Taylor mode (table 1); the ratio of = 0.52 is close. Similar arguments hold for peak memory of differentiable computation, stochastic approximation, and the other PDE operators (see table F2 for all comparisons).
D versus
Collapsed Taylor mode outperforms nested 1stoorder AD. For the exact and stochastic (weighted) Laplacians, collapsed Taylor mode is roughly twice as fast (consistent with the JAX results in fig. 1) while using only 40-50% memory. For the biharmonic operator, we also observe speed-ups; in the stochastic case up 9x in time, and 5x in memory (diffcrentiable).

---

Comparison with JAX. We also conducted experiments with JAX (+ jit) to rule out artifacts from choosing PyTorch, implementation mistakes in our Taylor mode library, or unexpected simplifications from the JIT compiler. We find that the choice of the MI. framework does not affect the results. E.g., when computing the exact Laplacian with nested first-order AD, Py PyTorch consumes 0.61 ms per datum (table 1), while JAX uses 0.58 ms (fig. and table G5). We find the same trend when comparing our collapsed Taylor mode and JAX's forward Laplacian. Interestingly, we noticed that JAX's Taylor mode was consistently slower than our Py Torch implementation, despite using jit. We hypothesize that this could stem from algorithmic differences the Taylor mode implementations and conclude from these results that (both ours, as well as the existing JAX) Taylor mode still has potential for improvements that may further increase the margin to nested first-order.
Conclusion
Computing differential operators is critical component in scientific machine learning, particularly for Physics-infen informed neural networks and variational Monte-Carlo. Our work introduces collapsed Taylor mode, simple yet effective optimization based on linearity in Faà di Bruno's o formula, that propagates the sum of highest-order Taylor coefficients, rather than propagating then summing. It contains recent advances in forward-mode schemes, recovering the forward Laplacian [21], while being applicable to stochastic Taylor mode [17, 27]. We demonstrated that collapsed Taylor mode is useful to compute general linear differential operators, leveraging Griewank et al. [14]'s interpolation formula. Empirically, we confirmed speed-ups and memory savings for computing (randomized) Laplacians and biharmonic operators after collapsing Taylor mode, in accordance with our theoretical analysis, and confirmed its superiority to nesting first-order automatic differentiation. As the optimizations are achieved through simple graph rewrites based on linearity, we believe they could be integrated into existing just-in- must-in-time compilers without requiring a new interface or burdening users.
Our work takes important step towards making Taylor mode practical alternative nested first-order differentiation in scientific machine learning, while maintaining ease use. Future work could focus integrating these optimizations directly into MIL compilers, broadening operator coverage of our Py PyToc implementation, and exploring additional graph optimizations for AD.
Acknowledgments and Disclosure of Funding
Resources used in preparing this research provided, in part, by the Province of Ontario, the Government of Canada through CIFAR, and companies sponsoring the Vector Institute. The research was funded partly by the DFG under Germany's Excellence Strategy The Berlin Mathematics Research Center MATH+ (EXC-2046/1, project ID:390685689, M.Z. acknowledger support from an ETH Postdoctoral Fellowship for the project "Reliable, Efficient, and Scalable Methods for Scientific Machine Learning".
References
[1] Arbogast, L. Du calcul des dérivations. 1800.
[2] Bettencourt, J., Johnson, M. J., and Duvenaud, D. Taylor-mode automatic differentiation for higher-order derivatives in JAX. In Advances in Neural Information Processing Systems (NeurIPS); Workhop on Program Transformations for ML, 2019.
[3] Bradbury, J., Frostig, R., Hawkins, P., Johnson, M. J., Leary, C., Maclaurin, D., and Wanderman- Milne, S. JAX: composable transformations of Python+NumPy programs, 2018.
[4] Carleo, G. and Troyer. M. Solving the quantum many-body problem with artificial neural networks. Science, 355(6325):002-66,, 2017.
[5] Dagréou, M., Ablin, P., Vaiter, S., and Moreau, How to compute Hessian- vector products? In International Conference Learning Representations (ICLR) Blogposts, 2024.
[6] Dangel, F., Müller, J., and Zeinhofer, M. Kronecker-factored approximate curvature for physics- informed neural networks. In Advances Neural Information Processing Systems (NeurIPS), 2024.
10

---

[7] Dwivedi, V. and Srinivasan, B. Solution of biharmonic equation in complicated geometries with physies informed extreme learning machine. Journal of Computing and Information Science in Engineering, 20(6), 05 2020. ISSN 1530-9827.
[8] Fa, K. S. Solution of Fokker-Planck equation for broad class of drift and diffusion coefficients. Phys. Rev. E, 2011.
[9] Faa Di Bruno, F. Note une nouvelle formule de caleul différentiel. Quarterly J. Pure Appl. Math, 1857.
[10] Foulkes, W. M., Mitas, L., Needs, R., and Rajagopal, G. Quantum Monte Carlo simulations of solids. Reviews of Modern Physics, 73(1):33, 2001.
[11] Fraenkel, L. Formulae for high derivatives of composite functions. In Mathematical Proceedings of the Cambridge Philosophical Society, 1978.
[12] Gao, N., Köhler, J., and Foster, A. folx forward Laplacian for JAX, 2023. URL http: //githul wwt./ooooo/eeeeeeeeeeee com///icoooooot//ooo...
[13] Griewank, A. and Walther, A. Evaluating derivatives: principles and techniques of algorithmic differentiation. SIAM, 2008.
[14] Griewank, A., Utke, J., and Walther, A. Evaluating higher derivative tensors by forward propagation of univariate Taylor series. Mathematics of Computation, 69, 1999.
[15] Hardy, M. Combinatorics of partial derivatives, 2006. arXiv.
[16] Hermann, J., Schätzle, Z., and Noé, F. Deep-neural-networrk solution of the electronic Schrödinger equation. Nature Chemistry, 12(10):91-887, 2020.
[17] Hu, Z., Shi, Z., Karniadakis, G. E., and Kawaguchi, K. Hutchinson trace estimation for high- dimensional and high-order physics-informed neural networks. Computer Methods in Applied Mechanics and Engineering, 2024.
[18] Hutchinson, M. A stochastic estimator of the trace of the influence matrix for laplacian smoothing splines. Communication in
and Computation, 1989.
[19] Karniadakis, G. E., Kevrekidis, I. G., Lu, L., Perdikaris, P., Wang, S., and Yang, L. Physics- informed machine learning. Nature Reviews Physics, 3(6):422-448) 2021.
[20] Li, R., Wang, C., Ye, H., He, D., and Wang, L. DOF: Accelerating high-order differential operators with forward propagation. In International Conference on Learning Representations (ICLR), Workshop on H4DDifierentiallquatie
In Science, 2024.
[21] Li, R., Ye. H., Jiang, D., Wen, X., Wang, C., Li, Z., Li, X.. He, D., Chen, Ren, W., al. A computational framework for neural network-based variational Monte Carlo with forward Laplacian. Nature Machine Intelligence, 2024.
[22] Paszke, A., Gross, S., Massa, F,, Lerer, A., Bradbury, J., Chanan, G., Killeen, T., Lin, Z., Gimelshein, N., Antiga, L., Desmaison, A., Kopf, A., Yang, E., DeVito, Z., Raison, M., Tejani, A.. Chilamkurthy, S.. Steiner, B.. Fang, L., Bai, J., and Chintala, S. PyTorch: An imperative style, nigh-perffrmance deep learning library. In Advances in Neural Information Processing Systems (NeurIPS). 2019.
[23] Pearlmutter, B. A. Fast exact multiplication by the Hessian. Neural Computation, 1994.
[24] Pfau, D.. Spencer, S., Matthews, A. G., and Foulkes, W. M. C. Ab initio solution of the many-electron Schrödinger equation with deep neural networks. Physical Review Research, 2020.
[25] Raissi, M., Perdikaris, P., and Karniadakis, G. E. 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.
11

---

[26] Reed, J., Ditito, Z., He, H., Ussery, A., and Ansel, J. torch.fx: Practical program capture and transformation for deep learning python. Proceedings of Machine Learning and Systems (MISys), 2022.
[27] Shi, Z., Hu, Z., Lin, M., and Kawaguchi, K. Stochastic Taylor derivative estimator: Efficient amortization for arbitrary differential operators. Advances in Neural Information Processing Systems (NeurIPS), 2024.
[28] Sun, J., Berner, J., Richter, L.,, Zeinhofer, M., Müller, J., Azizzadenesheli, K., and Anandku- mar, A. Dynamical measure transport and neural PDE solvers for sampling. arXiv preprint arXiv:2407.07873. 2024.
[29] Sun, R., Li, D., Liang, S., Ding, T., and Srikant, R. The global landscape of neural networks: An overview, 2020.
[30] Toulouse, J. and Umrigar, C. J. Optimization of quantum Monte Carlo wave functions by energy minimization. The Journal of Chemical Physies (JCP), 2007.
[31] Vahab, M., Haghighat, E., Khaleghi, M., and Khalili, N. A physics-informed neural network approach to solution and identification of biharmonic equations of clasticity. Journal of Engineering Mechanics, 148(2), 2022.
[32] Webber, R. J. and Lindsey, M. Bayleigh-Gauss-Nevve optimization with enhanced sampling for variational Monte Carlo. Physical Review Research, 2022.
12

---

NeurIPS Paper Checklist
1. Claims
Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
Answer: [Yes]
Justification: We propose an acceleration technique for Taylor mode automatic differentiation (AD), deseribe the scope of its applicability and how works, then empirically verify the claimed acceleration capabilities. This is clearly communicated in the abstract.
Guidelines:
The answer NA means that the abstract and introduction do not include the claims made in the paper.
The abstract and/or introduction should clearly state the claims made, including the contributions made the paper and important assumptions and limitations. A No or
NA answer to this question will not be perceived well by the reviewers.
The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings.
It is fine to include aspirational goals as motivation long it is clear that these goals are not attained by the paper.
2. Limitations
Question:
Does the paper discuss the limitations the work performed by the authors?
Answer: [Yes]
Justification: We specifically incorporate paragraph about limitations of the implementation in $4 and state the assumptions (linear differential operator) under which the proposed optimization technique applicable.
Guidelines:
The answer NA means that the paper has limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.
The authors are encouraged to create separate "Limitations" section in their paper.
The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g.. independence assumptions, noiseless settings, model well-specification, asymptotic approximations only holding locally). The authors should reflect on how these assumptions might be violated in practice and what the implications would be.
The authors should reflect on the scope of the claims made, e.g., if the approach was only tested few datasets or with few runs. In general, empirical results often depend on implicit assumptions, which should articulated.
The authors should reflect on the factors that influence the performance of the approach. For example, facial recognition algorithm may perform poorly when image resolution is low or images are taken in low lighting. Or speech-to-text system might not be used reliably provide closed captions for online lectures because it fails handle technical jargon.
The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size.
If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and fairesss.
While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, worse outcome might be that reviewers discover limitations that aren't acknowledged
the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an impor- tant role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations.
3. Theory assumptions and proofs
13

---

Question: For each theoretical result, does the paper provide the full set of assumptions and complete (and correct) proof?
Answer: [NA]
Justification: The paper does not present new theoretical results. All theoretical results used in this work (specifically, [14]) are cited, and we include self-contained introductions that are relevant to our approach in the appendix (specifically, ŞE...
Guidelines:
The answer NA means that the paper does not include theoretical results.
All the theorems, formulas, and proofs in the paper should be numbered and cross- referenced.
All assumptions should be clearly stated or referenced the statement of any theorems.
The proofs can either appear in the main paper or the supplemental material, but they appear in the supplemental material, the authors are encouraged provide short proof sketch to provide intuition.
Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix or supplemental material.
Theorems and Lemmas that the proof relies upon should be properly referenced.
4. Experimental result reproducibility
Question: Does the paper fully disclose all the information needed to reproduce the main ex- perimental results the paper to the extent that affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided or not)?
Answer: [Yes]
We describe the experimental procedure, including hardware and hyper- parameter details, in $4, together with a more detailed description in SF and G, where we also partially reproduce our results using a different ML library. We will open-source the code used to produce our results to further facilitate their reproducibility.
Justification:
Guidelines:
The answer NA means that the paper does not include experiments.
If the paper includes experiments, a No answer to this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.
If the contribution is dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable.
Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is a novel architecture, describing the architecture fully might suffice, or if the contribution is a specific model and empirical evaluation, it may be necessary to either make possible for others to replicate the model with the same dataset, or provide access to the model. In general. releasing code and data is often one good way to accomplish this, but reproducibility can also be provided via detailed instructions for how replicate the results, access to a hosted model (e.g., in the case of a large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed.
While NeurIPS does not require releasing code, the conference does require all submis- sions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example
(a) the contribution is primarily a new algorithm, the paper should make clear how reproduce that algorithm. (b) If the contribution is primarily the architecture clearly and fully.
model architecture, the paper should describe
(c) the contribution a new model (e.g., a large language model), then there should either be way to access this model for reproducing the results or way to reproduce the model (e.g., with an open-source dataset instructions for how to construct the dataset).
14

---

(d)
reproducible, they We recognize that reproducibility may be tricky in some cases, in which case authors are welcome describe the particular way they provide for reproducibility. In the case of closed-source models, it may be that access to the model is limited in some way (e.g., to registered users), but it should be possible for other researchers to have some path to reproducing or verifying the results.
5. Open access to data and code
Question: Does the paper provide open access to the data and code, with sufficient instruc- tions to faithfully reproduce the main experimental results, as described in supplemental material?
Answer: [Yes]
Justification: We will open-source the code for our PyTorch Taylor mode library (which is not the main contribution of our paper), as well as the code used to generate our results, including instructions how reproduce them. See the link the main text.
Guidelines:
The answer NA means that paper does not include experiments requiring code.
Please see the NeurlPS code and data submission guidelines (uttps://nips. cc/
for more details.
While we encourage the release of code and data, we understand that this might not be possible, "No" is an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (c.g., for new open-source benchmark).
The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurIPS code and data submission guidelines (https:
httpsscccccbcccccceeeeeeeeeccect /ttgg//ede//oeuumaaaiaeaeee //o///ooooee/eeeeeeee wssiisssssssioiooeoo oniooiccy for more details.
The authors should provide instructions on data access and preparation, including how access the raw data, preprocessed data, intermediate data, and generated data,
The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. If only subset of experiments should state which ones are omitted from the script and why.
At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).
Providing as much information as possible in supplemental material (appended to the paper) recommended, but including URLs to data and code permitted.
6. Experimental setting/details
Does the paper specify all the training and test details (e.g., data splits, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?
Question:
Answer: [Yes]
Justification: Our experiments consist of performance measurements. We clearly specify the protocol, used hardware, architectural details, and other haper-parameters in $4, with more details $F and
Guidelines:
The answer NA means that the paper does not include experiments.
The experimental setting should be presented in the core of the paper to a level of detail that is necessary to appreciate the results and make sense of them.
The full details can be provided either with the code, in appendix, or as supplemental material.
7. Experiment statistical significance
Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?
Question:
Answer: [Yes]
15

---

Justification: While we do not report enror bars, we choose the experimental procedure to minimize noise: To measure runtime, we report the smallest number from repeating run 50 times. Reporting mean and standard deviations is less conclusive because runtimes on hardware can vary due interference from other processes or warm-up effects during the first execution. Reporting the minimum time best approximation to the actual runtime is recommended, c.g. https: :/ http://raallpttttt cwm/pythom conm/pyttoo-ppofllinnn profiling! Guidelines:
The answer NA means that the paper does not include experiments.
The authors should answer "Yes" if the results are accompanied by error bars, confi- dence intervals, statistical significance tests, least for the experiments that support the main claims of the paper.
The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall run with given experimental conditions).
The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.)
The assumptions made should be given (e.g.: Normally distributed errors).
should be clear whether the error bar is the standard deviation the standard error of the mean.
It is OK to report 1-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have 96% CI, if the hypothesis of Normality of errors is not verified.
For asymmetric distributions, the authors should be careful figures symmetric error bars that would yield results that are out of range (c.g. negative eтror rates).
show in tables
error bars are reported in tables or plots, The authors should explain the text how they were calculated and reference the corresponding figures tables in the text.
8. Experiments compute resources
For each experiment, does the paper provide sufficient information on the com- puter resources (type of compute workers, memory, time execution) needed to reproduce the experiments?
Question:
Answer: [Yes]
Justification: We mention the exact type of GPU used in $4. Since we evaluate runtime and memory performances, we used a single GPU. The total compute time can estimated by multiplying our reported timings with the number of repetitions, then summing over all measurement points. totals less than 24 GPU hours.
Guidelines:
The answer NA means that the paper does not include experiments.
The paper should indicate the type of compute workers CPU or GPU, internal cluster, cloud provider, including relevant memory and storage.
The paper should provide the amount of compute required for each of the individual experimental runs as well as estimate the total compute.
The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary or failed experiments that didn't make it into the paper).
9. Code of ethics
Question: Does the research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics https: cc Couublic/Ethcsshiii ines?
Answer: [Yes]
Justification: After carefully reading the Code of Ethics, we believe our research conforms with it.
Guidelines:
16

---

The answer NA means that the authors have not reviewed the NeurlPS Code of Ethics.
If the authors answer No, they should explain the special circumstances that require deviation from the Code of Ethics.
The authors should make sure to preserve anonymity (e.g., if there is special consid- eration due to laws or regulations in their jurisdiction).
10. Broader impacts
Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?
Answer: [NA]
Justification: Our work introduces algorithmic improvements to Taylor mode automatic differentiation. We think they positively impact the research landscape of scientific machine learning, by reducing computational resources. We do not foresee any direct negative societal impacts of our work.
Guidelines:
The answer NA means that there is no societal impact of the work performed.
If the authors answer NA or No, they should explain why their work has no societal impact or why the paper does not address societal impact.
Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deploymeen of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.
The conference expects that many papers will be foundational research and tied to particular applications, let alone deployments. However, if there is direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster.
The authors should consider possible harms that could arise when the technology is being used intended and functioning correctly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional unintentional) misuse of the technology.
If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms to monitor how a system learns from feedback over time, improving the efficiency and accessibility of ML).
11. Safeguards
Question:
Does the paper describe safeguards that have been put in place for responsible release of data models that have a high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?
Answer: [NA]
Justification:
Our work neither uses data sets, nor trains models.
Guidelines:
The answer NA means that the paper poses such risks.
Released models that have a high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere to usage guidelines or restrictions to access the model or implementing safety filters.
Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.
We recognize that providing effective safeguards is challenging, and many papers do not require this, but we encourage authors to take this into account and make best faith effort.
17

---

12. Licenses for existing assets
Question: Are the creators or original owners of assets (e.g.. code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?
Answer: [Yes]
Justification: We give credit to code packages we use through citations, specifically JAX's Taylor mode [2] and forward Laplacian [12], as well torch. [26].
Guidelines:
The answer NA means that the paper does not use existing assets.
The authors should cite the original paper that produced the code package or dataset.
The authors should state which version of the asset is used and, if possible, include URL.
The name of the license (e.g., CC-BY 4.0) should be included for each asset.
For scraped data from a particular source (e.g., website). the copyright and terms of service of that source should be provided.
If assets are released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, paperswi thcode. com/ datasets has curated licenses for some datasets. Their licensing guide can help determine the license of dataset.
For existing datasets that are re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided.
this information not available online, the authors are encouraged to reach out to the asset's s creators.
13. New assets
Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?
Answer: [Yes]
Justification: Our PyTorch implementation of Taylor mode is fully fully-documented and tested. We are planning to open-source soon.
Guidelines:
The answer NA means that the paper does release assets.
Researchers should communicate the details the dataset/codeemmodll as part of their submissions via structured templates. This includes details about training, license, limitations, etc.
The paper should discuss whether and how consent was obtained from people whose asset used.
At submission time, remember to anonymize your assets (if applicable). You can either create anonymized URL or include an anonymized zip file.
14. Crowdsourcing and research with human subjects
For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well details about compensation (if any)?
Question:
Answer: [NA]
Justification:
The paper does not involve crowdsoureing research with human subjects.
Guidelines:
The answer NA means that the paper does not involve crowdsourcing research with human subjects.
Including this information in the supplemental material is fine, but if the main contribu- tion of the paper involves human subjects, then as much detail as possible should be included in the main paper.
18

---

According to the NeurlPS Code of Ethics, workers involved in data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector.
15. Institutional review board (IRB) approvals or equivalent for research with human subjects
Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or equivalent approval/teview based on the requirements of your country institution) were obtained?
Answer: [NA]
Justification:
The paper does not involve crowdsourcing research with human subjects.
Guidelines:
The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.
Depending on the country in which research is conducted, IRB approval (or equivalent) may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.
We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors adhere the NeurIPS Code of Ethics and the guidelines for their institution.
For initial submissions, do not include any information that would break anonymity (if applicable), such as the institution conducting the review.
16. Declaration of LLM usage
Question: Does the paper describe the usage of LLMs if it is important, original, or non-standard component of the core methods in this research? Note that if the LLM is used only for writing, editing, formatting purposes and does not impact the core methodology, scientific rigorousness, or originality of the research, declaration is not required.
Answer: [NA]
Justification:
Our work does not rely on LLM usage in any way.
Guidelines:
The answer NA means that the core method development in this research does not involve LLMs as any important, original, or non-standard components.
Please refer to our L.M policy (nttps://neuri)) wwweurrps.ce/oooeeeeeeeeeeeeeeee ere/Conneeeecce//225////// for what should or should not be described.
19

---

Collapsing Taylor Mode Automatic Differentiation (Supplementary Material)
Faà Di Bruno Formula Cheat Sheet
21
B Visual Tour: From Function to Collapsed Taylor Mode
22
Graph Simplifications
23
D Exploiting Linearity to Collapse Taylor Mode
26
D.1 Second-order Operators: Laplacian
27
(Collapsed) Taylor Mode for Arbitrary Mixed Partial Derivatives
28
E.1 Applied to the Biharmonic Operator
29
E.2 Pedagogical Approach for the Biharmonic Operator with 6-jets
30
E,3 Another Example: Mixed Third-order Derivatives
F
Py Torch Benchmark
32
F.1
Additional Analysis and Impact of torch. for.............................
32
F.2
Rank-deficient Weighted Laplacian
34
JAX Benchmark
34
H Numerical Complexity and Error Analysis
37
I Connections of Collapsed Taylor Mode to Existing Methods
39
L1 Connection to Randomized Laplacian via Hutchinson's 's Trace Estimator
39
1.2 Connection to the Forward Laplacian
40
20

---

A Faà Di Bruno Formula Cheat Sheet
To give some intuition on the Faa di Bruno formula, we illustrate eq. (4) for higher orders here:
h(wp) h1 (0h, 21)
a(bo) fa ( 91 (0g. h1) =/1 (0f. 21)
81
hg + (0h, 22)
#1
92
18 g. h
+ (3g. ha)
= (0 x
+0f ma)
(8 h. )
(8 g. h ;
10 f
hg + h.
+-388 hy
=/= = +3 f.
+000, ms!
(0g. hg)
(0 ma)
(8*h, at P4)
+0{88 ( A )1)
+0 %.
+6
9 h2?
m
24
14 +4(0 mi & ma?
94
+4 g. h1 & ha)
=14
+4:0 f ma)
3(0 m2
+-3 g. Ag
1.0g 00 (05.
(0h,
+ h4)
(0 g. e
+10 +10t0",, A.
m2)
+10 g. 98 h2)
+10 f. mm)
+10 h,, #1
+1000 11
ha)
+10
ma)
255
15
+1500 h. 2
95
+15(0" ty @ 2
=/
+15 f.m m
+ +5(0* h.
m4)
+5(0" g. h1 8 h4)
+5(0 aa m4!
10 h. 20h.
+10 h2 & ha) + +0g. ha)
+10 f.
+(ff.g))
a 06)
(0 a. 29,
( +15
+15 h.
+15 g. S b2)
+20 e
+2018 g. h
ha)
+20 an
+45(0 h, 8
+45 g. h h
+45 f.
+15 hh m1 (
+15(0" g. 822 h4)
+15 02 Sap
36
hg
96
+60(8" @ hg 9 hg)
=f +00000 f.ay
d Sag)
15 h. ml
+15 g. h
+15.8 f.
46 h.
ms)
+610
lgg)
+0(0² J mm)
15:00 h ma & m4)
+ 15 g. by & hu)
+ f. mo (4)
10000 n. 02)
1000 g. Ag
10:0 10:00 S.2 22)
+0h, za)
+ he)
+ +0f me
(8" g. k (
(0" F.
(2 h, m2)
22 a. 2h 8 hyt
12 20 my)
+35 h at mg) + +105 200 21
+3500 g. 84 lgg
ma) +105 f. @ m
+35 J.
10000000 h 8a @ 222
+35(0"h, m4)
+ +3518* a. 22 a ha)
435 J. 1 @ mi)
+210 22 @ h2 hi)
+210 f. a @ @ (ag))
+108:0 98;
+105,0* g.
+105 f. 2
h7
421 mg)
+21 g. 62 ha)
=t
421 J. 1 on)
+108(8 21 8
+105 g. hi @ ha & hi)
+10519 f.z 021
+70 (4)
+70,0" g. by k
+70.0 f. & 02
+105 h, @
+105(0 h @ hg)
+105 0
+ an @ mn)
+ g. bi @ ba)
+7 f. 8es)
+21 æg s
+21(0"g,
hg)
+21 6
+ab h. #1)
+ by ha)
+35 f.2
+ 27)
+ h7)
+05
(ath, 08)
(0 h
(0 a
+28 h. #00(0%h tos @ mz)
+ 28:87 a. h
hg)
+28 f de m2)
@ mn)
+ +o008 a. 25
+50 J.
mn
#210 h, @ 6
4210:3 h
+210;8 f.
+ro(0"h, m @ mg)
+ 770 a. h b4 ha) coona g. A @ ha ha)
+70 f
mn)
soata of wn 25) 4242 of e
+asu f. a ood
od es) +50,0 a @ ug)
4220 g, h @ h +5e ha leg
1420 f.a a
+ h. bo wg) 42008 h. of
e
420 a. h @ ha & hie
+ f. a @
+ #250 a1 6
#28013* g, h @ a 22
+280 J. a @ e 20
hs
+840 mj a ag)
+800 h k2 9 heg =f =
+040 f.a 8 02 @g)
+105 a
+ h. R.1
+1b 1.a
4 +2s(0"h, x mg)
28:0 g. h
hat
428 J. and)
+168 mg mg)
+108 g. h1 @ hgg he)
+168 f 0 ms))
+250 h
+280 g. hi ha h4)
+250 f. 2 88))
+210 h, m22 my)
+210 g. &
hg)
+210 f (
+280
27,
+280 g. b2 n 0a
+280 f.a2 2 6 22
+8:0 h. w0)
+s g, bi by)
+3 my)
+28 h. 2 es)
P0 9 he ho)
22188 f.22 6m)
+56 h. + A,, 92; (0h.
+56(0" g.
he)
+50 f (5)
+ +35 g. (22) (2g. ha)
+3510" 20)
+0f mn)
21

---

B
Visual Tour: From Function to Collapsed Taylor Mode
Derivative degree
Forward propagation
Steps:
1. Overload forward pass with Taylor arithmetic
2. Vectorize (illustration four directions)
.0g
he
go
3. Set anchor for all directions, eliminate redundancies
4. Collapse highest derivative
jet
Derivative degree
Forward propagation
propagation
degree
350
ho
1
vmap
(
{
h;
gi
K
t
hx
K
Derivative degree
Forward prupugation
(Mx--1))
(hx-r1)
igx-1a
jit
K
30g
ho
go
(
(bxcr
(mx>)
(
Derivative degree
Forward propagation
hg
go
K
(
(1x-1r3
jit
(3
(har)
to.r
K
(
(hR--)
(gke)
(xK-1*>)
(kg-1.)
1gx 1.
K
Σ.
Σ, hxg
Σ, 9K.
Figure B6: Step-by-step transformation of function into its collapsed Taylor mode. Collapsed Taylor mode can be implemented via the already existing jet (standard Taylor mode). vmap (vec- torization), and jit (graph transformation) interfaces. Therefore, users do not need learn new interface to benefit from our proposed method and can simply rely on standard Taylor mode and just- in-time compilation (after incorporating our proposed simplifications). Coloured boxes correspond to our Py Torch implementations of standard and collapsed Taylor mode from our experiments.
22

---

C Graph Simplifications
In this section, we illustrate the two graph simplifications that are required to collapse Taylor mode.
We will consider collapsing the 2-jet of f = sin as an example. Recall the propagation scheme eq. (D13) and assume that the Taylor coefficients are given by { æu), ( and {æ2rr) where r indexes the directions along which we evaluate the sum:
zz m1.r) (22r)
Fo.r 30 1.r
for = sin(2r0)
replicate wo
(D13)
fi, cos( @ æ1
22r
f2.
sin( 2 1r ① ir + cos(2)) C 22r
forr
sum highest component
fL.))
82,r
Here, sin applies element-wise and C denotes clement-wise multiplication. The computational graph for this procedure is displayed the following diagram, with input and output nodes highlighted in dark and light gray. The suffix rr means that all R corresponding tensors are stacked along their leading axis. replicate is a function that replicates tensor R times along a new leading axis, which is in PyTorch usually for free and without additional memory overhead (using torch. expand). All other functions refer to those of the PyTorch API:
x0
x0_r replicate(20)
f0.r sin stn(xx_rr)
x1_x
d2f u
-10_r
df cas rx
fi.r eineun >..." dr.r, xi.x)
tampl L = ainsun("........ -.... df.r, x1.r, xl_r)
x2_r
teap2..r einsua(", d2t_x, x2_r)
12_r tampi_r + terp2_r
f2 aun(f2_r, din=0)
Our simplification proceeds in two steps. First, propagate replicate nodes down the graph to remove repeated computations the same tensors. This is done in a forward traversal through the graph. Second, in a single backward traversal through the graph, we propagate the sum node up. After applying both steps, the graph looks follows:
xD
sin(x0)
10.r replicate(EO)
xi_rr
d2f -EO
df cus (x)(x0)
eiarux("r... Sr.." di, xl_r)
terp1 ainaun(" .........
"...", df, xi,r, xl.r)
x2_r
x2 sun (22_, din-0)
texp2 cineun d2e,
teap1
temp2
Two important properties of the new graph are (i) the replicate node moved to an output node, hence the corresponding redundant computation was successfully removed (ii) the highest component x2_r is immediately summed then propagated, i.e., we collapsed Taylor mode and avoid the separate propagation for all x2_r.
We will now illustrate the two simplification steps in full detail. The first stage starts from the original graph and pushes forward the replicate node, as illustrated step-by-step in fig. C7. The second stage starts from the graph produced by the replicate-push procedure, and propagates the final sum node up the graph, illustrated by fig. C8. This yields the final computation graph shown above.
23

---

x0
x0_r (x)
20_r ein(x0_r)
cinsum(" D...", df_r, x1_r)
xi_r
d2f_r ff0_r
df_r cas rx)_
templ_r einsun(".... ....." di_r, xi,r, xi_x)
x2_r
cemp2_r einsun(" " d2f_r, x2_r)
$2_= templ_r temp2_r
2 sun (f2_r, din-0)
x0
10 sin (x0)
10_r roplicate
x0_r replicate(x()
x1_r
d2f_r -f0_r
df_r cxo(xx)r))
f1_r cinsum ("..... df_r. x1_r)
templ_r einsun(", -...", dr,r, xI,r, xi_x)
tamp2_r cinsun("... .>..." d2f_r, x2_x)
x2_x
12_r campl_r tamp2_r
(f2_r, din-O)
x0
10 sin(x0)
_0_r replicate(fO)
cos (x0)
xi_r
d2f_r roplicata(-f0)
df_r (dplicata(dt))
fl_r eineum(". ("..an ...... dr_r, x1_r)
templ_r einsum(", ->..", dr_r, xl,r, xi_x)
x2_r
cemp2_.x einaun(". ....", d x2_r)
12_r templ_r temp2_r
(f2.r, din-O)
ein(x0)
10_r roplicata(tO)
f1_r cinsum ("... F... >r...., df, xi_r)
x1_r
d2f
-10
dr coa (x0)
templ_r cinsum ,...,.... ->r.. ", df, xi, xi_z)
x2_r
temp2_x einsun(" (". .r....", d2f, x2_r)
f2_r
templ_r temp2_r
(f2,r_ din=0)
Figure C7: Step-hy-step illustration of pushing replicate nodes down a computation graph.
24

---

sin(x))
replicate(tO)
flr einaun(""... T.... T............................... dt, xl_r)
x|_r
d2r
-10
dr (o2(x0)
templ_r ainsun ................................ .r..." ", di, xl_r, xi_r)
x2_r
temp2_r cinsum ->r..., d2t, x2_r)
f2.r templ_r tamp2_r
(f2_r, din-O)
x0
60 sin(x0)
10_r replicate (ro)
x1_r
d2f
-10
dr coe (xo)
1Lr einaun ..........
>r..., ", df, xi_r)
templ_r cinsum (*, .>r...", df, xi_r, xi_r)
x2_r
temp2_r einsun E'... >r.... d2f, x2_r)
teapl sun ( (teapl_r, din-0)
temp2 sum (tar222, din=0)
teap1 temp2
sin(x))
f0.r replicata
flr einamm(".... ..,....... df, x1_r)
x1_r
62f -ro
df = coa (x0)
temp1 cinsun(". ................ df, xl_r, x1_r)
tamp? sinaun(". ..,..... .>...", d2f, x2_r)
teapl temp2
x0
sin(x0)
fo.r replicate(1O)
fl.r cinsum("..... >r..", dr, x1_r)
x1_r
d2f
-f0
df coa cxa(x))
temp1 einaun ......,,...... >>.."' df, xi_r, xirr)
temp2 einsun (1) ....,,.... d2f,
x2_r
(x2__r, dim=o)
f2 templ
tenp2
Figure C8: Step-by-step illustration of propagating nodes up computation graph.
25

---

D Exploiting Linearity to Collapse Taylor Mode
Here, we illustrate the, idea behind propagating K K-jets through g h with input jets (J* (J*x), (t) Ek=o kt Thr The Taylor mode scheme results from inserting eq. (5) into eq. (4):
hg = h(()))
(hir) (0h( x1,r))
2D {air5}
{22rr = (20)) 2i,r @ x1.r) + (Oh(zo),
(x2.7)
Σ u(o) ( (Jloln(za)) 8
m
(
(KK,))
sEa
(Epart(K)
9o g(ho)
[9r
(0g(ho)) h1r1)
(g2,)
( g(o)) hi 8 h1s) (8g(ho), h2rr)
E v(o) (gllg((t)), 6 mear)
[g
(D13)
DEa
(Epart(K)
fo = f(zo)
[f1r]] = (0f 21rr)
(f2,r = @ f(xo) I1.r 2 xxrr) +
vlo) (ohl)((s)) f(x0), @ ....))
(SKr) =
Σ (Epart(K)
slice
-
{gKr}}
R
R
sump Σ
(2)
aa f(xo), @ _ v.)
gK2
Leveraging linearity in certain terms (in green) of the highest coefficient, as explained in eq. (6), instead leads to
ho
h(xp) (()))))
(hrr)
æp
(12,r
((Phh(z))) F1,r 8 a1,r) + (8h(xa), 22r1)
{
{ (v2.r
E2
Σ hgrr Σ
Σ
v(a) (/ (al1)((z)) 2 xx.r
R
Σ
o'pan(())
sEa
2 KKr
ah(wo)
DKKr
g(ho)
(g..)
(ho); hnrr))
{g2.r}
( (ho), hirr h)rr) + (a)))))
(D14)
Σ 9K
Σ
Σ
v(a) (gg)))))) & harr
()part(K))(G)
SEa
R
Og(ho),
Σ
hkrr
Pml
26

---

fo
f(x))
(fir)
(2)) x1,5)
(f2r)
(0") (z0) 21 8 20.1)) + (0f(2))) x2.1)
R
R
fk.
Σ
v(o) y10) f(zo), Zas sEa
Mp
@f(
Σ KKrr
slice
1
Ry 22
0)
f(x)), 1
9K.r
D.1
Second-order Operators: Laplacian
Here, we show details about the propagation schemes of standard Taylor mode and collapsed Taylor mode for the computation of the Laplacian of f. We consider the decomposition f o h.
Standard Taylor mode.
Using standard Taylor mode (eq. (D13)) compute the Laplacian yields
20
hn h(x))
[11,a] = (0h( x1,4)}
222,
[h2.d] { (02h1 Did @ x1.a) + (0h(xo), x2a))
90 00=g(ho) g(ho)
(g.d)) = (gg h2,d))
(92,a) { 00 (ho), hidd h1,)) + (8g(ho),
fo = f(z))
(D15)
[fi.a] (0 x1d)}
(f2,a 08 f(x) Fid 8 aa.a) (0f(zz).
slice
{g2
(m) [e
(92,)
A (xoo)
1=1
Collapsed Taylor Mode AD
Using our proposed collapsed Taylor mode, we get
To {x1.}}
ho = h(xo)
(hia) = (0h(20)) x1,)))
Σ 12 (n (oh(zo)) d=1
(l
Σ 22
22
d=1
d=1
90 g(ha)
(g) = ((0g(ho): k1,0)
D
D
D
Σ 92.d Σ
1 (ho)) hid & hid) + Og(ho): Σ
h2d
(D16)
d=1
d=1
fo = f(x))
[fid] = ((0f(x)))
Σ fe.d = Σ 0 (zx) Li,d x1d) @f((z)), Σ x2,d
slice
(g2d) ((1.==c==.-=0))
Af(
27

---

E (Collapsed) Taylor Mode for Arbitrary Mixed Partial Derivatives
Here, we introduce the notation of eq. (11). The right side of the formula sums over the set of all j E N' such that ||اا| = E.bik = K. For example, if I = 2 and |اllı 4, this set is {(4,0), 0), (0, 4), (3, 1), (1.3). (2,2)).
The coefficient Vi. is defined as
Ye.j
Σ UcmSi
(-1)lli-mli
[ m
|| !iaalli
(Imlll 1|4] ||:||
[
(E17)
The summation ranges over the set {m E N more, utilize the generalized binomial coefficient
[m] S | [m]r s [t]r, || > 0}. Further-
6-
:
allow the computation for all R and E N, which defined to be if = 0. The generalized binomial coefficient vectors the component-wise product of generalized binomial coefficients:
b
This notation also includes cases where the vector a has components of RR
Example computation.
Let us compute the coefficient (32.2).(31)- used by the biharmonic operator:
(322)
Σ
[m))
2
[m]
[m]] |mlli
|m
[m]]
1
| |1
4
mEN. mm] [m] [m
We have ((,)0,, (1,1), (2, 2), (1, (0. 1), (0, 2)}, which results in the terms
(-1)2-1+2-0
4
0
1
4
22
42
4.0
2
((-1)2-2+2-0
2
0
+(-1)2- 1
(2
2
4
2
4
1
+(-1)2 2+2-1 2-22211
2
(2
2
42
+(-1)2 2-222-2
2
1-112-1+2-
4
2(-1)2-0+2-1
(2
4
4.0
44
2
+(-1)2-0+2- 22222
4
28

---

The next step is to evaluate the binomial coefficients:
1
(-1).2-1.4.00
4
2
+1-1-4-0-
4
2 4
4
42.2.0-2
852
4
3
-1-2
963
4
+1-1-0-2-
41-2 8
(3) 4
-2.
96 3
4
-1.2-0.4.
2
+1-1-1-0.4
4
After simplification, the final result
852
4
3
4
2
8
3
4
(222). =(-1)
2
963
3
4
9
3
4
-40 81
128
81
-5
1
486
256
486
256
12
12
3
E.1 Applied to the Biharmonic Operator
To compute eq. (9) with eq. (II), we first select K = = =2,)= vd, = ed. and ed, Edy. Then we insert these parameters into the general equation eq. (11) and get
A2f(z0)=
Σ
1
Σ
Σ
(())) (ea, (il1 + ed, (5l2))1) 84)
7(2.2),5
4!
JEN
==
(22.))(4.0)
Σ
Σ
(())) (4ea))01)
d==1d d==1
D
+ (2.22)
Σ
f(x)), (4ex,)
d1 d==1
(E18)
D
0 f(0)), (3ea, + Edis) 04
+ (3(2))).331)
di =1= d==1
D
11 ()=)) (ed, +3edz) 181)
+ (1..))))(133
d1
d==1
D
Σ
(z()) (2ea, + 2eds) (11)
(22.2))(2.2)
d==1
29

---

Now, exploit the symmetry of the coefficients (222))((40) = (2,2)).(0.4) and (22.2) (122))
and the corresponding tensor basis expansion:
1
2D
(2.2)
Σ
0+ f(xo)), (4es,)
24
d1=1
D
(2.2).(33.
Σ
Σ
f(ro), (3ed, + )01)
(E19)
d=
D
D
Σ
((
Σ
f(ro)
(2ed; 20dz)) )01)
+ ((2.2)(222)
d==1d==1
Since the first sum captures all diagonal directions cd, = Cdz, we extract this from the second and third sums to further reduce the computational effort:
24 (
(2D + (7.2))..3. + (2.2)).
8 f(xo), (4es,) )04)
=1=1
D
D
Σ
Σ
( (3ea, + edg 101) )04)
+ (3.2)
(E20)
d==1
D
D
Σ
(xx))) + 2edz) )
+ (2222)
d==1
d2
Exploiting further symmetries in the last term's summation, we obtain
f2)()))
1
(2D + 27(22)).33.) + ((2.2)
(0'f(z)), )4es))
24
d1=1
D
+ (2222),
Σ
Σ
0 f(xo), (3ea. - eezz)
(E21)
d==1
da
D-1
Σ d=id;=dit1
Σ
(0)(( (2e,, 20 )01)
+ (2.22))
E.2 Pedagogical Approach for the Biharmonic Operator with 6-jets
different approach compute arbitrary-mixed derivatives was proposed in [27]. This approach relies, for the biharmonic operator, on the hand-selection of certain 6-jets to extract the required derivatives. The degree and directions for the jets are obtained by considering the Faà di Bruno formula for the 6-th coefficient fo (see ŞA). Selecting coefficients of the input 6-jet to ed, 22 = Cd2
and Ø3 24 25
36g = leads us to
82) fs f(xo), e 20 + 15 10" f(xo). ed, x eda))
(B22)
+ 45 (8+ (xo)) ea @ ed 15 (03 f(x)), e 83
Notice the blue term, which has the same structure the summands we want compute for the biharmonic operator. Therefore, a first 6-jet is computed as explained above. To cancel out the unwanted terms, we evaluate another 6-jet with the same input except 22 edg and adding the 6-th coefficient of this jet to eq. (E22) gives
2 (8 f (2o), e 006) + 90 (a f(x)). cd
(E23)
de
Finally, third 6-jet is computed with x2 0. The 6-th coefficient of this jet contains only
(0 f (x:p), (00)
(E24)
30

---

We obtain
90 (0* f(o)),
21 002 & x²2
(E25)
by subtracting twice of the 6-th coefficient of the third jet from eq. (E23).
To summarize the procedure, we evaluate the 6-jet three times. The first jet has the input æ1 e 22 Edg and Ø3 24 25 = 2e O, the second jet has the same input jet apart from 2g -dg and the third 6-jet takes æ² O. Then we add the 6-th coefficient of the first and the second and subtract twice of the 6-th coefficient of the third jet. Dividing by 90 provides the derivative corresponding the d1, d2 term of the biharmonic operator.
Standard Taylor mode would propagate + 18D vectors through every node, if we already exploit that all jets share zt- collapsed Taylor mode would pass 1+3 1572 vectors through every node of the compute graph. This more costly compared to our approach described before. In addition, until now, the selection of the jet degree and the input coefficients requires substantial human effort.
E.3 Another Example: Mixed Third-order Derivatives
As an additional example, consider computing Σ E 82 Ma)) This example is from [27, $F.2], which describes how compute these 3rd-order derivatives using 7-jets, The interpolation formula allows using multiple 3-jets instead. We expect to be favorable as Taylor mode scales polynomially in the derivative order.
Procedure.
The goal is to compute ΣP i=l
8² 8- (x) We proceed as follows:
j==1 Oazun,
1. Formulate the operator our notation:
ΣΣ f(z) Se,)
2. Compute the interpolation coefficients "p for E ((3,0), (2. (1,2)) and p =
(2, 1): (2.1)(033) - 7(2.1)(122) 16/27, 1221) -16/9, ((21))(3))))))))) 32/81.
3. Apply the interpolation equation (eq. (11)):
D ΣΣ
D
Σ qENZ, ||lll,
(L((x) + (gl2e,)
i=
Collapsed Taylor mode directly applied to these 4D2 3-jets. However, to exploit the full potential some further steps that leverage the structure are required.
4. Expand and manually simplify, using symmetries. The sums for (2.1)(30)) and ((21)) are similar; same for 7(2,1)(,11) and ((21))112) We only have 2D2 3 jets:
D. D
= (7(21)33.0) + 722.)0003)
ΣΣ
10' f(x), (3e,)
iml j=1
Σ 10t f(x), ((x))((e, (2e; + 83)
+ (1221)(21)) + ((2)))22)
We further observe that the first summation independent of j:
(T(21))33)) (22)) D
Σ
( f(x), (3e;)
+ +112.0 + +(200L2)
10' f(x), (2c; + 23)
31

---

Extracting the case i from the last term gives our final form
((72..))330)) ((2.))o,3)) + 7(2,1)(2.1) (211122))
Σ 10 f(x), (3e;)
D
+ (112,)2221) + ((21)(2)))
Σ
Σ
te,)
1=1 j=
This optimized version required D² 3-jets that can be collapsed.
F PyTorch Benchmark
Additional Analysis and Impact torch. compile
Here, compare the theoretically estimated performance improvements based counting the number proaaddpropaggted vectors with the empirically measured performance.
The number of propagated vectors is good empirical performance estimate. To estimate the performance ratio between standard and collapsed Taylor mode, we can use the number of additional vectors both modes propagate forward as we increase either the batch size or the number of Monte-Carla samples. This is a relatively simplistic proxy; e.g., it assumes that each vector adds the same computational load, which inaccurate as vectors corresponding to higher coefficients require more work and memory (as the Faà di Bruno formula contains more terms in general). Conversely, while incrementing the MC samples does add additional vectors that are propagated, does not introduce additional cost to compute store the derivatives, as they are already computed with just single sample. Table F2 summarizes the theoretical and empirical ratios. We find them to align quite well, despite the overly simplistic assumptions.
Concrete example. Consider the exact Laplacian. Adding one datum introduces + D versus 1 1+ vectors. For D 50, their ratio is (2+ (2+UU/2)2) 0.51. Empirically, measure that adding one datum adds 0.56 ms to standard, and 0.29 ms to collapsed, Taylor mode (table 1); the ratio 0.52 is close.
Compilation reduces memory, but not runtime. In table F3, repeat the benchmark from table using torch. compile. We observe that compiling can further reduce the memory footprint of all approaches for computing the Laplacian and weighted Laplacian, while the runtime remains roughly identical. For the biharmonic operator, we observe that compilation leaves runtime and memory footprint unchanged.
Table F2: Comparison of theoretical and empirical performance ratios between standard and collapsed Taylor mode. We list the number of additional vectors that are used when adding another data point (exact) or another Monte-Carlo sample (stochastic). The ratio of vectors offers a good
estimate of the empirically measured performance ratio.
Add datum
Laplacian (D 50)
Weighted Laplacian
Biharmonic
Mode
or MC sample
(D R 50)
(D 5)
# vectors (standard)
+ 2D
6D 2D + %/22² 3/2D + 4
vectors (collapsed)
2+ D
2 2 R
Exact
Thcoretical ratio #/#
0.51
0.51
0.77
Empirical time ratio
0.52
0.51
0.76
Empirical mem. ratio
0.45
0.45
0.63
vectors (standard)
2
vectors (collapsed)
1
3
Stochastic
Theoretical ratio #/#
0.5
0.5
0.75
Empirical time ratio
0.51
0.51
0.74
Empirical ratio
0.43
0.45
0.59
32

---

Laplacian
Weighted Laplacian ( 50)
Biharmonic
(D 50)
(D
Time [s]
1.0
0.5
0.0
0.0
0.0
Exact
Mem. [GiB]
5.0
2.5
0.0
0
1000
2000
1000
2000
500
1000
Nested -order
Standard Taylor
Batch size
Batch size
Batch size
Collapsed (ours)
(N 2048)
(N 2048)
(N = 256)
Differentiable
Non-diff.
Time [s]
1.0
1.0
0.5
0.5
Mem. [GiB]
Stochastic
5
20
0
40
20
Samples
Samples
Samples
Figure F9: Same as fig. 5, but using torch. compile.
Table F3: Same as table 1, but using torch. compile (i.e. fig, F9 in numbers).
Per-datum or -sample cost
Weighted Laplacian
Biharmonic via interpolation (E21)
Mode
Implementation
Laplacian
Nested 1 1"orrder
0.59 (1.0x)
0.58 (1.0x)
1.2 (1 (1.0x) 0.90 (0.72x)
Time [ms]
Standard Taylor
0.55 (0.93x)
0.55 (0.95x)
Collapsed (ours)
0.28 (0.48x)
0.29 (0.49x)
0.69 (0.55x)
Nested | 3-orrder
3.6 (1.0x)
3.4 (1.0x)
7.9 (1.x)
Exact
Mem. [MiB]
Standard Taylor
3.1 (0.88x)
3.1 (0.92x)
7.7 (0.98x)
(differentiable)
Collapsed (ours)
1.5 (0.43x)
1.5 (0.43x)
4.8 (0.61x)
Nested ("-order
1.2 (1.0x)
1.1 (1.0x)
3.5 (1.0x)
Mem. [MiB]
Standard Taylor
0.85 (0.69x)
0.84 (0. 73x)
1.2 (0.36x)
(non-diff.)
Collapsed (ours)
0.60 (0.49x)
0.60 (0.53x)
1.1 (0.33x)
Nested 1 -'-order
23 (1.0x)
23 (1.0x)
44 (1.0x)
Time [ms]
Standard Taylor
23 (0.98x)
23 (0.98x)
6.6 (0.15x)
Collapsed (ours)
12 (0.50x)
12 (0.50x)
4.9 (0.11x)
Nested 1 oroorerr
150 (1.0x)
150 (1.0x)
210 (1.0x)
Stochastic
Mem. [MiB]
Standard Taylor
(0.95x)
140 (0.95x)
(0.30x)
(differentiable)
Collapsed (ours)
73 (0.49x)
73 (0.49x)
38 (0.18x)
Nested f"-order
50 (1.0x)
50 (1.0x)
86 (1.0x)
Mem. [MiB]
Standard Taylor
33 (0.67x)
33 (0.67x)
15 (0.17x)
(non-diff.)
Collapsed (ours)
17 (0.33x)
17 (0.33x)
15 (0.17x)
33

---

Table F4: Exact weighted Laplacian for different ranks of the weightings. The lower the rank, the lower the memory and time consumption. Both scale approximately linear with the rank- stronger deviations for small ranks predicted by our theoretical analysis based on the number of prooaadproppggteed coefficients. The 'Full-rank' column is identical to the 'Weighted Laplacian" columns in Tables and F3.
Per-datum Implementation
Weighted Laplacian (D
50)
compile
sammple cost
Full-rank (D)
Half-full rank (P/2) Low-rank (P/10)
Nested -order
0.58 (1.0x)
0.30 (1.0x)
0.082 (1.0x)
Time [ms]
Standard Taylor
0.55 (0.95x)
0.28 (0.92x)
0.062 (0.76x)
Collapsed (ours)
0.29 (0.49x)
0.15 (0.49x)
0.040 (0.49x)
Nested 1rrorder
3.4 (1.0x)
1.8 (1 (xx)
0.43 (1.0x)
Mem. [MiB]
Standard Taylor
3.1 (0.92x)
1.6 (0.90x)
0.33 (0.78x)
(differentiable)
Collapsed (ours)
1.5 (0.43x)
0.80 (0.46x)
0.21 (0.50x)
Nested -order
1.1 (1.0x)
0.62 (1.0x)
0.15 (1.0x)
Mem. [MiB]
Standard Taylor
0.84 (0.73x)
0.47 (0.76x)
0.12 (0.82x)
(non-diff.)
Collapsed (ours)
0.60 (0.53x)
0.31 (0.49x)
0.073 (0.48x)
Nested 1rdorder
0.60 (1.0x)
0.31 (1.0x)
0.084 (1.0x)
Time [ms]
Standard Taylor
0.57 (0.94x)
0.29 (0.92x)
0.065 (0.78x)
Collapsed (ours)
0.29 (0.48x)
0.15 (0.49x)
0.042 (0.51x)
Nested -rddr
4.4 (1.0x)
2.2 (1.0x)
0.53 (1.0x)
Mem. [MiB]
Standard Taylor
4.6 (1.0x)
2.4 (1.0x)
0.60 (1.1x)
(differentiable)
Collapsed (ours)
2.1 (0.47x)
1.1 (0.50x)
0.39 (0.74x)
Nested -order
2.2 (1.0x)
1.1 (1.0x)
0.26 (1.0x)
Mem. [MiB]
Standard Taylor
1.2 (0.54x)
0.60 (0.54x)
0.14 (0.53x)
(non-diff.)
Collapsed (ours)
0.90 (0.41x)
0.46 (0.41x)
0.11 (0.41x)
F.2 Rank-deficient Weighted Laplacian
In the main text we use a weighted Laplacian with a full-rank weight matrix (i.e., R := r=rkk((D) = D). Since the weight matrix has full rank, the weighted Laplacian as expensive the unweighted Laplacian, and this is confirmed by our experiments. To show that the weight matrix's 's rank indeed affects the cost, ranks R E {P/2, P/10]. Table F4 contains the results of this analysis. We observe that going from full half-full rank roughly halves both the runtime and memory consumption for all implementations. For small ranks, this linear relationship weakens because the fraction of computations that do not scale with grows.
experiment with rank-deficient weight matrix this section and also consider the
G JAX Benchmark
This section presents experiments which show that the graph simplifications we propose to collapse standard Taylor mode are currently not applied by the jit compiler in JAX.
Comparing JAX implementations. Similar to our PyTorch experiment in $4, we compare three implementations of the Laplacian in JAX (all compiled with jax.jit):
1. Nested oororder AD computes the Hessian using jax. hessian, which relies on forward- over-reverse mode, then traces it.
2. Standard Taylor mode propagates multiple univariate Taylor polynomials, each which computes one element of the Hessian diagonal, then sums them to obtain the Laplacian. This implemented with jax. experimental
jet.jes and jax. vmap.
3. Collapsed Taylor mode relies on the forward Laplacian implementation in JAX provided by the folx library [12] and implements our proposed collapsed Taylor mode for the specific case of the Laplacian. folx also enables leveraging sparsity in the tensors, which is
34

---

Laplacian (D 50)
Biharmonic (D
(via nested Laplacians)
1.5
0.6
Time [s]
Time [s]
1.0
0.4
0.5
0.2
0.0
0.0
Exact
Mem. [GiB]
Mem. [GiB]
5
1000
1500
200
400
Batch size
Batch size
(N 2048)
(N 256)
1.5
0.6
Time [s]
Time [s]
1.0
0.4
0.5
0.0
0.0
Stochastic
Mem. [GiB]
10
Mem. [GiB]
2
0
10
20
30
40
50
0
10
20
30
Samples
Samples
Figure G10: JAX's jit compiler does not apply our graph simplifications to standard Taylor mode. Colors: Collapsed Taylor mode, standard Taylor mode, and nested first-order automatic differentiaion, opaque memory consumptions are for non-differentiable computations. Results GPU and we use a D - 768 + 768 + 512 + 512 + 1 MLP with tanh activations, varying the batch size. For each approach, relative speedup and memory reduction.
line the data and report the slope in table G5 quantify the
beneficial for architectures in VMC. To disentangle runtime improvements from sparsity detection versus collapsing Taylor coefficient, disable folx's sparsity detection.
For the biharmonic operator, we simply nest the Laplacian implementations.
We only investigate computing the exact Laplacian, as the forward Laplacian in folx currently does not support stochastic computation. We use the same neural network architecture as for our PyTorch experiments, fix the input dimension to = 50 and vary the batch size, recording the runtime and peak memory with the same protocol as described in the main text. JAX is purely functional and therefore does not have mechanism to build up a differentiable computational graph similar to evaluating a function in PoTorc where some Ieafs have requires_grad-TTee To approximate the
35

---

Table G5: JAX Benchmark from fig. G10 in numbers. We fit linear functions and report their slopes, i.e, how much runtime and memory increase when incrementing the batch size. All numbers shown with two significant digits and bold values are best according to parenthesized values.
Per-datum or sample cost
Biharmonic
Mode
Implementation
Laplacian
(via nested Laplacians)
Nested 1 -*-order
0.58 (1.0x)
0.87 (1.0x) 1.5 (1.7x)
Time [ms]
Standard Taylor
0.84 (1.5x)
Collapsed (ours)
0.29 (0.50x)
0.29 (0.33x)
Nested 1*-order
6.0 (1.0x)
7.0 (1.0x)
Exact
Mem. [MiB]
Standard Taylor
5.1 (0.85x)
12 (1.8x)
(differentiable)
Collapsed (ours)
1.9 (0.32x)
2.0 (0.29x)
Nested or-ord
1.4 (1.0x)
2.7 (1.0x)
Mem. [MiB]
Standard Taylor
1.7 (1.2x)
2.2 (0.83x)
(non-diff.)
Collapsed (ours)
1.4 (1.0x)
1.1 (0.39x)
Nested -rrore
24 (1.0x)
21 (1.0x)
Time [ms]
Standard Taylor
35 (1.5x)
9.5 (0.46x)
Collapsed (ours)
Not implemented
Not implemented
Nested -order
270 (1.0x)
190 (1.0x)
Stochastic
Mem. [MiB]
Standard Taylor
230 (0.87x)
77 (0.40x)
(differentiable)
Collapsed (ours)
Not implemented
Not implemented
Nested 1 orddr
58 (1.0x)
77 (1.0x) 7.9 (0.10x) Not implemented
Mem. [MiB]
Standard Taylor
71 (1.2x)
(non-diff.)
Collapsed (ours)
Not implemented
peak memory of computing a differentiable Laplacian in JAX, we measure the peak memory of first computing the Laplacian, then evaluating the gradient w.r.t the neural network's 's parameters which backpropagates through the same computation graph built by PyTorch.
Results (Laplacian). The left column of fig. G10 visualizes the performance of the three imple- mentations. We fit linear functions to each of them and report the cost incurred by adding one more datum to the batch in table G5. From them, we draw the following conclusions:
1. Performance is consistent between Py Poroh and JAX. Although our Py Torch implementa- tion does not leverage compilation, the values reported in tables I and G5 are consistent and only in rare cases differ by factor of more than two. This confirms that our Py PyTorch-based implementation of Taylor mode is reasonably efficient, and that the presented performance results in the main text are transferable to other frameworks like JAX.
2. Our implementation of collapsed Taylor mode based on graph rewrites in PyTorch achieves consistent speed-up with the Laplacian-specifie implementation in JAX. Specif- ically, we observe that collapsed Taylor mode/forward Laplacian use roughly half the runtime of nested 1d-order AD (compare tables I and G5). This supports our argument that our collapsed Taylor is indeed a generalization of the forward Laplacian, i.c., the latter does not employ additional tricks (leveraging sparsity could also be applied to our approach but are not aware of a drop-in implementation). It also illustrates that the savings we report in PyTorch carry over to other frameworks like JAX.
3. JAX's jit compiler is unable to apply the graph rewrites we propose in this work. If the JAX compiler was able to perform our proposed graph rewrites, then the j jit-compiled standard Taylor mode should yield similar performance than the forward Laplacian. How- ever, we observe a clear performance gap in runtime and memory, from which we conclude that the compilation did not collapse the Taylor coefficients. Our contribution is to point out that such rewrites could easily be added to the compiler's ability to unlock these performance gains at zero user overhead.
36

---

Results (biharmonic operator). For the biharmonic operator (right column of fig. G10 and table G5), we conclude that (i) the most efficient way compute biharmonics by nesting Laplacians (compare with table where Taylor mode uses the approach for general linear differential operators) and (ii) that nesting Taylor mode Laplacians more efficient than nesting orr-rr AD Laplacians, while also allowing to apply our collapsing technique.
H Numerical Complexity and Error Analysis
Setup. To illustrate the numerical properties of our proposed collapsed Taylor mode, consider a two-layered MLP with element-wise tanh activation IR + R The MLP denoted by f :- Φ h. The two linear layers are given as h R and g :R' + RCC g( = W200 + b2, with weights WI E IXD W₂ E RCx and bias compare the computational and storage complexity.
R h(x,) W1 + b1
b1 E ER'b, E RC Below for evaluating the sum of the second coefficients > ( f(x)), 2002) = >x=1 92.r (see eq. (5)) between collapsed and standard Taylor mode. For this toy example show (i) collapsing uses less operations and (ii) both methods are similarly stable based on our simplified error analysis.
well stability
Computational & storage complexity. Both vanilla and collapsed Taylor mode evaluate the function values (ho. Φo go) and the first derivatives ((hi,r; (1.r.91.r)) by propagating + R coefficients at each layer
ho Wi b, (m1r)) = ((11,
Φn x(ha) [ ((b()h)) h1.)))
go = Wi + b2 (g1,-) = (W2,21.)))
with @0(ho) ttanh(ho) diag(1 2 lx the tanh-activation layer's Jacobian. The propagation costs 1 + R matrix-vector multiplications with W1, 1 R matrix-vector multiplications with W2. . R Hadamard products with the derivative of Φ (since (diag(a), h1.r) ) = a O h1,r), and one Hadamard product to compute (o(ho). Additionally, there is one vector addition with the bias b1, one vector addition with b2, one vector subtraction in d<p(h)) (counted as vector addition), as well as the evaluation of @(hu). 3 + 3R vectors are stored.
For the second derivatives, vanilla Taylor mode propagates R vectors
((h2.r) =
(( ((0P)()))))) k + (8b(ho)) h2,)))
(192r) = [(2, 22,1)))
with activation Hessian @2(())) E R Ixxxxl of entries [8 Φ 1 j.j.k = 2220 ② (1 (1-402) 1.8.....
and contraction (02 (bo) h -200 O (1 00 h f.. These vectors summed up to >D= g This costs 2R matrix-vecte products with the weights, 1 + 3R Hadamard products, 2R - vector additions, and a single scalar multiplication.
get the result >m= (0² f(xa). ( v.
82
Table H6: Comparison of theoretical computational and storage complexity between standard Tay- lor mode and collapsed Taylor mode for a two-layer MLP computing the sum Σn (02 f(x)), 002
Computational Complexity
Operation
Standard Taylor
Collapsed (ours)
# Matrix-vecte products
4R +
2R + 4
Hadamard products
4R + 2
3R +
# Vector additions
2R + 2
2R +
# Scalar multiplications
1
1
# Activation evaluations
I
I
Storage Complexity
# Vectors stored
6R +
3R + 6
37

---

In contrast, collapsed Taylor mode propagates only single summed vector
n
R
R
R
n1
)
Σ
h
W
2,r
(3)
Σ
2
Σ
kg + app(oo)
Σ h2.r pm
r=l
R W2,
Σ
92,r
22.
This costs two matrix-vector vector products, 2 2R Hadamard products, 2R vector additions, and single scalar multiplication. Table H6 summarizes the accumulated costs.
Error analysis. For our numerical experiments. the result of all implementations (nested 1rdorder and standard/collapsed Taylor mode) was always checked to be close. To supplement this experimental error analysis, we sketch a simplified error analysis below. We assume that there are error-prone first- and second-order inputs {m + 1i.r). and {æ22 22r)r with errors ( )r. that can be seen as the error of previous propagation steps. error-prone ag would complicate our brief discussion too much and is ignored here. We consider again g e h. The error- error-iffluenced coefficients are denoted gē.r.
Using vanilla Taylor mode, the erroneous result is
R
((w2. (W1, + te,,) + (0b(hh)) (W1, æ2.r + E2.)) ))
F2
r=1
R
92.r
+
(0 (W1 @ (W1,
(W2, (ho)) (1, E1,)
( (1))),)))
+ (W2, (0 (ho). E1,, + (W2, (0@(ho), (W1, (2,)))
n
R W2 (a (W1, (2.))))
ge.r + Agen
All errors related to the first-order coefficients summarized in
Ag = Σ (W'o (W,, @ (W1, E1,))
+ (W2, (8o( (W (1, (W1, x1,
+ W2,, (W
The collapsed Taylor mode results in
n
n W2, Σ
Σ g2
((Po(ho), (W1, 2ir + (1,r) 82)
r=]
W2. (B0(to))
+
Wi, Σ
(I2,r + 52,r
r=l
38

---

R
92.
r=1
R
+
22
Σ
(oo))) (W1. @ (W1, + (0 (hoo), (11, @ (W1.21,))
r=1
+ ((o)) W1,,1,r1 )3,)))
W2
@p(ho),
W1. E 22,r
n
92. +
Agg,n
W2,
@p(ho),
W1. Σ E2.r
where the first-order errors are collected in
n
Agg,n = ( (2.,) Σ rml
(0o((h)) (1,, (W1,
+ 10 (hoo), (Wi, @ (1, + (ho), (W1, (1.))))))))
Error analysis (summary and discussion). Without considering soating--ping operations, the errors are equivalent. This is not surprising. since our collapsing method is mathematically equivalent to the standard Taylor mode on the same input coefficients.
Incorporating floating-pont point operations for the function evaluations, inner product, tensor product, and summations would greatly complicate the discussion, which not part of the paper. Still, the error could be split into the same three parts for both vanilla and collapsed Taylor mode. For the first-order errors Agī.n nt and AgÊ.R. however, even with floating-point operations, the errors are structurally similar, since apart from the most outer inner product (with W2) and the summation, all operations done in the same order. In practice, we would expect smaller errors for the collapsing method due to the reduced number of operations. The second error term, which collects the error of the second-order coefficients, could also reduce the accumulation of eтror terms. Of course, the actual condition and input, and output dimensions of the matrices are crucial. Theoretically, this could even lead similar error asymptotically. If inputs are small, one could argue that catastrophic cancellations are more likely to happen our case, since we sum first. But note that those cancellations are then also likely to happen in the standard Taylor mode, because weight matrices are often normalized, and the outputs of the activation functions are small if the input small.
We plan to investigate this more rigorously in the future.
I Connections of Collapsed Taylor Mode to Existing Methods
Here, we make the connection of collapsed Taylor mode to the forward Laplacian [21] and the randomized estimation of the Laplacian via Hutchinson's 's trace estimator [18] from [27] explicit.
I.1 Connection to Randomized Laplacian via Hutchinson's Trace Estimator
For simplicity, we consider a vector-tossclla function f RP + R, x + f(æ) (the general vector-to-vector case is straightforward but requires more notation) whose Laplacian is
D V V
Af(x) = Tr(V3
d=1
with (ax( E DXD the Hessian of evaluated at æ. Because the Laplacian can be expressed as trace of the Hessian, we can use Hutchinson's 's trace estimator [18] to estimate via Hessian-vector products with random vectors. Specifically, for any matrix A B DDD and a distribution p(v) over
39

---

vector v with unit covariance [Elvv' Ip) we can use the cyclic property of the trace and linearity of the expectation write
Tr(A) = Tr(AIb) = vv T]> B[Tr(Avv')] = T Av)]
Then, compute unbiased estimate of the right hand side by drawing S vectors 11 U2...., Ug ~ p(v) and evaluating the Monte-Carlo estimator
1 Σ A
(r(A) 0g Aus.
Applied to the Hessian, we can estimate the Laplacian of as
Af(x) S E fz))
Using our tensor notation, we can rewrite this into sum of terms involving (0)(()) " which can be computed with vanilla Taylor mode using S 2-jets (see eq. (4)):
8 D
2
1 Σ f(x). s
S
1,1=1
Instead of propagating then summing the 2-jets, we can also sum the vectors and then propagate the sum (assuming have composition, eq. (D16)), which is proposed collapsed Taylor mode:
8
1 S
8
(1)
Σ
82
1.2
Connection to the Forward Laplacian
We start by writing the propagation rules of the forward Laplacian (eqs. (5-7) in Li et al. [21]) for a function f = o h with : RC - R whose input we denote by hg E RC, The forward propagation consumes he h(xa),, the Jacobian Va, ho h(oo) IR DxC and the Laplacian A Aeoho ho = 3. E I:C:
ho
E C E DxC
9o (V 2. ho) (v)
ghaa)
20 ho
V 2x 90
R
Azg
hog
E RC
Am,90 = (V
Ahe Tr ((V mo ho) (V no V
E .R.
Let us rewrite this in terms of rows of the Jacobian denotes slice):
[ mo hola, E RC (where the colon subscript
ho {IV un hold. d=l Anp An
(
9o =g(h))
V 9old:
Va (V 9o)
Am..9 = (V ho )
Aho -d= V hold. V 9ol VV holT d.
In our tensor notation, this translates
ho {[V 2e ho d,; )am mo ho
)
90 g(to)
nV (0g (h)), (V ho hola.)a
Amn.9 (8g(ho), Aho) + E D
d==1 (ōg(ha), Va hol
To obtain the connection to Taylor mode, we define [V hold, = hi,d and Azo ho La h2,, and [V aoold, = I1,d and AwoHu -E 92 which allows us rewrite the forward Laplacian as
g(ho) fri, = (8g(ha), h1.e) -2 (gg(ho); h2,4) + Ed Edn (gg(ho). (Bg(ho) m13)
ho [h [ Ed h2,d
D
This yields our collapsed Taylor mode propagation: the first equation is simply the forward pass, the second equation propagates the first-order derivatives along /) directions, and the last equation propagates the collapsed second-order derivatives, as described by setting K = 2 in Equation (6).
40


Reviews:
[{"review_id": "sHbP6GngCj", "reviewer": "Reviewer_rNn6", "summary": "The paper proposes a modification of the Taylor mode method [1] for propagating higher-order derivatives to enable more (compute and memory) efficient computation of linear PDE operators. They note that PDE operators can be expressed as sums of contractions of derivative tensors with some coefficient tensor and use this to modify what is propagated. Specifically, instead of propagating separate tuples of derivatives (K-jets) for each term of the sum they note that one can propagate the sum of the highest-order derivative in the K-jet instead which could lead to a memory usage saving. They apply this to a Laplacian operator as well as a biharmonic operator going into detail as to how their approach would work for these operators. Their results show that they obtain improved memory and computational efficiency when compared to vanilla Taylor mode and the more naive nested first-order AD.\n\nRefs:\n[1] Griewank, A. and Walther, A. Evaluating derivatives: principles and techniques of algorithmic differentiation. SIAM, 2008", "questions": "- There seems to be a significant typo in eqn (6) where $\\bigotimes_{s \\in \\sigma} h_{K,r}$ should be $\\bigotimes_{s \\in \\sigma} h_{s,r}$\n- Adding an explanation/discussion for why collapsed Taylor mode has improved compute efficiency over vanilla mode would better help in understanding the full capabilities of their approach\n- The weighted Laplacian seems to be just the same experiment as the Laplacian (lines 241-242: \"For the weighted Laplacian’s coefficient matrix, we choose a full-rank diagonal matrix\") and so it isn't adding new information about the collapsed Taylor mode. Using a dense matrix would probably be the ideal setting for this experiment.", "limitations": "They have discussed limitations of their implementation but a wider discussion of the conceptual limitations of the collapsed Taylor mode is missing. Some suggestions:\n\n- Discussion of the compute/memory tradeoffs for different kinds of linear PDE operators would make it clear to the reader as to whether they can expect their approach to give them any significant gains over vanilla Taylor mode. \n- Related to the previous point, their experimental evaluation is a bit limited in that they only consider two linear PDEs that are very related (the biharmonic equation is just a Laplacian of a Laplacian) and use of the propagation of the sum of the highest-order Taylor coefficients (as highlighted in the \"Weaknesses\" above) isn't significantly utilized. Evaluating on a wider range of linear PDEs would give a clearer picture of their method.", "rating": 4, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 4, "significance": 3, "originality": 3, "strengths_and_weaknesses": "Strengths:\n- The explanation of their contribution is very thorough and clear and well elucidates the core ideas of their approach. Content in the appendix to aid in understanding of how their approach would work practically was very helpful. Also, the use of color coding and schematic diagrams are very useful additions to the paper which would otherwise make the details of the approach difficult to understand.\n- Making the connection between the forms of linear PDEs and the Faà Di Bruno formulae to derive a more efficient way to compute the PDEs is a notable contribution and is something that ought to be of interest to those working on PINNs or more broadly interested in AD. In addition, this connection could make it easier for other researchers to derive other efficiencies for operators not considered in this work.\n- The improved compute and memory performance over vanilla Taylor mode as shown in their results clearly demonstrates that their approach would be the preferable way to compute linear PDE operators.\n\nWeaknesses:\n- Their key contribution is a minor modification of vanilla Taylor mode and it doesn't seem like the key aspects of their approach i.e. propagating the sum of the highest-order Taylor coefficients instead of summing the separate propagations are used in their experiments. Specifically, for the Laplacian it is not clear as to what the difference in computation between standard and collapsed Taylor mode is (D15 and D16) other than the reduced memory saving that comes from storing a single vector instead of multiple. The same issue also looks to be present for the biharmonic equation as in their theory and experiments they don't seem to examine the case when non-zero high-order coefficients are propagated (they only consider the case when the zeroth and first-order coefficients are non-zero) which is when the sum of the highest-order Taylor coefficients would actually be used.", "was_revised": true, "final_justification": "So I have decided to leave my score as is because I feel that even though their approach can give compute and memory savings their lack of exploration of the applicability of the method makes the contributions of the paper not very significant. Specifically, they only apply it to Laplacian-type equations for moderately-sized dimensions which is a scenario of limited utility. Describing what use cases these computational savings unlock would definitely strengthen the paper", "rebuttal": "Dear Reviewer rNn6,\n\nthanks for your support and for catching the typo in Equation (6)!\nPlease find our responses to your questions below.\nLet us know if you have further questions; we would be happy to discuss.\n\n--- \n\n> for the Laplacian it is not clear as to what the difference in computation between standard and collapsed Taylor mode is\n\n> why collapsed Taylor mode has improved compute efficiency over vanilla mode\n\nThanks for insisting on further clarification.\n\nThe difference between standard and collapsed Taylor mode is that the collapsed version propagates fewer coefficients by computing the sum of the highest coefficients directly instead of propagating them separately, which reduces run time and memory.\n\n**Concrete example:** Let us illustrate this by looking at the *Laplacian propagation* at a point $\\mathbf{x}_0$ through a single linear layer $\\mathbf{h}_0 = \\mathbf{W} \\mathbf{x}_0 + \\mathbf{b}$ with weight matrix $\\mathbf{W}$ and bias $\\mathbf{b}$:\n\n- **Commonalities:** Both vanilla and collapsed Taylor mode propagate the function value $\\mathbf{h}\\_0$ and the first derivatives $\\mathbf{h}\\_{1,r}$ along directions $r=1, \\dots, R$ by computing\n\\\\begin{align}\n\\mathbf{h}\\_0 &= \\mathbf{W} \\mathbf{x}\\_0 + \\mathbf{b}\\\\,,\n\\\\\\\\\n\\mathbf{h}\\_{1,r} &= \\mathbf{W} \\mathbf{x}\\_{1,r}\\\\,.\n\\\\end{align} If we neglect the cost to add the bias, this costs $1 + R$ matrix-vector multiplications with the weight matrix, and storing the $1+R$ result vectors.\n- **Difference:** For the second derivatives, vanilla Taylor mode propagates $R$ vectors\n    $$\n    \\mathbf{h}\\_{2,r} = \\mathbf{W} \\mathbf{x}\\_{2,r}\\\\,.\n    $$ In contrast, collapsed Taylor mode propagates only a single vector\n    $$\n    \\underbrace{\\left(\\sum\\_r \\mathbf{h}\\_{2,r}\\right)}\\_{\\mathbf{\\tilde{h}}\\_2}\n    = \\mathbf{W} \\underbrace{\\left(\\sum\\_r \\mathbf{x}\\_{2,r}\\right)}\\_{\\mathbf{\\tilde{x}}\\_2}\\\\,.\n    $$\n\n**From this example, we can see that collapsing reduces both run time and memory:**\n- **Vanilla:** Store $1 + 2R$ vectors, apply $1 + 2R$ matmuls.\n- **Collapsed:** Store $2 + R$ vectors, apply $2 + R$ matmuls.\n\nWe will add this illustration to the appendix to improve the manuscript's clarity.\n\n---\n\n> it doesn't seem like the key aspects of their approach i.e. propagating the sum of the highest-order Taylor coefficients instead of summing the separate propagations are used in their experiments\n\n**We do propagate the sums (the key aspect) in all our experiments.**\nThis is reflected in the computational results:\nTake the linear layer example from above.\nFor large $R$, the ratio of matmuls approaches $(2 + R) / (1 + 2R) \\to 0.5$, and we thus expect the run time of collapsed Taylor mode to be half of vanilla Taylor mode.\nOur empirical results (Table 1) confirm that the time ratio is indeed close to $0.5$ for the Laplacian, and collapsed Taylor mode also uses less memory.\n\n--- \n\n> they don't seem to examine the case when non-zero high-order coefficients are propagated\n\nThis is a fair point, because we limit the discussion to linear PDE operators that are 'pure' in the sense that they are **linear combinations of derivatives that are all of the same order**.\nThis assumption leads to all propagated coefficients of order larger than 2 to be zero ($\\mathbf{x}_{s\\ge2, r} = \\mathbf{0}$) such that each $K$-jet computes pure $K$-th order derivatives.\nAs described in the paper, this covers a lot of practically relevant operators, like (weighted) Laplacians (e.g. Schrödinger and Fokker-Planck equations) and the biharmonic operator (e.g. plate equation).\n\nWe agree that generalizing the class of PDE operators (say towards linear combinations of mixed-order derivatives) is a promising direction for future work.\nThis may require some non-trivial intermediate steps though to formalize it mathematically.\n\n---\n\n> The weighted Laplacian seems to be just the same experiment as the Laplacian [...] Using a dense matrix would probably be the ideal setting for this experiment.\n\nYou are right that the weighted Laplacian we considered is exactly as expensive as the standard Laplacian.\nThe relevant quantity that determines this cost is not the sparsity of the weightings $\\mathbf{D}$, but their **rank** (Equation (8a)):\n- the standard Laplacian has $\\mathbf{D} = \\mathbf{I} = \\operatorname{diag}(1,1,1,\\dots)$ with $\\operatorname{rank}(\\mathbf{D}) = D$\n- our example weighted Laplacian has $\\mathbf{D} = \\operatorname{diag}((1,2,3,\\dots))$ which also has $\\operatorname{rank}(\\mathbf{D}) = D$. \n\nHence, they share the same cost.\nNote that we could have obtained the same result with any full-rank positive definite matrix $\\mathbf{D}$ that need not be diagonal/sparse.\nWhile we believe this is an interesting result in itself (dense and sparse combinations of second derivatives can be done at the same cost), we agree with your criticism that the results are to some extent redundant.\n\n**Additional experiment:** To address this, we experimented with weighted Laplacians with rank-deficient weightings ($\\operatorname{rank}(\\mathbf{D}) < D$) and will add this to the appendix.\nPlease find below a comparison between full-rank ($D$) and half-full rank ($D / 2$).\nBased on our theoretical analysis in the paper, we expect the performance ratio between vanilla and collapsed Taylor mode to be \n$(1 + 2 \\operatorname{rank}(\\mathbf{D})) / (2 + \\operatorname{rank}(\\mathbf{D})) \\approx 2$.\n**As you can see in the table, going from full-rank to half-full rank indeed approximately halves run time and memory.** (Table notation: *full-rank vs. half-full rank*)\n\n| Implementation | Per-datum time [ms] | Per-datum mem [MiB] (differentiable) | Per-datum mem [MiB] (non-diff.) | \n|-|-|-|-|\n| Nested 1st-order | 0.62 vs. 0.32 | 4.4 vs. 2.3 | 2.2 vs. 1.1 |\n| Standard Taylor  | 0.60 vs. 0.30 | 5.5 vs. 2.8 | 1.2 vs. 0.61 |\n| Collapsed Taylor | **0.31 vs. 0.15** | **2.4 vs. 1.3** | **0.92 vs. 0.47** |\n\nWe believe this addresses your concern regarding redundancy.\nLet us know if this makes sense to you.\nThanks again for your time and effort!"}, {"review_id": "hsEai7fGq2", "reviewer": "Reviewer_zJuU", "summary": "This paper proposes an optimized method to compute higher-order differential operators with Taylor-mode automatic differentiation. While typical automatic differentiation propagates only first-order derivatives (which requires nesting for higher orders), Taylor mode propagates a tuple of $K$ derivatives (a so-called “jet”). Leveraging the linearity of common PDE operators like the Laplacian, the authors show that the highest derivative term in the jet can be compressed when multiple perturbations are propagated simultaneously. This yields a theoretical speedup of roughly $K/(K-1)$, which is significant especially for low differentiation order $K$.  \nAs a side product, this method recovers previously-suggested optimizations, like the forward Laplacian, while being easy to automate. The authors demonstrate it with a PyTorch implementation based on computational graph rewriting, whose performance compares favorably to nested first order and standard Taylor mode.", "questions": "Here are the main additions which would improve the paper:\n\n- Explain the connection between collapsed Taylor mode and the individual optimizations it generalizes (especially the forward Laplacian).\n- Provide the experimental code to the reviewer, mainly assess the reusability of the underlying software for a wider audience (which will determine impact).\n- Explain why collapsed Taylor mode can lead to increased memory use for the biharmonic operator.\n\nHere are suggestions related to the presentation:\n\n- Since the key ingredient for collapsing to work is linearity in the last term of the Faa di Bruno formula, might it improve readability to push these details to the appendix and instead go with a simpified version in the main text? Maybe something like “the $K$\\-th derivative of $g \\circ h $ is $\\partial g$ times the $K$\\-th derivative of $h$ plus other lower-order terms in $h$”? I’m not sure this is a relevant suggestion to be honest.\n- In Equation (12), shouldn’t the two red sums be inside the inner product?\n- In Figure 5, what are “non-differentiable computations”, and should “opaque” be replaced by “transparent”?", "limitations": "Yes.", "rating": 5, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 4, "clarity": 3, "significance": 3, "originality": 3, "strengths_and_weaknesses": "## Strengths\n\n### Quality\n\n- The suggested approach is technically sound and insightful.\n- Experimental results on two differential operators support the use of collapsed Taylor mode, both for exact and stochastic computations.\n\n### Clarity\n\n- The paper demonstrates real effort to make theory and algorithms accessible, with the use of colors and diagrams to emphasize the salient points.\n\n### Significance\n\n- I am especially enthusiastic about the PyTorch library which automates graph rewriting based on torch.fx. Although the main contribution of the paper is algorithmic in nature, I assume that the open-sourcing of said library will be the actual driver of adoption for this otherwise complex method.\n\n### Originality\n\n- Collapsed Taylor mode gives a more generic perspective on previous isolated optimizations like the forward Laplacian, which is important for our understanding of automatic differentiation’s complexity.\n\n## Weaknesses\n\n### Quality\n\n- The main missing aspect is an explanation of how collapsed Taylor mode ties into existing methods like the forward Laplacian or the Hutchinson trace estimator. A short introduction to both of these techniques and a proof of the connection with collapsed Taylor mode would make the paper complete, but appendix D does not contain them.\n\n### Clarity\n\n- Through no fault of the authors, the entry cost into this paper is rather high due to the combination of tensor calculus and high-order derivatives with the Faa di Bruno formula. Perhaps some notational or exposition tricks could be used to facilitate the reader’s task (see below).\n\n### Significance\n\n- Without access to the code of the graph rewriting utility during review, it is hard to assess how generic and reusable this method will be for people unfamiliar with the underlying mathematics or implementation tricks.", "was_revised": true, "final_justification": "My score was already high and the rebuttal was entirely satisfactory, so I will keep it as is.", "rebuttal": "Dear Reviewer zJuU,\n\nthanks a lot for your strong support and constructive feedback!\nPlease find our responses below and let us know if you have follow-up questions.\nWe would be happy to discuss.\n\n---\n\n> missing [...] explanation of how collapsed Taylor mode ties into existing methods like the forward Laplacian or the Hutchinson trace estimator.\n\n> Explain the connection between collapsed Taylor mode and the individual optimizations it generalizes (especially the forward Laplacian).\n\nThanks for pointing this out. \nWe agree that making the connections of our method to existing approaches clear is essential for it to be adopted by the community.\nDue to the space limitations, these connections are currently hard to reconstruct because we frame everything in our tensor notation which focuses on generality.\n\nTo address your concern, **we will add a self-contained appendix that starts with equations** (5-7) from the forward Laplacian paper **and gradually transforms them into our tensor notation to show their equivalence** with collapsed Taylor mode.\nWe will do the same for the Hutchinson trace estimator, starting from its formulation via vector-Hessian-vector products.\n\nLet us know if this makes sense to you and whether you want all the details, and we will send a separate message about this.\n\n---\n\n> it is hard to assess how generic and reusable this method will be\n\nTo answer this, we summarize the core components of our code from a user and developer perspective (in our response to Reviewer GZnh, we also include a small code snippet that shows how these steps look like in code for the Laplacian):\n1. *(Start)* The entry point is a function `f` that implements the forward pass.\n  It **can use arbitrary PyTorch functions**, as long as `f` remains trace-able with `torch.fx`. \n  `f` could be an MLP implemented as `torch.nn.Sequential`, but can also contain `torch.nn.functional`s **and is therefore general**.\n2. *(Taylor mode)* To use Taylor mode, users can create the $K$-jet of `f` by applying the function transformation `jet` to `f` (implemented by our code).\n  **`jet` is extensible**: to add more operations, one has to implement their Taylor mode arithmetic according to the Faa-di-Bruno formula.\n3. *(Vectorize)* The PDE operators we discuss require evaluating $K$-jets along multiple directions. Going from evaluating one to multiple $K$-jets is done with `vmap`.\n4. *(Assemble PDE operator)* The vmap-ed $K$-jet can then be used to write a function that computes the desired PDE operator.\n   This function usually consists of three parts: (i) generate the ingoing Taylor coefficients\n   $\\\\{\\mathbf{\\mathbf{x}\\_{0,r}}\\\\}$, $\\\\{\\mathbf{\\mathbf{x}\\_{1,r}}\\\\}$, ..., along the directions $r$, \n   (ii) evaluate the $K$-jets along all directions using the vmap-ed $K$-jet, and (iii) combine the derivatives to obtain the final PDE operator.\n5. *(Collapse)* The last step is to simplify the PDE operator function via graph rewrites.\n   For this, our code provides a `simplify` function, which analyzes the passed function's compute graph for rewrites that lead to collapsed coefficients. \n   **`simplify` is also extensible**: it relies on rules that require implementing (i) detecting a pattern, and (ii) specifying its substitution.\n   \n\nPlease let us know if this clarifies your concerns regarding generality and re-usability.\nWe would be happy to follow up on this.\n\n---\n\n> Explain why collapsed Taylor mode can lead to increased memory use for the biharmonic operator.\n\nGreat question, thanks for asking!\n\n**Recap:** You are right that in our submission we found nested first-order AD to use less memory than vanilla/collapsed Taylor mode for computing the exact biharmonic operator (collapsing improved vanilla Taylor mode though, as expected).\n\n**Update:** We looked closer into this and identified an inefficiency in our implementation of Faa-di-Bruno (details below).\n**After fixing this inefficiency, collapsed Taylor mode now also uses less memory than nested first-order AD for the biharmonic operator.**\nSpecifically, collapsed Taylor mode **went from 9.3 MiB (submission) to 5.9 MiB (now)** per datum, outperforming nested 1st-order AD (7.9 MiB).\nThe inefficiency mostly affected higher derivatives, but we also got some small improvements for the Taylor mode Laplacians and weighted Laplacians.\n\n**Details:** In Faa-di-Bruno, we need to contract the derivative tensor against Taylor coefficients.\nOften, we contract multiple times against the *same* coefficient, but our submitted implementation did not leverage this.\nThink of multiplying multiple vectors `a * b * b * b`.\nWe originally did this with `einsum(a, b, b, b)`, which does not exploit that `b` appears 3 times and treats each operand as independent.\n`einsum`'s pair-wise evaluation then leads to a lot of intermediate tensors being saved for backpropagation.\nNow, our code instead computes `a * b ** 3`, which reduces the number of saved intermediates and therefore memory, while also giving a small speed-up.\n\n---\n\n> Maybe something like “the $K$-th derivative of $g \\circ h$ is  $\\partial g$ times the $K$-th derivative of $h$ plus other lower-order terms in $h$”? \n\nThanks for this suggestion.\nWe think your phrasing provides a nice complementary description of the math above Equation (6), which will help guide the reader, and will incorporate it.\n\n---\n\n> In Equation (12), shouldn’t the two red sums be inside the inner product?\n\nThe red sums are outside to highlight that Equation (12) has the same structure as Equation (5) for which we demonstrate in detail how the collapsing works.\nWe did not spell out the collapsed version of Equation (12) due to space limitations and will add this lengthier version to the next revision.\n\n---\n\n> In Figure 5, what are “non-differentiable computations”, and should “opaque” be replaced by “transparent”?\n\nRegarding opaque you are right, thanks, consider it fixed.\nWhat we mean by 'differentiable' versus 'non-differentiable' computations in code is the following:\n```python\nwith torch.no_grad(): # non-differentiable computation\n    laplacian = ... # no intermediate results need to be be stored\n    \nwith torch.enable_grad(): # differentiable computation\n    laplacian = ... # stores intermediates for later backpropagation\n```\nBoth modes of computation are relevant for applications.\nLet us know if this clarifies your question.\n\nThanks again for your time and effort!"}, {"review_id": "AYssci53mK", "reviewer": "Reviewer_SLnK", "summary": "The paper introduces a method to improve the efficiency of computing higher-order derivatives using Taylor-mode automatic differentiation (AD). The key idea is a graph-rewrite rule that exploits the linearity of highest-order derivative terms, enabling a collapsed representation that reduces memory usage and computational cost. Specifically, it is observed that in Taylor-mode AD, the propagation of the highest-degree coefficient is linear in the input, which allows the summation over multiple directions (used in computing operators like Laplacians) to be moved inside the AD graph. As a result, the system computes only one higher-order jet instead of many, saving time and memory.\nThe method is implemented as a graph-level transformation, and the rewrite rule identifies when a Taylor jet can be collapsed based on the operator structure. The  experimental results demonstrate improvements in speed and memory saving.", "questions": "Suggestions\n\n1. Add a formal analysis of time and space complexity that would clarify the cost savings to guide on when the collapsing method pays off.\n2. Broaden the experimental benchmarks to demonstrate robustness, scalability, and generality beyond toy MLPs.\n3. Compare with Reverse-Mode and Mixed-Mode AD.\n4.  Since high-order derivatives are often sensitive to numerical errors,  include a numerical stability analysis and show that collapsing is not unstable.", "limitations": "Yes", "rating": 3, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 2, "clarity": 3, "significance": 2, "originality": 3, "strengths_and_weaknesses": "Strengths \n1. Clear theoretical insight: The observation that in Faà di Bruno’s formula, the highest-order Taylor term is linear in its inputs leads to a simple graph-rewrite rule that reduces significantly the number of propagated jets.\n2.  A substantial computational savings is demonstrated in speedups and memory reductions.\n3. The collapsing rule is general-purpose, applicable to multiple higher-order operators, which makes it broadly applicable across SciML workflows.\n4. Simple integration with existing code, resulting in a minimal barrier to adoption.\n\n\nWeaknesses\n1. The collapsing method only applies when the higher-order operator being computed is a linear functional of the highest-order Taylor coefficient, hence it accelerates a narrow class of derivative computations.\n2. The speedups are measured empirically and the complexity reductions are not quantified formally, hence it lacks guarantees or limits on where collapsing breaks down or ceases to be beneficial.\n3. All benchmarks are performed on small MLPs with synthetic data, and there is no demonstration of  scalability or broad utility.\n4. The paper focuses on forward Taylor mode and compares primarily against unoptimized nested AD and not reverse mode or hybrid AD, thus it is not clear that this is the best method overall for higher-order AD.\n5. The paper does not analyze how collapsing affects numerical precision. Combining Taylor coefficients before propagation could change round-off behavior or accumulate truncation errors in high-order derivatives.", "was_revised": true, "final_justification": "Following the discussions, I raised my score to 4, I believe that the paper should be accepted.", "rebuttal": "Dear Reviewer SLnK,\n\nThanks for your comments and thoughtful feedback.\nIn the following, we answer to your raised concerns.\n\n---\n\n> accelerates a narrow class of derivative computations.\n\nWe respectfully disagree.\nWhile our method indeed relies on the assumption that the PDE operator is a linear functional of the highest-order Taylor coefficient, this structural requirement is satisfied by many important PDE operators:\n\n1. **The Laplacian.** Among its diverse applications, we stress the importance of computing Laplacians in high dimensions in quantum physics. This is a major bottleneck in the variational Monte Carlo method for the Schroedinger equation.\n2. **Weighted Laplacians.** These govern the diffusive terms in Fokker–Planck/Kolmogorov-type PDEs. Again, these are PDE operators in high dimensions and efficient schemes to compute them are paramount.\n3. **Higher-order elliptic operators:** This includes the biharmonic operator, which arises in elasticity and thin plate theory.\n\nAs these examples span diverse application domains, and efficient derivative computation is important for all of them, we do not consider our proposed methodology's scope to be narrow.\n\nClearly, PDE operators of different structure are worthwhile to explore as well.\nWe think our scheme can also apply to more PDE operators but leave this for future work as it will require more mathematical formalization.\n\n---\n\n> The speedups are measured empirically and the complexity reductions are not quantified formally\n\n> Add a formal analysis of time and space complexity that would clarify the cost savings to guide on when the collapsing method pays off.\n\nThank you for your insightful comment. You are right that a formal complexity analysis would strengthen the presentation. We will incorporate your suggestion in the following way (we make the below argument concrete for a linear layer in our response to Reviewer rNn6):\n\nThe key component for our analysis is computing sums of $K$th-order derivatives\n$$\\sum_{r=1}^R \\langle \\partial^K f(\\mathbf{x}), \\mathbf{v}_r^{\\otimes K}\\rangle\n$$ via Taylor mode, which must also compute the derivatives of lower degree.\n\n**Time complexity:** The cost of computing the lower-order coefficients ($0$ to $K - 1$) is identical for both collapsed and vanilla Taylor mode. The key difference lies in the highest coefficients ($K$th order), which can be directly understood from the chain rule:\n- In **vanilla Taylor mode**, for each of the $R$ directions, we compute at each layer\n$$  \n\\sum_{\\sigma \\in \\text{part($K$)}\\setminus \\\\{\\sigma_t\\\\}} \\nu(\\sigma)\\langle\\partial^K g(h_{0}), \\otimes_{s \\in \\sigma}h_{s, r} \\rangle + \\langle \\partial g(h_0), h_{K, r}\\rangle\n$$ followed by a vector summation over $r$ after the final layer. \n- In **collapsed Taylor mode** we instead compute *once* at each layer\n$$  \n\\sum_{r=1}^R\\left(\\sum_{\\sigma \\in \\text{part($K$)}\\setminus \\\\{\\sigma_t\\\\}} \\nu(\\sigma)\\langle\\partial^K g(h_{0}), \\otimes_{s \\in \\sigma}h_{s, r} \\rangle\\right) + \\langle \\partial g(h_0), \\sum_{r=1}^Rh_{K, r}\\rangle\n$$\nNote that the summation inside the second term is performed in the previous layer.\n\n**So in total, collapsed Taylor mode trades $R - 1$ Jacobian-vector products (matrix-vector multiplications!) for $R$ vector summations at each layer**.\n\n**Space Complexity:** The space complexity is similarly reduced. **Vanilla Taylor mode** stores the highest coefficient for every direction at each layer ($g_{K, r}$ for $r=1, \\dots, R$ in our example). Instead, **collapsed Taylor mode** stores the summed coefficient ($\\sum_{r} g_{K, r}$) at each layer.\n\n**This reduces space complexity from storing $R$ $K$th-order coefficients to storing $1$ $K$th-order coefficient at each layer.**\n\n---\n\n> All benchmarks are performed on small MLPs with synthetic data, and there is no demonstration of scalability or broad utility.\n\n**MLPs are common in the PINN literature**, which is one of the prime application fields for our developed method, explaining our choice and rendering the considered networks meaningful.\n\nBased on our theory, we expect the performance difference between vanilla and collapsed Taylor mode to carry over to a wide array of other architectures.\n**The methodology is completely architecture-agnostic.**\n\nIf you have suggestions for additional architectures we should show in the paper, do let us know and we would be happy to try to include them in a future revision.\nHowever, this will likely require supporting new operations in our Taylor mode library, and we might not be able to provide these results during the discussion.\n\n---\n\n> compares primarily against unoptimized nested AD and not reverse mode or hybrid AD, thus it is not clear that this is the best method overall for higher-order AD.\n\nCould you elaborate what you mean by hybrid AD?\n\nWe use forward-over-reverse mode for Hessian-vector products, which is the recommended way in PyTorch (see e.g. the docstring of `torch.func.hessian`).\nOther modes like reverse-over-reverse are less favourable and there exist empirical (Dagreou et al., 2023) and theoretical (Griewank et al., 2008) arguments for that.\n\nWe are not aware of any AD implementation for Python that allows going beyond combining forward and reverse mode.\nTherefore, **we believe that our 'nested first-order AD' baseline is the best currently achievable implementation in PyTorch.**\n\n---\n\n> The paper does not analyze how collapsing affects numerical precision.\n\n> include a numerical stability analysis and show that collapsing is not unstable.\n\nWe appreciate your request to address numerical stability and would like to answer it from a theoretical and empirical perspective:\n\n- **Theoretical:** While we are unaware of existing stability proofs for higher‑order Taylor mode AD, reference [1] showed **forward- and backward‑stability for first‑order AD**. Since the basic structure for propagating derivatives of any order remains the same, **we are strongly convinced that numerical stability can be proven also for the vanilla and the collapsed Taylor mode**. However, this is a very technical/mechanical task and we will not be able to provide general results during the discussion. To illustrate the idea for a corresponding proof, we will add a small example to the paper and highlight that collapsed Taylor mode is likely to reduce rounding errors due to its smaller number of floating point operations.\n- **Empirical:** In our experiments, we verify consistency between all implementations (nested first-order AD versus vanilla Taylor versus collapsed Taylor) by comparing the results  with `torch.allclose` with default arguments (`rtol=1e-5, atol=1e-8`).\nWe noticed **no discrepancies when computing in `float64`** (which is standard for applications like PINNs).\n\nWe hope this addresses your concern. Please don’t hesitate to let us know if further clarification would be helpful.\n\nThanks again for your time and effort!\n\n1. A. Griewank, K. Kulshreshtha und A. Walther: On the numerical stability of algorithmic differentiation. Computing 94:125--149 (2012)"}, {"review_id": "yrhGdg9miy", "reviewer": "Reviewer_GZnh", "summary": "Motivated by the evaluation of PDE operators, this paper introduces a collapsed Taylor mode forward algorithmic differentiation differentiation. The authors justify a form of chain rule for certain differential operators, represented as sums of high order derivatives, generalizing the Laplacian. The main idea is to propagate sums rather than high order tensors to be summed. This allows to avoid propagating the full highest order tensor and results in significant savings in terms of the number of arithmetic operations required to evaluate the operator.\n\nThe method is implemented in pytorch and tested on three differential operators. The authors report consistent improvements over the baseline nested AD both in terms of computation time and peak memory. The improvements are coherent with the arithmetic complexity estimates and the figures are confirmed using jax.\n\nThis is a convincing paper which adresses a relevant question and proposes an advanced automatic differentiation method which will likely reduce the training time of PINNs in the near future.", "questions": "(11), what does $[\\mathbf{j}]_i$ mean? It seems that the notation is not self contained in this respect. I interpret it as just a number, is it correct?\n\nRelated to the previous question, please indicate which part of [14] is refered to in Section 3.3. In its present form, the submission requires to read another full paper, with different notations.\n\nIn (12), the outer sum is not pulled inside. I guess this is because of non-linearity of the tensor product. Any hope to pull this inside?\n\nIn the present form (12) should be re-implemented for each setting? Could this be automatized as well?\n\nCould the authors provide a minimal working code example to reproduce some experiments? This would strengthen the submission.\n\nLine 89: \"composing f of atomic functions with known derivatives and the chain rule\" what does it mean?\n\nDoes Figure 4 contain all coefficients?", "limitations": "Yes", "rating": 5, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 3, "significance": 4, "originality": 4, "strengths_and_weaknesses": "### Strength \n\nThe reported improvements are very clean, consistent and convincing. \n\nThe proposed scheme can be directly combined with the stochastic estimator proposed in [26]. \n\nThe authors suggest that the proposed approach could be implemented with minimal modifications of existing Taylor mode AD frameworks, at least from the user point of view.\n\n\n### Weaknesses\nWhile the authors propose a solution for general differential operators in (10), the proposed solution is only illustrated using the biharmonic operator. The paper would benefit from more such examples. \n\nWhile equation (12) provides a one-fit-all solution for problems of the form (10), it requires to re-implement the equation in (12) for every new example. \n\nThe code is not part of the submission. The paper mentions that this will be open sourced, but the timeline and the actual content of this future library cannot be evaluated from the submission.", "was_revised": true, "final_justification": "See the \"comment after rebutal\".", "rebuttal": "Dear Reviewer GZnh,\n\nthanks for your strong support and constructive feedback.\nPlease find our answers below and let us know if you have any follow-up questions.\n\n---\n> (11), what does $[\\mathbf{j}]_i$ mean? It seems that the notation is not self contained in this respect. I interpret it as just a number, is it correct?\n\nThank you for pointing this out. As explained in the background section, we use square brackets to denote slicing.\nWe will revise the text to clarify the meaning of $[\\mathbf{j}]_i$ as $i$th entry of vector $\\mathbf{j}$ and ensure that the notation is clearly defined.\n\n---\n\n> The paper would benefit from more such examples.\n\nWe agree that the interpolation formula we use to present a fully-general approach can be daunting to parse at first.\n\nTo make this process more clear we further add an example in the appendix: computing $\\sum_{i=1}^D\\sum_{j=1}^D\\frac{\\partial^3}{\\partial x_i^2 x_j}f(x)$.\nThis example is inspired by Appendix F.1 of the stochastic Taylor derivative paper (Shi et al., 2024), which describes how to compute each mixed 3rd-order derivative using either one 7-jet, two 5-jets, or three 4-jets.\nLet us know if you would like to see all the details and we will follow up on this in a separate message.\n\n**TL;DR:** We show that computing this operator with the interpolation formula uses $D^2$ **3-jets (rather than multiples of $D^2$ jets of order $>3$)** that can also be collapsed.\nWe expect the interpolation formula to be favourable as Taylor mode scales quadratically in the derivative order, and we will experimentally verify that in a future version of the paper.\n\n---\n\n> In the present form (12) should be re-implemented for each setting? Could this be automatized as well?\n\nYou are right that Equation (12) must currently be re-implemented as it is the equation that interfaces with the (automated) Taylor mode and the (automated) collapsing.\n\nTwo things which are currently not automated are:\n\n1. We automated computing the coefficients $\\gamma_{\\mathbf{i},\\mathbf{j}}$, but did not design an interface that allows to specify a PDE operator by specifying the multi-index $\\mathbf{i} = (i_1, \\dots, i_I)$, i.e. the frequency of the directions. We believe this is possible.\n2. Using symmetries in the multi-index $\\mathbf{i}$ to simplify the number of jets. We believe fully automating this is non-trivial as it requires symbolic understanding of the math.\n\nWe will make sure to mention these aspects in the next revision.\n\n---\n\n> In (12), the outer sum is not pulled inside. [...] Any hope to pull this inside?\n\nYes, because Equation (12) has the same structure as Equation (5) for which we show how to perform the collapsing (pull the sum inside) in detail.\nWe will add the collapsed version of Equation (12) to the next revision, as we did not have space for it.\n\n---\n\n> Could the authors provide a minimal working code example [...]?\n\nHere is a minimal example (imports omitted) how to work with our code base to compute (collapsed) Laplacians.\nWe hope it provides further clarification.\n\n```python\n# 1) Define the function whose Laplacian we want to compute\ndef f(x: Tensor) -> Tensor:\n    return sin(x) # could also be a neural net\n    \n# 2) Apply Taylor-mode to compute 2-jets\njet2_f = jet(f, k=2) # signature: (Tensor, Tensor, Tensor) -> (Tensor, Tensor, Tensor)\n\n# 3) Vectorize to compute multiple 2-jets in parallel\nmultijet2_f = vmap(jet2_f) # signature: (BatchedTensor, BatchedTensor, BatchedTensor) -> (BatchedTensor, BatchedTensor, BatchedTensor)\n\n# 4) Use the multi-2-jet to compute the Laplacian\ndef laplacian_f(x: BatchedTensor) -> Tensor:\n    x0, x1, x2 = x, eye(...), zeros(...) # Equation (7b) in the paper\n    f0, f1, f2 = multijet2_f(x0, x1, x2)\n    return sum(f2, dim=0)\n\n# 5) Rewrite the compute graph (collapse)\ncollapsed_laplacian_f = simplify(laplacian_f)\n\n# ---\n\n# Sanity check\nx = rand(...)\nvanilla = laplacian_f(x)\ncollapsed = collapsed_laplacian_f(x) # same result, but less time and memory\nassert vanilla.allclose(collapsed)\n```\n\n---\n\n> Line 89: \"composing f of atomic functions with known derivatives and the chain rule\" what does it mean?\n\nBy this, we mean that implementing Taylor mode requires specifying the Faa-di-Bruno formula for basic operations (e.g. `+`, `sin`, `tanh`, ...) which users can then use to programmatically specify functions whose jets should be computed.\nThis is a common building principle for AD and makes it extensible, as new operations can be added one-by-one.\n\n---\n\n> Does Figure 4 contain all coefficients?\n\nYes. We will describe in more detail in the appendix how to manually compute these coefficients (our code also ships with a utility function to compute the interpolation coefficents $\\gamma_{\\mathbf{i},\\mathbf{j}}$).\n\nThanks again for your time and effort!"}]

Author rebuttal:

Meta-review: I recommend acceptance of this paper. All reviewers agree that the core idea, collapsing Taylor‐mode AD by rewriting the compute graph, leads to a clear practical benefits. In particular, reviewer rNn6 finds that "the explanation of their contribution is very thorough and clear and well elucidates the core ideas of their approach", reviewer zJuU finds that "the suggested approach is technically sound and insightful", reviewer GZnh finds that "the reported improvements are very clean, consistent and convincing". Even the more critical reviewers acknowledge the paper’s strengths. The paper makes a solid contribution.
Decision: Accept (poster)