Accelerating PDE-Constrained Optimization by the Derivative of Neural Operators

Thank you for your review with carefulness and details, and it is highly valuable and helpful to us!

1. Claims and Evidence

Nice catch on the mismatch of objective! As you pointed out, the objective of Microreactor is to minimize $J = -\int…$ in eq. (17) in appendix C.2. This objective is equivalent to maximize $J=\int…$ without the negative sign. We will revise the form of objectives to be consistent. In our implementation, we used maximization, please see Line 726 in utils.py. The blue line in Fig. 6 is obtained for illustration purpose. R-VF w/o ref. indeed runs without being given ground truth solution along optimization process. We check the quality of its output in the same way as R-VF with ref. The only difference is to drop the reference GT solution from input.

2. Methods And Evaluation Criteria

We respectfully disagree with the evaluation of marginal improvement. First, let us clarify that in Table 1, errors are shown by each component, and $L_{\cdot}$’s are sensitivity errors. We will modify the caption of Table 1 accordingly. We should emphasize that our main improvement is on derivative learning, which achieves error reduction 3.2%, (1.8%, 2%), (1.2%, 3.5%), 6.1% (or relatively 15.8%, (3.9%, 7.2%), (3.8%, 6.85%), 10.8%) from the second best. Second, the baseline models GNOT and Transolver are SOTA neural operators (NO), which leaves small room for improvement at operator learning (on physical fields). However, our model achieves Pareto front in the tradeoff between operator learning and derivative learning (See Fig. 5 and Section 4.1.1). Please see also Item 1 for Rebuttal for Reviewer 1 for the reason of the improvement of VF.

We appreciate the rigorousness of the reviewer and decide to supplement some experiment to address the question on the robustness of the improvement. Due to space limit, please check the table in Item 2 of the rebuttal for Reviewer 2.

$u_{gt}$ indicates the numerical solution corresponding to the current $\lambda$, NOT the numerical solution of PDECO.

It is a good question on how to make predicted $u$ and $J$ (blue arrow) exact, which is the key to our method. Given a GT solution ($\lambda_r, u_r$) obtained from numerical solver, let $\lambda_q =\lambda_r$ and thus $\varphi=0$. The input would be $(\lambda_q, u_r, 0)$. The output of reference neural operators (RNO) should be close to the exact reference solution $u_r$ due to our training algorithm. For each query, dataloader would randomly pick neighboring samples from the same optimization trajectory as references, including the queried data itself (See App. C.1). It forces RNO to learn an identity mapping when the deformation $\varphi$ vanishes.

3. Reference not discussed.

On geometry handling, [1] belongs to the framework of traditional NO, i.e., learning the mapping between geometry and solution. RNO is a generalization of NO to learn the change of solution according to deformation of geometry (and more general parameters, see Remark 3.2).

[2] is a relevant work, and we thank the reviewer for pointing this out. It has been cited in line 160-161 right. [2] is a general work on training NO in Sobolev space efficiently by exploiting the intrinsic low-dimensionality of the derivatives when handling the high-dimensional Jacobian. Our work learns the directional derivative of NO since the derivative is taken wrt some objective functions, which can be view as a special case of [2].

4. Other questions

We apologize for missing the information of dataset sizes. Please see Item 1 in Rebuttal 2 for details due to space limitation.

Please see Item 3 in Rebuttal 2 for discussion on wall-clock time for optimization.

It is a great question on if derivative learning would make NO less general. On the problem Inductor2D, we consider a flexible objective with a changing weight $\gamma$ between two terms (See eq. (19)). The dataset is drawn with random sampling of $\gamma$. Therefore, the NOs are able to predict different objectives for different $\gamma$. In principle, the dataset determines the generalization of NO, including objective functions.

Regarding RNO, $u$ and $\lambda$ are nodal functions on mesh. It is indeed possible to have $\varphi$ satisfying the relation between reference and query. To give a simple example, let $\varphi$ be a spatial rotation, and then the nodal value of $u$ remains unchanged, while coordinates $\lambda$ will be rotated accordingly. In the domain of shape optimization, the existence of $\varphi$ is the cornerstone for all applications. See Remark 3.2 for more details.

