Q9, what does the proposed program enable us to do ... is it whether the Euler-Lagrange equation is used?

Yes the direct approach is not new, in the experiment euler-lagrange approach seems better. What we have contributed here is the capability to use jax.grad(F) which automatically gives us the analytical gradient which is the same as derived from euler-lagrange. The advantage compared to explicit doing euler-lagrange are,

1. AutoFD handles tensor valued input and outputs automatically.
2. With AutoFD we do not need to write our code in the canonical form of euler-lagrange. We can compose functions more flexibly and use jax.grad for the functional gradient.
3. While euler-lagrange handles semi-local functional introduced in section 3. AutoFD handles non-local functionals via the integrate operator.

Q10, taking random grids in every iteration is often done in learning integrals

Taking random points will no doubt reduce the variance of both gradient estimators and make them converge to the same mean. We would like to emphasize that we are not developing a new method or trying to beat any baseline here. The purpose of this example to demonstrate that AutoFD enables us to use a very different gradient estimator via functional derivative.

We've also updated the figure in log scale.

Q11, confusions on DFT code

what Figure 4 (now Figure 5) means: Figure 4 shows the API of a library called libxc which provides intermediate values for the functional derivative. Traditionally, the functional derivative is achieved by composing the terms listed in Figure 4 using generalized euler-lagrange. Specifically, Vxc, Fxc, Kxc and Lxc corresponds to 1-4th order functional derivatives, right hand side of the colon shows the terms needed for each derivative. What we would like to convey here is that manual construction of functional derivatives of various order is tedious and laborious, while with AutoFD we can simple apply jax.grad on functionals.

Numerical results on DFT: This is a ongoing work, it is not available for the time being.

What is SCF: it is called the self consistent field, we updated it in the draft.

Q12, could you add more details of the experimental settings for reproducibility?

Yes we included the experimental settings for brachistochrone in Appendix G. The code is still being reviewed and will be directly open sourced in the near future.

Q13, PDE and FDE

Whether this work will benefit PDE and FDE would depend on whether an explicit function gradient is needed. I believe it is case by case. One potential use may be to regularize the functional gradient, like presented in this work: Neural Integral Functionals

Thanks for your time reviewing our draft, and the in depth questions on various technical points.

Q1, infinite dimensional arrays

This is brought up by more than one reviewers, and we admit that the previous presentation has not give enough background on it. In the updated PDF, we removed all quotes to "infinite dimension", which we find doesn't hurt the presentation. We think it will avoid confusion for the future readers.

Infinite-dimension is not implemented explicitly, we only extend the abstract shape of array to be infinite-dimensional. It is described in 4.1. In the updated PDF, we will stick to generalized arrays. Nevertheless let's explain how infinite dimensional is introduced in the shape.

For a plain array x of shape [5], we can retrieve the value at u by x[u], where u $\in \\{0,1,2,3,4\\}$. We have a finite number of u, therefore x is finite dimensional.

We extend the shape system to describe shape of a function, e.g. [f[5]] is the shape of a function taking an input of a 5 dimensional float vector. Mathematically f[5] represents a 5 dimensional vector space. Similarly, we can retrieve the value at u by calling the function at u, i.e. x(u), where u$\in R^5$. Since $R^5$ is a continuous space, we have infinite numbers of u.

Q2, could you add some reference papers about these three approaches for readers

Thank you for pointing it out, we have added references in the updated PDF. We removed the (3) because it doesn't explicitly calculate functional gradient.

Q3, does your proposed program support complex numbers

For complex valued function, definitely yes, we are already using it for density functional theory, which deals with wave functions.

For complex functions, it supports as good as JAX does. As many of the operators, linearize, nabla, linear_transpose are already implemented in JAX, we only provide the rules to perform functional differentiation for these operators. So they rely on JAX's capability to support complex functions. compose simply does lambda x: f(g(x)) so it trivially supports complex functions. As for integrate, I'm not very familiar with numerical integration for complex functions, but as long as it converts to summation over grids, I suppose there's no difference from real functions.

Q4, more definition

We have refactored Section 4.2 to connect the math and programming notations of JVP, VJP and transpose, including the Fréchet notation and our own Fréchet like notation for the linear transpose. We hope it is better this time!

Q5, does lack of inverse restrict the operations of the proposed AutoFD

