Accelerating PDE-Constrained Optimization by the Derivative of Neural Operators
Authors: Ze Cheng
Zhuoyu Li
Wang Xiaoqiang
Jianing Huang
Zhizhou Zhang
Zhongkai Hao
Hang Su
Venue: ICML 2025 poster (2025)
Keywords: neural operators
gradient-based optimization
PDE-constrained optimization
Primary area: deep_learning->algorithms
TL;DR: A framework of applying neural operators for PDE-constrained optimization

Abstract: PDE-Constrained Optimization (PDECO) problems can be accelerated significantly by employing gradient-based methods with surrogate models like neural operators compared to traditional numerical solvers. 
However, this approach faces two key challenges:
(1) **Data inefficiency**: Lack of efficient data sampling and effective training for neural operators, particularly for optimization purpose.
(2) **Instability**: High risk of optimization derailment due to inaccurate neural operator predictions and gradients.
To address these challenges, we propose a novel framework: (1) **Optimization-oriented training**: we leverage data from full steps of traditional optimization algorithms and employ a specialized training method for neural operators. (2) **Enhanced derivative learning**: We introduce a **Virtual-Fourier** layer to enhance derivative learning within the neural operator, a crucial aspect for gradient-based optimization. (3) **Hybrid optimization**: We implement a hybrid approach that integrates neural operators with numerical solvers, providing robust regularization for the optimization process.
Our extensive experimental results demonstrate the effectiveness of our model in accurately learning operators and their derivatives. Furthermore, our hybrid optimization approach exhibits robust convergence.

Accelerating PDE-Constrained Optimization by the Derivative of Neural Operators
Ze Cheng' Zhuoyu Li
Xiaoqiang Wang'
Jianing Huang' Zhizhou Zhang'
Zhongkai Hao
Hang Su"
Abstract
Predicted trajectory
PDE:Consttraind Optimization (PDECO) prob- lems can be accelerated significantly by employ- ing gradient-bard baseed methods with surrogate mod- els like neural operators compared to traditional numerical solvers. However, this approach faces two key challenges: (1) Data inefficiency: Lack of efficient data sampling and effective train- ing for neural operators, particularly for opti- mization purpose. (2) Instability: High risk of optimization derailment due to inaccurate neu- ral operator predictions and gradients. To ad- dress these challenges, propose novel frame- work: (1) Optimization-oriente training: leverage data from full steps of traditional op- timization algorithms and employ specialized training method for neural operators. (2) En- hanced derivative learning: We introduce a Fetual-Fourier layer to enhance derivative learn- ing within the neural operator, crucial aspect for gradient-based optimization. (3) Hybrid opti- mization: We implement hybrid approach that integrates neural operators with numerical solvers, providing robust regularization for the optimiza- tion process. Our extensive experimental results demonstrate the effectiveness of our model in ac- curately learning operators and their derivatives. Furthermore, hybrid optimization approach exhibits robust convergence.
Ground truth trajectory
dkp
u
Ar
Figure 1: Illustration of Error Accumulation in Gradient- Based Optimization with Neural Operators. Blue Trajectory: Represents the optimization path (Ae, ût), where ün is pre- dicted by the neural operator. Red Trajectory: Represents the trajectory of the numerical solutions u, corresponding to X₁- While the optimization process proceeds along the blue trajectory guided by the neural operator, the red trajectory of numerical solution may deviate significantly, highlighting the potential for error accumulation and suboptimal conver- gence.
parameter space A,
min XEA
J(A) ((u,))=0 =0
(1)
1. Introduction
const.
In this paper, consider solving PDE-constrained opti- mization problems (PDECO) in high-dimension design
where == A) is the objective function, and C is a differential operator that describes a PDE with design parameter and solution u. For clarity, we will henceforth omit the tilde from J where the meaning unambiguous. While gradient-free optimization methods like Bayesian optimization (Snoek et al., 2012) and particle swarm op- timization (Kennedy & Eberhart, 1995) are effective for low-dimensional design spaces, gradient-based methods are generally more suitable for high-dimensional PDECO problems, for example, shape optimization (Sokolowski & Zolésio, 1992) and topology optimization (Bendsoe
Bosch (China) Invest Lid., Shanghai, China 'Dept. of Comp. Sci. & Techn., Institute for AI, BNRist Center, Tsinghua-Bosee Joint ML Center, Tsinghua University. Correspondence to: Ze Cheng <ze.cheng bo.bosch.com....
Proceedings of the 4rnd International Conference on Machine Learning, Vancouver, Canada. PMLR 267, 2025. Copyright 2025 by the author(s).
'Our code is available https//giiubb.co//rchnng ai/Opt RNO.

---

Accelcrating PDECO by the Derivative of Neural Operators
Noural Openatr
AU
Model
Openatr
Norrl
d
Figure 2: (Left) The overview of training neural operators with derivative learning, where curved arrow indicates differentia- tion by autodiff to obtain 0J/0x. (Right) The iterative process of gradient- optimization with neural operators of fixed parameters 0*
Sigmund, 2013). Some problems, such as those arising in density functional theory (DFT) (Kohn & Sham, 1965), may even involve combination of both approaches at different stages.
volving intricate geometries, boundary conditions, and ma- terial properties. Moreover, effectively sampling data near the optimal region, characterized by sophisticated patterns and complex structures, using random methods is highly improbable, akin to a "monkey typing Shakespeare." (Borel, 1913)
This paper focuses on solving high-dimensional PDECO using gradient-basee methods. These methods rely on the gradient of the objective function with respect to the design parameters, commonly referred to as sensitivity:
To address these limitations, we leverage the data generated during the optimization process itself. Optimization algo- rithms, such as the adjoint method, produce sequences of so- lutions and sensitivities as they converge towards the optimal region. These sequences, often discarded as "byproducts" in traditional optimization, form valuable trajectories in the design-solution space. However, directly training neural operators on these trajectories can be inefficient and prone to overfitting due to the high correlation between successive data points. To mitigate this, we propose training method that leverages the relative change between data points to achieve higher learning efficacy. This approach exploits the inherent structure of the optimization process, where successive iterations provide valuable information about the search direction and proximity to the optimal solution.
aJ JJ du OJ ax
Du ax
Ox
The design parameters are then iteratively updated according to A4+1 f A - "7 ax where 17 represents the learning rate. Among the various methods for gradient computation, the adjoint method (Pontryagin, 2018) is well-established and efficient approach.
aJ
Despite the efficiency of the adjoint method, optimization processes can be computationally expensive due to the high cost of solving large-scale PDEs. This often results in long convergence times, spanning hours or even days. To ad- dress this, various surrogate models (Forrester et al., 2008; Zhu & Zabaras, 2018; Fahl & Sachs, 2003; Audet et al., 2000) are employed to replace computationally expensive simulations, leading to significant speedups, often reducing computation times to seconds or milliseconds. While some surrogate models directly map the design space to objective values, their applicability limited by their dependence on the specific objective function. In contrast, surrogate mod- els capable of predicting underlying physical fields (e.g., displacement, stress, temperature) offer greater flexibility. These models can be used to compute any desired quantity of interest, including the objective function. Neural opera- tors (Kovachki et al., 2023; Lu et al., 2021; Hwang et al., 2022; Lu et al., 2019) exemplify this approach and show great promise for accelerating gradient-based optimization PDECO problems.
The second major difficulty derailing optimization. Sup- pose have trained neural operator
Go A U
A + u
(4)
where 0 is the learned parameter, A and U denote design space and solution space. The gradient (2) can then be cal- culated by aj RJ &Ca 8J ax (u ax ax Fig. 2 (right) illustrates ej how to compute using neural operators. This approach, however, introduces significant issues: (i) There is a dis- crepancy between Ge(A) and the true solution u, leading to differences between J(Go) and J()); (ii) similarly, there exists an error between 192 and the true derivative du, (iii) ax
through iterations of gradient-basee optimization, these er- rors accumulate, which often results in the optimization trajectory entering
However, training neural operators for PDECO presents significant challenges. Firstly, constructing a representative training dataset is crucial. Traditional sampling methods like Gaussian random fields (Neal, 2012) and uniform sam- pling often fail to capture the complex and diverse scenarios encountered in optimization problems, particularly those in-
(OoD) regions of the neural operator, causing unreliable predictions and erratic behavior. As illustrated in Fig. 1, the error of solution prediction grows out of control from initial point (Ap, un) to the final optimized point (Az, ur). Hence, maintaining in-distribution is critical for reliable optimization.
2

---