The training pipeline of RNO is summarized in Algorithm 1 in Appendix D. To represent and construct $\varphi$ would be rather simple with the data from optimization. One can use the shift of mesh nodes between ref. and query for shape optimization, or use the difference between control variable $\lambda$ of ref. and query for topology optimization. See line 256-265 right in the paper.

Thank you for reviewing our paper, your feedback is in-detail and comprehensive. Among your comments we notice that you particularly expressed “having a hard time following virtual-Fourier (VF) layers and distinguishing VF from FNO”. We would like to clarify this point first since it lies in the core of our motivation of this work.

FNO and its variants require a uniform-grid domain due to the need of applying fast Fourier transform. geo-FNO uses a 1-to-1 mapping to deform a uniform grid to an arbitrary shape, and the computation essentially still happens on uniform grids. To simply put, the input shape (H, W, C) must be fixed. VF on the other hand can be applied to arbitrary number of mesh grids (T, C), where the length of input T can change. This is similar to transformer-based neural operators and preferable since we are dealing with data from shape/topology optimization, where mesh grids can be arbitrary and uniform grids are inapplicable.

The way of handling arbitrary length input with VF layer consists of, 1. Project the input onto a fixed number of virtual sensors; 2. Process the virtual signals by Fourier layers; 3. Project the processed signals back to physical space (See Figure 3). The projection technique is entirely independent from FNO and inspired by Transolver instead. Additionally, our projection is different from Transolver due to our consideration of learning derivatives (See Remark 3.1).

We emphasize that, to our best knowledge, this is the first work that extends Fourier layers to arbitrary mesh grids. The benefit is its advantage in learning derivatives due to its simple structure compared to transformer counterparts (See items 2&3 below).

To address your other questions:

1. Sensitivity is defined on line 69 left by equation (2), $\frac{\partial \tilde{J}}{\partial \lambda} = \frac{\partial J}{\partial u} \frac{\partial u}{\partial \lambda} + \frac{\partial J}{\partial \lambda}$. We will bold the word sensitivity here.

2. We meant Fourier transform is a linear operator, which has a simple derivative, i.e., $\frac{d}{dx}\mathcal{F}(z)= \mathcal{F}(\frac{dz}{dx})$. This motivates us to integrate Fourier layer into our design to learn derivatives. We will modify this part to improve clarity.

3. Consider the derivative of the attention of a transformer, $\frac{d}{dx} softmax(QK)V$. The derivative must be taken wrt to Q, K, V, hence involving product rule. Also, the derivative of softmax is $s_i(\delta_{ij}-s_j)$ (See line 189-191). Thus the derivative of attention would be roughly and informally $s_i(\delta_{ij}-s_j)(Q’K+QK’)V + S(QK)V’$. Hence, we pointed out that the derivative of transformer is assigned with strong inductive bias (probably undesired). Imagine learning derivatives of operators with layers constructed as above. The expressiveness is therefore restricted.

On the other hand, the derivative of Fourier transform is as simple as $\frac{d}{dx}\mathcal{F}(z)= \mathcal{F}(\frac{dz}{dx})$, and hence has “no additional bias”. We conducted more analysis of the derivatives of VF in Remark 3.1 and Appendix A.

4. Indeed, VF is the building block of RNO. Namely, according to the general architecture of neural operators $\mathcal{G}_{\theta}:= \mathcal{Q}\circ\mathcal{L}\circ\cdots\circ\mathcal{L}\circ\mathcal{P}$ (On line 149 right, in the beginning of Section 3), VF instantiates the integral operator $\mathcal{L}$. To smoothly transit between sections, we will clarify the connection in the beginning of section 3.3.

5. Claims and Evidence