It is technically feasible to support inversion in JAX, we can add inversion rules for each primitive w.r.t each input. However, this seems an independent project because function inversion has much wider use outside of functional differentiation, for example, used in normalizing flows.

One further question would be, what if the function is not invertibe? It then has to be worked out case by case. For example, piecewise invertible functions like abs can be separated into invertible pieces to work out the contribution from each piece. For more complex functions, there is no general way of deriving the transpose rule, in my opinion, it is more of a mathematical limitation than a programming one.

Q6, grid points

In the integration operator, we provide user the interface for supplying their own numerical integration grids. Our main contribution is on the automatic function differentiation software, there is no numerical/algorithmic contributions here.

There are various APIs for obtaining different grids in existing packages, which can be used with our integrate operator. e.g. legendre grid in numpy. However, the users need to case by case decide which grid to use depending on the functions they are integrating, because there is no universal method.

Q7, font

We fixed it in the updated PDF.

Q8, for most of the statements in Section 4.4, I would like to see quantitative results.

We include a case study in Appendix D where the redanduncy is exponential to the depth of function composition, with function call caching we achieve a linear time cost, while the naive implementation suffers a exponential cost.

Thank you for the reply, revision, and additional experiment. I also read the other reviews and understand more about the contributions of the paper. They addressed some of my concerns, and the clarity is improved; therefore, I changed Soundness (1 to 2), Presentation (1 to 2), and Rating (3 to 5) accordingly. Still, there is a critical problem on clarity. Let me make some more comments.

[Comment (major)]

• Some mathematical terms, such as "cotangent" and "Frechet derivative", may still confuse some readers. They have more rigorous, high-level definitions in math and the terms used in the paper are special cases of them. The abstract, intuitive definitions written in the paper may help readers understand the meaning, but may still keep readers questioning what they are anyway.

[Comment (major)]

• Please add more details of the experiment; e.g., the number of steps in training, optimizer, learning rate, weight decay, and other hyperparameters, if any.

[Comment]

• Please add the label of the horizontal axis of Figure 2.

[Comment]

• Applications in Section 5 are limited to integral functionals. Adding an example of non-integral functionals would be helpful.

[Additional question]

• Do you have an idea to deal with the functionals whose functional derivative includes the Dirac delta function, which I understand is difficult in general?

[Typo]

• In Appendix E, "... that is not dependent on x, therefore ..." should be "... that is not dependent on x; therefore ..."

Thank you for the changes to the paper.

rigorous, high-level definitions

I understood your situation. Then, I think adding reference papers/texts to immediately after the math terms is a simple way to go. How about it?

non-integral functionals

I literally meant non-integral functionals, e.g., $F([f]) = \frac{\int dt \sqrt{1 + {f^\prime}^2(t)} }{\int dt \sqrt{1 + f(t)} }$ (even a toy problem is OK), but anyway, the extra experiment in Appendix H is also a nice, practical example, and it made the paper more convincing and helpful for users.

Thanks for your time reviewing our draft and the in depth comments on various technical points.

Weakness 0, hard to follow mathematics of functional analysis, Fréchet derivatives

We refactored Section 4.2 to connect the math and programming notations of JVP, VJP and transpose, including the Fréchet notation and our own Fréchet like notation for the linear transpose. We hope it is better this time!

Weakness 1, Confusion on infinite dimensional array

This is brought up by more than one reviewers, and we admit that the previous presentation has not give enough background on it. You're right that the infinite-dimensional array is only abstract. In the updated PDF, we removed all quotes to "infinite dimension", which we find doesn't hurt the presentation. We think it will avoid confusion for the future readers.

Nevertheless, let's also explain why the generalized array described in 4.1 is infinite-dimensional. For a plain array x of shape [5], we can retrieve the value at u by x[u], where u $\in \\{0,1,2,3,4\\}$. We have a finite number of u, therefore x is finite dimensional.

We extend the shape system to describe shape of a function, e.g. [f[5]] is the shape of a function taking an input of a 5 dimensional float vector. Mathematically f[5] represents a 5 dimensional vector space. Similarly, we can retrieve the value at u by calling the function at u, i.e. x(u), where u$\in R^5$. Since $R^5$ is a continuous space, we have infinite numbers of u.

Weakness 2, introduce Radul et al

Thanks a lot for this suggestion, In our refactored Section 4.2 we show how JVP can be converted to VJP via a linear transpose.