Accelcrating PDECO by the Derivative of Neural Operators
Thus, address the two fundamental questions, i.e., (i) how to effectively sample and utilize data near optimality, and (ii) how to control the growing error of neural operators during optimization, we proposed the following solutions:
explored, the impact of architecture on the smoothness of the learned operator, crucial aspect for gradient-based op- timization, remains relatively understudied. PINO (Li et al., 2024) investigates the derivatives of neural operators but focuses mainly on enforcing residual losses from governing equations, Our work delves deeper into how the architecture of neural operators influences their derivatives, essentially examining their smoothness properties.
1. We utilize the full data generated from optimization trajectories as the training set. In order to effectively achieve operator learning from such data, we adopt reference neural operators (RNO) (Cheng et al., 2024) to learn extended solution operators.
Sobolev training (Czarnecki et al., 2017; Tsay, 2021) en- hances the smoothness of neural networks, critical prop- erty for gradient-based optimization. Sensitivity supervision can be viewed as a specific instance of Sobolev training. In- stead of randomly sampling derivatives in all directions, sen- sitivity supervision focuses on derivatives with respect to the optimization variables (e.g.. design parameters). DeePMD (Zhang al., 2018) exemplifies this approach by impos- ing supervision on the forces acting on atoms, which are essentially the derivatives of the predicted energy with re- spect to the atomic positions. Our method follows a similar principle, utilizing ground truth sensitivities obtained from classical numerical methods to supervise the sensitivities predicted by the derivative of neural operators.
2. We analyze the effect of different neural operator archi- tectures on their derivative performance and introduce a Virtual-Fourier layer to enhance derivative learning.
3. We propose a hybrid inference method that combines neural operators with numerical solvers to suppress errors and regularize optimizing process.
2. Related works
Surrogate models have received significant attention in re- cent years, particularly within the realm of PDECO. Early ef- forts such as (Audet et al., 2000) employed Kriging models (Gaussian processes) coupled with gradient-free optimiza- tion techniques. (Fahl & Sachs, 2003) applied Reduced- Order Models (ROMs) leveraging proper orthogonal de- composition (POD) alongside sophisticated optimization strategies. More recently, operator learning has emerged as powerful paradigm, with (Hwang et al., 2022) outlining a general framework and formulating the optimization as a variational approximation. (Wang et al., 2021) explored the combination of DeepONet (Lu et al., 2019) with physics- informed loss on regular domains, while (Xue et al., 2020) demonstrated the effectiveness and versatility of amortized approaches for linear PDECO problems.
Hybrid methods combining neural networks with tradi- tional numerical methods have shown promise in solving PDE-related problems. Notably, Hsieh et al. (2019) have demonstrated convergence guarantees for their hybrid ap- proach, a significant advantage currently lacking in many neural operator frameworks. However, their method is lim- ited to linear PDEs. List et al. (2022); Um et al. (2020) integrate neural networks within the numerical solution pro- cess, often refining solutions iteratively from coarse to fine grids. These existing methods focus on solving PDEs. In contrast, (Allen et al., 2022) presents a line of works on dif- ferentiable simulators. Our method is different from these works. During inference, our model cooperates with tradi- tional solvers enhance the accuracy and stability of the optimization process. Importantly, the neural operator is trained independently (off-line), similar to traditional neural operator approaches.
Neural Operators, particularly instantiated by integral oper- ators (Kovachki et al., 2023). provide a flexible and powerful framework for modeling operators. Two prominent architec- tures have emerged: Fourier- based neural operator (FNO) (Li et al., 2020; Tran et al., 2021; Liu-Schiaffini et al., 2024), and transformer-basee neural operators (Li et al., 2022; Cao, 2021; Hao et al., 2023; Wu et al., 2024; Xiao et al., 2024). FNOs excel on data with fixed, uniform grids by leveraging the Fast Fourier Transform. Geo-FNO (Li et al., 2023) can predict on meshes deformed from uniform grids. In con- trast, transformer-based neural operators effectively handle arbitrary meshes. To enhance scalability, researchers have introduced linear transformers (Cao, 2021; Hao et al., 2023), reducing computational complexity to O(N) where N is the number of mesh points. Projection strategy (Wu et al., 2024) has also been proposed to decouple the complexity of attention from the number of mesh points.
3. Methods
Let da, d1, h represent the dimension of input space, output space, and embedded space, respectively. A neural operator can then be formulated as Go- = Q L o.... L P, where P de + R is lifting operator, Q : R2 -> RRd is
a projection operator, and L Rah + R² is the integral operator.
3.1. Sensitivity Loss
Since the gradient of neural operators is crucial based optimization, we impose derivative learning to neural
to gradient-
Although various neural operator architectures have been
3

---

Accelcrating PDECO by the Derivative of Neural Operators
transformers leads limited expressiveness due unde- sired inductive bias. On the other hand, the derivative of Fourier-based layers with discrete Fourier transform FF is simply dz F(2) = F( da
N:NxC
N
M
M
d
dz Therefore, the expressiveness of
the derivative of Fourier-based layers is untempered.
Nevertheless, Fourier- layer has a key limitation that it can only deal with uniform mesh points of fixed num- bers, while many applications, particularly optimization problems rely on irregular meshes with varying numbers of points. To address this limitation, inspired by Wu et al. (2024), we propose Virtual- Frtual-Fourier layer (see Fig. 3 for an overview), which transforms irregular meshes to a virtual physical space with a fixed number of sensors. This transfor- mation is achieved by point-wise projection and weighted aggregation.
11
Figure 3: Overview of structure of Firtull-Furres (VF) layer. The VF layer is composed by three steps: I. Weighted aggregating mesh points to a virtual physical space by equa- tion (8); II. Signal processing in the virtual physical space by a Fourier-basee based layer (9); III. Projection of the processed signals back to physical space by equation (10).
operators. Derivative learning in this paper means imposing supervision on the sensitivity of the neural operators defined as (2). Previous studies (Czarnecki et al., 2017; Tsay, 2021; et al., 2024) have shown that derivative learning is beneficial for neural networks.
First, let N denote the number of points in the physical space, M the number of sensors in the virtual physical space, and C' the number of channels. While N can vary across different cases, M remains fixed and is pre-determined. Given input æ E N and point-wise transformation Project: RIXC -> RIXM let
Given the sensitivity LJ obtained from numerical solvers, X
along with differentiable method of calculating J given u and X, we impose sensitivity loss the neural operators,
1: Project(2;,) VC E ERIX
OJ
Lgens = J A)
ax ||-
(5)
In our implementation, Project is a simple linear trans- form, though could potentially be implemented as non- linear transformation if needed. Here. hi represents the logit of the i-th point of being lin- early classified to the j-th sensor in the virtual physical space. Then we apply softmax(() along the first dimen- sion of I to obtain the weights w(()) E ER &RNXA xM specifically, (I) w w
Fig. (left) provides an overview of the computational graph for training neural operators wtih sensitivity loss. Consequently, the total loss is
RNX and
L= (A) u|l a. Laens:
where a >> is hyperparameter controlling the weight of the sensitivity loss.
exp(l,)) exp(1i,), 1, ,N. Let
3.2. Virtual-Fourier Layers
Zj
11) x E gIxc
j=1, M.
(8)
The inffuence of neural operator architecture on derivative learning has received comparatively less attention in the literature. We find that traditional transformer-based archi- tectures have limited expressiveness for this task. Here, we introduce a novel operator layer that exploits the simple derivative properties of Fourier layers. Notably, our layer also generalizes Fourier layers to arbitrary grids, a technical contribution that may be valuable in its own right.
Here, E RAxxC can be intuitively interpreted I-D signal in a virtual physical space sampled by M sensors, each with C channels.
Remark 3.1. Equation (8) is similar yet different from its counterpart of Transolver (Wu et al., 2024), which let w; = Softmax((roject((z))) E RIxM (softmax(.) applied along the second dimension of t) and 2j
In the following, use bold letters to indicate ten- sors of different dimensions. Let us consider commonly- used transformer-based and bausee neural operators. For a transformer, the derivative of the nonlinear atten- tion unit consists of (i) the derivative of Softmax of at- tention mechanism, let 8 = Softmax(2), z E RAN then del s; (ōtr- - sz); (ii) the chain rule and the product dz,
ΣN Wa, E RI Both treatments preserve the nature > wey of the whole layer as an integral operator since they only introduce additional measures the hidden space RCC, and the weight w's ''s are probability density functions. The main difference lies in their derivatives. The derivative of equa- tion (8) introduces less bias due to its simpler structure, which motivates us to adopt this forn. See Appendix A for detailed analysis of the derivatives.
rule of derivatives with respect to Q. K, and V. See Ap- pendix A for more detailed analysis. The derivative of

---

Accelcrating PDECO by the Derivative of Neural Operators
These signals z's are then processed by a Fourier-based layer:
primary purposes: efficient learning from optimization tra- jectory data and enabling hybrid optimization approach to mitigate error accumulation.
z' =F=F (Fz)) E R.MxC.
Reference Neural Operators (RNO) (Cheng et al., 2024), initially designed for learning smooth solution dependen- cies on domain deformations in shape optimization, can be extended a broader range of gradient- based optimization problems. Our core idea is that RNO leverages "nearby" data points within optimization trajectory to learn high- quality predictions. We find that the inherent relationships among these trajectory points are crucial for effective opera- tor learning, allowing RNO to capture fine-grained solution changes resulting from small input perturbations. Specifi- cally, our goal is to learn an operator,
The Fourier series Fz is truncated finite number of modes k E & such that Fz E CkxC and the weight parameter Rg ClxCxC, Fourier-based layers assume transitional symmetry on kernel integral operator, and due to convolution theorem (Proakis, 2001), Rg element-wise multiplies with frequency coefficient of z. In the origi- nal treatment of FNO (Li et al., 2020), the interaction be- tween channels is allowed. In this paper, with slightly abused notation, we take < i S k, 1 S j,l S C, and define the operation in (Rg- (Fz)) from equation (9) to be (Rg- (Fz))u = Eラ= (=1,((=)).. To see the simplicity of the derivatives of Fourier-based layer, notice dz )), which does not change the struc- ture of the layer. This is highly favorable for derivative learning.
A:Axu +u
(11)
dz' = -Ro (F
(1g, ur 4) + Uig
(12)
where 24 and A are Banach spaces, and T = C*(A), 8 2 1 represents a smooth transformation space on A. Also, we assume that there exist sp T such that (),)) = X, and Up 0 Ugg For example, in shape optimization, A denotes the geometry space which can be substantiated as a space of signed distance function, and T denotes a collec- tion of smooth deformation defined on A. Particularly, due to the bijection between reference mesh and query mesh the deformation X can be discretized as a shift vector field between them. In topology optimization problem, A rep- resents the space of input function, T can be assumed as any smooth, invertible transformation space, and to can be substantiated as the difference X, - -Ay-
Finally, the processed signals are transformed from the vir- tual physical space back to the physical space with the same logits from equation (7). However, this time softmax(-) is applied on the second dimension of L. Namely, let M exp(4.,), w (2) E RN and wh = exp(l..)/Ef= j = =,,... M and we get
M
x
w 2
2j ERICC
(10)
Although Fourier-basse based layer is sensitive to the order of input sequence, the Virtual-Fourier layer is permutation equivariant similar to Transformer-base positional embedding). To see this, one needs to notice that 4 only depends on o; due to the pointwise projection and the fact that under permutation equation (8) is invariant and equation (10) is equivariant.
Remark 3.2. The poos-posss of the mapping (11), specifically the existence, uniqueness and smoothness of yp, hinges on the property of the operator G' : A + U. These properties, such continuity and differenniiaiility, are non- trivial and inherently dependent on the specific PDE under consideration. For instance, the shape holomorphy of steady Stokes and Navier-Stokes equations, as demonstrated by Cohen et al. (2018), establishes the smooth dependence of the solution on small shape perturbations. This implies that given a reference shape A,- and a slightly perturbed shape Ag. the difference between the corresponding solutions ur and u, is bounded by the norm of perturbation. For com- prehensive overview of shape optimization theory, we refer readers to Sokolowski & Zolésio (1992). In the context of topology optimization, analogous results of fluid flows regarding the smooth dependence of the solution on the density coefficient have been established by Haubner et al. (2023); Evgrafov (2005; 2006). These results impose certain limitations on the applicability of our method. For example, significant shape deformations that alter the topology should be avoided, and the governing PDEs must exhibit sufficient regularity with respect to the design parameters.
layers (without
Following (Kovachki et al., 2023), the final output of a Virtual-Fourier layer is o( x ++x++++x)) b+z"), where W E RCxC b E RC and a is a nonlinear activation, such as GeLU (Hendrycks & Gimpel, 2016) in our implementation.
The overall computational complexity of the layer is ONNCC+ + 22 By grouping channels into heads, the computation cost can be reduced to O(NMC + M²C² /n), and the number of parameters in Ra decreases from OKC22) to O(KC2/1 /n). The derivative of equation (8) and (10) do introduce some additional bias to the layer, and we provide the analysis in Appendix A.
3.3. Training and Optimization with RNO
Finally, a key component of our neural operator-based gra- dient optimization method RNO. Its design serves two
Remark 3.3. In Cheng et al. (2024), the target of RNO
5

