Automatic Functional Differentiation in JAX
Authors: Min Lin
Venue: ICLR 2024 poster (2024)
Keywords: JAX
automatic differentiation
calculus of variation
higher-order functions
functionals
operators
JVP rules
transpose rules
Primary area: infrastructure, software libraries, hardware, etc.
TL;DR: Describes a augmentation to JAX that enables automatic functional differentiation.

Abstract: We extend JAX with the capability to automatically differentiate higher-order functions (functionals and operators). By representing functions as infinite dimensional generalization of arrays, we seamlessly use JAX's existing primitive system to implement higher-order functions. We present a set of primitive operators that serve as foundational building blocks for constructing several key types of functionals. For every introduced primitive operator, we derive and implement both linearization and transposition rules, aligning with JAX's internal protocols for forward and reverse mode automatic differentiation. This enhancement allows for functional differentiation in the same syntax traditionally use for functions. The resulting functional gradients are themselves functions ready to be invoked in python. We showcase this tool's efficacy and simplicity through applications where functional derivatives are indispensable.

Published as a conference paper at ICLR 2024
AUTOMATIC FUNCTIONAL DIFFERENTIATION IN JAX
Min Lin
Sea Al Lab
linmin@sea.com
ABSTRACT
We extend JAX with the capability to automatically differentiate higher-order functions (functionals and operators). By representing functions as a generaliza- tion of arrays, we seamlessly use JAX's existing primitive system to implement higher-order functions. We present a set of primitive operators that serve as foun- dational building blocks for constructing several key types of functionals. For every introduced primitive operator, we derive and implement both linearization and transposition rules, aligning with JAX's s internal protocols for forward and reverse mode automatic differentiation. This enhancement allows for functional differentiation in the same syntax traditionally ing functional gradients are themselves functions ready to be invoked in python. We showcase this tool's efficacy and simplicity through applications where func- tional derivatives are indispensable. The source code of this work is released at https: gi thub. wwt sai sg/ autofd.
for functions. The result-
INTRODUCTION
In functional analysis, functions can be characterized as continuous generalizations of vectors. Correspondingly, the principles of linear algebra, derivatives, and calculus for vectors be extended to the spaces of functions. Based on this generalization, we can build higher-order functions whose domain/codomain are function spaces. These higher-order functions are traditionally called functionals or operators. For example, definite integral is a functional that maps function to scalar, and the differential operator is an operator that maps a function to its derivative function. The generalization of gradient descent to functionals is particularly interesting to the machine learning audiences. Like we optimize function by moving the input in the gradient direction, functional can be optimized by varying the input function using the functional derivative. Gradient descent for functionals has been used for a long history in science and engineering, named the calculus of variations. Well known examples include the least action principle in Lagrangian mechanics, and the variational method quantum mechanics.
While the theoretical framework for functional differentiation is established, there remains notable gap in computational tools tailored to automatically compute functional derivatives. To this date, functional differentiation with respect functions are usually done in one of the following ways: (1) manually derived by human and explicitly implemented (Zahariev et al., 2013; Gonis, 2014); (2) for semi-local functionals, convert it to the canonical Euler-Lagrange form, which can be implemented generally using existing AD tools (Hu et al., 2023; Cranmer et al., 2020). In stark contrast, the domain of automatic differentiation ( (AD) for functions that maps real tensors has seen significant advancement recently and widespread application in various domains, e.g. robotic simulation (Freeman et al., 2021), cosmology (Piras Mancini, 2023) and molecular dynamics (Schoenholz & Cubuk, 2021). Like AD has transformed several fields and industries, we believe automatic functional differentiation (AutoFD) holds the potential to catalyze a similar wave of innovation.
Based on the idea that functions can be represented as a generalization of arrays, we propose that AutoFD can be implemented in the same vein as AD. Building AutoFD in this way avoids the complexity of symbolic approaches that depends on analytical integrals. JAX provides elegant machinery for supporting both forward mode and reverse mode AD without redundant code. The forward mode differentiation is implemented via linearization rules, also called jacobian vector product (JVP) rules for each primitive computation. The reverse mode differentiation consists of three steps (Radul et al., 2023), 1. perform forward mode AD on function; 2. unzip the linear and

---

Published as a conference paper at ICLR 2024
nonlinear part of the function; 3. transposition of the linear part. Therefore, to make a primitive computation differentiable in JAX, it is crucial that the primitive is implemented with a JVP rule and transpose rule. Similarly, AutoFD for an operator relies higher-order generalizations of JVP and transpose rules, which are well defined in mathematics. The Fréchet derivative extends the idea of derivative to higher-order functions whose inputs and outputs are functions. And the adjoint operator generalizes the concept of transposition from function to higher order functions.
Having introduced the fundamental mechanisms for AutoFD, we now delve into the specific operators we include in this work. We motivate the choice of operators from the types of operators and functionals that are commonly used, which we summarize as non-linear semi-local functionals and linear non-local operators. These two types cover the most used functionals in various applications, for instance, the majority of exchange-correlation functionals in density functional theory fall under semi-local or linear non-local functional categories (Perdew & Schmidt, 2001). It turns out that to build any functional described in Section 2, there is a set of five essential operators, namely, compose, V. linearize, linear transpose and integrate. These operators, along with their JVP and transpose rules, are detailed in Section 3.2. To ensure the completeness of the operator set, we also introduce some auxiliary operators in Section 3.3. We discuss the applications of AutoFD in Section 4, presenting both opportunities for new methods and improvements on the coding style.
2 OPERATORS AND FUNCTIONALS
By convention, we define operators as mappings whose domain and codomain are both functions. For example, the V operator when applied on function returns the derivative of that function. Functionals are mappings from functions to scalars. To distinguish them from plain functions that map real values, we denote both functionals and operators with capital letters, with an extra hat for operators. We use round brackets to denote the application of operators on functions as well as the application of functions on real values. e.g. (z)) means application of operator Ô on function f. and apply the resulting function on x. Without loss of generality. we write integrals without specifying the domain. For simplicity, we present the results on scalar functions when there no ambiguity. The actual implementation supports functions that map arbitrarily shaped tensors or nested data structure of them. The operators we aim to support can be categorized into three different types based on their property.
Local operator generalizes element-wis wise transformation in finite dimensional vector spaces. Consider the input function f X + Y, with any function ØL : X Y Z, local operator d needs to satisfy
i : X + Z; + 01 ( (x)).
Semilocal operator extends local operator by introducing extra dependencies on the derivatives of the input function up to a finite maximum order n.
$s())
X 2: + ss(x. ((x), Vs(x). (()) (()))
This type of operator is called semilocal because & (f)(z) not only depend on the value of f(x) at the point but also on the function values in the infinitesimal neighborhood, i.e. f(x + dz), via the knowledge of the derivatives.
Nonlocal operators are the operators that are neither local nor semilocal. Although they do not need to assume any specific forms, one of the most interesting nonlocal operator is the integral transform,
The function Φ : U x X + Y is called the kernel function. Integral transform generalizes finite dimensional matrix vector product; therefore it is a linear mapping of f. The Schwartz kernel theorem (Ehrenpreis, 1956) states that any linear operators can be expressed this form.
Integral functionals are functionals that conform to the following form:
F F:F-RR
(x)(
2

---

Published as a conference paper at ICLR 2024
Correspondingly, integral functionals are called local, semilocal or nonlocal depending on the property of operator क The components used to build functional approximators in existing works belong to one of the above types. For example, the fourier neural operator (Li et al., 2020) is a direct generalization of multilayer perceptron, the linear layers nonlocal integral transforms, while the nonlinear activation layers are pointwise transformation that follows the form of local operator. The lagrangian neural networks (Cranmer et al., 2020) is implicitly composition of integral functional with learnable semilocal operator. In applications like density functional theory, most empirically developed functionals belongs to the semi-local local family. Recently, non-local functionals are also introduced for better treatments of electron correlation (Margraf & Reuter, 2021).
3
IMPLEMENTATION
3.1 GENERALIZED ARRAY
In JAX, array is fundamental data structure that corresponds to vector, matrix and tensor in mathematics. Primitives are algebraic operations for transforming the arrass. Following the notation of JAX, we use f[3, 5] to describe float array of shape by 5. To represent functions as generalied arrays, we first generalize the notations. The shape of a function is represented as a list of its return value and each arguments in the format of F[ret, argo, .. For example, F[f[], f[3], f[2, 3]] describes function: I R2x3 + R. In AutoFD, this abstract information is implemented as custom class and registered via jax. core. py pytty_eaave aval _aal_mappings types. FunctionType Type] to jax.core.......oo.............
make JAX recognize and convert python functions into generalized arrays when tracing higher order function.
3.2 PRIMITIVE OPERATORS
The primitive operators considered in this work focused to realize the most used types of operators and functionals described in Section 2. To enable functional derivatives for these primitives, as mentioned in the introduction, need define JVP rule and tranpose rule for each of them. Here we only make connections between the programming notation and math notation of the JVP and tranpose rules, which should be enough for understanding the implementation of AutoFD, For formal definitions, we refer the readers to Section II.5 in Sternberg (1999) for tangent/cotangent vectors, and Definition 3.6.5 in Balakrishnan (2012) for Fréchet derivative.
The JVP function in JAX is defined as
jax. jxx.j
8x 5z))
= = (ox)).
(5)
primal tangent
Fréchet
Where (z) = dz x the jacobian of function f at point æ. The JVP function also called the forward mode gradient, it evaluates change of function f at point æ along the direction Sr. However, backward mode gradient is what we usually use in backpropagation. In JAX, the function associated with backward gradient is vector jacobian product (VJP), VJP means vector is on the left hand side of the jacobian, but we usually transpose it instead to obtain column vector.
0f
jax.vjp(f
8y = J 58
primal
cotangent
Usually, forward mode and backward mode are implemented separately. However, Radul et al. (2023) introduces that idea that the VJP is simply the linear transposition of JVP. Here we introduce the notations of linear transpose in programming and in math. With f(x) = Mæ; f (v) = M y.
jax.linear _ttneeerrttaapppeeee_e
/ prinal
y
(v)
(()
(6)
cotangent
our own Fréchact like notation
It can be easily seen that when apply linear transposition to JVP: óx ++ (z(2) it results in the VJP: dy + J (x) 8y. Therefore, in JAX, each primitive is given a JVP rule and a transpose rule, together they are used to derive the backward mode gradient in JAX.
The JVP and transpose rules be generalized higher order functions. We use the Fréchet derivatives D and 0 denote JVP rules and partial JVP rules. (())) denotes the JVP of
3

---

Published as a conference paper at ICLR 2024
operator Ô with the primal value f and tangent value öf. We add ó prefix for functions from the tangent and cotangent space. Similarly to the Fréchet notation, we introduce the symbol T and T; to denote transpose and partial transpose rules. Note that although transpose rules generally do not require a primal value, partial transposes do sometimes need the primal values when the operator is not jointly linear on the inputs. Therefore, we explicitly write out the primal values for transpose rules for consistency; i.e. T ()) (ōh) denotes transposition of the operator Ô with respect to f. with the primal values f. 9 and cotangent value 8h. The generalization of transposition to operators is the adjoint of the operator T(Ô) = O", satisfying (Óf,g) = (f. O"
To support high order derivative, the JVP and transpose rules need be implemented using primitives exclusively, which means any operator used in the right hand side of the JVP and transpose rules needs to be implemented as a primitive themselves. For example, the JVP rule (7) of the compose operator uses the linearize operator L. which is described in Section 3.2.3. To save space, we leave the proof of these rules to the Appendix.
3.2.1 THE COMPOSE OPERATOR
Function composition is a frequently used operator, with g: x -Yf: The compose operator defined as f o X + Z: x + fg(())) Here the compose operator written as an infix operator, alternatively, we can use the symbol C' to represent the compose operator, and Ö(f. g) to denote the composition of f and g. The compose operator can be generalized to more than two functions. For example, Č(f, 91, 92) describes the function f(o (1))(z), (2())) Implementation of compose in python is simple
def Compose_impl() g): return lambda *args: f(g(*args))
The JVP and transpose rules of compose are derived as follows:
a,())().)) 0g C(L(I)
(7)
())))) 6f C(ōS.g).
(8)
fōh f g det ((g if g is invertible.
()()))
[undefined
otherwise.
(5h- C(T()) 8h) if is linear. ((())))))) undefiine
(10)
otherwise.
As be seen, the implementation of the compose operator requires auxiliary operators [llinerrize) and (transpose), which will be defined independently in Section 3.2.3 and 3.2.4. The function inverse the right hand side of the (C()) rule is not implemented because currently general mechanism for checking the invertibility and for defining inversion rules of the primitives are not yet available in JAX. The compose operator the basis for all local operators defined in (1). Replacing the dependency on with an identity function 1, all local operators can be composed with Č(Ør, I, r + (I ff))) A large number of common operators can be defined via Ć; for example, we overload the infix operators in python to support the syntax + y which converts to Ō((x,)) v) + æ + y, f.g). All direct algebraic operations on functions in our code examples are supported in a similar way.
3.2.2 THE V (NABLA) OPERATOR
The V operator converts a function to its derivative function. In JAX, the corresponding transforma- tions are jax. jacfwd and jax. jacrev. Its JVP and transpose rules are:
D(V)()= of + + V(5/)-
(11)
T(V)()+ 8h - 5h.
(12)
The implementation of the V operator is readily available in JAX as jax. jacffw and jacrev. The JVP rule is the same the primal computation because V is linear operator. The transpose
4