Weakness 3, how the math map to core programming language

We provide a minimal but self contained code snippet for the nabla operator on scalar functions in Appendix A, in the hope to provide a better overall understanding of the system.

As for the comments on the point free style, I'm not sure I fully understand, I guess

def nabla(f):
  return jax.grad(f)

is somehow point free because we don't pass any argument to f? However, when we apply nabla on a function, we pass that function to nabla, thus it is not point free. Therefore, the operators are implemented in the non-point free plain python code, but the functions passed into the operator are used in a point free style.

how a function is defined: We registered a custom function to abstract value converter at jax.core.pytype_aval_mappings[types.FunctionType], which enables JAX to convert a function into a abstract array that we defined during tracing. This is added to section 4.1

Weakness 4, invertibility seems a severe limitation

It is technically feasible to support inversion in JAX, we can add inversion rules for each primitive w.r.t each input. To invert a function, we first trace the function into JAXPR and transform the JAXPR graph by applying the inverting rule on each primitive. However, this seems an independent project because function inversion has much wider use outside of functional differentiation, for example, used in normalizing flows.

One further question would be, what if the function is not invertibe? It then has to be worked out case by case. For example, piecewise invertible functions like abs can be separated into invertible pieces to work out the contribution from each piece. For more complex functions, there is no general way of deriving the transpose rule, in my opinion, it is more of a mathematical limitation than a programming one.

Weakness 5, are there problems not solvable with ordinary JAX

Yes, this work is originally motivated by density functional theory (DFT) in quantum chemistry. In DFT, it often happens that we need the functional derivative explicitly, for example, the loss function could be in the form of $L(\theta) = \int \frac{\delta F}{\delta f}(x) g^\theta(x)dx$. Given the functional $F$, it is not very trivial how to write this loss directly in JAX. It is usually hand derived and implemented, like how bp is done without automatic differentiation in the old days. By AutoFD, we would like to offer the "What You See is What You Get" (WYSIWYG) style of programming for problems where functions are first class citizens, e.g. quantum physics deals with wave functions.

Weakness 6, Tradeoffs between two differentiation methods

We bring up the different methods to demonstrate that functional derivatives could provide a different angle when we solve the same problem. In the brachistochrone problem it happens to be better via functional derivative, hypothetically, it brings us closer to the minimal in the space of function than in the space of parameters. However, further theoretical investigation is needed to reach a conclusion whether it is general to all problems. My feeling is that most likely it is case by case depending on the variance of the estimator, therefore we didn’t expand further on this topic.

Thanks for the further comments. I think we're aligned on the form of "point free" programming. On the comment about the connections to Elliott's work, we indeed drew a connection in the related work section,

"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 a compositional manner when the code contains high-order functional transformations, which is related to but distinct from our work."

In short, we believe they are related but distinct topics, let me explain why using the composition operator.

As you mentioned we can describe $f$ in lambda calculus f = lambda x: ..., where x is a point in the domain of f. In the same vein, we can describe the $\circ$ as $\circ=\text{lambda}\\;f,g: f\circ g$, where $f$ is a point in the domain of $\circ$. In the paper, we call $\circ$ an operator, or higher order function. Here let's be more specific and call $f$ a first order function and $\circ$ a second order function, where the inputs to the second order function are first order functions. I believe in the context of your discussion, point free is relative to the first order function, because $x$ never appears, but it is not point free to the second order $\circ$ because $f$ and $g$ are its points.

Now let's expand on the distinction from Elliott's work, let's use $D(f)(x)$ to denote differentiating $f$ at the point $x$. In short, Elliott's studies $D(f\circ g)(x)$ while ours studies $D(\circ)(f)$.

• Elliott's work studies how to perform AD for first order function more efficiently when your program has second order functions. Concretely in the example of composition, how to calculate $D(f \circ g)(x)$. As you mentioned, although we can naively trace $f\circ g$, it would be more efficient if we know the algebraic relation $D(f\circ g)(x) = D(f)(g(x))\circ D(g)(x)$ (Theorem 1 in Elliott's).

• Our work studies how to perform AD for second order function in the sense of variational calculus. Again in the composition example, we study how $D(\circ)(f)$ and $D(\circ)(g)$ can be computed automatically. Let's call it Automatic Functional Differentiation (AFD) instead of AD.