---

Accelcrating PDECO by the Derivative of Neural Operators
14
Predicted J trgecpor Cround Inth / tojedayy
ctedicted hajesto Cround tuth Sojectry
Ag
Ax
Ag Ag
Figure 4: Optimization Process with RNO. This figure illustrates the optimization trajectory using the RNO model. Red dashed trajectory illustrates the true solutions and objectives without access. Black solid trajectory stands for optimization with neural operators. Blue arrows indicate updates to the ground truth solution based on the current input design parameters X. The optimization trajectory can be recalibrated whenever a numerical solver called, ensuring robustness and accuracy. (Left) The optimization trajectory is projected onto the objective function J design parameter A plane. (Right) The optimization trajectory is projected onto the solution space U design parameter A plane. Note that both panels represent the same optimization trajectory.
defined to be the difference du = u, op ug. inspired by material derivatives. We modify the target to be simply ug- With skip connections outside of the integral operator layers, setting My as the target would implicitly learn ōu. There are two benefits with this modification. Firstly, it simplifies implementation. Secondly, reference neural operators can be re-interpreted as an extension of vanilla neural operator (3), which unifies RNO and NO. In practice, it enables the flexibility of utilizing RNO as NO, which is helpful if reference is expensive or less worthy acquiring.
Pefeerence-Quert pairs. To train RNO, we require a pair of data points for each query. Since our data originates from optimization trajectories, we pair each data point with its nearest neighbor within the same trajectory, excluding neighbors that exceed a predefined distance threshold. This pairing process is implemented within custom dataloader. For further details, please refer to Appendix C.I.
Reference dropout ratio. During training, we randomly drop reference inputs (ur; s) so that RNO reduces a vanilla neural operator that maps from to u. Enabling this feature is preferable because it allows RNO to perform inference flexibly without external references. This is partic- ularly beneficial at the beginning of the optimization process, as RNO can operate without relying on any ground truth so- lution. Consequently, we can avoid unnecessary numerical computations during the initial stages of the design process. The reference dropout ratio is set 0.3 to balance training performance.
Remark 3.4. Although RNO an extension of neural oper- ators that maps from a larger space A x l x T to solution space U1 empirically observe that it achieves higher learning efficacy compared to vanilla neural operators. See Section 4. A possible underlining reason is that the train- ing procedure of RNO is analogue to contrastive learning, which takes advantage of the inner relation between data samples, thereby improving data utilization. Given N data samples, each data with k neighbors, the number of pairwise data samples is OKNN). Since data samples are collected from optimization processes using traditional methods, they are naturally closely related. Based on this, RNO is able to effectively leamn the changes of solutions that result from small changes in input.
3.3.2. OPTIMIZATION WITH RNO
To mitigate the accumulation of errors generated by neural operators during optimization, a straightforward approach involves validating intermediate optimization results and restarting the optimization process from the validated point. Then neural operators must effectively utilize newly ac- quired ground truth solutions. Traditional neural operators typically require retraining with the newly acquired data, which is computationally expensive in single optimization run. In contrast, since RNO accepts a reference solution as input to predict the response for a query, it can recalibrate its predictions and adjust the optimization direction whenever new ground-truth solution becomes available. See Fig. 4 for an illustration of the method.
3.3.1. CIAIMMZAATOO-RREENNTI TRAINING RNO
We continue to use the notation of neural operator and rep- resent RNO by Go = Q L... Lo Σ P:, where P; 's are lifting operators corresponding to each element of triplet (Ny, Ur, Ur,p)) :p. The training objective function is defined as (6). We summarize the training algorithm in Algorithm D. This training paradigm is specifically tailored for optimization data, thus inherently optimization-orientee
6

---