The reason why we did not compare with traditional sampling method is the following. For example, consider random sampling in parametric space of the drone problem (Fig. 10). The parameter space is $[0,1]^{30000}$, given 3E4 mesh nodes. The likelihood of sampling an optimal design is almost surely zero. Therefore, an optimal design would be an out-of-distribution (OoD) data for random sampling method. Since neural operators would fail on predicting OoD data, they are inapplicable for optimization tasks. Thus, in order to make data near optimality in-distribution, in our opinion, it is necessary to adopt optimization data sampling. See the discussion on line 97-109 left in our paper.

The data efficiency of our approach is exemplified in Table 1. Due to space limit, please see discussion in Item 1 of Rebuttal 2.

6. Methods And Evaluation Criteria

We apologize for missing the detail of data size in our submission. Please also see Item 1 of Rebuttal 2 for details. The dataset size we use is around 1E3, which is similar to the setup of many traditional neural operators. Therefore, the cost of dataset is NOT significantly more than traditional datasets.

Regarding runtime, please see discussion in Item 3 of Rebuttal 2.

Thank you for reviewing our paper, your summarization is precise and your suggestion on additional references are well received. Particularly, we thank you for suggesting the 2nd work that we'd like to discuss the distinction with our work. The first distinction is that their work is targeting time-dependent fluid problems. Our work is targeting static problems, e.g., steady flow, electromagnetics on frequency domain, solid mechanics. Second, their approach intends to replace numerical simulation with surrogate models. Our work is complementary to surrogate approach that hybrids with numerical solvers. It is an interesting question on how to apply our hybrid method to problems with rollout operations, which can be explored in the future.

Below are some frequently asked points we’d like to address:

1. Dataset sizes and Table 1

Microreactor, Fuelcell, Inductor, Drone have 100x12, 30x20, 100x10, 10x20 samples respectively as in num_trajectory x num_steps. We added the information in appendix.

In Table 1, all errors are listed by components, and $L_{\cdot}$’s are sensitivity errors (we will clarify this in caption). Naïve supervised learning paradigm for neural operators perform poorly due to overfitting highly correlated optimization data. The data from optimization steps are quite similar to each other. Our solution to learn by reference neural operators (RNO) is much more efficient in terms of low error rate (Comparing models with and without “R-”).

To see the improvement due to virtual-Fourier (VF), we observe that the sensitivity errors are reduced by 3.2%, (1.8%, 2%), (1.2%, 3.5%), 6.1% (or relatively 15.8%, (3.9%, 7.2%), (3.8%, 6.85%), 10.8%) from the second best. See also the balance of tradeoff between error and sensitivity error in section 4.1.1, VF achieves the best Pareto front.

2. Robustness of improvement of R-VF

The following results are obtained on 4 runs on Microreactor with different random seeds. R-VF consistently outperforms the baselines. Note that, LA: Linear Attention, PA: Physics Attention, and R-: reference neural operators.

Also, we doubled the size of Microreactor dataset to 200 trajectories (each of 12 steps and effectively 2400 samples). 40 traj. are kept as test set. We train models on 80, 160 traj. respectively. We observe that with larger dataset, R-VF kept advantage over baselines on predicting both physical fields and sensitivities. The main improvement is still on sensitivities.

3. Wall-clock time for optimization with R-VF

The runtime of our hybrid method consists of time on optimizing with NO and time on calling numerical solvers. The latter dominates the runtime, and optimization with NO is much cheaper. A potential cost saving of hybrid method lies in fewer calls of numerical solvers.

One crucial aspect of optimization with neural operators is the method of optimization. In our implementation we adopted GD and MMA, but there are many other choices such as Adam, L-BFGS, SIMP, etc. A good optimization method can reach higher objective values and save both the numbers of iterations and function calls. Thus, the optimization method significantly affects wall-clock time. Besides, there are some important techniques in optimization, e.g., gradient filtering, projection, clipping, etc. A comprehensive and fair comparison requires all these techniques. Therefore, in our opinion it is premature to compare the runtime in this work.