To this point, I hope it is clear how AFD can be implemented using the same mechanism in JAX, as one can find the resemblance between $D(\circ)(f)$ and $D(f)(x)$. More specifically, in python AFD does the following

def F(f, g):
  return integrate(compose(f, g))
dFdf = jax.grad(F, argnums=0)(f, g)

Although the line integrate(compose(f, g)) is point free in your definition because there is no x fed to f or g, it is not point free for the second order function F. Appendix A provides more example on how we implement an operator in python, again the code is point free in first order, but not point free in second order.

Completeness of primitives

We explained what we meant by completeness in 4.3

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.

Namely, because the JVP and transpose rules of operator A could depend operator B, we checked and implemented a closed set such that any operator in the set only depend on the other operators in this set. We didn't mean that the set of operators we introduced in the paper is complete to implement any functionals. For example, the functional $f\mapsto\max_x f(x)$ would be nontrivial to implement and differentiate, and can't be expressed with the operator primitives we introduced.

I tried to find duplication from Elliott's work, it seems to me that $dup: \lambda a\rightarrow(a,a)$ is a first order function instead of an operator (second order function).

Inversion of function

Inversion of a function is not necessary for AD (differentiating first order functions), however when we generalize AD to AFD (differentiating second order functions), the summations is generalized to integration. For example, the matrix vector product $Mv$ is generalized to $\int M(x,y)v(y) dy$.

In first order functions, the transposition of a linear function $x\mapsto Mx$ is $y\mapsto M^\top y$, we can see the relation that $\langle Mx, y\rangle=\langle x, M^\top y\rangle=\sum_i (Mx)_iy_i$. When generalized to operators, we use $\hat{O}^*$ to denote the transposition (adjoint) of an operator $\hat{O}$, it also needs to satisfy $\langle \hat{O}f, g\rangle=\langle f, \hat{O}^*g\rangle=\int (\hat{O}f)(x)g(x) dx$. Notice that the former involves a $\sum$ which is generalized to $\int$ in the latter. The function inverse comes out when we change the coordinate of integration, see Appendix C1 for more details.

These are both excellent explanations, thank you. Is there any way you can work them into the text of the paper?

In particular, other readers who are familiar with Elliott's work (as I am) might also become confused when they see AD over what /seem/ to be a very similar set of operators being defined in very different ways.

With respect to dup, it is indeed a first-order operator, but it's important because (1) it is not usually treated as an operator in graph-style AD, but (2) its derivative becomes (+) in reverse-mode, which is one of the main sources of complexity and symmetry-breaking in the graph-style AD algorithms, so its inclusion is absolutely required.

I am wondering if something similar might not also be true in your setup. It might also help with the common-subexpression-elimination problem that you discuss; AFAIK, the example you give could be easily resolved by using dup in an appropriate way.

Thank you for your time spent reviewing this draft, the positive acknowledgement of our contribution, and the suggestions. We address each point as below.

Weakness 1 & Q1, function inversion not implemented

It is technically feasible to support inversion in JAX, we can add inversion rules for each primitive w.r.t each input. To invert a function, we first trace the function into JAXPR and transform the JAXPR graph by applying the inverting rule on each primitive. This seems an independent project because function inversion has much wider use outside of functional differentiation, for example, used in normalizing flows.

Q2, how to supply numerical grid

To our knowledge, the Fourier, Legendre grids mentioned in Section 4.2.5 are two of the most used numerical grids. For example, they are used in the recent state space model for time series. They are exact when used with linear sums of planewaves / polynomials respectively. However, they are also used for other functions, which is equivalent to expanding the functions in planewaves/polynomials with a truncation error.

Q3, release of code

The source code will be open sourced under Apache license, it will be provided as a standalone python package. Documentation and tutorials will be provided too.

Thank you for your time spent reviewing this draft, the positive acknowledgement of our contribution, and the suggestions. We address the individually as follows

Weakness 1, introduce JAXPR instead of XLA:

We bring up XLA mainly in section 4.4 to mention that optimizations like common subexpression elimination (CSE) is only performed at the XLA level, and the python level lacks this mechanism. We think it is necessary for motivating function call caching. Although the primitives we implement has strong ties to JAXPR, the actual implementation (We added in Appendix A) does not require the knowledge of JAXPR. We think it would be too heavy to introduce to this level.