Accelcrating PDECO by the Derivative of Neural Operators
Dataset
Component
LA
PA
VF
R-VF w/ ref
R-LA
R-PA
R-VF (Ours)
p
7.12e-2
7.67e-2
6.36e-2
5.56e-2
1,45e-2
1.38e-2
1.33e-2
1.6le-1
1.76e-1
2.16e-1
1.25e-1
2.85e-2
2.53e-2
2.31e-2
Microreactor2D
6,87e-1
6.86c-1
6.77e-1
5.35e-1
7.34e-2
7.02c-2
6.84e-2
8.79e-2
9.33e-2
1.04c-1
7.33e-2
2.64e-2
2.56e-2
2.42e-2
Ls
2.05e-1
2.02e-1
1.70e-1
1.63e-1
1.18e-1
1.5le11
1.05e-1
1,37e-2
1.34e-2
1.30e-2
u
2.99e-1
2.19e-1
2.97e-1
1.84e-1
1.14e-2
1.14e-2
1.10e-2
Fuelcel12D
3
2.60e-1
2.36e-1
2.10e-1
1,42e-2
1,40e-2
1.38e-2
L1
4.90e-
4.62e-1
4.44c-1
L2
3.18e-1
2.77e-1
2.57e-1
Br
5,43e-
4.62c-1
4.57e-1
3.93e-1
9.65e-3
1.14e-2
8.66e-3
B=
3.76e-1
3.24e-1
3.91e-1
3.74e-1
1.04e-2
1.05e-2
9.55e-3
Inductor2D
Ly
3.12e-
3.64e-1
3.00e-1
L
5.11e-1
5.77e-1
4.76e-1
5,51e-
2.42c0
4.27e-1
4.42e-1
1.08e-1
9.60e-2
7.64e-2
Drone3D
U2
5.51e-1
6.60e-1
3.88e-1
3.84e-1
7.01e-2
7.25e-2
5.60e-2
U3 7.50e-1
9.91e-1
1.17e0
5.95e-1
1.32e-1
1.18e-1
1.20e-1
Ls
6.48e-1
5.66e-1
5.05e-1
Table 1: The prefix "R-" indicates models being modified into RNO framework. See definition of baseline models in Section 4. All metrics relative-l2 error. Errors shown by each component. L. 's sensitivity errors marked in blue. Bold is the best and underline is the second best. Red values are the best among all models without reference.
Smoothen the gradient. During inference, we adopt two tricks to further smoothen gradient while proceeding opti- mization: (1) Inject noise into the inputs and take average on the outputs. Let WN(0,a), o), where is set to 1 1% of the standard deviation of the input,
Algorithm Optimization with RNO
Input: RNO Ge. random initial input λg. learning rate 7 > 0, buffer list B with size N2, Ground truth solution Ugt None with empty initialization, worm_up_steee for optimization before the first validation, radius around Ugt a criterion to trigger validation.
8J
vxJ(u(A+ei).)) N
(13)
for = to T - do
and Dist( (4 ugt) > r then Update ugt with numerical solver Reset buffer [uge]
>> ayarm_up_shee
(2) We buffer the most recent optimization steps with certain size M, and then make all steps as reference and take average on the outputs.
end if
Compute de by Eq. (14) and 5J by Eq. (13)
Store in B
N2 Σ 3=1
Update det1 + de 1 --
Upredd
N₂
(14)
end for
The total number of forward passes for single optimization step is Ni Since the derivative of the summation of J with respect to all inputs À + E; are independent to each other, the derivative in (13) can be computed by Autograd in one pass. Also, (14) can be computed parallel. The optimization algorithm is summarized as Algorithm
optimization problem that minimizes both pressure drop and flow velocity variance between channels. For eletro- magnetics we have Inductor2D: A 3D axial-symmetric shape optimization problem that aims to reduce material cost while maximizing inductance. For solid mechanics we have Drone3D: 3D topology optimization problem that minimizes elastic strain energy under distinct load cases: vertical acceleration and motor torque. All datasets are organized as num. Iraj nurn slep, where raum Iraj and num numn-step are the number of trajectories and the num- ber optimization steps. Train and test sets are split in 8:2 by trajectories. Please check more details of each dataset in Appendix. C.
4. Experiments
We benchmark four challenging optimization datasets from various physics backgrounds, all simulated using COM- SOL 6.0. For fluid dynamics we have Microreactor2D: 2D topology optimization problem that maximizes the reaction rate of channel design; FuelcelI2D: A 2D shape

---

Accelcrating PDECO by the Derivative of Neural Operators
infinity
Model R-PA R.VF
0.45
0.04
Alpha
0.4
0.0
Sensriiity
0.35
0.5
Objective
Numerical
Model
Hybrid R-V
0.3
0.02
RVVF reference
0.25
0.2
0.00
0.026
0.028
Rel2 Error
Num, of Calls
Figure 5: Trade-off between operator and derivative accu- racy. This plot illustrates the trade-off for three models trained with different values of the hyperparameter on the Microreactor2D dataset. R-VF consistently achieves points on the Pareto front, demonstrating a strong balance between accurate operator predictions and accurate derivative esti- mations.
Figure 6: Objective Function Values during Optimization on Microreactor2D. The horizontal axis represents the number of calls to the numerical solver. "Hybrid R- GFV refers to the optimization algorithm (Algorithm 1) using the trained R- VF model. "R-VE w/o reference" represents the optimiza- tion process relying solely on the R-VF model without any calls to the numerical solver.
Computing the sensitivity loss (Equation 6) requires the objective function value, ., for each data point. To ensure differentiabilityy we employ simplified approximations for the objective function. For example, we approximate the average value on irregular meshes by averaging node values and incorporate constraints additional objectives using Lagrange multipliers. These simplifications can introduce discrepancies between the true objective function used in data generation and the approximated objective used for computing neural operator sensitivities. To address this discrepancy, we utilize an equivalent form the cosine sim- ilarity for the sensitivity loss for all datasets except Microre- actor2D, where the original formulation is used. Details of this equivalent formulation are provided in Appendix B.
Effectiveness of RNO: All RNO variants demonstrate effec- tive learning across all datasets. Notably, R-VF consistently outperforms baselines in both operator and derivative learn- ing. We highlight that the error reduction of sensitivity loss from the second best is between 3.8% to 17.5%, suggesting that the Virtual-Fourier (VF) layer effectively reduces bias in derivative learning.
Interestingly, the RNO training procedure itself appears to mitigate overfitting. When op- erating without a reference solution (by setting the refer- ence dropout ratio to 1), RNO effectively degenerates into vanilla neural operator. However, even in this degenerate state, its performance often surpasses that of all traditional neural operators, as observed in the "R-VF w/o reference" column.
Overfitting suppression:
We compare Virtual-Fourier (VF) layer to two baselines: linear attention (LA) from GNOT (Hao et al., 2023) and physics attention (PA) from Transolver (Wu et al., 2024). Except for the layer structure, all models share the same architecture, c.g., lifting and projection operators, and the similar. We add prefix "R-" to
4.1.1. THE TRADEOFF DUE TO SENSITIVITY Loss
The sensitivity loss (6) introduces a tradeoff between op- erator accuracy and derivative accuracy. To investigate this tradc-off, we analyze the Pareto front of each model on Microreactor2D for different values of the coefficient a = 0.1, 0.5, 0.9). Figure 5 presents the results, where the relative error represents the average error across all components. R-VF outperforms both baselines and demon- strates strong performance on all a'
number of parameters models that are modified into RNO framework.
4.1. Learning Results
Table I reveals several key observations. Limitations of tra- ditional training: Vanilla neural operators trained with stan- dard methods exhibit poor performance across all datasets. This is likely due to the inherent similarity of data points within optimization trajectories near the optimal region, leading to a lack of diversity compared to traditional training datasets. Limited diversity can easily result in overfitting.
1.2. SCALING ABILITY OF LIRTUAL-FOURIER
In order to study the scaling ability of VE, we doubled the size of Microreactor dataset to 200 trajectories (each of
8

---

Accelcrating PDECO by the Derivative of Neural Operators
12 steps and effectively 2400 samples). 40 trajectories are kept test set. We train models 80, 160 trajectories respectively. In Table 3 of Appendix E, we observe that with larger dataset, R-VF keeps advantage over baselines on predicting both physical fields and sensitivities.
Microreactor2D
Drone3D
Mesh nodes
4.9e3
2.1e4
R-VF (GPU)
0.53
0.62
R-VF (CPU)
1.89
3.20
Numerical (CPU)
3.20
49.40
4.2. Optimization Results
Table 2: Wall-clock runtime comparison in second/ itter.
To evaluate our hybrid approach (Algorithm 1), we ran- domly selected a test case from the Microreactor2D dataset and performed optimization from scratch using a fully trained R-VF model. We employed the Method of Mov- ing Asymptotes (MMA) (Svanberg. 1987) and Gaussian smoothing for sensitivity updates. We compared our hybrid approach with traditional numerical gradient-based opti- mization. Additionally, we benchmarked the R-VF model's performance within Algorithm against a baseline where the R-VF model operated without feedback from the numer- ical solver. In this "w/o reference" scenario, we validated intermediate objective values every 20 optimization itera- tions.
affects the overall wall-clock time. Moreover, there are some important techniques in optimization, such as gradient filtering, projection, and clipping. A comprehensive and fair comparison requires all these techniques. We provide a primitive reference for wall-clock time in Table 2.
Experiments were conducted on both a laptop with CPU 2.5 GHz (1 11th Gen Intel i7) and a GPU Nvidia V100. The reported time is the average value of 10 iterations, measured in seconds per iteration. The runtime of traditional numeri- cal methods increases sharply as the problem scales, since numerical methods suffer from the curse of dimensionality. In contrast, R- VF runtime increases moderately due to its linear computational complexity w.r.t. the number of mesh nodes.
As shown in Figure 6, the hybrid approach (R-VT with Algo- rithm 1) demonstrates rapid and stable convergence. In con- trast, optimization relying solely on the R-VF model quickly derails without numerical solver feedback. Although the hybrid approach may converge to a suboptimal solution due to factors such as derivative accuracy errors, objective func- tion discrepancies, constraint handling, density projection, etc., it exhibits robust convergence behavior.
5. Conclusion
This work addresses the challenges of applying neural opera- tors to PDECO problems using gradient-baie based methods. We propose a novel framework that incorporates a specialized training procedure that leverages data generated from opti- mization process, a novel Virtual-Fourier layer to improve the accuracy of derivative predictions, and a hybrid approach that integrates neural operators with traditional numerical solvers. Through extensive experiments, we demonstrate the effectiveness of our proposed framework in accurately learn- ing operators and their derivatives, leading to significant improvements in the robustness of gradicnt- optimiza- tion with neural operators.
Importantly, within the first four calls to the numerical solver, the hybrid R-VF optimizer achieves objective values comparable to, or even exceeding. those obtained by the numerical optimizer, suggesting potential for computational cost savings. Furthermore, each call to the numerical solver in our hybrid approach costs approximately half much as one adjoint method step (due to the expense of solving the adjoint equations), further amplifying the potential savings.
1.2.1. WALL-CLOCK TIME COMPARISON
The runtime of our hybrid method consists of the time spent optimizing with RNO and the time calling numerical solvers. The latter dominates the runtime, and optimization with RNO is much cheaper. Thus, a major potential cost saving of the hybrid method lies in fewer calls for numerical solvers.
Impact Statement
This paper presents work whose goal is to advance the field of Machine Learning. There are many potential societal consequences of our work, none which we feel must be specifically highlighted here.
A 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, and SIMP (Bendsøe, 1989). A well-chosen optimization method can achieve higher objective values and reduce both the number of iterations and the number of function calls. Thus, the optimization method significantly
A Acknowlegement
This research is a result of the collaboration within Bosch- Tsinghua Machine Learning Center. This work was sup- ported by the National Key Research and Development Program of China (No. 2020AA40106302), NSFC Projects (Nos. 92370124, 62350080, 92248303, 62276149).
9

---