Thank you for reviewing our paper in such depth and being very informative, especially on reference. It is highly valuable to us!

1. Experimental Designs Or Analyses

We respectfully disagree that Virtual Fourier (VF) Layer is least important. It reflects our thoughts on designing architecture for learning derivatives. To this sense, it is as crucial as inference with reference in our framework. In Table 1, all errors are listed by components, and $L_{cdot}$’s are sensitivity errors (we will clarify this in caption). We observe that the sensitivity errors are reduced by 3.2%, (1.8%, 2%), (1.2%, 3.5%), 6.1% (or relatively 15.8%, (3.9%, 7.2%), (3.8%, 6.85%), 10.8%) from the second best. This improvement is due to the architecture of VF. See also the balance of tradeoff between error and sensitivity error in section 4.1.1, VF achieves the best Pareto front.

Specifically, the derivative of VF layer has a simpler form than baseline models such as transformers. Consider the derivative of the attention of a transformer, $\frac{d}{dx} softmax(QK)V$. The derivative must be taken wrt to Q, K, V, hence involving product rule. Also, the derivative of softmax is $s_i(\delta_{ij}-s_j)$ (See line 189-191). Thus the derivative of attention would be roughly and informally $s_i(\delta_{ij}-s_j)(Q’K+QK’)V + S(QK)V’$. Hence, we pointed out that the derivative of transformer is assigned with strong inductive bias (probably undesired). Imagine learning derivatives of operators with layers constructed as above. The expressiveness is therefore restricted.

On the other hand, the derivative of Fourier transform is as simple as $\frac{d}{dx}\mathcal{F}(z)= \mathcal{F}(\frac{dz}{dx})$, and hence has “no additional bias”. We conducted more analysis of the derivatives of VF at Remark 3.1 and Appendix A.

2. Questions For Authors

Following the analysis from last question, we hope that it is clearer to see the derivative of VF having more expressiveness compared to transformer counterparts. We will modify the paper to reduce the ambiguity on “complex derivatives” and “additional bias”.

In our paper, there is another type of example of $\varphi$ besides topology optimization, namely shape optimization. Let $\varphi$ be a deformation between two shapes, which is instantiated as mesh grids warping and finally discretized as mesh grids shift. See Fuelcell and Inductor examples in our experiment. Thus, in this paper, for both topology and shape optimization, the discrete $\varphi$ can be represented as $\lambda_r-\lambda_q$.

It is definitely an interesting question to consider other form of $\varphi$. In fact, the deformation $\varphi$ is not unique. For example, given two shapes of the same topology, there are infinite many ways of deforming one to another. How does $\varphi$ affect the operator learning? Is it possible to take integral of RNO along a path of warping a shape to a much more different shape beyond perturbing? We believe there is a vast field to explore in this area.

3. Weakness

We appreciate your precise judgement on recognizing the reliance of our method on classic solver and the nature of data as optimization trajectories. However, we’d like to discuss with you as these topics do not present as weakness in our opinion.

Hybrid methods have great potential in scientific computing area. One of the main reasons is that modern computing algorithms are developed with profound theories and abundant ecosystems, e.g., open/closed source software. The guaranteed accuracy and reliability are still indispensable in the era of data-driven approach. We are glad you mentioned the work and effort of surrogate models. Advocates embrace them for the dominating advantage in cost. We absolutely agree and would love to contribute to this community. However, we hold a slightly different view that AI will not replace numerical computing but rather play as a complementary part, which will be different from what happens to CV or NLP. It is because of this reason hybrid method is getting popular in AI4PDE, e.g., Multigrid method and other hybrid methods mentioned in our paper.

The cost of optimization dataset is NOT larger than traditional neural operators. We apologize for missing this information of dataset sizes and added to appendix. Please see item 1 in Rebuttal 2 for details due to space limit. Also, in our opinion the derivative information would be necessary if we tend to obtain surrogate models with good gradient property.

4. Missing metric of wall-clock time

Please see Item 3 in Rebuttal 2 due to space limit.