Summary
The paper proposes a new method, STDE, to compute high-order differential operators of (high-dimensional) neural network-represented functions. Two key ingredients include a generalization of Taylor-mode AD and randomness in the algorithm. This work shows impressive performance of STDE, in terms of both speed and memory.
Strengths
- This work addresses an important and challenging task in physics-informed machine learning.
- The illustration and analysis of previous AD algorithms and STDE are clear and comprehensive.
- STDE displays noticeable potential to improve the computation of arbitrary differential operators in both memory consumption and efficiency.
- The work provides new insight into introducing randomness other than SDGD or HTE whose basic ideas are employing Monte Carlo over dimensionality.
Weaknesses
1. Compared with original back-propagation or forward-Laplacian(FL), STDE is a random algorithm that does not give **precise** derivatives. In precision-demanding tasks, e.g. solving many-electron Schrödinger equations, the inaccurate value of differential operators will possibly result in an unreliable solution. To make the result more convincing, a cost-accuracy curve (STDE with different number of random samplings, i.e. $|J|$ in Eqn.(17)) should be added.
Furthermore, in light of this, the comparison between STDE and FL in super high dimensions (in table 1,2) is unfair. A possible remedy is to compare STDE with modified FL, which employs randomness to be more scalable, for instance, Monte Carlo over dimensions. The author can refer to [1] for more technical and implementation details of FL.
2. All the test cases of high-dimensional high-order operators are Laplacians (including $\Delta^2$). The high-order diagonal differential operator is mentioned in L216 but does not appear in the experiments. A comparison between STDE and direct back-propagation is expected.
The applications are all coming from the small area of PINN. To demonstrate the broad impact of this work, the author should target more tasks of general interest, e.g. second order derivative of NN parameter, which plays an important role in accelerating optimizations.
3. Some of the highly related existing works are missing. E.g.,
- Regarding L63, the forward rule for general different operators is done in [1].
- Regarding the geometric interpretation of Taylor-mode AD in sec 3.4, although not explicitly mentioned in Ref. [6] in the paper, is discussed in section 2 of [2]. The author should cite [6] or the first paper presenting similar ideas, and clarify their novelty if they still hope to claim generalizing univariate Taylor mode AD as their contribution.
4. As is discussed in Section 4.4, STDE can not be applied to *arbitrary* high-order operators. Obviously, this restricts the potential application scenarios of STDE.
5. The author didn’t discuss how large $l$ would be for large $k$ (notations from L165), which is necessary to assess the effectiveness of STDE for general high-order derivatives.
6. Some parts of the writing are confusing. See **Questions**.
Also, there are several misuses of terminologies. E.g.,
- Fréchet derivative, not Frechet.
- The form of second order PDE in Eq.(18) is not ‘general’. For instance, you are assuming ellipticity in the second-order term. Please refer to Eq.(10) or Eq.(1) in [1] for *general* second-order operator.
- In Appendix I.1 and the main text, ‘inseparable and effectively high-dimensional’, I guess what the authors intend to say is ‘elliptic’. The linear part of all the test equations is merely a Laplacian, casting doubt on the effectiveness of STDE for general elliptic equations.
- In Appendix I.2 and the main text, these PDEs are not nonlinear Fokker-Planck Equations, I think the correct name would be ‘semilinear parabolic equation’.
[1] Li R, Wang C, Ye H, He D, Wang L. DOF: Accelerating High-order Differential Operators with Forward Propagation. ICLR 2024 Workshop on AI4Differential Equations In Science
[2]Tan, Songchen. Higher-Order Automatic Differentiation and Its Applications. Diss. Massachusetts Institute of Technology, 2023.
Questions
1. Which function does ‘univariate’ refer to in Sec 3.4, $g$ or $F$? What is the first contribution (in the introduction), generalizing Taylor mode AD to the case where $g$ is multivariate, or $F$? I am skeptical of the authors’ claiming it as a contribution/novelty because it seems like this is merely about replacing scalar Faà di Bruno's formula with that of tensors.
2. In Appendix D, where $g$ defines an $n$-dimensional manifold, is $n$ always set to 1 in this work? If not, please point me to where the method with general $n$ is discussed.
3. Although the authors claim to remove the exponential complexity w.r.t. the order of derivatives, for general operators the size of coefficient tensor **${C}$** in Eqn(10) grows exponentially and the sampling procedure in Eqn(15) will become hard.
4. Could the authors think of any real application scenarios with high-order high-dim operators? As for the case of low-dim high-order operators, the authors need to check if the differential operator takes up a considerable proportion of the whole process, e.g. training PINN. If not, the benefit of accelerating differential operators might be very limited.
5. Does STDE have the potential to make use of sparsity in the neural network function, as shown in FL and [1]? For the specific case of second-order operators, how would STDE perform when $C$ is dense but low-rank?
6. Instead of tackling a scalar function (e.g. Laplacian of a function), could STDE help with the computation of tensor-value functions, e.g. Hessian or curvature? Hessian appears in Monge-Ampère equation and Ricci curvature appears in Einstein equation.
On writing and notations:
1. Can you clarify on ‘$[2]$’ in L214 with more details or concrete examples? For me, $d^2u_\theta$ in Eqn.(16) is a (0,2)-tensor (Hessian), but why does it take three input $(a,e_j,0)$? Please also clarify on ‘$[3]$’ in Eqn(44).
2. In Eqn.(19), it should be $\sum_{i=1}^d$.
3. In L247, do you mean $v\sim p$,, instead of $(a,v,0)\sim p$?
4. In the ‘100D’ column of table 2, 539 is not the smallest number. Why is it in bold?
5. There needs to be more description in section F.1.
It is good to remind the readers that Einstein notation is employed in Eq.(45) and that the indices $n,\ n’,\ n’’$ are omitted.
The author should inform the readers that in the discussion $n=1$ for $n$ in Eq(40) and notations like $\partial^4u(x)(e_i,e_j,0,0)$ refers to setting $v^{(1)}=e_i,\ v^{(2)}=e_j$ instead of the taking $e_i,e_j,0,0$ as inputs for a (0-4) tensor.
Overall, this paper has the potential to be a good work. Hope the concerns above can be well addressed and then I will upgrade my assessment accordingly.