Accelcrating PDECO by the Derivative of Neural Operators
References
Forrester, A., Sobester, A., and Keane, A. Engineering design via surrogate modelling: a practical guide. John Wiley & Sons, 2008.
Allen, K. R., Lopez-Guevara, T., Stachenfeld, K., Sanchez- Gonzalez, A., Battaglia, P., Hamrick, J., and Pfaff, T. Physical design using differentiable learned simulators. arXiv preprint arXiv:2202.00722 007288, 2022.
Hao, Z. Wang, Z., Su, H., Ying, C., Dong, Y., Liu, S., Cheng, Z., Song, J., and Zhu, Gnot: general neu- ral operator transformer for operator learning. In Inter- national Conference Machine Learning, pp. 12556- 12569. PMLR, 2023.
Audet, C., Denni, J., Moore, D., Booker, A., and Frank, P. A surrogaee-moeelbaale method for constrained optimiza- tion. In 8th symposium multidisciplinary analysis and optimization, pp. 4891, 2000.
Haubner, J., Neumann, F,, and Ulbrich, M. A novel den- sity based approach topology optimization of stokes flow. SIAM Journal on Scientific Computing, 45(2):A338- A368, 2023.
Bendspe, M. P. Optimal shape design as a material dis- tribution problem. Structural optimization, 1:193-202, 1989.
Hendrycks, D. and Gimpel, K. Gaussian error linear units (gelus). arXiv preprint wtt.:16606.84411 2016.
Bendsoe, M. P. and Sigmund, O. Topology optimization: theory, methods, and applicationn. Springer Science & Business Media, 2013.
Hsieh, J.-T., Zhao, S., Eismann, S., Mirabella, L., and Er- mon, S. Learning neural pde solvers with convergence guarantees. arXiv preprint arXiv:1966011200, 2019.
Borel, É. La mécanique statique et I'irréversibilité. Phys. Theor, Appl., 3(1):189-1966 1913.
Hwang, R., Lee, J. Y., Shin, Y., and Hwang, H. Solving pde-constrained control problems using operator learn- ing. In Proceedings the AAAI Conference on Artificial Intelligence, volume 36, pp. 4504 4512, 2022.
Bower, A. Applied mechanics of solids, 2009.
Cao, S. Choose a transformer: Fourier or galerkin. Ad- vances in neural information processing systems, 34: 24924 24940, 2021.
Kennedy, J. and Eberhart, R. Particle swarm optimization. Proceedings of -iitte9ninteanationnl conference on neural networks, volume 4, pp. 1942-1948. ieee, 1995.
Cheng, Z., Hao, Z., Wang, X., Huang, J., Wu, Y., Liu, X., Zhao, Y., Liu, S., and Su, H. Reference neural operators: Learning the smooth dependence of solutions of pdes on geometric deformations. In Forty-first International Conference on Machine Learning, 2024.
Kohn, W. and Sham, L. J. Self-consistent equations includ- ing exchange and correlation effects. Physical review, 140(4A):A1133, 1965.
Cohen, A., Schwab, C., and Zech, J. Shape holomorphy of the stationary navier-stokes equations. SIAM Journal on Mathematical Analysis, 50(2): 1752, 2018.
Kovachki, N., Li, Z.. Liu, B., Azizzadenesheli, K., Bhat- tacharya, K., Stuart, A., and Anandkumar, A. Neural operator: Learning maps between function spaces with applications to pdes. Journal of Machine Learning Re- search, 24(89): 2023.
Czarnecki, W. M., Osindero, S., Jaderberg, M., Swirszez, G., and Pascanu, R. Sobolev training for neural networks. Advances in neural information processing systems, 30, 2017.
Li, Z., Kovachki, N., Arizzadenesheli, K., Liu, B., Bhat- tacharya, K., Stuart, A., and Anandkumar, A. Fourier neural operator for parametric partial differential equa- tions. arXiv preprint arXiv:20/0.0899) 2020.
Evgrafov, A. The limits of porous materials in the topology optimization of stokes flows. Applied Mathematics and Optimization, 52(3):263-271) 2005.
Li, Z., Meidani, K., and Farimani. A. Transformer for partial differential equations' operator learning. arXiv preprint arXiv: 220.223671.....1. 2022.
Evgrafov, A. Topology optimization of slightly compress- ible fluids. ZAMM- Journal of Applied Mathematics and Mechanies/Zeitschrij für Angewandte Mathematik und Mechanik: Applied Mathematics and Mechanics, 86(1): 46-62 2006.
Li, Z., Huang, D. Z., Liu, B., and Anandkumar, A. Fourier neural operator with learned deformations for pdes on general geometries. Journal of Machine Learning Re- search, 24(388):1-26 2023.
Fahl, M. and Sachs, E. W. Reduced order modelling approaches to pde-constrained optimization based on proper orthogonal decomposition. Large-scale PDE- constrained optimization, pp. 268 280. Springer, 2003.
Li, Z., Zheng, H., Kovachki, N., Jin, D., Chen, H., Liu, B., Azizzadenesheli, K., and Anandkumar, A. Physics- informed neural operator for learning partial differential
10

---

Accelcrating PDECO by the Derivative of Neural Operators
equations. ACM/JMS Journal of Data Science, 1(3):-227, 2024.
Um, K., Brand, R., Fei, Y. R., Holl, P., and Thuerey, N. Solver-in-the-lore Learning from differentiable physics to interact with iterative pde-solvers. Advances in Neural Information Processing Systems, 33:6111-6122 2020.
List, B., Chen, L.-W., and Thuerey, N. Learned turbulence modelling with differentiable fluid solvers: physics-basee loss functions and optimisation horizons. Journal of Fluid Mechamics, 949:A25, 2022.
Wang. S., Bhouri, M. A., and Perdikaris, Fast pdd- constrained optimization via self-supervised operator learning. arXiv preprint arXiv:2110. 2021.
Liu-Schiafini, M., Berner, J., Bonev, B., Kurth, T., Aziz- zadenesheli, K., and Anandkumar, A. Neural operators with localized integral and differential kernels. arXiv preprint arXiv:2402.1653 2024.
Wu, H., Luo, H., Wang, H., Wang. J., and Long, M. Tran- solver: A fast transformer solver for pdes on general geometries. arXiv preprint arXiv:2402.02336, 2024.
Xiao, Z.., Hao, Z., Lin, B., Deng, Z., and Su, H. Improved operator learning by orthogonal attention. In Proceed- ings of the 41st International Conference on Machine Learning, pp. 54288-54299, 2024.
Lu, L., Jin, P., and Karniadakis, G. E. Deeponet: Learning nonlinear operators for identifying differential equations based on the universal approximation theorem of opera- tors. arXiv preprint 1xxvv1910003113, 2019.
Xue, T., Beatson, A., Adriaenssens, S., and Adams, R. Amortized finite element analysis for fast pde-constrained optimization. In International Conference on Machine Learning, pp. 10638-10647. PMLR, 2020.
Lu, L., Jin, P., Pang, G., Zhang, Z., and Karniadakis, G. E. Learning nonlinear operators via deeponet based on the universal approximation theorem of operators. Nature machine intelligence, 3(3):21--23) 2021.
Zhang, I.., Han, J., Wang, H., Car, R., and E, W. Deep potential molecular dynamics: scalable model with the accuracy of quantum mechanics. Physical review letters, 120(14): 1430001, 2018.
Neal, R. M. Bayesian learning for neural networks, volume 118. Springer Science & Business Media, 2012.
O' O'Leary- Roseye-rseeerry, T., Chen, P., Villa, U., and Ghattas, O. Derivative-infone informed neural operator: an efficient frame- work for high-dimensional parametric derivative learning. Journal of Computational Physics, 496: 12555, 2024.
Zhu, Y. and Zabaras, N. Bayesian deep convolutional encoder-decoder networks for surrogate modeling and uncertainty quantification. Journal of Computational Physies, 366:415 447, 2018.
Pontryagin, L. S. Mathematical theory optimal processes. Routledge, 2018.
Proakis, J. G. Digital signal processing: principles algo- rithms and applications. Pearson Education India, 2001.
Snoek, J., Larochelle, H., and Adams, R. P. Practical bayesian optimization of machine learning algorithms. Advances in neural information processing systems, 25, 2012.
Sokolowski, J. and Zolésio, J.PP. Introduction to shape optimization. Springer, 1992.
Svanberg, K. The method of moving asymptotes-a new method for structural optimization. International journal for numerical methods in engineering. 24(2):359-373, 1987.
Tran, A., Mathews, A., Xie, L., and Ong. C. S. Factorized fourier neural operators. arXiv preprint arXiv:2111.13802 2021.
Tsay, C. Sobolev trained neural network surrogate models for optimization. Computers & Chemical Engineering, 153:107411 2021.

---

