Paper: Automatic Functional Differentiation in JAX (ICLR-2024)

Reviewer comment (Reproducibility):
In Abstract (and Introduction), the authors state: "By representing functions as infinite dimensional generalization of arrays, we seamlessly use JAX’s existing primitive system to implement higher-order functions." How was it realized in the proposed method? Could you point the part of the manuscript that concretely explains it?

Author response:
**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`.

Reviewer comment (Writing):
In Introduction, the authors state: "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." Could you add some reference papers about these three approaches for readers who are not that familiar with functional differentiation?

Author response:
**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.

Reviewer comment (Writing):
In Section 3, the authors state: "The Schwartz kernel theorem states that any linear operators can be expressed in this form." Could you add a reference here? It is more reader-friendly.

Author response:
**reference to the formal definitions**

Good point, we updated accordingly in Section 3.2.

Reviewer comment (Reproducibility):
In Section 4.2, the authors state: "The primitive operators considered in this work are focused to realize the most used types of operators and functionals described in Section 3." Does your proposed program support complex numbers?

Author response:
**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](https://jax.readthedocs.io/en/latest/notebooks/autodiff_cookbook.html#complex-numbers-and-differentiation
) 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.

Reviewer comment (Writing):
In Section 4, it would improve the clarity of the 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 paper, because not all readers are familiar with both functional analysis and computer programming. Changing the order of explanations can also be an option (e.g., the beginning of Section 4.2.3 can be an intuitive explanation of the Frechet derivative as a generalized directional derivative).

Author response:
**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!

Reviewer comment (Reproducibility):
In Section 4.2.1, the authors state: "The function inverse on the right-hand side of the $T_f (\hat{C})$ rule is not implemented..." Does it restrict the operations of the proposed AutoFD? Could you take some examples?

Author response:
**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.

Reviewer comment (Presentation):
In Section 4.3 and Eq. (7--8), the authors state: "We denote them as undefined because..." Changing the font of "undefined" would be good, e.g., mathtt.

Author response:
**Q7, font**

We fixed it in the updated PDF.

Reviewer comment (Evaluation):
For most of the statements in Section 4.4, I would like to see quantitative results.

Author response:
**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.

Reviewer comment (Novelty):
In Section 5.1, the authors performed an experiment to solve the brachistochrone problem. The authors simplify the problem to a parametric fitting of $y^{\theta}(x)$. This is how conventional methods in the numerical analysis of functionals work, as stated in Sections 1 and 2. What does the proposed program enable us to do, or what is the difference from the conventional methods?

Author response:
**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.

Reviewer comment (Theory):
In Section 5.1, 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.

Author response:
**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.

Reviewer comment (Experiments):
In Section 5.1, it is noted that directly minimizing in the parameter space is limited by numerical integration, and it easily overfits to the numerical grid. Taking random grids in every iteration is often done in learning integrals, which would fill the gap between the red and the other curves. I would like to see the performance of such a more rational baseline.

Author response:
**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.

Reviewer comment (Presentation):
In Figure 3 and 4, 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?

Author response:
**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.

Reviewer comment (Experiments):
I would like to see numerical results about DFT using AutoFD.

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

Reviewer comment (Writing):
In the SCF loop of the DFT calculation, what is SCF?

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

Reviewer comment (Reproducibility):
Could you add more details of the experimental settings for reproducibility?

Author response:
**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.

Reviewer comment (Reproducibility):
Could you show the code for reproducibility?

Author response:
**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.

Reviewer comment (Experiments):
Can we use the proposed AutoFD to solve a functional differential equation that includes functional derivatives (the task is to get the functional that satisfies the given equation that includes functional derivatives)?

Author response:
**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](https://openreview.net/forum?id=aQuqw6eVKP)

Reviewer comment (Experiments):
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.

Author response:
**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.

Reviewer comment (Presentation):
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. It is suggested to add a diagram of the pieces of JAX's architecture relied upon, and remove mentions of XLA-mapping of ops etc. from the latter parts of the paper.

Author response:
**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.

Reviewer comment (Reproducibility):
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.

Author response:
**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.

Reviewer comment (Experiments):
While neural operators, and most specifically Fourier Neural Operators, are presented as a clear motivation for the work, they are 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 be much more succinct with `autofd` while matching the performance of the JAX-native implementation.

Author response:
**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.

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.

Reviewer comment (Evaluation):
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?

Author response:
**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.

Reviewer comment (Theory):
Is there an implicit trade-off computationally or conceptually to the chosen representation as an infinite-dimensional array?

Author response:
**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`.

Reviewer comment (Reproducibility):
Why did you choose a purely Python-based implementation of AutoFD, as compared to a version of AutoFD acting directly on the JAXPR?

Author response:
**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.

Reviewer comment (Evaluation):
For the _Efficiency of Implementation_, have you considered tracing your implementation with perfetto?

Author response:
**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.

Reviewer comment (Reproducibility):
How do you anticipate an inversion operator in JAX to be feasible? Would it be possible to implement such an operator if you were writing the operators at the level of the JAXPR?

Author response:
**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.