For the overall connection with JAX, we refined Section 4.2 our updated PDF, where we first introduced the mechanism of JAX and how we extend to higher order functions.

Weakness 2, provide code implementation

We’ll make the code open source in the near future. We also include a minimal but self contained example in Appendix A in our updated PDF to provide better clarity on how the math are mapped to code.

Weakness 3, FNO implementation with AutoFD

Thanks for bringing this up, along the way we actually thought of whether FNO is a good demonstration. It is true that we can implement FNO with a more succinct syntax, e.g. directly algebraic operations on functions. However, this syntactic simplicity does not relate directly to functional derivative. In the case of FNO, once we construct the forward computation, we then differentiate only with respect to the parameters, which deviates from our major contribution of differentiation with respect to functions. Therefore we focus on examples where functional derivatives makes a difference.

Q1, analysis of computational properties:

Yes, since our primitive operators are in python, they lack the equivalence of common subexpression elimination (CSE) in XLA, we therefore introduced function call caching for the remedy. It is discussed in section 4.4. In our updated PDF, we include a case study in Appendix D where the redanduncy is exponential to the depth of function composition, with function call caching we achieve a linear time cost, while the naive implementation suffers a exponential cost.

Q2, infinite-dimensional array:

The infinite-dimensional array is only abstract, based on the reviews, we decided to remove the quotes to "infinite dimension" in our rebuttal pdf, which doesn't hurt the presentation. We think it will avoid confusion for the future readers.

Nevertheless, let's also explain why the generalized array described in 4.1 is infinite-dimensional. For a plain array x of shape [5], we can retrieve the value at u by x[u], where u $\in \\{0,1,2,3,4\\}$. We have a finite number of u, therefore x is finite dimensional.

We extend the shape system to describe shape of a function, e.g. [f[5]] is the shape of a function taking an input of a 5 dimensional float vector. Mathematically f[5] represents a 5 dimensional vector space. Similarly, we can retrieve the value at u by calling the function at u, i.e. x(u), where u$\in R^5$. Since $R^5$ is a continuous space, we have infinite numbers of u.

Q3, why pure python instead of JAXPR

We need to clarify that we do capture higher order functions as JAXPR, but unlike ordinary JAXPR whose primitives consumes arrays (and thus can be lowered and compiled), in our JAXPR the primitives are operators that consumes python functions, therefore the operators themselves have to be python functions. It would be easier to understand by the minimal code example in Appendix A.

Q4, tracing with perfetto

The profiling tools and visualization tools like perfetto is more suitable for analysing the compiled program. As mentioned above, we have a case study that measures the efficiency in Appendix D.

Q5, inversion operator in JAX

It is technically feasible to support inversion in JAX, we can add inversion rules for each primitive w.r.t each input. We first trace the function into JAXPR and transform the JAXPR graph by applying the inverting rule on each primitive. This seems an independent project because function inversion has much wider use outside of functional differentiation, for example, used in normalizing flows.

In the new revision PDF, we updated our experiment section with an additional section that tests AutoFD on an MLP-like neural operator. For simplicity we choose a uniform grid for integration but does not exactly follow the details of FNO. We give the operator a target function to fit, and use the $L_2$ distance between the prediction and the target as the functional loss.

As discussed in the weakness 3, usually neural operator is trained by gradient descent directly in parameter space. In our experiment, we perform gradient descent in the function space, by subtracting the kernel and bias functions with their functional gradient. For example, we start with a very simple kernel function $k(x,y)$, then we perform $k(x,y)\leftarrow k(x,y)-\eta\frac{\delta L}{\delta k}(x,y)$, now $k(x,y)$ gets more complex, causing the next $\frac{\delta L}{\delta k}(x,y)$ to be even more complex. In Figure 9 we show how the JAXPR graph of $k(x,y)$ grow over time. Due to this compounding effect, we can only perform 4 step of gradient descent. However, in these 4 steps, we can see the loss are decreasing, and the prediction gets closer to the target function.

This new experiment is mainly design to show that AutoFD handles nonlocal functionals very well, and can correctly deal with procedures like above that leads to highly complex functions. By no means do we suggest training neural operator in this way, as discussed in weakness 3, FNO may not be the most suitable problem for AutoFD.

In the new revision, we include experiment on neural operator which is a generalization of MLP to higher order. We have trained it using gradient descent directly in the function space, and show results in Appendix H.