Accelcrating PDECO by the Derivative of Neural Operators
A. Derivative of Virtual-Fourier Layer
First, let us take a closer look at the derivative of nonlinear attention, di softmax(QK]V ignoring constant coefficient and the transpose of matrix. Suppose t E RR and softmax(t). Then Øs, ăt; 5z(őı; s.). Thus, the derivative would be roughly and informally 85(6, (:)(Q' + QK") + softmax(QK)V' Hence, the derivative of such layers is assigned with strong inductive bias (probably undesired), and the expressiveness would be severely restricted. Even for linear transformers, e.g., GNOT (Hao et al., 2023), the derivative of normalization for linear attention has a complicated form due to the quotient rule.
In contrast, the computation of the derivative of Virtual-Fourier layers essentially occurs on the Softmax function. For equation (8), its derivative is
0zj = w 0x;
N
Tliss x; 0x:
+
W R.j (6
(15)
Additional bias
Proof. Notice that 4,, only depends on x; by (7). / = exp(l.,)/ ixp(4,),. 1,..., ,N, which is softmax())
along the first dimension of I. Then it directly follows from the derivative of softmax() )), product rule and chain rule. D
(2) In equation (10) since the dependence on zj is linear, we mainly focus on the derivative of "k." Similarly, we have
Ow k.j (5,, WKgs) Ox:
(2)
(16)
8x:
Proof. Again, due to (7) and w exp(4.))/E exp(4.)). = 1, M, which is softmax(-) along the second dimension of lig. It directly follows from the derivative of softmax(-), product rule and chain rule.
Both the derivative of (8) and (10) introduce some bias to the derivative of operators due to the derivative of softmax(-),, but compared to transformer-base neural operators, the derivative of Virtual-Fourice layer is considerably simpler.
B. Sensitivity Loss
Two technical challenges arise in computing the sensitivity loss (equation 5):
Sensitivity Representation: For shape optimization datasets, COMSOL provides sensitivities with respect to curve parameters (e.g., coefficients of Bernstein polynomials), not directly with respect to the coordinates of mesh nodes. To simplify the comparison, we obtain the shift in node coordinates and maximize the cosine similarity between this shift and the sensitivity of the neural operator with respect to the node coordinates. Equivalently, we normalize both vectors and minimize the la loss between them.
The optimization method used within COMSOL handle constraints introduces complexities in sensitivity computation. To simplify this aspect, we incorporate constraints as additional objectives in the objective function J using Lagrange multipliers. This simplification may introduce errors in the computed sensitivities. To mitigate this, we again prioritize maximizing the cosine similarity (minimize the L2 loss between normalized sensitivities), while also choosing a small value for the coefficient a in equation 6.
Constraint Handling:
C. Dataset and Dataloader
C.1. Reference Dataloader
To train RNO with pairwise data query and reference, we custom dataloader to select reference point from the same optimization trajectory for each query point. Given a trajectory ((Ai, w.)).N where N, is the number of steps. the reference for the i-th data point is chosen from the index range [1, where / = max(0, i d) and r = min(N [N. i + d). Note that the reference index j can be equal to i. Additionally, the relative difference between the reference solution u, and the query ug must be less than a threshold d,.. The dataloader randomly selects a reference point that satisfies both conditions. In our implementation, d = 2 and d, = 0.5.
12

---

Accelcrating PDECO by the Derivative of Neural Operators
Because both the query and reference selected from the same trajectory, they share one-to-one correspondence between mesh grids. This is inherently true for data generated from topology optimization, where all meshes identical. For shape optimization, this one-to-one correspondence is preserved after shape deformation. This consistent mesh correspondence is a key advantage of using optimization trajectories as training data. Without it, constructing appropriate transformations and implementing interpolations between different meshes, as required in (Cheng et al., 2024), would be necessary.
C.2. Microreactor
(a)
(b)
Figure 7: (a) The optimized density of porous media. (b) The corresponding concentration c.
We modified a model from COMSOL optimization gallery that designs a catalytic microreactor by topology optimization.
The objective function is
1 max J = volume(()) / (1 - =)eds?
(17)
where e is concentration, kn = 1 is the rate constant and s()) E [0, 1] denotes the control variable, the volume fraction of solid catalyst. Let D be diffusion coefficient and u be the velocity, and the overall physics is governed by Navier-Stokes and diffusion equations,
p(u Vju -V V Vuu+ (Vu)²)- a())u
V v.u
V DVe) = Cc.
The data is structured as À = (Ii, 22,6, m) is a 4-tensor with first two dimensions of spatial coordinates and e the control variable. The last dimension is mask m E {0, 1 for active domain where E is defined, and N is the number of mesh nodes. u = (p, u, v, c, s). which are pressure, velocity components, concentration and the sensitivity s = de = E, - Eg be the transformation between reference and query.
dd Finally, let
We randomly draw the widths and heights of the reaction region, as well as the heights of inlet and outlet. The size of the dataset is numjrrjj num_steps 100 x 12.
C.3. Fuelcell Bipolar Plate
We model a simple distribution area for a fuel cell bipolar plate, the governing equation is Navier-Stoos equation and the objective function consists of two parts: the total pressure drop and the variance flow among all channels,
min J = (pl intet Ploutiet) + (I ū) chonnet
(18)
where is the control variable of the internal walls in distribution areas close inlet and outlet. The objective can be further simplified if we set Plouttet 0 and assume that flow velocity in each channel is constant and hence we only need to measure the velocity end of arm area (or the inlets of channels). In COMSOL, X is the control variable of shapes which is modeled by the coefficients of Bernstein polynomials. In our implementation of optimizing shape with neural operators, we simply optimize the coordinates of the movable mesh nodes in arm area so that À stands for the coordinates of those points,
The data is structured as = (I1, F2, w, m), a 4-tensor with first two dimensions of spatial coordinates and w, m E (0, 1] N are mask for internal walls and mask for arm area (free shape domain), where N is the number of mesh nodes.
13

---

Accelcrating PDECO by the Derivative of Neural Operators
(a)
(b)
Figure 8: (a) The pressure of optimized design. (b) The velocity component u of optimized design.
(p, u, v, 1æ drz), where (dz1, dr.2) is the shift of coordinates. ip (I1, 22), (2), is the difference of spatial coordinates of mesh nodes between reference and query.
We randomly draw the widths and heights of channels, as well as the widths and heights of inlet and outlet. The size of the dataset is numjtraj x ntm_steps = 30 x 20.
C.4. Inductor
||BI|
Axial symmetry
T.
Core
0
Coil
Infinite boundary
Figure 9: (Left) The 2D inductor model, consisting of core, coil, air, and infinite boundary. (Right) The complete 3D domain after revolving the domain around the symmetry axis.
We model the magnetic field response of an inductor under axial symmetry assumption as seen Fig. 9. The objective function consists of competing components: maximizing the inductance (magnetic energy storage capability) while minimizing the cross sectional area of the inductor.
min J Sq (B H)2mrdrdz (1 (1 1drdz
/
(19)
12
where H stands for magnetic field intensity, B for magnetic flux density, / for current, and À for the control variable for the cross section shape. Particularly, we randomly draw 7 E [0.7, 0.9] to enable diversity on objective functions. For simplicity, kHz by solving the magnetic wave equation the
we investigate the steady state frequency domain response at f 2t 2D inductor cross section Ω using COMSOL,
V H + 2 Eull
X
(20)
14

---

Accelcrating PDECO by the Derivative of Neural Operators
where J stands for current density, is the angular frequency, μ and e permeability and permittivity. The excitation current is set to have small amplitude so that the induced magnetic field can be assumed stay within the linear regime (u 600 Tm/A) of the core material without hysteresis effect. Therefore, the steady state field response takes the form (B, H) = (B, H) at identical phase. The objective function in Eq. 19 can then be simplified by replacing B H with the corresponding amplitude jH' and performing integral only within the core region, as is only significant within the core region.
The data is structured À = (r; 2, Y, my, m2), a 5-tensor with first two dimensions of axial-symmetrical spatial coordinates (r,2) and my, m2 E {0, 1] N are masks for coils and magnetic core, where N is the number of mesh nodes. (B,, Bz, dr, dz), where (dr, dz) are shifts of coordinates. &p (), 2)g is the difference of spatial coordinates of mesh nodes between reference and query.
We randomly draw the radius and heights of core, the radius and heights of inner hole, the widths and heights of coil and the parameter y in objective function. The size of the dataset is num Jraj x num_steps 100 10.
C.5. Drone
(a)
(b)
Figure 10: (a) The optimized density e filtered by a threshold. (b) The magnitude of corresponding displacement u of the loadcase of vertical acceleration. We modified a model from COMSOL optimization gallery that designs the arms of drone with linear clastic material. The problem minimizes elastic strain energy with physical variables such as x 3 strain, stress tensors and displacement. Our modification on the objetive is, with Einstein sum notation,
min J = J beu,VV + f. L,u,A,
(21)
where u, is the displacement, by is force, t4 is traction on surface òR and æ denotes the density of the plastic material. This equivalent to elastic strain energy due to the principle of minimum potential energy (Bower, 2009). Also, the governing equation will be simplified as the Navier equations of elasticity (Chapter 5. Bower (2009)). The only field variable of interest is displacement. The problem is optimized under two load cases, a vertical acceleration on the whole body and torque on the motor surface. The optimization process is performed twice to obtain the complete drone body. See Fig. 10.
single arm, and the final result is then mirrored
The data is structured as À = (æ1, T2, F3. 0, f1, J2, fs, my, m2, ma), a 10-tensor with first spatial coordinates and E [0, 1|.i the control variable of density model for topology optimization where N is the number of mesh nodes. f's are force components and m's are masks for the surface where torque is applied, surfaces for symmetry conditions and volume aj
assigned with 0. Displacement u (i,, u2 U3 s), sensitivity s 00 and transformation yp 0, 0.- 0g-
We randomly draw the widths, heights and thickness of the arm of drones, as well as the torque and accelerations of of two loadcases. The size of the dataset is num Iraj x num sleps 10 x 20.
D. Algorithms
15

---

Accelcrating PDECO by the Derivative of Neural Operators
Algorithm 2 Optimizatiion-oriented training of RNO
Input: RNO Go. dataloader that loads data in pairs. (Ay. Uy) and (Ar: u,), flag Sens for training with sensitivity loss, dropout ratio rdrop that drops reference. for epoch 0 to N₂ do
Draw 01 1]
if < Tdrop then
Drop u, and +
end if
Upred = ( Ur; (p)
Compute loss Upred Ugl if Sens then
Compute J and aJ by autodiff ah Update loss L by (6). end if
optimizer.zero() L.backward() oplimizer.updulee end for
E. More Experiments
P
v
La
R-LA-80
1.64e-2
3.25e-2
7.83e-2
2.65e-2
2. 69c-2
R-PA-80
1.5le2
2.74e-2
7.46e-2
2.48e-2
2.54e-2
R- VF-80
1.48e-2
2.57e-2
7.56e-2
2.42e-2
2.30e-2
R-LA-166
1.40e-2
2,42e-2
7.16e-2
2.29c-2
2.37e-2
R-PA-160
1.32e-2
2.18e-2
6.81e-2
2.15e-2
2.30e-2
R-VF-160
1.31e-2
2.17e-2
6.95e-2
2.15e-2
2.18e-2
Table 3: The results of scaling dataset Microreactor21D. Errors are shown by each component. I., is sensitivity emoo. Bold the best and underline is the second best.
Models
p
L, 2.16e-2 + 3.55e-4
R-LA
1.29e-2 + 3.3e-4
2.33e-2 1.04e-3
6.85e-2 + 1.94e-3
2.32e-2 + 9.3e-4
R-PA
1.25e-2 4.1e-4
2.22e-2 1 6.6e-4
6.6e-2 + 1.23e-3
2.28e-2 +: 2.1e-4
2.02e-2 7.8e-4
R-VF
1.2e-2 + 3.2c-4
2.06e-2 = 6.1e-4
6.42e-2 + 1.42e-3
2.14e-2 + 6.89e-4
1.76e-2 + 3.81e-4
Table Results of Microreactor21) with random seeds. Showing consistent advantage of VE.
16


Reviews:
[{"review_id": "LYM9gpZDoJ", "reviewer": "Reviewer_wq4G", "summary": "Authors propose a general architecture-agnostic framework that can be used for PDE-constrained optimization and a Virtual Fourier Layer suitable for data on irregular grids.\n\nThe framework consists of three essential parts: (i) a particular structure of inputs and outputs for selected neural networks, (ii) a way to generate training data and train the model, (iii) an inference strategy that utilises classical solver.\n\n(i) Roughly, the inputs to the neural networks are (solution for current parameters, desired parameters parameters, difference between current and desired parameters) and the output is a solution for desired parameters.\n\n(ii) Data is generated along the trajectories of the optimization process and the training is performed with intermittent masking of two inputs: (solution for current parameters, difference between current and desired parameters). The latter encourages neural networks to be more robust and learn mapping from parameters to solution.\n\n(iii) At the inference stage authors propose to sparingly correct optimisation trajectory using classical solver. This is possible because of the special structure of the neural network described in (i).\n\n## update after rebuttal\n\nSummarised in https://openreview.net/forum?id=LFF7kUQ5Rp&noteId=gR2VRvNh9Z", "questions": "", "limitations": "", "rating": 4, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for reviewing our paper in such depth and being very informative, especially on reference. It is highly valuable to us!\n\n1.\t**Experimental Designs Or Analyses**\n\nWe 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.\n\nSpecifically, 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.\n\nOn 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.\n\n2.\t**Questions For Authors**\n\nFollowing 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”.\n\nIn 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$. \n\nIt 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.\n\n3.\t**Weakness**\n\nWe 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.\n\nHybrid 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.\n\nThe 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.\n\n4.\tMissing metric of wall-clock time\n\nPlease see **Item 3 in Rebuttal 2** due to space limit.", "claims_and_evidence": "In my view most claims made by authors are well-supported.\n\nThe one missing metric is a wall-clock time, required by different methods to reach optimal solution or approximation to thereof. Without this data it is not possible to claim that the proposed method leads to faster PDE-constrained optimization as stated in the title.", "methods_and_evaluation_criteria": "Authors consider several challenging PDE-constrained optimisation problems and perform reasonable ablation study. I believe that both methods and evaluation criteria are adequate.", "theoretical_claims": "Authors do not provide novel theoretical claims.", "experimental_designs_or_analyses": "The most important part of experiments is the ablation study done by authors.\n\nThe proposed approach consists of three techniques: (i) training strategy and the overall design of architecture (inputs and output), (ii) Virtual Fourier Layer, (iii) inference strategy. Given that, it is reasonable to evaluate the effect of each component.\n\nQuite appropriately authors perform such ablation tests and present the results in Table 1. From this table I can conclude that corrections with reference at inference play the most significant role, the effect of training strategy is the second most important factor and the improvement from the use of Virtual Fourier Layer is the least important but still pronounced.", "relation_to_broader_scientific_literature": "PDE-constrained optimization is a mature field with a large number of classical techniques available https://link.springer.com/book/10.1007/978-3-319-13395-9.\n\nThe approach by authors is within a scope of surrogate modelling https://arc.aiaa.org/doi/10.2514/6.2000-4891. In the PDE-constrained optimization a repeated solution of PDE is required, so the idea is to replace accurate but numerically expensive classical solver with numerically cheap surrogate model. Various surrogate models were developed including ROM https://link.springer.com/chapter/10.1007/978-3-642-55508-4_16 and neural networks https://arxiv.org/abs/2111.04941, https://arxiv.org/abs/2110.13297.\n\nThe main innovation of the current contribution is a design of a special training scheme and network structure that allows one to use a numerical solver in conjunction with a surrogate model.", "essential_references_not_discussed": "Authors do not discuss development of surrogate modelling at all. In place of that they provide references for related problems: prediction with neural operators, training that promotes smoothness, hybrid solvers that combine neural networks and classical techniques. \n\nI suggest authors provide a brief review of PDE-constrained optimization methods that incorporate neural networks for surrogate modelling. Several references are provided above, but more thorough review will improve the overall quality of the contribution.", "other_strengths_and_weaknesses": "**Strengths:**\n1. The approach by authors is clearly more robust than alternative. This robustness is achieved by learning small corrections to the reference solution and by the use of classical solvers that can provide unbiased reference.\n2. The developed framework is general and can be used applied with arbitrary architectures\n\n**Weaknesses:**\n1. Classical solver is still required at the inference stage.\n2. More challenging collection of training data: (i) data is collected along optimisation trajectories, (ii) in addition to input-output pairs the derivative of the loss with respect to parameters is collected.", "other_comments_or_suggestions": "The article is generally well-written, so I have only a few minor points:\n1. The need for Figure 1 is not evident, since it does not explain much beyond the fact that optimization trajectory with approximate solver diverges from trajectory with more accurate solver. Figure 4 seems to demonstrate the same effect.\n2. Line 408, right column. \"without numerical solver feedback quickly derails.\" -> \"Without numerical solver feedback quickly derails.\"", "questions_for_authors": "1. Section 3.2. contains several fragments on bias and complexity of derivatives. For example:\n   1. Lines 187-188 \"For a transformer, the nonlinearity of attention unit causes complex calculation of derivatives ...\"\n   2. Lines 196-197 \"Fourier-based layers is practically linear, and therefore does not introduce additional biases ...\"\n   3. Lines 190-192, right column \"The derivative of equation (8) introduces less bias due to its simpler structure, which motivates us to adopt this form.\" \n   4. Lines 210-213 \"Since Fourier-based layer is practically linear, its derivative does not add additional bias to the layer. This is highly favorable for derivative learning.\"\n   I do not understand these claims. Can the authors please clarify what they mean by the \"complex calculations\" of derivatives? Do they refer to numerical complexity? Isn't it the case that numerical complexity of automatic differentiation (AD) is roughly the same as forward pass through the architecture?\n   Next, what do authors mean by \"bias\" of derivative? Again, since AD is used the derivative is computed up to machine precision for all but pathological cases. Since authors use fairly stable architectures and compute derivatives with respect to input, I would not suspect to see instabilities in AD.\n2. In equation (12) authors use $\\varphi$ and explain that for topological optimization this is $\\lambda_r - \\lambda_q$. From Appendix one can find that this is the case for all considered examples. Can the authors please provide an example where $\\varphi$ has a different form or reflect in the main text that $\\varphi$ is a difference of controlled variables for most of the problems.", "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": {"ethical_review_concerns": "na"}}, {"review_id": "UFzPA0Nmv2", "reviewer": "Reviewer_V7mT", "summary": "This paper proposed a novel framework to enhance PDE-constrained optimization (PDECO) with neural operators, addressing data efficiency and robustness. Key innovations include data-driven training, a Virtual-Fourier layer for improved derivative learning, and a hybrid optimization approach integrating neural operators with numerical solvers. Experiments show accurate operator learning, robust convergence, and improved optimization performance.", "questions": "", "limitations": "", "rating": 4, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "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.\n\nBelow are some frequently asked points we’d like to address:\n\n1.\tDataset sizes and Table 1\n\nMicroreactor, Fuelcell, Inductor, Drone have 100x12, 30x20, 100x10, 10x20 samples respectively as in num_trajectory x num_steps. We added the information in appendix.\n\nIn 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-”). \n\nTo 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.\n\n2.\tRobustness of improvement of R-VF\n\nThe 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.\n| Models | p | u | v | c| $L_s$ |\n|---|---|---|---|---|---|\n|R-LA| $1.29e-2 \\pm 3.3e-4 $ | $2.33e-2  \\pm 1.04e-3 $ | $6.85e-2 \\pm 1.94e-3$ | $2.32e-2 \\pm 9.3e-4$ | $2.16e-2 \\pm 3.55e-4$ |\n| R-PA | $1.25e-2 \\pm 4.1e-4 $ | $2.22e-2  \\pm 6.6e-4 $ | $6.6e-2 \\pm 1.23e-3$ | $2.28e-2 \\pm 2.1e-4$ | $2.02e-2 \\pm 7.8e-4$ |\n| R-VF | $\\textbf{1.2e-2} \\pm 3.2e-4 $ | $\\textbf{2.06e-2}  \\pm 6.1e-4 $ | $\\textbf{6.42e-2} \\pm 1.42e-3$ | $\\textbf{2.14e-2} \\pm 6.89e-4$ | $\\textbf{1.76e-2} \\pm 3.81e-4$ |\n\nAlso, 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.\n|     | p | u | v | c| $L_s$|\n|---|---|---|---|---|---|\n|R-LA-80| 1.64e-2 | 3.25e-2 | 7.83e-2 | 2.65e-2 | 2. 69e-2 |\n| R-PA-80 | 1.51e-2 | 2.74e-2 | $\\textbf{7.46e-2}$ | 2.48e-2 | 2.54e-2 |\n| R-VF-80 | $\\textbf{1.48e-2}$ | $\\textbf{2.57e-2}$ | 7.56e-2 | $\\textbf{2.42e-2}$ | $\\textbf{2.30e-2}$ |\n|R-LA-160| 1.40e-2 | 2.42e-2 | 7.16e-2 | 2.29e-2 | 2.37e-2 |\n| R-PA-160 | 1.32e-2 | 2.18e-2 | $\\textbf{6.81e-2}$ | $\\textbf{2.15e-2}$ | 2.30e-2 |\n| R-VF-160 | $\\textbf{1.31e-2}$ | $\\textbf{2.17e-2}$ | 6.95e-2 | $\\textbf{2.15e-2}$ | $\\textbf{2.18e-2}$ |\n\n\n3.\tWall-clock time for optimization with R-VF\n\nThe 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. \n\nOne 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.", "claims_and_evidence": "Yes.", "methods_and_evaluation_criteria": "Looks good to me.", "theoretical_claims": "No such", "experimental_designs_or_analyses": "All Look good to me.", "relation_to_broader_scientific_literature": "We need such optimization works using operator learning.", "essential_references_not_discussed": "The authors may add some reviews of papers from 'AI-aided geometric design of anti-infection catheters' and 'Physical Design using Differentiable Learned Simulators'", "other_strengths_and_weaknesses": "No such.", "other_comments_or_suggestions": "No such.", "questions_for_authors": "No such.", "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": "iQpFNnRMtp", "reviewer": "Reviewer_1Qkw", "summary": "The authors propose a neural-operator based approach for solving PDE-constrained optimization problems. While this approach is not new in the literature, the authors come up with several innovations to improve existing approaches: (i) data-driven training by using trajectories generated by traditional optimization algorithms, (ii) enhanced derivative learning through virtual-Fourier layers, and (iii) hybrid optimization integrating neural operators with numerical solvers. The authors show the potential of their approach on several datasets from shape and topology optimization.\n\n## update after rebuttal\nI thank the authors for responding to my comments. I have raised my score from 1 -> 2. I am hesitant to raise my score further because of the initial quality of presentation in the paper.", "questions": "", "limitations": "", "rating": 2, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "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.\n\nFNO 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.**\n\nThe 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). \n\nWe 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). \n\nTo address your other questions:\n\n1.\tSensitivity 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.\n\n2.\tWe 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.\n\n3.\tConsider 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.\n\nOn 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.\n\n4.\tIndeed, 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.\n\n5.\t**Claims and Evidence** \n\nThe 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.\n\nThe data efficiency of our approach is exemplified in Table 1. Due to space limit, please see discussion in **Item 1 of Rebuttal 2**.\n\n6.\t**Methods And Evaluation Criteria** \n\nWe 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.\n\nRegarding runtime, please see discussion in **Item 3 of Rebuttal 2**.", "claims_and_evidence": "I do not believe that the claims made in the submission are properly supported.\n\nThe authors claim in the abstract that their “extensive experimental results demonstrate the effectiveness of our model in accurately learning operators and their derivatives”. However, what is the threshold for accuracy? I am also skeptical of this claim because their method is outperformed by the numerical method in Figure 6. To me this suggests that the derivatives are not being learned to sufficient accuracy. Furthermore, I’m not fully convinced that experiments of four different datasets are enough to be considered extensive.\n\nAlso, how exactly do the authors address the data efficiency issues that they describe in the abstract? I think the authors should compare to other data sampling approaches (other than generating data from traditional optimization algorithms) to show that their method yields an improvement.", "methods_and_evaluation_criteria": "How much data is needed to train the neural operator in this framework? Is this data easy to generate? If the data generation itself is very expensive, would this method be useful for PDE-constrained optimization?\n\nAnother thing the authors should address is the runtime (in seconds) of their method, especially since it seems like the classic numerical method outperforms their approach in Figure 6. Does Hybrid R-VF at least have a faster runtime than the numerical method?", "theoretical_claims": "There are no theoretical claims made in the paper.", "experimental_designs_or_analyses": "The experimental design and analysis seem fine.", "relation_to_broader_scientific_literature": "To the best of my knowledge, the paper makes an interesting contribution to the broader literature on PDE-constrained optimization. In particular, the virtual-Fourier layer seems to be novel, and could be used in future work to improve the derivative accuracy in neural operators.", "essential_references_not_discussed": "The paper is missing references to DeepONet, which is also popular in the operator learning literature.\n\nThere is also work within the last two years that generalizes FNO to irregular meshes, e.g., geo-FNO (https://www.jmlr.org/papers/volume24/23-0064/23-0064.pdf) — the authors should address this in the related work section, where they claim that the applicability of FNO to irregular meshes is limited.", "other_strengths_and_weaknesses": "Combining neural operators with PDE-constrained optimization is an interesting idea. However, there is a lot of room for improvement in the presentation (see “Other Comments Or Suggestions” and “Questions For Authors”). I found section 3.2, which introduces virtual-Fourier layers, particularly hard to follow. I had a hard time distinguishing what aspects of the virtual-Fourier layer were being introduced by the work, and what aspects were just based off of the FNO.", "other_comments_or_suggestions": "Using $\\hat \\cdot$ to denote variables with ground-truth values could be confusing to readers skimming the paper. Typically $\\hat \\cdot$ is used to denote an estimate of a ground-truth value.\n\nI’d recommend adding a notation subsection at the end of the introduction. For example, the bold notation for tensors and the use of $\\mathcal F$ for the Fourier transform could be part of this subsection.\n\n“Softmax” and “Project” are italicized in Remark 3.1, but they are not italicized in the rest of the paper.", "questions_for_authors": "1. Page 2: The paper refers to “sequences of solutions and sensitivities” that are computed by the adjoint method. However, the mathematical expressions for the sensitivities are never shown in the paper. Presenting the mathematical expressions for these sensitivities (and perhaps, a brief overview of the adjoint method) will improve the quality of presentation in the paper.\n\n2. Page 4: What does it mean for the Fourier transform of Fourier-based layers to be “practically linear”. Does this mean that the Fourier transform of these layers is a linear operator, or that it has linear time complexity. This statement should be clarified in the paper. \n\n3. Page 4: What does it mean for the derivative to “not add additional bias to the layer”? Is this supposed to motivate why the authors use FNO rather than transformer-based architectures?\n\n4. Page 5: After the discussion of virtual-Fourier layers in section 3.2, the paper abruptly jumps to training and optimization with RNO. How are the virtual-Fourier layers related to the RNO? Are they the building block of RNO?", "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": "oZTEXbw9ky", "reviewer": "Reviewer_kQVE", "summary": "The paper improve on the reference neural operator (RNO) for PDE-constrained optimization (PDECO) by (1)  using the full trajectory data of PDECO and the sensitivity loss (2) introducing the virtual Fourier layer (3) using numerical solution to correct the neural operator solution.", "questions": "", "limitations": "", "rating": 2, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your review with carefulness and details, and it is highly valuable and helpful to us! \n\n1.\t**Claims and Evidence**\n\nNice 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`.\nThe 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.\n\n2. **Methods And Evaluation Criteria** \n\nWe 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.\n\nWe 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**.\n\n$u_{gt}$ indicates the numerical solution corresponding to the current $\\lambda$, NOT the numerical solution of PDECO. \n\nIt 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.\n\n3.\tReference not discussed.\n\nOn 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).\n\n[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].\n\n4.\tOther questions\n\nWe apologize for missing the information of dataset sizes. Please see **Item 1 in Rebuttal 2** for details due to space limitation.\n\nPlease see **Item 3 in Rebuttal 2** for discussion on wall-clock time for optimization.\n\nIt 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.\n\nRegarding 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.\n\nThe 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.", "claims_and_evidence": "The claims are generally supported by the evidence. \n\nSome clarification is needed regarding Figure 6:\n- Why is the objective increasing as the number of calls increases? From the appendix, this example (Microreactor2D) is a minimization problem. \n-  As mentioned in the caption,\"R-VF w/o reference represents the optimization process relying solely on the R-VF model without any calls to the numerical solver\". I think that means the num. of calls should be 0 for R-VF w/o reference? How do we get the blue line", "methods_and_evaluation_criteria": "The proposed methods make sense. While the evaluation require improvement. \n\n- Table 1 seems to come from 1 testing data set. The performance should be evaluated on multiple testing problems and report the mean/std.\n- From table 1, the improvement from VF also seems to be marginal. For example, comparing the best and the second best, the objectives (L) decreases by roughly 0.03, 0.02, 0.02, 0.01 etc. The improvement might not be robust when the number of test case increases.\n\nAnd some clarification are needed:\n- Algorithm 2, what is $u_{gt}$, is it the numerical solution corresponding to the current $\\lambda$? Or is it the \"ground truth solution to the PDECO\"?\n- Figure 4, how to make the predicted $u$ and $J$ exact (blue arrows) ? From my understanding of algorithm 2, we don't have access to the red trajectory.", "theoretical_claims": "There is no theoretical claim.", "experimental_designs_or_analyses": "See Claims And Evidence.", "relation_to_broader_scientific_literature": "The work improve on an existing method (RNO) for PDECO. While effective, the broader impact might be limited.", "essential_references_not_discussed": "The author should discuss how does their method of handling different geometry differs from [1]\n\nAnd how does their use of derivative differs from [2]\n\n[1] Li, Z., Kovachki, N., Choy, C., Li, B., Kossaifi, J., Otta, S., Nabian, M.A., Stadler, M., Hundt, C., Azizzadenesheli, K., Anandkumar, A., 2023. Geometry-Informed Neural Operator for Large-Scale 3D PDEs.\n\n[2] O’Leary-Roseberry, T., Chen, P., Villa, U., and Ghattas, O. Derivative-informed neural operator: an efficient framework for high-dimensional parametric derivative learning.", "other_strengths_and_weaknesses": "Strength:\n- the example problems are challenging.\n\nWeakness:\n- see other boxes.", "other_comments_or_suggestions": "No other comments.", "questions_for_authors": "Since numerical solutions are used to assist RNO, it's unclear that the hybrid approach is more effective than using numerical solver alone. In particular, the key questions for practitioner will be:\n- How many sample/trajectories are needed to generate the training data? \n- What is the total time for pre-training and what is the time for each solve?\n\nFrom figure 4, it seems that the performance stays flat after 4 function calls, this suggest that the accuracy can not be improved with more function calls for Hybrid approach. It's totally reasonable that the accuracy is lower than the numerical solver, but the author need to show that the speed makes up for the loss of accuracy.\n\nAnother important question is, does using the sensitivity make the Neural operator less general? An important advantage of NO is that, once pre-trained, it can be used for different objective function. With proposed method, it seems modifying the objective function would require retraining the whole neural operator. \n\n\nSince RNO is an essential part of the work, I think the author should review RNO with more detail, in addition to the high level introduction in section 3.3. Some related questions: \n- Are both $u$ and $\\lambda$ are nodal function on a mesh, or $\\Omega\\rightarrow \\mathbb{R}$ Is it really possible to have $\\varphi(\\lambda_r)=\\lambda_q$ and $u_r \\circ \\varphi ^{-1} = u_q$? \n- What is the training pipeline for RNO? How to represent and sample $\\varphi$?", "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:

Meta-review: This submission considers a novel neural-operator approach (improving on the reference neural operator approach) for PDE-constrained optimization (PDECO). The main ideas are a virtual-Fourier layer for learning derivatives and improved training by modifying the optimization scheme. In my opinion, the authors have sufficiently addressed the majority of the concerns brought up by reviewers during the rebuttal period (and scores have been raised as a result, giving final scores 4,4,2,2). Although there are still some concerns about presentation (e.g., choices of notation, convoluted explanations), overall the reviewers agree that the work is novel and interesting to the community. I therefore recommend to accept this paper and strongly encourage the authors to work on the presentation aspects of the paper.
Decision: Accept (poster)