---

Published as a conference paper at ICLR 2024
rule uses V as the divergence operator, which can be implemented C(trace, V(ōh)). We provide quick proof of the transpose rule here for scalar variable functions, the full derivation is provided in Appendix C.2. To see the transpose rule, write the V operator in the Schwartz kernel form, i.c. V 5 (x - y)- Like finite dimensional linear transposition swaps the axis that is contracted, the T(V) = ff dx (' - y) simply changes the variable under integration, using integral by parts,
verify that
f dx 8' -
-( f d (
((V)()())
The V operator can be used for constructing semilocal operators together with the Ô operator, i.. semilocal operator in the form of (2) be implemented as C(os, J. Vf, v(*)), f), where 1 + is the identity function, and (*) is the repeated application of V for times.
3.2.3 THE LINEARIZE OPERATOR
The linearize operator has the same meaning as the Fréchet derivative D. Except that can be used with any higher-order functions, while we reserve the symbol L(ʃ) : dx + (z)) for the linearization of functions. The linearize operator is necessary for the completeness of local operators, as it is used in the JVP rule (7) of the compose operator. The JVP and transpose rules of the linearize operator are defined as following:
D(L( 8f - if
(13)
T(L)() 6h -Vi((yhh(..))
(14)
The iy symbol used in the transpose rule represents integration over the variable The corresponding function transformation implemented in JAX is jax.
3.2.4 THE LINEAR TRANSPOSE OPERATOR
The linear transpose operator can be applied only on linear functions. It is easy to see that the operator itself also linear. We only manage derive general form for the adjoint of the transpose operator when the cotangent function is invertible.
D(T) 8f f(ff).
(15)
0 5h 1 det Voh- if 8h is invertible. T())): undefined
(16)
otherwise.
The corresponding transformation defined in JAX for linear transposition is jax www.:innar_traaspooo.
3.2.5 THE INTEGRATE FUNCTIONAL
The integrate functional indispensable component for building functionals. It is required in the integral functionals (4), for constructing nonlocal operators (3), and used in the transpose rule (14) of the linearize operator. The integrate functional is called an integrate operator when we perform integration over only part of the variables. We use for integrate operator, with subscript for the variable under integration. For function of two variables, i=()) + f f(x, y)... The JVP rule and transpose rules are:
(17)
T(L(() 8h (x,y St((y))
(18)
Only limited subset of functions admit analytical integrals, the Risch algorithm provides a robust mechanism for determining the existence of an elementary integral and for computing it. In practical applications, numerical integration techniques are more commonly employed. These techniques are rooted in the concept that function can be expanded in series of integrable functions. Methods of numerical integration often hinge on a designated grid of integration points associated with respective weights. For example, the Fourier series converts to summation over uniform grid, and the Gaussian quadrature induces points that are the roots of the Legendre/Laguere polynomials. With the grids and corresponding weights, integration of function converts to a weighted summation over finite set of points. We implement the integrate operator by supplying a numerical grid for the integrand function.
5

---

Published as a conference paper at ICLR 2024
3.3 COMPLETENESS OF THE PRIMITIVES
Besides the above listed operators, few other operators are needed for completeness. For example, the operator to permute the order that arguments are passed to function: PermuteArgsal (f) T, + [(y, the operator to zip together two functions: Zip(J, (g)) : 2; y + f(x) g((). We can inspect whether these rules forms a closed set by examining whether the right hand side of the rules (7) to (18) can be implemented with these primitives themselves. Most of the rules are readily implementable with the primitives, with few exceptions:
1. Rule (9) and rule (16) requires the inverse of a function, which is not implemented because the general mechanism to invert a function is not available in JAX. When the function is a non- invertible map, the transpose rules can still exist mathematically (Equation 25). However, it has to be implemented case by case based on specifie properties of the mapping.
2. Rule (14) involves an integral, as discussed Section 3.2.5, there is no general rule for integration for any functions. Therefore, the accuracy of the integral is limited by the chosen numerical grid.
3.4 EFFICIENCY OF THE IMPLEMENTATION
There are a few practical issues considered in our system. Firstly, chain function composition is expensive, if we have function h, and we compose with the expression foo h + g h. In python g( (h())) without proper treatment, h(x) will be invoked twice.
code this is lambda x: fxh(x)) The situation is exacerbated when the resulting function is again composed with other functions multiple times. Although in JAX, common sub-expression elimination (CSE) is performed to remove redundant computations during the JIT compilation, we find practice that the computation graph could grow to prohibitively large size before we could JIT it. To resolve this problem, we add a cache for the function calls. When the same function invoked several times with the same argument, the cached output is directly returned. With a large amount of redundant computations, this optimization greatly reduced the computation time of the resulting function (Appendix D).
Another performance issue hides in the V operator. For one, the reverse mode derivative approxi- mately doubles the nodes in the computation graph, the size of the graph grows quickly when we perform higher order derivatives. Moreover, it is quite common in applications that we need to per- form various orders of differentiation on different variables, e.g. (f)((() y) and (f( y). It is worth noting that the inner differentiation of the first expression (f( shares computation graph with the second expression. Again, it could be prohibitively expensive for the tracing phase. This issue is specific our system but in general exists in JAX. We see an opportunity for optimization using our system because we build graph made of V operators and functions without actually executing it. It is desirable compile mixed order differentiation of functions into efficient programs, for example, Faà di Bruno mode differentiation (Bettencourt et al., 2019). We leave this optimization as a future work.
APPLICATIONS
In numerous application scenarios, function approximation problems are often transformed into parameter fitting tasks, such as using neural networks or linear combinations of function bases. Gradients in the function space may appear unnecessary, as our ultimate optimization takes place in the parameter space. In this section, we focus on application cases where functional derivative plays crucial role.
4.1 SOLVING VARIATIONAL PROBLEM
Variational problems are problems that search for the optima of a functional. We consider local functional of the form Fly) = (x), Vy²(x))dr Conceptually, we just need to compute 5f
the functional derivative 8y and perform gradient descent for functionals. In practice, to make it implementable, we usually introduce a parametric function y (x). For optimization, we can directly compute the gradient w.r.t. 0.
θ a8 F(g0)=
a [l- (z),
(19)
6

---

Published as a conference paper at ICLR 2024
Optimized
Runctional Gradient
0.0
groutruth
101
min
mirimize
-0.4
-10-
-0.8
-167
-107
-1.2
0.5
0.0
0.4
0.5
Figure 1: Brachistochrone curve fitted with different loss
While the above is the most straightforward method for this problem, with the introduction of functional gradient, enable few more options. One of the options is to perform chain rule through the functional gradient, namely, we compute the functional derivative 6T(y) pointwisely backpropagate the gradient to the parameter,
first, then
by
08 F(y²)= - 8T((")
(z) dr 00
8y" (x)
(20)
Notice that in Equation (20) the functional derivative can be explicitly expanded using the Euler- Lagrange formula. With some derivation (Appendix E), we can see that Equation (20) and (19) are two different estimators for the gradient of 0, though they share the same expectation.
One more choice we have is to directly minimize the magnitude of the functional gradient based on the fact that the optimal, functional gradient should be zero.
/
8T(x") 5y""
min 6
(x)
(21)
We apply the above methods on the classic brachistochrone problem. The brachistochrone represents the optimal path that minimizes the travel time for a particle moving between two points in the gravitational field. Assume the initial position x, y 0, 0, and the end position 2, = 1, 1. The travel time can be represented as functional of the curve, denoted as F(y) = So VI+ Vy(a)² /V(a))//V((zzda The numerator calculates the length of the curve, the denominator is the speed of particle. We use (x) nn sin(TT)) - the approximator, where nn is an MLP and the other terms are used to constrain the initial and end positions. More details of the experimental settings can be found in Appendix G
Instead of implementing Equation (20) explicitly, we illustrated below the simplicity of using AutoFD to compose the functional and compute functional gradient. The results are summarized in Figure 1, the methods (19), (20) and (21) are named minimize F, minimize F via FD and minimize FD correspondingly. It worth highlighting that the directly minimize in the parameter space limited by the numerical integration, it easily overfits to the numerical grid. Meanwhile, the functional derivative involved in the other two methods seems to provide strong regularization effect, leading to curves close to the groundtruth. Better functional optimal are found as can be seen in Figure (right) that the functional gradient closer to zero. Whether Equation (20) is better than Equation (19) in general context needs further investigation.
from autofd. operators import integrate, sqrt, nabla
def y (params, x):
def Fly: Callable) -> Callable:
return integrate (sqr (streettttttttabaaarl (1+nabla(y, argnums=1)*x2) (esp=)/x/////(((((() /sqrt(-))) fd (y)

---

Published as a conference paper at ICLR 2024
4.2 DENSITY FUNCTIONAL THEORY
Density functional theory (DFT) is widely used for calculating electronic structures. Although the DFT theory proves the existence of functional that maps electron density the ground state energy. the accuracy of DFT in practice relies heavily on approximated exchange-correlatin (XC) energy functionals with the general form
Exc(p) = / p(T))xc(P)(T)dr
(22)
Where p the electron density function, and (x) the XC energy density which is dependent on p. A comprehensive list of XC functionals are maintained in LibXC (Lehtola et al., 2018), its JAX translation introduced in JAX-XC (Zheng & Lin, 2023). In the Self-Consistent Field (SCF) update of the DFT calculation, the effective potential induced by the XC functional is computed as its functional derivative Vxe (e) Traditionally, the Vxo is derived via the Euler-Lagrange equation. 8p. Its implementation is often done in the parameter space of p and entangled with the numerical integral scheme. With automatic functional differentiation, the computation of Vxe can be decoupled from the choice of parameterization and integral schemes. In the code example below, gga_x_pba is implemented as an semilocal operator using the primitives defined in this work. Vxc can be obtained by calling ggrd on Exe.
from autofd. operators import integrate i point in 3D
r jip.. arpaarray([O.1 1, 0.2, 0.3])
# Exc, exchange-correlation energy functional
def exc(rho: Callable) -> Float:
epsilon_xc sga_x_pbe(rho)
return integrate rho)
# Vxc
functional derivative Exc
vxc = jax (ex.g) (rho)
# Vxc is itself a function callable with as input
vxc(r)
4.3 DIFERRNNIAATIG NONLOCAL FUNCTIONALS
Both the brachistochrone and DFT shown in previous sections are examples of integral functionals with semi-local functionals that conforms 2. In this section, we explore more interesting case where a nonlocal operator is involved. We consider the higher order version of multilayer perceptron (MLP), where the input is a IR - R function. The ith linear layers follows f((x) I k(x,, f) + b(x),, where & is the kernel function resembling the weights in MLP, and b acts as the bias. Activation functions are applied pointwise after the linear layer. For the last layer, we don't 't apply any activation, but compares the the output with a target function with Le loss. This is often called neural operator or neural functional, or which the Fourier Neural Operator (FNO) is special case where the integral is done via fourier series.
To learn a neural operator, traditionally we use parameterized neural networks as k k (x.y) y) and (x) where f, & are network parameters. The gradient descent is directly done in the parameter space, seems like we can't apply AutoFD here. There is, however, an expensive way to do it, we can directly subtract the functional gradient from the old function, i.e. k <- k 77 8L(K), where is the learning
rate and both and 5L(k) 5k are functions. is expensive because the new I expands in the computation graph. The purpose of this experiment is only show that AutoFD is capable of handling such complex computations. We show in Appendix H that our proposed update moves in the direction of smaller losses, though each step gets more expensive and yields more complex k.
RELATED WORKS
AD has long history of wide applications across various domains in science and engineering. More recently, the success of deep learning has brought automatic differentiation focal point in machine
8

---

Published as a conference paper at ICLR 2024
learning. In the world of machine learning, several differentiable programming frameworks backed by highly optimized vectorized implementations are proposed in the past few years, for example, Theano (Bastien et al., 2012), Mxnet (Chen et al., 2015). Tensorflow (Abadi et al., 2016), Jax (Frostig et al., 2018) and Torch (Paszke et al., 2019). All of these frameworks focus on AD of functions that map real values. Differentiation of functionals and operators, however, has been mostly derived analytically and implemented case by case. There are a number of works that studies AD in the context of higher-order functionals (Pearlmutter & Siskind, 2008; Elliott, 2018; Shaikhha et al., 2019; Wang et al., 2019; Huot et al., 2020: Sherman et al., 2021). They mainly focus on how AD can be implemented more efficiently and in compositional manner when the code contains high-order functional transformations, which is related to but distinct from this work. For example, Elliott (2018) studies how DJ o (x) can implemented more efficiently using the algebraic relation of DJ (a)) = (())(((a)) (g(=:) (ar); while in this work, we study the differentiation of the higher-order function itself, namely (f))
The most closely related works are Birkisson & Driscoll (2012) and Di Gianantonio et al. (2022). Birkisson & Driscoll (2012) computes the Fréchet derivative of operators by representing functions as linear combination of Chebyshev basis functions. Di Gianantonio et al. (2022) describes language called Dual PCF which is capable of evaluating forward mode directional derivatives of functionals. The method is based on dual numbers and there is reverse mode support mentioned. There are also implementations of functional derivatives in symbolic math packages. For example, Both Mathmematica (Wolfram Research) and SymPy (Meurer al.. 2017) implement the semi-local functional derivative in the canonical Euler-Lagrange form. In Maple (Maplesoft), there is also the FunDiff command, which relies on the calculus of Dirac delta function and its derivatives based on the provided information. Their implementation differs from the AD approach we take in this work, and they do not support vectorized execution on modern hardwares.
DISCUSSION
In this work, we introduce AutoFD, a system to perform automatic functional differentiation. We take a novel approach to functional differentiation by directly reusing the AD machinery for higher order functions. We implemented a core set of operators that covers various useful types of functionals. We discuss several limitations here potential directions for future work.
Completeness: As discussed Section 3.3, in several of the rules, the inversion operator is required. It would rely on systematic mechanism to register the invertibility and inverse function for the primitives, at the time of writing, such mechanism is not implemented in JAX.
Analytical integration: It is desirable in applications like quantum chemistry to analytical functions and integrate them analytically. While integrating symbolic packages like SymPy to JAX (Kidger) could provide this functionality, it is limited to scalar functions. Automatically de- vmapping the vectorized primitive to scalar functions could be one potential path to generally bring analytical integral to JAX.
Static shape: AutoFD requires accurate annotation of functions using jaxtyping. This is a design choice allow early raising of errors as it more informative than delayed the execution of resulting functions. However, this not only adds extra work but also limits the flexibility of using AutoFD. Further exploration is required for a better trade-off.
Programming in mixed order: For example, the partial transformation in python is a mixed order operator that binds an argument a function (both considered inputs the partial operator). While is possible to support gradients for both the argument and the function. Complications emerge during the just in time compilation, jitting a mixed computation graph not possible because the operator primitives are pure python and do not support lowering. Ultimately we would like to remove this constraint and program differentiably for real values as well as any order of functions.
9

---

Published as a conference paper at ICLR 2024
REFERENCES
Martín Abadi, Paul Barham, Jianmin Chen, Zhifeng Chen, Andy Davis, Jeffrey Dean, Matthieu Devin, Sanjay Ghemawat, Geoffrey Irving, Michael Isard, al. (TensorFlow): system for [Large-Scale] machine learning. In 12th USENIX symposium on operating systems design and implementation (OSDI 16), pp. 265- 283, 2016.
Alampallam V Balakrishnan. Applied Functional Analysis: A, volume 3. Springer Science & Business Media, 2012.
Frédéric Bastien, Pascal Lamblin, Razvan Pascanu, James Bergstra, Ian Goodfellow, Arnaud Bergeron. Nicolas Bouchard, David Warde-Farley, and Yoshua Bengio. Theano: new features and speed improvements. arXiv preprint arXiv:1211.5500 2012.
Jesse Bettencourt, Matthew J Johnson, and David Duvenaud. Taylor-mode automatic differentiation for higher-order derivatives in jax. In Program Transformations for ML Workshop at NeurIPS 2019.
Asgeir Birkisson and Tobin A Driscoll. Automatic fréchet differentiation for the numerical solution of boundary-value problems. ACM Transactions on Mathematical Software (TOMS), 38(4):1-29, 2012.
Tianqi Chen, Mu Li, Yutian Li, Min Lin, Naiyan Wang, Minjie Wang, Tianjun Xiao, Bing Xu, Chiyuan Zhang, and Zheng Zhang. Mxnet: A flexible and efficient machine learning library for heterogeneous distributed systems. arXiv preprint arXiv:1512.0122 2015.
Miles Cranmer, Sam Greydanus, Stephan Hoyer, Peter Battaglia, David Spergel, and Shirley Ho. Lagrangian neural networks. arXiv preprint arXiv:2003.044 04630,, 2020.
Pictro Di Gianantonio, Abbas Edalat, and Ran Gutin. A language for evaluating derivatives of functionals using automatic differentiation. arXiv preprint arXiv:2210.06095. 2022.
Leon Ehrenpreis. On the theory kernels of schwartz. Proceedings of the American Mathematical Society, 7(4):713-718, 1956.
Conal Elliott. The simple essence of automatic differentiation. Proceedings of the ACM on Program- ming Languages, 2(CCF)):1-29, 2018
C. Daniel Freeman, Erik Frey, Anton Raichuk. Sertan Girgin, Igor Mordatch, and Olivier Bachem. Brax a differentiable physics engine for large scale rigid body simulation, 2021. URL http: //gi thub. www.coo///oooo gwwg///ooll......oo
Roy Frostig, Matthew James Johnson, and Chris Leary. Compiling machine learning programs via high-level tracing. Systems for Machine Learning, 4(9), 2018.
A Gonis. Functionals and functional derivatives of wave functions and densities. World Journal of Condensed Matter Physics, 2014, 2014.
Zheyuan Hu, Tianbo Li, Zekun Shi, Kunhao Zheng, Giovanni Vignale, Kenji Kawaguchi, YAN Shuicheng, and Min Lin. Neural integral functionals. ICLR 2023 Workshop on Physics for Machine Learning, 2023.
Mathieu Huot, Sam Staton, and Matthijs Vākár. Correctness of automatic differentiation via diffcolo- gies and categorical gluing. FoSSaCS, pp. 319-338, 2020.
Patrick Kidger. sympy2jax. https://githii com/ gowgl//olllleeoooo sttps//mmpyyjaaa.
Susi Lehtola, Conrad Steigemann, Micael JT Oliveira, and Miguel AL Marques. Recent developments in libxe comprehensive library of functionals for density functional theory. SoftwareX, 7:1-5, 2018.
Zongyi Li, Nikola Kovachki, Kamyar Azizzadenesheli, Burigede Liu, Kaushik Bhattacharya, Andrew Stuart, and Anima Anandkumar. Fourier neural operator for parametric partial differential equations. arXiv preprint arXiv:2010.08895 2020.
10

---

Published as a conference paper at ICLR 2024
Maplesoft. functional differentiation. URL http:///wwww..mmtttttttto//////s v wwt.::pppppptithhhsssssss/aaaai
Johannes T Margraf and Karsten Reuter. Pure non-local machine-learned density functional theory
for electron correlation. Nature communications, 12(1):344, 2021.
Aaron Meurer, Christopher Smith, Mateusz Paprocki, Ondřej Čertík, Sergey B. Kirpichev, Matthew Rocklin, AMIT Kumar, Sergiu Ivanov, Jason K. Moore, Sartaj Singh, Thilina Rathnayake, Sean Vig. Brian E. Granger, Richard Muller, Francesco Bonazzi, Harsh Gupta, Shivam Vats, Fredrik Johansson, Fabian Pedregosa, Matthew J. Curry, Andy R. Terrel, Stěpán Roučka, Ashutosh Saboo, Isuru Fernando, Sumith Kulal, Robert Cimrman, and Anthony Scopatz. Sympy: symbolic computing in python. PeerJ Computer Science, 3:, January 2017. ISSN 2376-5992. doi: 10.7717/peerj-cs. 103. URL https://doi org 16. 7717 j 103.
Adam Paszke, Sam Gross, Francisco Massa, Adam Lerer, James Bradbury, Gregory Chanan, Trevor Killeen, Zeming Lin, Natalia Gimelshein, Luca Antiga, et al. Pytorch: An imperative style, nigh-performance deep learning library. Advances in neural information processing systems, 32, 2019.
Barak A Pearlmutter and Jeffrey Mark Siskind. Reverse-mode ad in a functional framework: Lambda the ultimate backpropagator. ACM Transactions on Programming Languages and Systems (TOPLAS), 30(2):1-36, 2008.
John P Perdew and Karla Schmidt. Jacob's ladder of density functional approximations for the energy. In AIP Conference Proceedings, volume 577, pp. 1-20. American Institute of Physics, 2001.
Davide Piras and A Spurio Mancini. Cosmopower-jax: high-dimensiional bayesian inference with differentiable cosmological emulators. arXiv preprini arXiv:2305.06347, 2023.
Alexey Radul, Adam Paszke, Roy Frostig, Matthew Johnson, and Dougal Maclaurin. You only
linearize once: Tangents transpose to gradients. Proceedings of the ACM on Programming Languages, 7(POPL): 1246-1274 2023.
Samuel S Schoenholz and Ekin D Cubuk. Jax, md framework for differentiable physics. Journal of 2021.
Statistical Mechanics: Theory and Experiment,
Amir Shaikhha, Andrew Fitzgibbon, Dimitrios Vytiniotis, and Simon Peyton Jones. Efficient differentiable programming functional proay-proessiing language. Proceedings of the ACM on Programming Languages, 3(ICFP):1330 2019.
Benjamin Sherman, Jesse Michel, and Michael Carbin. Computable semantics for differentiable programming with higher-order functions and datatypes. Proceedings of the ACM on Programming Languages, 5(P0PL):133,, 2021.
Shlomo Sternberg. Lectures differential geometry, volume 316. American Mathematical Soc., 1999.
Fei Wang, Daniel Zheng, James Decker, Xilun Wu, Grégory M Essertel, and Tiark Rompf. Demysti- fying differentiable programming: Shift/reset the penultimate backpropagator. Proceedings of the ACM on Programming Languages, 3(IFF))1131, 2019.
Inc. Wolfram Research. Variationald. URL https: httt::://reeeeence... www.ww..cco/wwaaaaaammmmmmmmmma lamm/l/gggage///// Bbrraaaiionlllettodu/uuuorri//// hwtp://awaaaaiiiaaaeeeaeaaeoaaoe html.
Federico Zahariev, Sarom Sok Leang, and Mark S Gordon. Functional derivatives of meta-generalized gradient approximation (meta-gga) type sichange-correlation density functionals. The Journal of Chemical Physics, 138(24), 2013.
Kunhao Zheng and Min Lin. Jax-xc: Exchange correlation functionals library in jax. In Workshop on" Machine Learning for MItLr 2023, 2023.
11

---

Published as a conference paper at ICLR 2024
A EXAMPLE OF IMPLEMENTATION
To clarify how the math are correspondingly implemented as extension to JAX, we show a minimal implementation of the operator V. We restrict the implementation take only scalar function, that the divergence V 8h is equal to the gradient Vôh. With this simplification, the JVP and transpose rules are
((V)(f) 88 + (6ff). (f)) 6h + Vöh.
(23)
(24)
Here's list of mappings between math symbols and the code.
f:
of: df
ôh: dh
V: nabla
(&f): nabla_.vo__ruee((.)), (df,)
(oh): mabla_ttaanspos_rulleddda f)
We first implement V as a JAX primitive.
nabla_p = core. Primitive (ivabbaal"""))
def nabla(f):
return (f)
mmabla_p.def_ite def nabla_impl(f) (f):
return jax. gf)
@nabla_p. lecpllappdde__assteeeee def def_abstract_eval
# has scalar input and output fax.grad(ff has signature as return shape
def nabla is a linear operator
tangents):
f, df = primals[0], tangents[0] return nabla(f), nabla(df)
def gabla_transposs cunppose_rull primal):
# According to the transpose rule in math dh = cotangent
# we assume here negation function # is already implemented by the compose operator return "nabla(cotangent)
# we register the jvp and transpose rules.
jax. interpreters. ad. primittve_jivssieee jax. interpreters ad. primi tive transposes nabla_transsnneneeeee transpose_rule
= nabla_ivp_rule
12

---

Published as a conference paper at ICLR 2024
Now define some random functions as primal, tangent and cotangent values.
def f(x):
return sin(x)
def df(x):
return x*
def dh(x):
return jnp. (x)
Finally, show how the operator can be invoked, and how to perform automatic functional differen- tiation on the operator.
# We can use nabla directly on f,
nf = nabla(f)
# triggers nabla_impl,
jnp.cos
# Or, we can compute the jvp of nabla, remember nabla is an operator
# jap here is computing the forward mode gradient for an operator!
nf, ndf = jax. jox primals=(f,), tangents=(df,)) # triggers nabla_jvp_rule, nf: jnp.cos, ndf: lambda 2x
# Since nabla is linear operator, can transpose it. tnabla jax. traneerrrtrasssseeaaaaa pr primals=(f,)) # linear transpose of an operator still an operator, # we apply this new operator tnabla on the function dh. tndh tnabla(dh) # triggers tabla__transpoe__rkke tndh: lambda -jnp.
# Or, can do the backward mode gradient on nabla primal_out, vjp_function jax. vjp(nabla, f) # invoke the vjp function the cotangent dh
vip_dh (d)_f
# this triggers both nabla_jvp_rule and mabla_transpose.c
# vjp_dh: lambda x: -jnp.exp(C
B PROOFS OF JVP RULES
JVP rules are trivial for linear operator, for a linear operator Ô, the JVP rule are always simply applying the same operator on the tangent function,
:f
For set of operators, V. L and T are all linear operators that need no extra proof for the JVP rules. We give here step by step derivation for the JVP rules of the Ć' operator.
(,(O)(f.))(6)(e) lim f(g(x) + r5g(x)) f(((x))
T-0
T
d
di f(g(x) + rdg(z)) ---= v(f(o(z)) 8g(x))
-((
13

---

Published as a conference paper at ICLR 2024
[((())) + +5f(((())) f((()))
((( lim T-+
=(((()))
= -((
C
PROOFS OF TRANSPOSE RULES
Given operator Ô the adjoint of an operator Ó" satisfies (Óu,o) v) (u, On
C.1
COMPOSE
The Schwartz kernel form of compose operator is C(S.)) = Ō( ((2))) (x))d..
/ drō(2 9())f((z)
(c(f.)).5h)
dyōh(v)
[ dx dx / dyd(o dyd(xx99())dmk g(v)))h(())
(S. dyō(s ((()))(()))
(25)
- daf(()) / dsl
f drf(z) ( (1) (
(J. (U,C det det \)g
Therefore, transposition Ô w.r.t. fss
())))))))(h)) = C(sh, g det V(y
In the case where f is linear, we can write f(x) as JJæ where Jy is the jacobian matrix of f. We omit the case where f is nonlinear, as transposition only defined and implemented for linear operators. Transposition w.r.t. is simple
/ (Faly)) 8h(y))d
(C({,)),hh)
/ g (T) 8h(y)dy
(g. (ff((f) 8h))
Therefore, transpose of Ò w.r.t. iis,
T,(C)(f.))(hh) = C(f()),
C.2 NABLA
A brief proof for single variable function given in the main text, here expand to the multi- variate case.
(Vf.g) Σ / (yg](t) / do0'(2; v,)f
1j
Σ / -v)o(()))) / ((V--),2))de
(J. -V.
14

---

Published as a conference paper at ICLR 2024
C.3 LINEARIZE
(L(f),8h) /Iʃ d26' - =x)f(2)ōr 8h(x, oz dxdddd
/def(e) / dōd 6x I dad' 2)8h(2, 6z)
/ (zff()) (-Jozvo8hh,.oo 62))
(f. - VzVohl, (x)))((
Therefore,
(()))))())== = (on))oh((,,,aa
C.4 LINEAR TRANSPOSE
Linear transpose of linear transpose operator sounds interesting, Since linear transpose only be applied to linear functions, we write the function being transposed as x + Jyz, transpose of T(f) y + JJ y. We're only able to derive the adjoint of T when 8h is an invertible mapping.
(f)
(JF v) 5h(y)dy
y Joöh(y)dy
i 6h (x)) Jjædöh
det V(Sh v(6h-1)x)) (x)) Jyædæ
(f. det ((8h (/6hh-)|687 |1hh
Therefore,
((7)()(()) = (5h) = det V(ōh ((0h-1)16 1[8h"" -1
C.5 INTEGRAL OPERATOR
(4()),8h) = dmws / do: dæif Zm.)
/ f())0h(z)dr (f.
Where 8h is F., Ini 8h(4x)). Therefore the adjoint the integral operator is simply augmenting the cotangent function öh with unused arguments to match the domain of the f.
15

---

Published as a conference paper at ICLR 2024
D COMMON SUBEXPRESSION ELIMINATION VIA CACHING
In Section 3.4, we introduced the composition foh g goh. We can recursively nest this composition by setting h4 = f h + 9 h4-1- The redundant computations are then exponential to the depth of the composition. Specifically, we use the following code for measuring the execution cost for different nested depth, with and without function call caching.
# sin, exp, tanh for h, f, g respectively. def F(h):
for in range
f(h) + g(h)
return
# Fh is h f(h) g(h) nested depth times
Fh = F(h)
# time the execution
t1 time. time()
jax.jit (0.)
t2 = time. time ()
cost t2 t1
When the above nested function composition implemented naively, the time cost of computation grows exponentially with the nest depth, because there are two branches f f(h) and (h) at each composition. However, when function call caching is used, at each level one branch cached result of the other branch, resulting in linear time cost with respect to the nest depth. We plot the time cost vs the depth in Figure
reuse the
Execution time vs nest depth
w/o cache
w/ cache
101
Execution time (s)
109
10-1
0
5
10
15
20
25
30
Nest depth
Figure 2: Function call caching significantly improves execution efficiency.
16

---

Published as a conference paper at ICLR 2024
E
COMPARING GRADIENT ESTIMATORS
Subtracting Equation (19) and (20) we get:
OT
(v)(2)
01
0x((x)
dr
OV (x)
V (④) 00
/ OV (x)
OT
(
OT
0yf(x)
00 dr
88
(x)
/ (v)
01
Oy (a)
dx
00
= / (v)
Oy (x)
dx
(26)
00
In Equation (26), the integral evaluate to a constant that is not dependent on æ;; therefore taking V on the integral yields O. This proves Equation (19) and (20), when the integrals are discretized, are different estimators of the same quantity.
17

---

Published as a conference paper at ICLR 2024
F HIGHER ORDER FUNCTIONAL DERIVATIVE IN LIBXC
def (setf Inp, cutputatione do_ excexee do_vxc=True, fo_ db_caceaalie, do to.taeet)):);
grid.
Parsmeters
Inp 1 np op-ndarray on diot of np np.ndarnay
input cictionery NusPy arrey-Like structures that provide the density grid and its derivaties. These LabTed: density grid
signz the contracted donsity gradients
lapt laplacian
density
energy density
Each femily functionals requires different derivetives:
LDK:
GBA: rho, sigra
(aptionat),
output dist of ndaandyyyyy Costionsl, Ronne)
Contains dictionary Function sill build outsut space you. outout dictionary depends derivates requested. A compretansive List is provided beloe for each functionsl, fomily.
array-Like structures
supplied this
LDK:
EXC:
VXG: vrho
vžrho2
LXC: várhod
0BA:
EXC:
VXE: vrho, veigmn
v2rho2,
v2signa2
RXC: várho3, v3rho2siges, 98rhosigna2, vásigna3
LXC: várhoG, várholsiges, várho2signs2, väroosignat, vásigra4
EXC:
VXE: vrto, veigen, vlapt. (optionat), vtau vžrho2,
vžrholapl, varhota., v2sigea2, v2signntapt, väsignstuu, v2tap12, v2lasttau, v2tau2 IXE: várhos, várhozziges, värnoztapt, vārhožtau, virtosigna2, v3rholag12, värhelapltau, #ärhosignlapp, sisigna2aap,, väsigra2ta, varhoteu2, vásigea3, vicignolap12, väsignstapttau, vāsigratzu2, v3tapt3, vStap12tau, vSlapltaud, v3tau3
LXc: virbod, värholsiges, vérhoStapt, värhottaw, vörho2sigra2, vérhozesgeatast, vimozsigeatau v4rhožtap12, virho2lapttau, vérho2tau2, vérhosigre3, virhosigre2led vértosigastou, virbosigestaptz, vimesigntaaits, vörtosignetau,, vérholasts, Varholapt2tau, wirholapltac2, véreetau3, visigraSlast, visigm/tapttuu visigna2tau2, vésignstap13, vidignstap12tuu vecigeelapitas2, vesignataus, y4lap14, v4laplitau, válap12tau2, vélapitau3,
wisiges2tap2,,
Figure 3: Many terms are involved in constructing higher order functional derivatives.
18

---

Published as a conference paper at ICLR 2024
G EXPERIMENTAL SETTINGS FOR BRACHISTOCHRONE
For the brachistochrone experiment, the initial position is at (Io, yo) = (0, 0) and the end position is (#T, (1, -1). The functional we minimize
Fly) L'' V1++/(3)//-oo111dd
We aim to find
y" arg min Fly)
We use (()) (LP((r) sin( x to ensure that it passes through (0,0) and (1, -1). The MLP is a multi-layer perceptron that maps R - R, with hidden dimensions as {128, 128, 128, 1J. All layers uses the sigmoid d function as the activation function, except for the last layer which has activation function.
For the integration, the ending point is at 1. The reason we choose a non zero starting point is because a singular point for the integrand, i.e. the denominator (0) = 0. For the optimization process, since is toy problem, we use out the box adam optimizer from optax, with fixed learning rate le and optimize for 10000 steps.
uniformly sampled grid of 50 points, the starting point is at 0.01 and
is worth noting that this is not the optimal setting for doing integration, one can either use a more densely sampled grid use Monte Carlo integration for better fitted brachistochrone curve. We choose this setting to show that Equation (19) and Equation (20) are very different estimators, and that Eqation (20) could have regularizing effect the scale of functional derivative (Figure 1) due to the fact that functional derivative is explicitly used.
19

---

Published as a conference paper at ICLR 2024
H NONLOCAL NEURAL FUNCTIONAL WITH FUNCTIONAL GRADIENT DESCENT
We describe our precedure of optimizing the nonlocal neural functional in more detail here. The neural functional we use has two linear operator layers, the first layer uses a 'tanh' activation function, while the second layer uses activation. The final output function is compared with the target function via L2 loss. We take a learning rate of 0.1 for steps. The code of this experiment is presented below.
import operators
def f(x:
Float32[Array, ""]) Float32[Array, ""]:
return jnp.sin(( x * inp.pi)
def b(x: Float32[Array, ""]) >> Float32[Array, ""]:
return jnp. sin(x jnp.pi)
def y(:: Float32[Array, ""]) Float32[Array, ""]: return jnp. co jnp.pi)
def k(y: Float32[Array, ""], Float32[Array, ""]) -> Float32[Array, ""]: return jnp. sin(y) jnp. cos(x)
def layer b, f, activate=True):
# here @ f is syntatic sugar for o.integrate(k * broadcast(f), argnums=1)
g k f if activate:
o. numpy. numpy.tanh(e,, tgnh(g)
return
else:
return
def losrams, f, t): two layer mlp k1, b1, k2, b2 params
h1 = layer b1, f, activation=True)
h2 = layer b2, h1, activation-False)
return o. ih12 +)**2)
# initialize both k1, k2 to k, b1, b2 to b param (k, b, k,
# perform gradient steps
1 = lossmp,ram, f, t) ofinttt"iittial loss: (1)")
for in range(3):
grad jax. (loss) (param, f, t)
param = jax. tree_uteett wtwl. tree_map (lambda x, dx: 0.1 dx, param, grad) loss (param, f, t)
print at step (i): (1)")
As can be seen, in the limited number of steps we take, the loss steadily goes smaller (Figure 4). We visualize the prediction from the neural functional the target function in Figure 5, and the kernel k(x,, y) from the first layer of the neural functional Figure 6. The reason we only takes 4 steps of descent is because the learned kernel function gets prohibitively large. We show the JAXPR graphs of k( y) from the first neural functional layer for each step in Figure 7, notice that the graph for step 4 is too big that we failed to render it.
20

---

Published as a conference paper at ICLR 2024
2.5
1.0
Figure 4: Loss vs Step for the nonlocal functional training
114
612
332
330
0.3-
620
4.37
-45
Saya O0
Figure 5: Predicted function vs target function at each step.
02-
64
34
3.4
6A
15
21
14
300
67
84
04
50
04
26
10
18
02
24
66
3.0
Figure 6: Visualize kernel k(x,,) from the first layer each step.
fow
(a) step 0
(b) step 1
(c) step 2
(d) step
Figure 7: Jaxpr graph of k(x y) at each step.
21


Reviews:
[{"review_id": "z4aOXPpvt4", "reviewer": "Reviewer_Mbhg", "summary": "In this work the authors introduce a package for the machine learning framework JAX, to enhance JAX with functional differentiation, stemming from functional calculus / the calculus of variations. This enables the easier expression of e.g. neural operators, as they implicitly rely on the functional framework. To this end, the authors introduce 5 new language primitives, namely `compose`, `nabla`, `linearize`, `linear transpose`, and `integrate`. All implemented in pure Python, and complemented by the theoretical derivation of the individual operators.\n\nThe introduced extension is subsequently validated on examples from particle dynamics with the brachistochrone problem, and the exchange-correlation functional stemming from density functional theory.", "questions": "A number of questions arose while reading through the paper:\n\n- Did you perform an analysis of the computational properties of AutoFD, and most specifically the way JAX compiles AutoFD code if it is not being mapped to XLA-ops?\n- Is there an implicit trade-off computationally or conceptually to the chosen representation as an infinite-dimensional array?\n- Why did you choose a purely Python-based implementation of AutoFD, as compared to a version of AutoFD acting directly on the JAXPR?\n- For the _Efficiency of Implementation_, have you considered tracing your implementation with perfetto\n- How do you anticipate an inversion operator in JAX to be feasible? Would it be possible to implement such operator, if you were writing the operators at the level of the JAXPR?", "limitations": "", "rating": 8, "confidence": 4, "soundness": 4, "presentation": 2, "contribution": 3, "strengths": "Where this paper shines is its connection with a very strong theoretical background stemming from functional calculus, and how it derives its proposed extensions to JAX from said theoretical motivation. The two chosen examples only underline this further.\n\nIt removes previous constraints from a machine learning framework, hence acting in large part as an enable of future work using functional calculus, such as for learned operators.", "weaknesses": "At the same time, the work suffers from a number of unclear treatments of JAX, and a failure to establish the usage of the proposed framework for neural operators, such as the Fourier Neural Operator.\n\nJAX as a framework:\n- The authors explain the tracing into a DAG, and the mapping of primitives to XLA, the compilation backend underpinning JAX. What they miss in this instance though is the intermediate layer of JAXPR, JAX's internal representation, and the mentioning of operation-tracing, and XLA-compilation, only leads to a lack of clarity. I'd suggest to add a diagram of the pieces of JAX's architecture you rely on, and remove mentions of XLA-mapping of ops etc. from the latter parts of the paper.\n- While provided for some introduced operations, a number of operations do not have their implementation code attached to them. The addition of the code in the main paper, or the appendix would contribute greatly to further the clarity of the exposition.\n\nNeural Operators:\n- While neural operators, and most specifically Fourier Neural Operators, are presented as a clear motivation for the work, they are sadly not used in the experimental evaluation of the presented extension. Addition of a functioning Fourier neural operator based on `autofd` would significantly strengthen the paper's claims. Evaluation of a Fourier neural operator could for example be in the form of a normal JAX-based implementation, and an `autofd`-using implementation, where the code could for example be much more succinct with `autofd` while matching the performance of the JAX-native implementation.\n\nIn addition the draft has a number of minor typos, the addressing of which would improve the legibility greatly. For example:\n- Page 1, last paragraph: \"Base\" -> Based\n- Page 2, last line: \"maps arbitrary\" -> map arbitrarily", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "hpboRV2z2h", "reviewer": "Reviewer_4ztS", "summary": "This paper introduces functional differentiation (i.e., derivatives of functionals, which are functions that take other functions rather than values as inputs). The package directly builds on JAX's automatic differentiation machinery, introducing a new datatype for functions (section 4.1) and then implementing the appropriate linearization and transposition rules for each primitive operation. The primitive operations implemented are differentiation (nabla), linearization (Fréchet derivative), linear transposition, and integration, as well as some utility operators (composition, permutation, zipping). Caching is used to perform common subexpression elimination while the computation graph is being built. Finally, some experiments using functional differentiation are performed (brachistochrone problem, exchange-correlation functionals) as showcase.\n\n---\n\nUpdated my rating from 6 to 8 in response to the changes during the discussion period.", "questions": "Could the machinery used to register derivatives not also be used to register function inverses?\n\nI would also love to see a more thorough discussion of the integration trade-offs: How should the user know what numerical grid to provide? And how sensitive should they expect the outcome to be to the grid provided?\n\nAs this is a software package, I would also like to know some of the details regarding public release: Will the code be on GitHub (or some other platform)? Under what license? Will there be documentation, tutorials, or notebooks? Does the code take the form of a separate Python package, or is it a fork of JAX?", "limitations": "", "rating": 8, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "I think this paper is quite straightforward, but not in a bad way: The text is clearly structured, and a good balance is found between the theory behind and the implementation of the framework. I appreciate how the authors were able to re-use JAX's machinery, which allows them to benefit from a lot of JAX's strength (compilation backends, debugging tools, etc.) and the simplicity of the brachistochrone in code is quite compelling.", "weaknesses": "The main weaknesses of this paper are some of the limitations discussed in section 6. In particular, not having any approach for function inversion seems like a shortcoming.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "HHBILQGVha", "reviewer": "Reviewer_RHkE", "summary": "In this paper, the authors extend the JAX automatic differentiation system to support higher-order derivatives -- i.e. derivatives over functions (often called functionals), rather than arrays.\n\nSome background: the core JAX autodiff system is based on a first order functional programming language.  The only data type that it supports is the Array (or tuples/dictionaries of Arrays), and all of the built-in differentiable primitive operations are first-order functions over arrays.  Consequently, JAX autodiff can only compute gradients with respect to Arrays. \n\nIt is important to note that the larger JAX library has plenty of higher-order API calls (e.g. vmap, linearize, jvp etc.), and python itself supports higher-order functions.  However, the core autodiff system works by *tracing* python functions to a first-order computation graph, and then computing gradients on that graph.\n\nIn this paper, the authors extend JAX type system to support functions as first-class citizens in the computation graph, and they introduce a set of differentiable primitive operations that are higher-order (e.g. function composition).  The extended autodiff \nsystem uses the same JAX autodiff machinery, but it can compute gradients with respect to functions/functionals as well as Arrays.\n\nThe authors describe the differentiation rules for a core set of higher-order primitive operations, and give a few examples, primarily drawn from engineering and physics, for doing autodiff over functionals.", "questions": "Please see weaknesses, above.", "limitations": "", "rating": 6, "confidence": 2, "soundness": 3, "presentation": 2, "contribution": 3, "strengths": "The paper is reasonably well-written, and the idea is mathematically interesting.", "weaknesses": "Although the paper is reasonably well-written, I had a very hard time following it.  Part of the problem is my fault.  Although I have deep knowledge of both JAX and automatic differentiation systems in general (having implemented several of them myself), I am not \nparticularly familiar with the underlying mathematics of functional analysis, Fréchet derivatives, and so on, which are used in this paper.   I have not checked the math, and I am happy to defer to other reviewers who have a deeper understanding.\n\nSince this paper is being submitted to ICLR, rather than a computational mathematics conference, I had expected the authors to provide a gentle introduction to some of the underlying concepts, suitable for ML practitioners.  Sadly, they do not, and I suspect that most ICLR readers will have the same problems understanding it that I did.\n\nThe authors claim that \"functions are represented as infinite dimensional generalizations of arrays\", but they do not explain how, nor do they even cite any source for this claim.  (Perhaps a functional analysis textbook?)  Moreover, I assume the \"infinite dimensional array\" is simply a mathematical abstraction, and it is unclear to me why it is even relevant.  In the body of the paper, functions actually seem to be represented in the usual way as symbolic programs.  As one would expect from a practical algorithm, infinite dimensions do not appear.\n\nThe actual differentiation rules are written in terms of forward derivatives and *transpose* operations.  The use of transpose is due to recent work by Radul et al., but will likely be unfamiliar to ML practitioners who are used to traditional backpropagation.  It would have been helpful if the authors had given an example of how these two operators work in a conventional (non-higher-order) setting, before diving into the higher-order case.  \n\nEven after doing my best to read this paper carefully, I am still unsure about how higher-order functions (functionals) are actually defined and represented in the core language. The differentiation rules given here do not seem to constitute what I would consider to \nbe a core programming language.  E.g. the authors provide differentiation rules for function composition, but not for function application or function definition.  That leads to me believe that the code for functionals would have to be defined in a point-free functional programming style, as is used by some other autodiff systems in the literature.  However, the examples in the paper just show ordinary python.  Is the python code traced, and then translated into the primitive operations, as is usual for Jax?  What are the details of this translation, since going from python to a point-free representation is not necessarily trivial?\n\nThe differentiation rules for function composition $f \\circ g$ require that $g$ is invertible.  If composition is the basic mechanism used to build complex programs (as is usually the case in point-free languages), then this would seem to be a very severe limitation.\n\nMy final point of confusion is that in most cases of practical interest, the functional that we are computing a derivative for is really just an ordinary symbolic function that is parameterized by some array $A$.  The value that we actually want to solve for is $A$ -- the higher-order derivative is just an intermediary step.  In this situation, ordinary Jax works just fine without higher-order automatic differentiation; the higher-order operations are eliminated by tracing and partial evaluation, and the first-order autodiff system then solves for the gradient of $A$.\n\nThe authors actually allude to this fact, but do not provide a detailed discussion of the tradeoffs between solving directly for $A$ using tracing and first-order autodiff, and doing something more complicated with higher-order autodiff.  Given the various limitations and restrictions on the higher-order methods, the former seems decidedly simpler.  Why would I want to use this system?  Can you explain why some problems can't be solved with ordinary Jax?\n\nFinally there are some typos which added to my confusion.  E.g. page 5 second paragraph $C(x, y \\to x + y, f, g)$ needs extra parens: $C((x, y) \\to x + y, f, g)$ otherwise it makes no sense.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "SJ2qy5yLDB", "reviewer": "Reviewer_VvN5", "summary": "The authors show that automatic functional differentiation (AutoFD) can be implemented in the same vein as automatic differentiation (AD) in JAX.\nThe authors introduce operators, namely, compose, ∇, linearize, linear transpose, integrate, JVP, and transpose rules.\nThe authors provide two applications of AutoFD: Solving the brachistochrone problem and density functional theory.", "questions": "# [Question (major)] Infinite dimensional generalization of arrays\nIn Abstract (and Introduction),\n> By representing functions as infinite dimensional generalization of arrays, we seamlessly use JAX’s existing primitive system to implement higher-order functions.\n- How was it realized in the proposed method? Could you point the part of the manuscript that concretely explain it?\n\n# [Comment (minor)] Reference \nIn Introduction,\n> To this date, functional differentiation with respect to functions are usually done in one of the following ways: (1) manually derived by human and explicitly implemented; (2) for semi-local functionals, convert it to the canonical Euler-Lagrange form, which can be implemented generally using existing AD tools; (3) use parametric functions and convert the problem to parameter space.\n- Could you add some reference papers about these three approaches for readers who are not that familiar with functional differentiation?\n\n# [Comment (minor)] Reference\nIn Section 3,\n> The Schwartz kernel theorem states that any linear operators can be expressed in this form.\n- Could you add a reference here? It is more reader-friendly.\n\n# [Question] Complex numbers\nIn Section 4.2,\n> The primitive operators considered in this work are focused to realize the most used types\nof operators and functionals described in Section 3.\n- [Question] Does your proposed program support complex numbers?\n\n# [Comment (major)] More definitions\nIn Section 4,\n- It would improve the clarity of the present paper to add a rigorous and/or intuitive definition of the Frechet derivative, cotangent space, and transpose rule and also an illustration of the primitive operations used in the present paper, because not all of the readers are familiar with both functional analysis and computer programming. Changing the order of explanations can also be an option (e.g., the begging of Section 4.2.3 can be an intuitive explanation of the Frechet derivative as a generalized directional derivative).\n\n# [Question] \nIn Section 4.2.1,\n> The function inverse on the right hand side of the $T_f (\\hat{C})$ rule is not implemented...\n- Does it restrict the operations of the proposed AutoFD? Could you take some examples?\n\n# [Question (major)] Grid points\nIn Section 4.2.5 and Experiment\n> We implement the integrate operator by supplying a numerical grid for the integrand function.\n- How did you choose the grid points?\n- How critical is the numerical error?\n- The proposed integral scheme looks like nothing but the conventional numerical integral. Is there any difference?\n\n# [Comment (minor)] Font\nIn Section 4.3 and Eq. (7--8),\n> We denote them as undefined because...\n- Changing the font of \"undefined\" would be good, e.g., mathtt.\n\n# [Comment (major)] Efficienty\n- For most of the statements in Section 4.4, I would like to see quantitative results.\n\n# [Questions and Comments (major)] Experiment: Solving variational problem\nIn Section 5.1, the authors performed an experiment to solve the brachistochrone problem.\n- The authors simplify the problem to a parametric fitting of $y^{\\theta}(x)$. This is how conventional methods in the numerical analysis of functionals does, as is stated in Section 1 and 2. What does the proposed program enable us to do, or what is the difference from the conventional methods?\n- What is the difference between Eq. (17) and (18--19)? Is it whether the Euler-Lagrange equation is used?  If yes, the performance gap given in Figure 2 may come from it.\n\n> It is worth highlighting that the directly minimize in the parameter space is limited by the numerical integration, it easily overfits to the numerical grid.\n- Taking random grids in every iteration is often done in learning integrals, which would fill the gap between red and the other curves. I would like to see the performance of such a more rational baseline.\n\n> Better functional optimal are found as can be seen in Figure 2 (right) that the functional gradient is closer to zero.\n- The authors should use log scales.\n\n# [Question and Comment (major)] Experiment: Density functional theory\nIn Figure 3 and 4,\n- I could not understand what Figure 4 means, potentially because Figure 4 is not a complete code like Figure 3. Could you provide more details of it? \n- Could you clarify how difficult it is to implement higher order derivatives? \n- Is Figure 3 simply a wrapper of Figure 4?\n- I would like to see numerical results about DFT using AutoFD.\n\n> In the SCF loop of the DFT calculation, ...\n- What is SCF?\n\n# [Comment (major)] Detailed experimental settings\n- Could you add more details of the experimental settings for reproducibility?\n\n# [Comment (major)] Code submission\n- Could you show the code for reproducibility?\n\n# [Question] PDE and FDE\n- Can we use the proposed AutoFD to solve a partial differential equations that include functional derivatives?\n- Can we use the proposed AutoFD to solve a functional differential equation that include functional derivatives (the task is to get the functional that satisfies the given equation that include functional derivatives)?\n\n# [Comment] Other journals and conferences \n- Automatic differentiation is actively discussed in, e.g., MFPS, TOMS, ICFP, TOPLA, POPL, and FoSSaCS. These community might have much more interest in the present paper.\n\n\n# [Comment (major)] Typos\nPlease proofread the manuscript before submission.\n- In Section 4.1, \"...one level of generalization, the term function in...\" should be. e.g., \"...one level of generalization. That is,  the term function in...\" \n- In Section 4.2, \"... for consistency. i.e. ...\" should be \"... for consistency; i.e. ...\"\n- In Section 4.2, \". Which\" should be \", which\"", "limitations": "", "rating": 5, "confidence": 3, "soundness": 2, "presentation": 2, "contribution": 2, "strengths": "- Automatic functional differentiation is required in many research areas, such as physics, chemistry, mathematics, and machine learning. The topic is very relevant.", "weaknesses": "- Presentation is poor (see Questions and Comments below for details) and much more clarification is needed.\n\n- The experiments are not reproducible.\n\n- The baseline of the experiments is too weak.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (poster)