PACE: Pacing Operator Learning to Accurate Optical Field Simulation for Complicated Photonic Devices
Authors: Hanqing Zhu
Wenyan Cong
Guojin Chen
Shupeng Ning
Ray Chen
Jiaqi Gu
David Z. Pan
Venue: NeurIPS 2024 poster (2024)
Keywords: AI for Scienece
Optical simulation
Neural opeartor
AI for PDE
Primary area: machine_learning_for_physical_sciences

Abstract: Electromagnetic field simulation is central to designing, optimizing, and validating photonic devices and circuits. 
However, costly computation associated with numerical simulation poses a significant bottleneck, hindering scalability and turnaround time in the photonic circuit design process.
Neural operators offer a promising alternative, but existing SOTA approaches, Neurolight, struggle with predicting high-fidelity fields for real-world complicated photonic devices, with the best reported 0.38 normalized mean absolute error in Neurolight.
The interplays of highly complex light-matter interaction, e.g., scattering and resonance, sensitivity to local structure details, non-uniform learning complexity for full-domain simulation, and rich frequency information, contribute to the failure of existing neural PDE solvers.
In this work, we boost the prediction fidelity to an unprecedented level for simulating complex photonic devices with a novel operator design driven by the above challenges.
We propose a novel cross-axis factorized PACE operator with a strong long-distance modeling capacity to connect the full-domain complex field pattern with local device structures.
Inspired by human learning, we further divide and conquer the simulation task for extremely hard cases into two progressively easy tasks, with a first-stage model learning an initial solution refined by a second model.
On various complicated photonic device benchmarks, we demonstrate one sole PACE model is capable of achieving 73% lower error with 50% fewer parameters compared with various recent ML for PDE solvers.
The two-stage setup further advances high-fidelity simulation for even more intricate cases.
In terms of runtime, 
PACE demonstrates 154-577x and 11.8-12x simulation speedup over numerical solver using scipy or highly-optimized pardiso solver, respectively.
We open-sourced the code and *complicated* optical device dataset at [PACE-Light](https://github.com/zhuhanqing/PACE-Light).

PACE:
Pacing Operator Learning to Accurate Optical Field Simulation for Complicated Photonic Devices
Hanqing Zhu Wenyan Cong Guojin Chen ; Shupeng Ning",
Ray T. Chen Jiaqi Gu', David Z. Pan
Arizona State University
The University of Texas at Austin
aagahu@utexas.edu, 'dpanßece. utexas. edu
Abstract
Electromagnetic field simulation central to designing, optimizing, and vali- dating photonic devices and circuits. However, costly computation associated with numerical simulation poses a significant bottleneck, hindering scalability and turnaround time in the photonic circuit design process. Neural operators offer promising alternative, but existing SOTA Serroache,,NeuOOigge. predicting high-fidelig fields for real-world complicated photonic devices, with the best reported 0.38 normalized mean absolute error in NeurOLight The inter- plays of highly complex light-matter interaction, c.g., scattering and resonance, sensitivity to local structure details, non-uniform learning complexity for full- domain simulation, and rich frequency information, contribute to the failure of existing neural PDE solvers. In this work, we boost the prediction fidelity to an unprecedented level for simulating complex photonic devices with novel oper- ator design driven by the above challenges. We propose a novel cross-axis fac- torized PACE operator with a strong long-distance modeling capacity connect the full-domain complex field pattern with local device structures. Inspired by human learning, we further divide and conquer the simulation task for extremely hard cases into two progressively easy tasks, with a first-stage model learning an initial solution refined by second model. On various complicated photonic de- vice benchmarks, we demonstrate one sole PACE model is capable of achieving 73% lower error with 50% fewer parameters compared with various recent MI. for PDE solvers. The two-stage setup further advances high-fideligy simulation for even more intricate cases. terms of runtime, PACE demonstrates 154-577x and 11.8-12x simulation speedup over numerical solver using scipy or highly- optimized pardiso solver, respectively. We open sourced the code and compli- cated optical device dataset at PACE-Light.
struggle with
Introduction
With advances in integrated photonics, photonic structures capable of transmitting processing in- formation are gathering increasing interest, fueled by the optical communication [24] and the recent resurgence of photonic analog computing [5. 20, 23, 36, 39, 40]. Light engpwerep communication and computing offer a promising pathway for reshaping future Al systems, prompting the optical community to discover compact, customized devices [7, 32, 38] overcome the limitations of bulky optical components. In this optical design process, numerical simulators, e.g., the popular finite dif- ference frequency domain (FDFD) algorithm [11], is heavily used to obtain accurate optical fields for characterizing and optimizing device behavior. However, the significant time and computational
"This work was done when Guojin Chen was visiting scholar at UT Austin.
38th Conference on Neural Information Processing Systems (NcurIPS 2024).

---

(e) Complicated Right-matter Interaction
(6) Non-uniform loarning complexity
(e) Framework (f. P..P.
A'
Local Rescoramcs
Long missing fill with oiverse inner structure
Permittivity
Input source H
(b) change Largely
(d) Frequency Info
Energy Spechum
Weve prior (Px (P..Ps)
104
ML-based Operator
Siiphilyy changed port
Highty- difigrent hote
Hyfr m))
Wavenumber
law treq
high freq
Figure I: Challenges of complicated optical device simulation: (a-d) and learning framework (c).
costs associated with Maxwell partial differential equation (PDE) simulations, exacerbated by the need for finely tailored meshes and numerous simulation runs for iterative optimization, pose sub- stantial bottlenecks in the design loop.
Recently, neural PDE solvers [3, 8, 9, 15, 16, 27, 33] have emerged as promising surrogate models for fast and accurate PDE solving. NeurOLight [8] represents the state-of-the-art (SOTA), tending neural operators to parametric photonic device simulations in a physics-agnostic manner. However, it still exhibits large emrors in simulating real-world complicated optical devices, reporting a 0.38 normalized mean absolute error on the etched multi-mode interference (MMI) device [26]. One may wonder what the major challenges are, given the successes of neural operators in many scientific PDEs. Firstly, for complicated devices, the permittivity distribution is discrete and highly contrasting, transforming the Maxwell PDE into multi-scale problem [1], further leading com- plex light-matter interactions such as scattering and resonance, as illustrated in Fig. (a)) Secondly, their optical fields highly sensitive local structural changes; even minor alterations can signif- icantly impact the field, as depicted in Fig. (bb). Moreover, with diversifying field patterns along the light propagation path, it shows non-unine non-uniform learning complexity especially in regions distant from the input light source. Finally, a spectral analysis provides insights into the frequency-doman challenges, illustrated in Fig. I(d). Unlike simpler systems where low frequencies dominate (e.g., Darcy flow shown in Fig. 10), complicated devices exhibit rich frequency spectra with high- frequency components. This diversity underpins the difficulty faced by previous neural PDE solvers in accurately simulating complicated photonic devices, supporting the assertion in [15] that no single model can universally solve all types of PDEs.
In this work, we tackle the challenging real-world complicated optical device simulation problem. We vastly boost prediction fidelity and keep 154-577 and 11.8-12x speedup over traditional nu- merical solver [1 [1]] on a 20-core CPU with scipy or highly-optimized pardiso solver, respectively.
Overall, make the following key contributions:
We introduce a novel cross-axis factorized PACE operator backbone, effectively capturing complex physical phenomena across the full domain parameter-efficient manner.
We employ a divide-and-conquer approach inspired by human learning for extremely chal- lenging cases, with a first-stage PACE-I to learn a rough approximation of the optical field, refined by second-stage PICE-II.
On various complicated device benchmarks, one sole PACE significantly outperforms base- lines, achieving 73% lower error with 50% % fewer parameters. Even compared to the best baseline, it lowers prediction error by over 39% with 17% fewer parameters. Our two- stage method further advances high-fidelity simulation for extremely challenging cases.
We open-source the complicated optical device datatset and code at PACE-Light to facili- tate AI for PDE community.
2 Preliminaries
2.1 Neural Operators for PDE
Recently, neural operators have emerged as a novel approach for developing machine learning models aimed at solving partial differential equations (PDEs). These models focus on learning

---

the mapping between the function spaces in a purely data-driven fashion. This holds the general- ization capability within a family of PDEs and can potentially adapted different discretiza- tions. Various function bases are utilized to build the operator learning model, such as the Fourier bases [16, 29, 2, 8), wavelet bases [9], spectral method [33], and attention layer [15, 3, 14]. These models have demonstrated remarkable performance and efficiency in solving specific types of prob- lems, often achieving record-breaking results in certain applications. Despite their successes, it's important to recognize that the field of PDEs encompasses a wide variety of equations, each with its own unique properties and characteristics, As pointed out in recent research [15], there is no guarantee that single type of data-driven model can effectively address all types of PDEs.
2.2 Optical Field Simulation with Machine Learning
Analyzing the propagation of light through optical devices is crucial for the optimization and design of photonic circuits. For a linear isotropic optical device, with a time-harmonic continuous- wavee light beam shining on its input port, we can obtain the steady-state electromagnetic field distributions E(r) = xEx + EE, + EE, and H(r) = xH, + H + 2H, by solving the steady-state frequency- domain curl-of-curi Maxwell PDE under absorptive boundary conditions
[11],
((15 V x Vx) cor, E(r) jwJ,(r), (V x (c, (r)Vx) - )o H(r) = juJ (r) (1)
where V x the curl operator, Pu is the vacuum magnetic permeability, en is the vacuum electric permittivity, e, is the relative electric permittivity, and Jm and J. the magnetic and electric cur- rent sources, respectively. The finite difference frequency domain (FDFD) method, a widely adopted numerical technique detailed in [11], is used to discretize these continuous domain equations into an M x N mesh grid. This transforms the Maxwell PDEs into a linear system AX = Solving this system with large sparse matrix A E M is computationally expensive and challenging to scale. Although improvements have been made, such as replacing the scipy solver with the more efficient pardiso solver, the process remains prohibitively costly for large-scale applications.
Building neural networks (NNs) to accelerate this tine-consuming simulation process has been in- vestigated in predicting some key design parameters [26] or the entire optical field [30, 17, 4, 8]. NeurOLight ght extends the neural operator to optical field simulation, enabling learning a physics- agnostic parametric Maxwell PDE solver and achieving SOTA accuracy, while its performance real-world complicated photonic device still not satisfactory.
3 Understand the Problem Setup and Challenge
In this study, we aim to build physics-agnostic neural operator V for parametric photonic device simulation in a data-driven fashion to approximate the ground-truth Maxwell PDE solver V A + U described in Eq. (1), Here, 2/ represents the solution space for the optical field in COxds and A (2, er, 4r,4, J) represents the observation space of the Maxwell PDE, both defined over the continuous 2-D physical solving domain Q (le, 42). We follow NeurOLight [8] to discretize the simulation domain Ω as Ω (M, Alz, Als) with adaptive mesh granularity, i.e., with grid steps Al, //M and Al, = L,/N. Moreover, (2, Ey, w) the raw observation A is encoded as informative wave priors, P3 e ""x"rr 1:" Al, and Pa ed "*x"rr =1 Al, where 2 (0, 1, ... .N 1). reflecting the propagation behaviors of light through different media. The input light source J further modeled as a masked light source field H.
1, M 1) and
Therefore, as illustrated in Fig. I(e), the overarching objective formulated as learning operator V that maps At (c. H Pay P.) to the target field 2/ by optimizing the empirical error,
0" min Ea [[Va(()) u)],
3.1 Challenges in Predicting the Light Field of Complicated Photonic Devices
NeurOLight [8] delivers a pioneering effort extending neural operators to the simulation pho- tonic devices, achieving SOTA accuracy. However, it still yields significant errors, particularly for real-world complicated devices, with a reported 0.38 normalized mean absolute error for etched MMI device [26, 12]. This Ieads to interesting reflection: despite the successes of neural
3

---

operators in solving scientific PDEs, why do they still fall short in complicated photonic device sim- ulation? Below, we provide detailed analysis that highlights the underlying learning challenges.
Complicated light-matter interaction in the optical field of real-world photonic device. Permittivity cr, a critical parameter in photonic devices, greatly impacts how light propagates through media. Designing new devices often involves manipulating the e, distribution across the domain. However, due to manufacturing limitations, e, changes are discrete rather than smooth. Moreover, researchers explore patterning materials with highly contrast permittivity to design compact devices [26, 31J. This discrete and highly contrasting permittivity transforms the Maxwell PDE into multiscale PDE problem [1], with complicated light-matter interactions such as scattering resonance happening, shown in Fig. 1 (a), which has been shown difficult to predict from both scientific computing and operator learning perspectives 35J.
Significant prediction field variations from minor structural changes. Due to the complex hight-matted interactions within the field, even a slight change in the photonic structure can result in drastically different optical fields under the same input conditions, as shown in Fig. 1(b). This calls for a powerful backbone model that is capable of building the relationship between local rival changes with the global optical field transition.
Non-uniform learning difficulty along the spatial domain. As shown in Fig. 1(c), with light shining in from a specific position and direction, it propagates through the media, resulting in non-uniform learning difficulties along the spatial domain. Due to the vast diversity of potential internal structures along the light propagation path, the light patterns are becoming highly diverse. Consequently, the data collected for training also incorporates the same phenomenon where many similar patterns are seen during training near the input sources, whereas the model faces more diverse patterns at greater distances. This makes it hard for the model learn how to predict further regions, especially when the domain is elongated. This issue is analogous to the roll-out eтror encountered in temporal PDE modeling at the large time steps.
Rich frequency information lies in the predicted field. We show the energy spectrum of the optical field in the frequency domain in Fig. (d). The field, characterized by complex interac- tions such as scattering and resonance, exhibits rich frequency information, unveiling the learning complexity from frequency-domain analysis. This confirms the usage of high-frequee frequency modes in NeurOLight, underscoring the need for a parameter-efficien robust, and powerful backbone model to resolve the parameter efficiency and overfitting issue with large modes.
4 Proposed PACE Methods
In this paper, we follow the standard operator learning model architecture as
a'(r) + (o(r) + 01 (r) + -r.. - u(r), E
(3)
We start with the convolutional stem used in [8] to project the PDE observation at (r) into higher- dimensional feature space of dimension C. This is followed by a sequence of K cascaded neural operator blocks, which gradually reconstruct the complex optical field within the C' dimensional space. At last, a head with two poset-wise convolutional layers projects the (r) to the optical field space u(r). Fig. 2(a) shows the proposed PACE neural operator block structure, formulated as,
V (r) == PPN ((CC)))) + vn) Uk, Vr 92 e (r) pre-norm (r)),,
(4)
where K. is the our proposed PACE operator and FFN(-) is feedword network used in [8]. To stabilize the model performance when scaling to deeper layers, we add pre- pre-normalization [34] and follow [13] to add double skip. In this work, we consistently use the NeurOLigh operator in the first two blocks to align our model with the horizontal and vertical wave prior encoding method adopted from NegrtL which we found slightly improves our accuracy.
4.1 Parameter-efficient and Effective Cross-axis Factorized PACE Operator
The neural operator design is key to obtaining satisfactory accuracy on a given PDE task. With the well-discussed challenges in Sec. 3.1, we derive key insights that have guided the development of our PACE operator in Fig. 2(b): (1) Long-distance full-domain modeling capacity, especially effectively modeling how local features impact the whole domain; (2) Isotropic model architecture with no

---

(x(+) E R
Self-weigthed Path
Pre-norm
High-freq
Projection
PACE
Mnit
Cross-axis factorizad intagial kormel
groue
R Chelted
Ckofet
FFN
(x)1(() RR
PACE Operator
Figure 2: (a) PACE block with double skip and pre-normalization; (b) Our cross-axis factorized PACE operator.
down-sampling/ patching without losing local details; (3) Parameter efficiency under the needs of capturing high-frequency features.
Given the isotropic requirements, an operator based on Fourier bases is an ideal candidate as achieves full-domain attention in the ((dogn) time complexity. However, the rich frequency
information lying in the optical field requires the
of large frequency modes, making the FNO [16] with huge parameters and severe overfitting issues. NeurOLight [8] and Factorized FNO [29] propose to decompose the FNO block with independent I-D FNO blocks in the full N-dimensional domain & (see Fig. 3), therefore, solving the parameter concern when utilizing high-frequenc, modes and serving as a regu- larization for overfitting. The only difference between NeurOLi ght and Factorized FNO [29] is whether they chunk the input copy the input to the independent 1-D FNO block. We argue that their theoretical success is attributed to the implicit full-domain integration Corollary 4.1.
E
R E ChaGG
R" Chop,
Frr
Figure 3: Factorized FNO [27, 8].
Corollary 4.1. The factorized Fourier integral operator K. [29, 8] factorizes the original Fourier integral operator [ [16] along each dimension n in the N-dimension domain Ω
(ICoe)(r1) = FF (r2)))((r)), Vri E
(5)
where each item explicitly computes a 1-D kernel integral, Sa, s(r1, r2)" plicitly implements full-domain kernel integration in Q by stacking K. i.e. Ko o K1 ...,,
It im-
However, the reliance on implementing full-domaie integration with multi-layers makes them weak operator candidates to achieve our first requirement, i.e., a strong model that is capable of building full-domain modeling between local structures with the global fields.
Proposed cross-axis 2-D factorized integral kernel. Aware of the above shortcomings of previous factorized FNO variants, in our 2-D domain, we propose to factorize the full domain integral in cross-axis way along the horizontal (h) and vertical (v) axis:
(Kz))(r1) (())) Vr: & Ja.
(6)
M((1,)2)* Sa,
(ra))
Vra e Q
This factorization enables an explicit factorized full-domain integration. provides a strong way to capture the relationship between points in the domain Ω. building the relationship between local structure with the complicated field pattern. The implementation the above cross-axis integral can be efficiently implemented by Fourier Transform F(-) when the kernel k(T1, k(rr r2),
as follows.
(Kox)(ri) (no))())=FF FF (F. (r(r)))))r1,, Vr. E 2
in nlogm complexity (n MN in our 2-D cases with Ω CMXN).
Group-wise cross-axis integration. For input r with a channel dimension C, can be viewed as the sampling of set of functions {r) on grid point in the 2-D discretized domain ? Ω2 x
5

---

Cross-stage feature distillation
Sten
PACE
Hoad
Siem
PACE 9
A'= = (er, Pa,
((Vata),
(6., Vai(a))
A(9n(s. Va.), P())")
The proposed cascaded learning flow with two stages. The first stage learns an initial and rough solution. followed by the second stage to revise it further. A cross-stage distillation path is used to transfer the learned knowledge from the first stage the second stage.
Figure 4:
Og. The Iearnable integral kernel intrinsically performs information exchange along different grid points in f2 Similar to the multi-head design in Transformer, which assumes different heads extract different information, we can also partition the C. basis functions into g disjoint sub-groups and feed each sub-group through our cross-axis factorized kernel. This grouping further reduces the number of parameters to (KA + k.) x C.C. showing significant parameter reduction compared to FNO (Kh, Ky X C,C,) and Factorized FNO ((KA + Ky) x C.C,), showing excellent parameter efficiency when utilizing large frequency modes is a must. We do an ablation study in Appendix A.4 to investigate the choices of different group g. where we find g 4 strikes the best between parameter efficiency and model performance.
Explicit projection unit & for extracting high frequency information. The optical field shows rich information in the frequency spectrum, reciting a special care of high-frequency information. Besides utilizing high-frequency modes, we propose to add an explicit projection module before the cross-axis integral, which is very simple as one linear layer followed by non-linear activation, given non-linear activation is known to help generate high-frequency features [22].
Self-weight weighted path for enhanced instance-based local feature attention.
The optical field's re- sponse is intricately linked to the minute variations in different photonic device structures. A self- weighted path introduced ensure the model can pay different attention to regions of significant influences for varying device structures. An instance- instance-based weight is generated by passing the fea- ture map after the projection unit through a linear layer and Sigmoid unit, and then multiplied with the results after the cross-axis integral unit to provide instance- based attention.
Overall, the above ingredients assembled together Fig. 2 (b), which implements a self-weighted 2-D cross-axis factorized integral transform.
proposed PACE operator, shown in
4.2 Cascaded Learning from Rough to Clear
With the effective PACE operator design, the prediction fidelity can be largely improved by only us- ing 12-layer PACE model (see Section. 5.2.1). But for some complicated benchmarks (e.g., etched MMI 3x3/5x5), it still yields ~ 10% mean squared error, which is not satisfying. straightforward solution might involve scaling up the model size, expecting additional layers would enhance perfor- mance. However, as demonstrated [27], scaling to deep layers shows saturated performance after exceeding a specific number.
Existing ML for PDE solving work typically learns model a one-shot way by directly learning the underlying relationship from input-output pairs Unlike AI systems, humans don't learn new and difficult tasks in a one-shot manner; instead, they learn skills progressively, starting with casier tasks and gradually moving harder ones. For example, instead of directly learning how to solve equations, students first learn basic operations, such as addition and multiplication, and then move on to solving complex equations.
Hence, inspired by this human learning process, unlike previous work that directly learns one-stage model, we propose to divide the challenging optical field prediction problem into two sequential la- tent tasks. The first task, undergoing the same problem setup as discussed Sec. 3, could predict an initial, rough optical field based on the less informative raw PDE observation (we only have the light source and device permittivity distribution). Then, the successive second task could refine the rough prediction further by capturing more details and nuances, by accepting the predicted field Va and

---

device permittivity e, as the imput. Therefore, we assign higher Fourier modes to enable sufficient capacity. The divide-and-conquer way results in a cascaded two-stage model architecture, as shown in Fig. 4. The cascaded learning model is trained jointly (PACE-I PACE-II) with the optimization target as the sum of two losses o(Vor 01 (a), u) [(V oo (a): fy), u), where the first L(Vor (a); serves as intermediate supervision that enfores the first stage model condensate the learned knowl- edge. To better connect the two-stage model, we propose a cross-stage feature distillation path to distill learned feature from the previous stage to the last by using simple Linear -Siggo path.
5 Experimental Results
5.1 Experimental Setup
Benchmarks: We evaluate our methods on real-world complicated photonic devices that pose sig- nificant simulation challenges for ML surrogate models. This includes the Etched MMI with ran- domly placed rectangular cavities, used in [8], and the metaline device [37, 19] featuring two layers of randomly dimensioned meta-atoms. These devices present a highly discrete and contrast per- mittivitty distribution and complex light-matter interactions, making them ideal for testing the effectiveness of our model. We generate our datasets using the open-source 2-D FDFD simulator, Angler [11], with generation details in Appendix A.1.
We evaluate the proposed PACE model against a range of baselines, including the SOTA neural operator work, NeurOLight [8], for optical simulation. We also include representative op- erator learning models for scientific PDEs based on Fourier bases(FNO [16], Factorized FNO (F- FNO) [28, 291, U-NO [2], tensorized FNO (TFNO) [13]), attention kernels [15], and the latent spectral method (LSM) [33]. We also incorporate UNet [17, 4] and Dilated ResNet (Dil-ResNet) [25]. For fair comparison, we keep model size budget of under/near 4 million (M) parameters for baselines, except LSM [33] where the original implementation is adopted. Details on model configurations are in the Appendix A.3.
Baselines:
Training setting and metric: All models undergo training for 100 epochs using the AdamW opti- mizer with weight decay of 1e-5 in batch size of 4. To balance the optimization among different fields, we use normalized mean squared error (N-MSE) as the learning objective,
(a)) & (a)) (E(a) v \/Il" (a))"
(8)
We don't t use the previously-used mean absolute error (MAE) [8] as the metric given for complex- valued optical fields; argue that L2 distance is accurate metric to evaluate the distance in the complex plane with a detailed analysis in Appendix A.6. We adopt the superposition- mix-up technique [8] to generate input light combinations randomly to augment training data.
5.2
Main Results
5.2.1
Prediction Quality of Single PACE Model
In Tab. 1, we compare our 12-layer PACE model with various baselines on multiple real-world de- vice benchmarks, showing significant 73.85% % smaller test emror with 51.67 % fewer parameters average. Notably, even when compared to the best baseline, 16-layer NeurOLight we show over 39% lower test error with over 17% fewer parameters. Given the challenge @ that trial structure change can totally change the optical field, model relying on downsampling or patching fails to capture the local details, confirming the failure of the UNet and Transformer model. Moreover, the challenge 1 and challenge @ call for a powerful model with long-distance modeling capability. Al- though Dil-ResNet utilizes a dilated block to enlarge the receptive field, it is insufficient for large domain, validated by the result that it shows much better accuracy on the small Metaline than the etched MMI3x3. Capturing long-range dependency with the Fourier operator provides an efficient way to the isotropic model without any downsampling, therefore making the Fourier-operator type model show consistently better accuracy than other baseline methods. However, due to the chal- lenge O that there is rich frequency information in the predicted field, FNO-2d falls short due to the impediment of utilizing large modes given the large parameter count. We also compared it with the tensorized FNO 2d. However, we find the general tensor decomposition hurt the accuracy of this challenging task. NeurOLight shares a similar insight of Factorized FNO by factoring Fourier ker- nel with several independent 1-D Fourier kernels; however, as we argued before, it fails to establish
7

---

Table Comparison # parameters, training error (last epoch), and test error on three benchmarks among our PACE and various baselines. We use geo-means to report overall improvements across different benchmarks.
Benchmarks
Model
#Params (M) $
Train Hir (10-
Test Hir (10-2)
UNet [17.4]
3.88
63.03
65.32
Nil-ResNet [25]
4.17
51.79
Riched MMI 3x3
Attention-based model [15]
3.75
70.05
69.85
U-NO [2]
4.38
34.22
42.86
Latent-speetral method [33]
4.81
55.07
55.16
FNO-2d
3.99
32.51
38.71
'Tensnrized FNO 2d [13]]
2.25
35.52
36.61
Factorized FNO-2d [29]
4.02
24.2
32.81
NeurOLight [8]
2.11
15.58
17.21
PACE
1.71
9.51
10.59
UNet [17.4]
3.88
65.73
66.01
Etched MMI 5x5
Altention-bused model [15]
3.75
74.16
74.20
U-NO [2]
4.38
37.92
42.24
-ppeett method [33]
4.81
53.9
54.01
FNO-2d [16]
3.99
33.12
36.49
Tensorized FNO-2d [13]
2.25
39.11
39.45
Factorized FNO-2d [29]
4.02
22.18
26.06
NeurOLight [28]
2.11
18.04
17.41
PACE
1.71
11.66
11.91
UNet [17,4]
3.88
39.12
39.61
Dil- Dil-ResNNt [25]
4.17
12.37
13.20
Metaline 3x3
Attention-based model [15]
3.75
63.99
64.10
U-NO [2]
4.38
19.27
22.09
Latent-spectral method [33]
4.81
31,60
31.94
FNO-2d [16]
3.21
19.73
20.88
Tensorized FNO-20 2d0dddir [13]
1.58
30.60
31.04
Factorized FNO-2d [29]
2.68
9.28
YeurOLight [8]
1.49
6.76
6.09
PACE
1.24
3.32
2.82
Improvement over best baseline NeurOLight [8]
-17.70%
-1.23% -72.57% %
-39.03%
Improvement over all baselines
-51.67%
-73.85%
Table Comparison between our two-stage model and simply scaling more layers. All models use the same Fourier modes sctup.
Benchmarks
Model
Cross-stage dist.
#Params (M)
Train Err (10
Test Err (10-?)
PACE- 12 layer
1.73
9.51
10.59
PACE-20 layer
3.135
6.46
7.04
Etehed MMI 3x3
PACE-I PACE-II
3.151
4.66
5.83
PACE-I + PACE- II
V
3.151
4.14
5.32
PACE- 12 layer
1,73
11.66
11.91
PACE-20layer
3.135
7.74
7.88
Etehed MMI 5x5
PACE-I PACE-II PACE-I + PACE-IL II
X
3.151
6.17
6.78
3.151
5.43
6.15
strong full-domain modeling capacity by linking local details the global complex field. Overall, our PACE block benefits from a physically meaningful cross-axis Fourier kernel factorization, equip- ping the capacity to capture full-domain dependency in parameter-efficient way. Visualization of predicted results is in Appendix A0.0.
5.2.2 Quality Improvement with Two-stage Model
We further compare the proposed cascaded two-stage model with the common practice of solely increasing # layers. We set the PACE-I 12-layer PACE model with Fourier modes(#Mode =70, RMMode=4), and PACE-IP II as 8-layer PACE model with larger Fourier modes (#Mode =100, #Mode =40). As shown in Tab. 2, the two-stage setup introduces slight overhead for one extra set of stem and head but shows clear margin over only increasing the number of layers in terms of both train emror and test error. The cross-stage feature distillation further provides meaningful guidance by transferring learned features to the second-stage model, leading to the best accuracy for the two- stage setup- Appendix A.7, we also show that the cross-stage distillation trick can improve model accuracy, similar to more costly training setup, by training the two-stage models sequentially.

---

5.2.3 Speedup over Numerical Tools
To develop a fast surrogate ML model that can replace the Maxwell PDE solver, it's crucial to evaluate the speed-up of our PACE model compared to
the FDFD numerical simulator Angler [[]]. We vary the
10P
simulation domain size and set the grid step to 0.05 nm,
577.7 x
sealing the discretized size pardiso linear solvers, re-
5.1
spectively and number of frequency modes to ensure the model has sufficient capacity capture the entire sim- ulation domain. For comparison, we employ a 20-layer joint PACE model. As shown in Fig. our PACE model achieves a speed-up of 150-577 x and 12x over Angler on 20-core Intel i7-12700 CPU using the scipy and We further set larger simulation granularity, 0.075 nm, to check speedup if we tolerate simulation quality loss in commercial tools. However, we find that setting a larger granularity results in a significantly different field, as qualitatively shown in reb-Fig.3, with a corresponding
Kumiiee
10a
154.2x
12.0x
10.6
8
10
PACE
Sim- (0075)
Sim.-P -0
10
Domain Size (jinr)
Figure Speedup PACE over angler [11] using scipy (SV pardiso (P) with simula- tion granularity (0.05mm) and (0.075mm).
N-MSE error of 1.2. Even though this case, PACE still shows 5.1.1..6x speedup over pardiso- based Angler with much better fidelity.
5.3 Discussion
Cross-axis PACE block design choices. within the PACE operator to assess their effectiveness. The self-weighted path, which provides instance-specific weights, significantly improves model accuracy across various photonic de- vice patterns. Removing this compo- nent results in a 17% increase error, highlighting its importance. Similarly, eliminating the high-frequency projec- tion unit leads to 23% worse error, phasizing its crucial role in capturing higg-fryy features. To further illustrate this, we visualize the feature maps in the frequency domain before and after applying the nonlinear activation in the high-frequency projection unit. As shown in Fig. 11, the nonlinear activation effectively amplifies high- high-frequency components, supporting our claim and validating the design decision incorporate an additional high-frequency projection path. Lastly, we replace our cross-axis Factorized integral kernel with a recent tensorized FNO (TFNO) [13] (tucker decomposition with rank 0.02). While TFNO effectively models long-range dependencies, matching our parameter count required aggres- sive decomposition, which significantly degraded performance. This comparison underscores the advantage of our physically grounded cross-axis factorized kernel.
In Tab. 3, we independently alter individual components
Table 3: Model design ablation on Metaline dataset.
#Params (M)
#Train Err (10-3),
#Test Err (10-3)4
Variants
8-layer PACE
0.82
5.65
4.82
No self-weighted path
0.8
6.33
5.66 (40.84)
No projection unit
0.8
6.58
5.97 (+1.15)
Use TENO
1.06
10.80
9.51 (+4.69)
Generalization to -nttrf-distribution testing. As operator model that is parameter-agnostic,, important to test the generalization for out- of-distribution data with unseen parameters. We re-generate photonic devices with different device configurations (size, etched region, etc.)
Test N-M
0.21
Sean A training Unooen A training
0.16
and unseen frequencies in our interested wave- length range (1.53-1.565 pam), i.c., C-band. As shown in Fig. 6, our PACE model generalizes well on unseen simulation frequency and new devices. It is vital test to prove the useful- ness PACE in helping device design within an interested wavelength range. We also test the accuracy outside the C-band, where PACE shows good accuracy on neighboring wave- lengths while holding a 10- 15% error at further range. This is expected since wave propagation sensitive to frequency. It can be mitigated by incorporating sampled wavelengths into training.
0.11
Interssted G-band rarge
0.06
0.01
1.50
1.52
1.54
1.56
1.58
1.60
Wavelength 2 (fI))
Figure 6: Generalize to unseen wavelength in inter- ested C-band (1.53-1,565) and outside C-band.

---

10a
NeurOLight
10*
PACE
10
10"
10
Large error in
free
targed Fall track high-
output
10a
terget
output
50
100
150
200
50
100
150
200
Figure 8: The radial energy spectrum of predicted fields from NeurOLight and PACE. NeurOLight fails to align precisely with the targeted field in both low-frequency and high-frequene parts.
Are PACE general enhancer module for Fourier-type operator? We further investigate whether our new PACE operator is general enhancer for other Fourier operators, rather than dedicated mod- ule for our own model architecture. We randomly insert four PACE blocks into Facztoried FNO [27] and test the error on Metaline3x3 and Etched MMI 3x3 benchmarks, showing up 28% error duction as shown Fig. with much fewer parameters.
Comparison with operator for multi-scale PDE. Notic- ing that our problem shares similar complexities in solv- ing multi-scale PDEs with neural operator [18, 35], we further compare our approach with the recent method [35] that alternates Fourier operator with dilated convolution layer to better capture local details. On the etched MMI 3x3 dataset, we implement 14-layer model with alter- nating NeurOLight block and dilated convolution layer. It yields a 1.73 M parameter count similar to our PACE but shows a 17.4 N-MSE error, much worse than ours (10.59).
40.0
4.0
Test Em. (le-2)
30.0
3.0
Param (M)
48%
20.0
2.0
10.0
0.0
0.0
Figure 7: Insert 4 PACE module (0: g=2; g=4) randomly in Factorized O
Spectrum of the predicted field: The predicted field spectrums of PACE and NeurOLight are in Fig. 8, Although NeurOLight uses the same frequency modes, it fails to align well with both the low-frequency and compared to NeurOLight,
regions. PACE excellently aligns with the baseline spectrum
6 Conclusion
In this work, we pace the simulation fidelity on highly challenging complicated photonic devices to unprecedented level. Our novel cross-axis factorized PACE operator enables the neural PDE solver to capture complex relationships between local device structures and the resulting complex optical field across the entire simulation domain. Furthermore, learning paradigm to further enhance the prediction quality when one sole PACE is not sufficient, demonstrating better quality enhancement than simply adding more layers. Experiments demon- strate that PACE achieves remarkable 73% reduction in error with 50% fewer parameters compared previous methods. Our method also offers significant speedup (11.8x 577x) over traditional numerical solvers. Looking forward, we aim to integrate our model into the design optimization loop for photonic devices and circuits. Moreover, we want emphasize that our proposed operator and learning strategy are not dedicated to photonic cases but generally applied to challenging PDE problems with similar problem characteristics, e.g., multi-scale PDE problems.
introduce cascaded two-stage
Limitations and Broader Impact. This work focuses on steady-state optical field solutions using the FDFD method. Exploring the effectiveness of operator learning for the Domain (FDTD) can be an interesting direction. Moreover, the FFT kernels on GPU are not fully op- timized [6]. Employing specialized, optimized FFT kernels can unlock even greater computational efficiency on GPUs, further accelerating the neural PDE solver.
Time
7 Acknowledgments and Disclosure of Funding
We acknowledge NVIDIA for donating its A100 GPU workstations and the support from TILOS, NSF-funded National Artificial Intelligence Research Institute. Additionally, this work was sup- ported by the Air Force Office of Scientific Research (AFOSR) through the AFOSR project, con- tract FA9550-23- 1-0422, and the Multidisciplinary University Research Initiative (MURI) program under contract No. FA9550- 17-1-0071.
10

---

References
[1] Habib Ammari, Yves Capdeboscq, and Hyconbae Kang. Multi-scale and High-contrast PDE; from Modelling, to Mathematical Analysis, Inversion, volume 577. American Mathematical Society, 2012.
[2] Md Ashiqur Rahman, Zachary E Ross, and Kamyar Azizzadenesheli. U-no: U-shaped neural operators. arXiv e-prints, pages arXiv-220M4 2022.
[3] Shuhao Cao. Choose transformer: Fourier or galerkin. Advances in neural information processing systems, $4:24924-24944 2021.
[4] Mingkun Chen, Robert Lupoiu, Chenkai Mao, Der-Han Huang, Jiaqi Jiang, Philippe Lalanne, and Jonathan Fan. Physics-augmented deep learning for high-speed electromagnetic simula- tion and optimization. Nature, 2021.
[5] Johannes Feldmann, Nathan Youngblood. Maxim Karpov, Helge Gehring, Xuan Li, Maik Stappers, Manuel Le Gallo, Xin Fu, Anton Lukashchuk, Arslan Raja, Junqiu Liu, David Wright, Abu Sebastian, Tobias Kippenberg, Wolfram Pernice, and Harish Bhaskaran. Parallel convolutional processing using integrated photonic tensor core. Nature, 2021.
[6] Daniel Y. Fu, Hermann Kumbong, Eric Nguyen, and Christopher Ré. FlashFFTCCnv: Efficient convolutions for long sequences with tensor cores. 2023.
[7] Jiaqi Gu, Chenghao Feng, Hanqing Zhu, Zheng Zhao, Zhoufeng Ying, Mingjie Liu, Ray T Chen, and David Z Pan. Squeezelight: A multi-operand ring-based optical neural network with cross-layer scalability. IEEE Transactions on Computer-Aided Design of Integrated Circuits and Systems, 42(3):807-819. 2022.
[8] Jiaqi Gu, Zhengqi Gao, Chenghao Feng, Hanqing Zhu, Ray Chen, Duane Boning, and David Pan. Neurolight: A physics-agnostic neural operator enabling parametric photonic device simulation. Advances Neural Information Processing Systems, 35:14623-14636 2022.
[9] Gaurav Gupta, Xiongye Xiao, and Paul Bogdan. Multiwavelet-based operator learning for differential equations. In Proc. NeurIPS, 2021.
[10] Jayesh K Gupta and Johannes Brandstetter. Towards muliispatiooeemppral-ssallll generalized pde modeling. arXiv preprint arXiv:2209.156h 2022.
[11] Tyler W. Hughes, Momchil Minkov, lan D. Williamson, and Shanhui Fan. Adjoint method and inverse design for nonlinear nanophotonic devices. ACS Photonics, 2018.
[12] Junhyeong Kim, Berkay Neseli, Jae yong Kim, Jinhyeong Yoon, Hyeonho Yoon, Hyo hoon Park, and Hamza Kurt. Inverse design of an on-chip optical response predictor enabled by deep neural network. Opt. Express, 2023.
[13] Jean Kossaifi, Nikola Kovachki, Kamyar Azizzadenesheli, and Anima Anandkumar. Multi-grid tensorized fourier neural operator for high-resolution pdes. arXiv preprint arXiv:2310.001/0, 2023.
[14] Zijie Li, Kazem Meidani, and Amir Barati Farimani. Transformer for partial differential equa- tions' operator learning. Transactions on Machine Learning Research, 2023.
[15] Zijie Li, Dule Shu, and Amir Barati Farimani. Scalable transformer for pde surrogate modeling. Advances in Neural Information Processing Systems, 36, 2024.
[16] Zongyi Li, Nikola Kovachki, Kamyar Azizzadenesheli, Burigede Liu, Kaushik Bhattacharya, Andrew Stuart, and Anima Anandkumar. Fourier neural operator for parametric partial differ- ential equations. In International Conference on Learning Representations (ICLR), 2021.
[17] Joowon Lim and Demetri Psaltis. Maxwellnet: Physics-driven deep neural network training based on maxwell's 's equations. Appl. Phys. Lett., 2022.
[18] Xinliang I Liu, Bo Xu, and Lei Zhang. Ht-net: Hierarchical transformer based operator learning model for multiscale pdes. 2022.

---

Factorized fourier [19] Nina Meinzer, William L Barnes, and Ian R Hooper. Plasmonic meta-atoms and metasurfaces. Nature photonics, 8(12):889- 898, 2014.
[20] Shupeng Ning, Hanqing Zhu, Chenghao Feng, Jiaqi Gu, Zhixing Jiang, Zhoufeng Ying, Jason Midkiff, Sourabh Jain, May H Hlaing, David Z Pan, et al. Photonic-electronic integrated circuits for high-performance computing and ai accelerators. Jourmal of Lightwave Technology, 2024.
[21] Mario Ohlberger and Barbara Verfurth. new heterogeneous multiscale method for the helmholtz equation with high contrast. Multiscale Modeling & Siimulation, 16(1):385-411 2018.
[22] Bogdan Raonic, Roberto Molinaro, Tim De Ryck, Tobias Rohner, Francesca Bartolucci, Rima Alaifari, Siddhartha Mishra, and Emmanuel de Bézenac. Convolutional neural operators for robust and accurate learning of pdes, Advances in Neural Information Processing Systems, 36, 2024.
[23] Bhavin J. Shastri, Alexander N. Tait, T. Ferreira de Lima, Wolfram H. P. Pernice, Harish Bhaskaran, C. D. Wright, and Paul R. Prucnal. Photonics for Artificial Intelligence and Neu- romorphic Computing. Nature Photonics, 2021.
[24] Yaocheng Shi, Yong Zhang, Yating Wan, Yu Yu, Yuguang Zhang, Xiao Hu, Xi Xiao, Hongnan Xu, Long Zhang, and Bingcheng Pan. Silicon photonics for high-capacity data communica- tions. Photonics Research, 10(9):A106AA11 2022.
[25] Kim Stachenfeld, Drummond Buschman Fielding, Dmitrii Kochkov, Miles Cranmer, Tobias Pfaff, Jonathan Godwin, Can Cui, Shirley Ho, Peter Battaglia, and Alvaro Learned simulators for turbulence. In International conference on learning representations, 2021.
[26] Mohammad H. Tahersima, Keisuke Kojima, Toshiaki Koike-Akino, Devesh Jha, Bingnan- Wang, and Chungwei Lin. Deep neural network inverse design of integrated photonic power splitters. Sci. Rep., 2019.
[27] Alasdair Tran, Alexander Mathews, Lexing Xie, and Cheng Soon Ong. neural operators. arXiv preprint arXiv:2111.13802, 2021.
[28] Alasdair Tran, Alexander Mathews, Lexing Xie, and Cheng Soon Ong. Factorized fourier neural operators. In NeurlPS Workshop, 2021.
[29] Alasdair Tran, Alexander Mathews, Lexing Xie, and Cheng Soon Ong. Factorized fourier neu- ral operators. In The Eleventh International Conference on Learning Representations, 2023.
[30] Rahul Trivedi, Logan Su, Jesse Lu, Martin F. Schubert, and Vuckovi.. Data-driven acceleration of photonic simulations. Sci. Rep., 2019.
[31] Barbara Verfürth. Heterogeneous multiscale method for the maxwell equations with high con- trast. ESAIM: Mathematical Modelling and Numerical Analysis, 53(1):35-61, 2019.
[32] Zi Wang, Lorry Chang, Feifan Wang. Tiantian Li, and Tingyi Gu, Integrated photonic meta- system for image classifications at telecommunication wavelength. Nature communications, 13(1):2131, 2022.
[33] Haixu Wu, Tengge Hu, Huakun Luo, Jiammin Wang, and Mingsheng Long, Solving high- dimensional pdes with latent spectral models. In Proceedings of the 40th International Con- ference on Machine Learning. ICML'2S '23 JMLR.org. 2023.
[34] Ruibin Xiong, Yunchang Yang, Di He, Kai Zheng, Shuxin Zheng, Chen Xing, Huishuai Zhang, Yanyan Lan, Liwei Wang, and Tieyan Liu. On layer normalization in the transformer architec- ture. In International Conference Machine Learning, pages 10524-10533. PMLR, 2020.
[35] Bo Xu and Lei Zhang. Dilated convolution neural operator for multiscale partial differential equations. 2023.
12

---

[36] Xingyuan Xu, Mengxi Tan, Bill Corcoran. Jiayang Wu, Andreas Boes, Thach G. Nguyen, Sai T. Chu, Brent E. Little, Damien G. Hicks, Roberto Morandotti, Arnan Mitchell, and David J. Moss. 11 TOPS photonic convolutional accelerator for optical neural networks. Na- ture, 2021.
[37] Sanaz Zarei, Mahmood-reza Marzban, and Amin Khavasi. Integrated photonic neural network based on silicon metalines. Optics Express, (8(24):3668--36666 2020.
[38] H. H. Zhu, J. Zou, H. Zhang, Y. Z. Shi, S. B. Luo, al. Space-efficient optical computing with an integrated chip diffractive neural network. Nature Communications, 2022.
[39] Hanqing Zhu, Jiaqi Gu, Hanrui Wang, Zixuan Jiang, Zhekai Zhang, Rongxing Tang, Cheng- dynamically- photonic transformer accelerator. In 2024 IEEE Interna- tional Symposium on Performnree Computer Architecture (HPCA), pages 686-703. IEEE, 2024.
hao Feng, Song Han, Ray T Chen, and David Pan. Seghtening-transtoroee operated
[40]
Hanqing Zhu, Keren Zhu, Jiaqi Gu, Harrison Jin, Ray T Chen, Jean Anne Incorvia, and David Z Pan. Fuse and mix: Macam-enabled analog activation for energy -efficienn neural acceleration. In Proceedings of the 41st IEEE/ACM International Conference Computer-Aided Design, pages 1-9, 2022.
13

---

A Appendix
A.1 Dataset Generation
We generate our customized etched MMI and Metaline dataset using the open-source FDFD simu- lator angler [11]. For each type of device, random sample 5. 12 K device configuration following the Tab. 4, and generate single-source data by sweeping the input light over the input ports. We ran- domly sample the device's 's physical dimension, input/output waveguide width and input light source frequencies. For etched MMIs, we randomly sample etched cavity sizes, ratios (which determine the number of cavities in the MMIs), and permittivities in the controlling region. For Metaline, randomly sample the metaatom physical dimension with a fixed total number of 20.
We discretize the domain of etched MMI by 80 x 384, and the domain of metaline by 128 144.
Table 4: Summary of etched MMI device design variable's sampling range, distribution, and unit.
Valle/Disttbuutinn
Variables
Unit
Metaline 3 x
Etched MMI 3
Etched MMI
Length
2488 10)
2/220, 30)
L((25, 35) 147755
pum
Width
4(10, 12)
(45.5, 7)
Port Length
1.5
1.5
1,5
p07
Width
24(0.5, 0.8) 3
24(0.8, 1.1)
24(0.8,
Taper Length
4.5
4.5
Taper Width
1.3
1.3
pm
Border Width
0.25
0.25
0.25
pm
PML Width
1.5
1.5
1.5
pam
Wavelengths A
U(155, 1.565)
((1553, I. 1.565) 2/(0.05, 0.1)
U((1,, 1.565)
pm
Cavity Ratio
24(0.05, 0.1)
Note
2-layer meta atoms (2.07, 12.11]
random
random slats
Relative Permitiivtty
(2.07, 12.11]
[2.07. 12.11)
A.2 Training Settings
We implement all models and training logic in PyTorch 2.3. We use A100 and A6000 to train our models and report the latency running on single A100 GPU with torch. compile. For Bench- marking FDFD simulator performance, we use the Intel 12th Gen Intel(R) Core(TM) i7-12700 with 20 CPU cores. We split single-source examples into 72% training data, 8% validation data, and 20% test data.
For training, we set the number of epochs to 100 with an initial learning rate of 0.002, cosine learning rate decay, and a mini-batch size of 4. We use adam W as the optimizer with the weight decay le-5 to avoid over-fitting, Moreover, we apply stochastic network depth with linear scaling strategy.
A.3
Model Designs
To ensure comprehensive evaluation, compare our proposed model against recently published and available SOTA baselines, encompassing various architectural paradigms such as the Fourier- operator models, attention-based models, and latent space methods. To maintain fairness in com- parison, we constrain the parameter count of all models to be under 4 M in most cases and use open-sourced implementations.
Here, we report the model details for baselines for the etched MMI dataset.
UNet. We construct a 4-level convolutional UNet with a base channel number of 36, following the open-sourced implementation'. The total parameter count is 3.88 M.
Dil-ResNet [25]. We use the implementation in open-sourced pdearena [10]', with a channel num- ber of 128 and enabled normalization. The total parameter count for Dil-ResNet is 4.17 million.
Phapps//githu..com///eee http://hhhhhh.ooooocmmmmmmenaaan
14

---

FNO-2d [16]. We use 6 2-D ENO layers, and the Fourier modes are set to (#Mode =32, #Mode =10) for the etched MMI dataset and (#Mode =16, #Mode =16) for the Metaline dataset, resulting in the total parameter count as 3.99 M or 3.21 M. We use the implementation"
Tensorized FNO-2d [13]. One obvious advantage that model features low parameters. Hence, we will compare it with tensorized FNO, which compresses the model weights with the tensor de- composition method. We adopt the implementation in and use the model designed for the darcy flow problem. We use 5 2-D FNO layers, and the Fourier modes are set to (#Mode =40, #Mode =20) for the etched MMI dataset and (#Mode =24, #Mode =24) for the Metaline dataset. Tucker decom- position is used with rank of 0.42. The total parameter count is 2.25 M and 1 1.58 M, respectively, for the two types of datasets.
F-FNO-2d. For factorized Fourier neural operator (F-FNO), we use 13 F-FNO layers with a channel number of 52. The Fourier modes set to (#Mode =70, #Mode =40) for etched MMI dataset and (#Mode =36, #Mode =36) for Metaline dataset, leading to the total parameter count as 4.02 M and 2.68M. The FNO-2d implementation is referred to We use the same projection head as ours.
U-NO-2d [2]. For U-shaped neural operators, we follow the implementation". We use their 11-layer UNO with base channel 24. The total parameter count is 4.38 M.
Attention-based operator [15]. For the attention- neural operator, we choose the most recent Transformer- type model [15] and use its official implementation" " We use 3 layer-attention with 12 heads. The total dimension 128. The total parameter count is 3.75 M.
Latent Spectral Method [33]. For the latent spectral method, we use the original implementation in The number of bases is set 12, and the channel number is 32. The patch size is set to 4x4. The total parameter count is 4.8 M.
NeuroLight [8]. We use the same implementation" in the original paper with 16 layers. For the etched MMI dataset, the Fourier modes are set to (#Mode =70, #Mode =40). For the Metaline dataset, Fourier modes set to (#Mode =36, #Mode =36). The total number of parameters is 2. and 1.49 M for the two cases.
PACE, For our proposed PACE, we use 12 layers, with the first two being the same factorized layers in [8], since we found it is important first to generate some meaningful wave patterns and then do global information swapping. The Fourier modes are set to (#Mode =70, #Mode =40). We use the same convolution stem [8] to extract information before going through the feature propagator and the same projection head. The total number of parameters is 1.73M. For the second stage PACE-II model, we use layers with all being PACE operators, where Fourier modes are set to (#Mode =100, #Mode =40) For the Metaline dataset, we solely use 12-layer PACE with Fouier modes being (#Mode =36, #Mode =36).
A.4 Ablation study on group number choices
We run an 8-layer PACE model on the Metaline dataset by setting the group size to 1, 2, 4, and show the train and test error in Tab. We use #group=4 in our paper, which balances between parameter efficiency and test error.
A.5 Ablation Study of Double Skip and Pre-Normalization
We further investigate whether the observed improvements in accuracy are attributed to the incorpo- ration of double skip connections and pre-normalization, which were incorporated into model to stabilize it in deeper layers with better generalization. We add these two techniques to NeurOLight ght and compare them with ours PACE in Tab.. 6. The double skip and pre-normalization model generalize well for test data, while the training error is slightly improved as normalization
make the
15

---

Table 5: Ablation on # group on 8-layer PACE model on the Metaline dataset.
#Params
#Train Err
#Test Err
Group
(M)
(10 2
(10-2)
1
2.15
4.89
4.55
2
1.27
5.33
4.80
4
0.825
5.65
4.82
Table 6: Ablation on the comparison between PACE and NeurOLight with the adopted double skip and pre-normalization.
#Params
#Train Err
#Test Err (10 3)
Model
Double skip & Pre-norm.
+
(10 2)
NeurOLight-16 layer
2108258
15.58
17.21
NeurOLight-16 layer
2110306
15.26
15.87
PACE-12 layer
X
1709026
10.32
11.06
PACE-12 12 layer
1710562
9.60
10.60
can be seen some linear affine. However, it still shows much worse accuracy than model, especially given the context our model is shallower with fewer parameters.
A.6 L2 distance is a more informative metric compared distance for distance evaluation
Corollary A.1. Consider two complex numbers in polar form, 21 Ti and z2 = r2L02. Their mean square error is rotation invariant, as shown by:
31, 22 1 In cos r2 cos Φ2 |ri sin Φ1 sin dzl
2ri cos(Ør 62).
This distance metric depends solely on the difference in signal norm and angle between and 22. However, their mean absolute error is the rotation variant:
21, #2/1 [r cos Φ1 cos sb2) + [ri sin d1 sin db).
(10)
In the complex plane, optimizing MAE equates to minimizing the summed LI distances the real and imaginary components. However, the LI distance is variion-variant. A simple rotation of the two complex numbers on the plane results in changing L1 distance, as shown Fig. 9, while the true distance does not alter. Therefore, it is not appropriate metric it cannot accurately measure proximity in the complex plane. In this way, we use L2 distance the loss (mean squared loss) that rotation invariant, proved in corollary A. which exactly captures the distance in the complex plane.
Im
stanoo
Re
Figure 9: L1 and L.2 distance in the complex plane.
16

---

Energy Spectrum
100
Wavenumber
Figure 10: Radial energy spectrum for one solution of Darey flow problem.
Table 7: Results of train two-stage model sequentially.
Benchmarks
Model
Cross-stage dist.
#Params (M) +
Train Eir (10-3)
Test Eir (10-²)
PACE- 12 layer
1.73
9.51
10.59
Etched MMI 3x3
PACE-I PACE-II
3.151
3.91
5.06
PACE: 12 layer
1.73
11.66
11.91
Etched MMI 5x5
PACE-I + PACE-II
3.151
5.51
6.15
A.7
Train two-stage model sequentially
We also show the error by training proposed two-stage model sequentially in Tab. 7, which shows a similar error to our joint training approach when equipping with our proposed cross-stage feature distillation.
Training sequentially is more costly than joint training, as second-stage training requires first infer- encing with the first stage get the predicted results.
A.8 Visualization of energy spectrum
We generate the prediction field's radial energy spectrum by first transferring the image from the spatial domain to the spectral domain and then shifting the transferred image to the center. Then, the wavenumber computed as the distance with respect to the center.
We sum the squared magnitude of the Fourier coefficients that fall into the specific number, which implemented following open-sourced code 11,
We also visualize one example of darcy flow problem, as shown Fig. 10. It shows highly distant characteristics compared to optical field, with most information concentrating low- cowffreueeccy parts.
A.9 Visualization feature map before/after non-linear activation in our explicitly designed high-frequency projection path
We visualize the first channels of feature maps before and after the nonlinear activation in last PACE layer by showing them in the frequency domain. As shown in Fig. 11, the nonlinear activation can ignite high-frequu fighfrreuency features, which confirms our claim and validates our design choice of injecting extra high- figh-frequency projection path.
A.10 Visualization of prediction
We provide visualization figures on etched MMI 3x3 devices in Fig. and metaline devices in
Fig. 13. We provide the predicted fields (a). the groud-truth field (a)" and the residual er- ror (a)" - (a) of Dil-ResNet, Facztoried FNO, NeurOLight and our PACE. For etched MMI test cases, show both the single 12-layer PACE model and the joint 20-layer model PACE-I PACE-II. Our PACE shows much better prediction results with a near- black error map compared to other bascline methods.
http://wtmhooooolooooooooooo/oer utils/utils, spectrum.pyy
17

---

Sefore nonlinear
After nortinear
Figure 1I: Frequency-done domennn visualization of feature map before and after non linear activation in the PACE block(The center represents low frequency). The pattern is shifted to the center to understand the fre- quency content hetter.
(a)(
000
C0C
Di-ReeNNet
(4))
(P)(a) (4)
(a)
Factorized FND
(a))
(a) (a)
(a)
iNeurolight
(+))
()(()) - [()
(a)
Pace
RVV (a))
(9) (a)
(()(a))
Pacc Pace II
(a))
(0'(a) v(a)
Figure 12: Visualization of test cases on etched MMI 3x3 devices with random input sources.
18

---

R(V(a))
DI ResNot
R()*
(a) EI
(a)
R(V(a))
Factorized
(a))
FND
(a) (a))
R()(a)))
Neurolight
(())
(a) (a)
(()(a))
I. AAB
Pace
(())
00
000
- 2
(a) (a)
Figure 13: Visualization of several test cases on Metaline devices with random input sources.
19

---

NeurIPS Paper Checklist
1. Claims
Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
Answer: [Yes]
2. Limitations
Question:
Does the paper discuss the limitations of the work performed by the authors?
Answer: [Yes]
3. Theory Assumptions and Proofs
Question: For each theoretical result, does the paper provide the full set of assumptions and complete (and correct) proof?
Answer:[Yoe
4. Experimental Result Reproducibility
Question: Does the paper fully disclose all the information needed to reproduce the main experimental results of the paper to the extent that it affects the main claims and/or conclu- sions of the paper (regardless of whether the code and data are provided or not)?
Answer:[Yes]
5. Open access to data and code
Question: Does the paper provide open access to data and code, with sufficient instruc- tions to faithfully reproduce the main experimental results, as described in supplemental material?
Answer: [Ye]]
Justification:
We open-source the dataset and code.
Guidelines:
6. Experimental Setting/Details
Does the paper specify all the training and test details (e.g., data splits, hyper- parameters, how they chosen, type of optimizer, etc.) necessary understand the results? Answer:[Yes]
Question:
7. Experiment Statistical Significance
Question:
Does the paper report error bars suitably and correctly defined or other appropri- ate information about the statistical significance of the experiments?
Answer:[Yes)
Justification: We report the emror bar when we test the performance on generalization ex- periments.
8. Experiments Compute Resources
For each experiment, does the paper provide sufficient information on the com- puter resources (type of compute workers, memory, time of execution) needed reproduce the experiments?
Question:
Answer:[Yes,
9. Code Of Ethics
Question: Does the research conducted in the paper conform, in every respect, with the wwterrisscc.ppuuieee
NeurlPS Code of Ethics
Answer:[Yes:
10. Broader Impacts
Question: Does the paper discuss both potential positive societal impacts and negative socictal impacts of the work performed?
20

---

Answer:[Yes Yes]
11. Safeguards
Does the paper describe safeguards that have been put in place for responsible release of data models that have high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?
Question:
[nsw
12. Licenses for existing assets
Are the creators or original owners of assets (e.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?
Question:
Answer:[NA]
13. New Assets
Question: Are new assets introduced in the paper well documented and the documenta- tion provided alongside the assets?
Answer:[Yes]
Guidelines:
Crowdsourcing and Research with Human Subjects
Question: For crowdsourcing experiments and research with human subjects, does the pa- per include the full text instructions given to participants and screenshots, applicable, as well as details about compensation (if any)?
Answer:[NA]
15. Institutional Review Board (IRB) Approvals or Equivalent for Research with Human Subjects
Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or equivalent approval/review based the requirements of your country or institution) were obtained?
Answer:[NA]
21


Reviews:
[{"review_id": "Z8AwAWYoha", "reviewer": "Reviewer_ZkVJ", "summary": "This paper develops a computationally efficient learning based PDE simulator for modeling optical fields within photonic devices. The proposed method builds upon NeurOLight [7] and adds a novel factorization of the integral kernel and a two-stage approach to prediction---produce a rough initial estimate that is then refined. The proposed method is noticeably more accurate than [7] and still considerably faster than conventional, non-learning based PDE solvers. An ablation study validates the importance of both the factorization adn the two-stage approach.", "questions": "Are the \"real-world\" devices actually manufactured and measured or merely simulated? I found section 5.1 unclear.", "limitations": "Yes", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 2, "contribution": 3, "strengths": "Efficiently and accurately simulating photonic devices is an important problem with broad applications.\n\nTHe proposed method is more accurate than existing learning based solvers and 12x faster than conventional methods", "weaknesses": "The paper is written for a somewhat narrow audience and contains a number of typos. E.g., \"boarder impact\". \n\nThe 577x comparison number is somewhat misleading and should be dropped from the abstract. The proposed method may be much faster than scipy, but its only 12x faster than paradiso.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you so much for your valuable feedback on our work! Please see our responses below:\n\n**Q.1 Somewhat narrow audience**\n\nWe would like to stress that our paper aligns well with NeurIPS’s interest and has essential machine learning contributions to AI for the PDE community. Please check the detailed answer in our common response Q.1.\n\n**Q.2 Confusing description in speedup**\n\nThank you so much for pointing out the description is misleading. We will modify it to “In terms of runtime, our PACE model demonstrates 154-577x and 11.8-12x simulation speedup over numerical solver using scipy or highly-optimized pardiso solver, respectively”.\n\n**Q.3 Unclear description of “real-world” devices**\n\nThe devices used are claimed to be in the real world since they have real-world geometry, permittivity distribution, input light sources, and wavelength ranges. (check detailed response in common response Q.2) \n\nAs we cannot measure the optical field after the device is manufactured, photonic designers rely on simulation tools to obtain it to aid the design loop. Therefore, in this paper, we also use the simulation tool to generate the ground-truth optical field to assess the effectiveness of our model.\n\n**Q.4 Typos**\n\nThank you so much for your careful reading, which is crucial to improving our paper quality. Since we cannot upload the revised manuscript now, we will fix all typos and update the final manuscript.\n\nThank you so much for reviewing our paper and this response. We look forward to further discussions with you if you have any other questions!", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "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": "sy3vrIql2c", "reviewer": "Reviewer_9WdN", "summary": "This paper propose a novel cross-axis factorized PACE operator with a strong long-distance modeling capacity to connect the full-domain complex field pattern with local device structures. Additionally, the authors proposed a two-stage model learning method for extremely hard cases, with a first-stage model learning an initial solution refined by a second model. Finally, experimental results show that the proposed PACE model can achieve good performance with fewer parameters.", "questions": "1. Training details in Sec. 5.1 and A.2 should be integrated.\n\n2. The effectiveness of pre-normalization and double skip is not well verified. In Table 6 of the appendix, you should list the results of PACE-12 layer without double skip and pre-norm.\n\n3. At line 211, the authors claim that non-linear activation is known to help generate high-frequency features. Therefore, some visualization of the high-frequency features is needed to support this statement.", "limitations": "NA", "rating": 6, "confidence": 4, "soundness": 4, "presentation": 3, "contribution": 4, "strengths": "1. Major challenges for simulating real-world complicated optical devices is well analyzed.\n\n2. Their proposed PACE model significantly boost prediction fidelity to an unprecedented level with much faster running speed than traditional numerical solver on a 20-core CPU.\n\n3. A two-stage model learning method is used for extremely challenging cases.\n\n4. Experimental results show that on the complicated device benchmarks, their proposed PACE model significantly outperforms previous baseline models.\n\n5. A dataset on for the complicated optical device is released to facilitate AI for PDE community.", "weaknesses": "1. I do not agree with the authors that the proposed two-stage learning method belongs to a divide-and-conquer approach. From my perspective, it is called coarse to fine.\n\n2. The writing should be improved as some sentences are not precise. For example, at line 19&20, “one sole PACE model …fewer parameters” (compare with which method?)\n\n3. It is common that deep learning based method runs much faster than traditional numerical solvers. The testing error of the traditional numerical solvers may plays as a groundtruth for the upper bound of precision. Thus, the authors should list the testing error of the traditional numerical solvers.\n\n4. In Table 1, except for the number of Params, you should list other metrics like Flops and testing time to show the efficiency performance of the proposed PACE model.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you so much for your appreciation of our contributions and for providing valuable feedback on our work! Please see our responses below:\n\n**Q.1 Use “coarse to fine” instead of “divide and conquer”**\n\nThank you for your suggestion. We agree that our learning flow could be more precisely described as learning the PDE solution in a coarse-to-fine manner, given that we divide the challenging cases into two progressively latent tasks and generate high-fidelity PDE solutions from rough to clear. We will update our description.\n\n**Q.2 Testing errors of traditional numerical solver**\n\nSince the true optical field is not measurable, photonic designers typically rely on ***fine-grained*** numerical simulations to obtain the optical field. The simulated optical field is treated as a reliable reference (ground truth) to guide the photonic design optimization.\n\nIn this paper, we set the simulation granularity to 0.05nm, sufficient to obtain a reliable referenced optical field. We treat this fine-grained simulation as the ground truth (assuming 0 error or target) and test the error of the predicted field from ML models against it. Therefore, we don’t list the tested error of numerical solvers.\n\n**Q.3 Include run-time in Tab. 1**\n\n\nWe report the run-time for etched MMI3x3 (same for etched MMI5x5) here for your reference, and the final Table 1 will be updated. We tested the run-time on a single A100 GPU using` torch.utils.benchmark`, averaging multiple runs. Models were compiled with `torch.compile`, except for TFNO, which is not compatible with `torch.compile`.\n\nOur PACE shows excellent accuracy, superior parameter efficiency, and better run-time compared to the previous SOTA NeurOLight. \n\nWe’d like to emphasize that **accuracy** is the highest priority metric when using an ML surrogate model to guide the photonic device design and optimization loop. It is worthwhile to sacrifice speed for much better fidelity. As shown in the table, although several models show faster inference speed, they fail to deliver good accuracy. And model like FNO shows saturated model accuracy when further scaling to deeper models[1]. \n\nMoreover, PACE can be further accelerated with customized acceleration (e.g., grouping operation) and a better 1-D Fourier kernel. Even with the current implementation, PACE shows a much-accelerated runtime with a 12x speedup compared to the highly optimized Pardiso-based numerical solver. Therefore, PACE is highly effective in serving as an **accurate** and **fast** ML surrogate to accelerate the photonic design loop.\n\n**Parameter count** is also a key metric for our problem. The presence of high-frequency features necessitates the use of high-frequency modes, which is known to greatly increase the parameter count in FNO and cause overfitting, motivating the development of factorized neural operator work (e.g., Factorized FNO, Tensorized FNO). As a novel neural operator, we utilize a physically meaningful cross-axis factorization and outperform previous factorized FNO work with a significant accuracy margin, while maintaining comparable speed with Factorized FNO.\n\n| Model | #params (M) | Test error (1e-2) | Test time($\\pm$σ) (ms) |\n|------------------------|-------------|-------------------|----------------|\n| UNet  | 3.88 | 63.03 | 1.79($\\pm$0.02) |\n| Dil-ResNet| 4.17| 51.34 | 5.89($\\pm$0.06) |\n| Attention-based module | 3.75 |70.05 |  10.41($\\pm$0.01) |\n| UNO | 4.38 | 34.22 | 5.34($\\pm$0.02) |\n| Latent-spectral method | 4.81| 55.07| 3.84($\\pm$0.06)   |\n| FNO-2d  | 3.99| 32.51| 4.24($\\pm$0.03)   |\n| TFNO-2d  | 2.25| 35.52| 4.92($\\pm$0.01)   |\n| Factorized FNO-2d| 4.02| 24.2 | 9.93($\\pm$0.03)   |\n| NeurOLight | 2.11 | 15.58| 12.98($\\pm$0.14)  |\n| PACE  | 1.71  | 9.51| 11.44($\\pm$0.03)  |\n\n\n**Q.4 Incomplete ablation study on pre-normalization and double skip:**\n\nWe include the training results of the PACE-12-layer model without double skip and pre-normalization. The training error, test error, and parameter counts are listed as follows:\n| Model| Double skip & Pre-normalization | Params  | Train error | Test error |\n|------------|---------------------------------|---------|-------------|------------|\n| NeurOLight-16layer | w/o| 2108258 |15.58 |17.21 |\n| NeurOLight-16layer | w/| 2110306 |15.26 |15.87 |\n| PACE-12layer|w/o| 1709026 |10.32 |11.06 |\n| PACE-12layer|w| 1710562 |9.51|10.59 |\n\nAs stated in the paper, double skip and pre-normalization are not the primary sources of accuracy improvement but are beneficial for stabilizing the model when scaling to deeper layers.\n\n**Q.5 Visualization of  feature map before/after non-linear activation in our explicitly designed high-frequency projection path**\n\nWe visualize the first 6 channels of feature maps before and after the nonlinear activation in the last PACE layer by showing them in the frequency domain. As shown in reb-Fig. 1, the nonlinear activation can ignite high-frequency features, which confirms our claim and validates our design choice of injecting an extra high-frequency projection path. Furthermore, in Sec. 5.3, we show that removing the projection path results in a considerable accuracy loss.\n\n**Q.6 Writing**\n\nThank you so much for pointing out that some parts of our description are not precise, which is of great importance to improve our paper quality further. We will modify the wording in the final version, e.g., one sole PACE model is compared with various recent ML for PDE solvers.\n\nWe agree that the training details in Sec. 5.1 and A.2 should be integrated, and we will do it in the revised version with one more page budget.\n\n[1] Tran, Alasdair, et al. \"Factorized fourier neural operators.\"ICLR 2023.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "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": "725VVypa1Y", "reviewer": "Reviewer_TfAG", "summary": "The authors propose PACE operators, which perform much better than existing methods on predicting optical fields for real-world complicated optical devices. The work builds up on and is compared to the NeurOLight framework which represents the SOTA in this field. The benefits of physics-agnostic PACE operators stem from their cross-axis factorized structure. This enables modeling of complicated devices in a parameter-efficient manner. The authors also propose a two-stage learning process which further improves the performance of PACE operators. Finally, the authors open-source this dataset consisting of optical field simulations for complicated devices. The authors demonstrate strong improvements over SOTA in solving the Maxwell PDE for complicated devices through novel operator design and learning workflows.", "questions": "1. The etymology of \"PACE\" is not clear. The all-caps spelling suggests it is an acronym but it is not expanded anywhere in the paper. If there is another contextual meaning of the word \"pace\" that is applicable to the model, it is not clear from the paper. Why is it called PACE?\n2. The works cited in the paper (reference 25, 12) seem to use the word \"complex\" to describe non-trivial optical devices. Would this be a better word to use than \"complicated\" used in the title and throughout the paper?\n3. In \"Generalization to ood testing\" (line 322), the results are only shown for interpolation, i.e., for wavelength ranges within the training domain. It is great that the model interpolates well in this range. However, for true ood generalization, the authors should also show results for extrapolation. How does the model behave for wavelengths outside the training range (outside 1.53-1.565 micrometers).\n4. Related to point of weaknesses, what is the contribution of using the different N-MSE metric as the learning objective? Does the performance of NeurOLight also increase when trained using N-MSE?\n5. In section 5.2.3 Speedup over numerical tools: FNOs are resolution invariant by design. Does PACE also exhibit this property? In the speed up section, were multiple PACE models trained for each discretization or are the results shown for one model with inference over different discretization? \n6. The authors qualitatively state in section A.9 that the prediction results show \"near-black error map\" (line 582). However it seems that the colorbars have different scales across different models. There are no colorbars in Figure 11.  Also, since in cases like these with very low errors, it would be interesting to see the absolute errors in log space. This might shed further light into structural sources of errors. How would this look in log space?", "limitations": "The authors have a brief limitations section where they state that exploring operator learning on FDTD methods and general optimizations to FFT kernels on GPUs are interesting future directions.\n\nWhile this is already a parameter efficient model, it is a known limitation of Neural Operators that they don't scale well with increasing dimensions. It might be helpful to address this limitation here. How would the PACE operator scale in 3D and beyond?\n\nI don't think that there are any negative societal implications for this work.", "rating": 7, "confidence": 4, "soundness": 4, "presentation": 3, "contribution": 3, "strengths": "Strengths:\n1. The paper is well-structured with clear background and motivation. The authors highlight the issues that make optical field prediction for complicated devices challenging for existing methods.\n2. The authors introduce a novel PACE operator that effectively utilizes the identified challenges to mitigate the difficulties in optical field prediction for complex devices.\n3. The authors demonstrate the strong performance of PACE operators compared to the existing NeurOLight framework as well as other commonly used ML PDE solvers such UNets, FNO variants and attention based methods. The proposed method achieves significantly lower errors with comparable or lower number of parameters, while also being faster than numerical methods.\n4. The authors also validate all the proposed improvements individually with well-designed ablation studies.\n5. The curation of the dataset for optical field simulations of complicated devices will also be helpful to the community for future benchmarking.\n\nOverall the paper serves as a good example of identifying a clear problem in a specific domain,  thoroughly analyzing the shortcomings of existing methods, and leveraging domain knowledge to enhance neural operators. The results are strong and showcase the applicability of this method for optical device simulation.", "weaknesses": "1. While demonstrating strong results on a challenging problem, my main concern for this paper is that it might be too specific of an application for a broader ML conference like NeurIPS. The paper might be better suited to a photonics journal. This can be mitigated by demonstrating results on other standard PDE benchmarks and showing improvements over existing methods. The authors allude to this in line 362 Conclusion but concrete results would greatly improve the paper.\n2. The interpretation of errors is not clear. First the normalized absolute error (relative error) of NeurOLight is repeatedly mentioned. Then the authors use the normalized mean squared error as the learning objective (which is rightly justified in A.6). However it is not clear what these errors mean. The paper would also be clearer if it is explicitly stated what error metrics are used in the various table (I assume it is N-MSE). What is the tolerance for acceptable errors for downstream tasks? What magnitude of error makes it comparable to existing numerical solvers used in practice?\n3. It is common in ML-PDE papers to show speed ups over numerical methods while the ML models have less accuracy. If you were to achieve this same lower accuracy using the numerical methods, it might greatly reduce the runtime of numerical methods as well. What is the runtime of Angler/scipy/pardiso that achieves comparable accuracy to the PACE operator?\n5. The authors mention multiple times about inspiration from \"human learning\". However there is no description of what this means. Section 4.2 on cascaded learning can be improved by expanding on how this related to human learning with appropriate citations.\n6. It has been mentioned that the model preserves the spectrum of the solution field. The author show the spectrum of the reference field in Figure 1d and Figure 9 (for Darcy Flow). However there is no discussion on the spectrum of the predicted solution fields. The authors should visualize the spectrum of predicted fields and compare it to the reference.\n\nThe paper would benefit from a thorough proof-reading. There are multiple typographical errors and some sentences are not clear. Examples are:\n- Line 5: space after comma \", NeurOLight,\"\n- Line 162: satisfactory instead of satisfied\n- Line 188: shortcomings instead of shortness\n- Line 196: O(nlogn)\n- Line 205: \"is a must\" is extra/\n- Line 224: satisfactory instead of satisfying\n- Line 238: What does condensate mean here?\n- Line 252: \"neural\"\n- Line 255: Establish LSM acronym here\n- Line 269: I think the authors mean to say that there is a 73.85% reduction in error. The way this sentence is phrased right now implies that there is a 73.85% error.\n- Line 284: Fourier\n- Line 294: Sorely seems to not fit here. Maybe \"just\" or \"only\"?\n- Line 297: Appendix\n- Lines 314 and 315: It is a bit confusing. The sentence can be made clearer by avoiding the use of worse.\n- Line 319: Decomposition doesn't seem to fit here.\n- Line 336: Factorized\n- Line 362: Restricted might be a better alternative to dedicated.\n\n\nSome figures are hard to decipher:\n- Figure 1 a,b,c should be made larger so that the readers can clearly see the issues mentioned in the Challenges section.\n- Figures 5, 6, 7 show important comparisons between different methods and should be made larger.\n- Figure 9 should have subfigures. Also 9a (Darcy flow field) should have axis labels and colorbar.\n- Figures 10 and 11 should be made bigger so that the light fields and errors are clearly visible. The absolute error is not closed with \"|\" in the blue boxes in Figure 10. Figure 11 is missing colorbars that show the magnitudes of the optical fields and errors. \n\nRight now the dataset URL is behind a hyperref \"URL\". It might be helpful to mention the full address somewhere in the paper for accessibility.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you so much for your valuable feedback on our work! Please see our responses below:\n\n**Q.1 Narrow audience and other benchmarks**\n\nWe'd like to emphasize that our paper aligns with NeurIPS’s interests and offers novel ML contributions to the AI for PDE community (check common response Q.1)\n\nThe standard PDE benchmarks typically use smoothly changing coefficients and, therefore, don’t share the same PDE characteristics with ours (e.g., our permittivity is highly discrete and contrastive), making them unsuitable for evaluating our method.\n\nWe notice that our problem shares a similar pattern with multi-scale PDE, where PDE coefficients are generated contrastively. Recent works [1][2] also show the failure of existing AI-based PDE solvers in multi-scale PDE. Unfortunately, [1][2] do not open-source their datasets. Thus, we compare SOTA [2] with PACE on our benchmark (See Sec. 5.3). PACE significantly outperforms [2], with a 10.59 N-MSE error over 17.4.\n\n**Q.2 Used errors:**\n\nWe consistently use normalized mean squared error (l2 loss) in training and testing. We argue that mean squared error is a more informative/right metric for evaluating complex variables’ distance in A.6 (not a contribution to boost accuracy).\n\n Normalized mean absolute error (l1 loss) is to show NeurOLight's indecent accuracy in complicated cases. We intentionally report the same metric in NeurOLight to avoid unmatched numbers.\n\n**Why normalized MSE, not regular MSE?**\nWe randomly sample device configuration and light combination. The resulting optical field yields distinct statistics, e.g., total energy. We need to normalize them to ensure a balance in optimization/evaluation on different optical fields.\n\nAs for acceptable error for downstream tasks, PACE achieves ~5% N-MSE error, sufficient to guide device optimization in the early device design and exploration stage (better fidelity than a larger simulation grid shown in the next question).\n\n**Q.3 Speedup over numerical methods with less accuracy:**\n\nWe can’t set a targeted error in the simulator as the angler uses a direct linear solver, not an iterative solver. However, we can reduce the simulation granularity to speed up the process with sacrificed accuracy. In our experiment, we use a commonly used 0.05nm to obtain ground truth. A larger granularity (e.g., 0.075 nm) leads to a very different field (qualitatively compared in reb-Fig.3; quantitive N-MSE error 1.2).\n\nAs you suggested, we added PACE's speedup over the 0.075nm simulation in reb-Fig.5. PACE still shows a 5.1-10.6x speedup with much better fidelity.\n\n**Q.4 Unclear human learning**\n\nPlease see common response Q.3.\n\n**Q.5 Spectrum of the predicted field**\n\nThe predicted field spectrums of PACE and NeurOLight are in reb-Fig.3. PACE excellently aligns with the baseline spectrum compared to NeurOLight, while NeurOLight uses the same frequency modes.\n\n**Q.6 Etymology of PACE:**\n\nSorry for not annotating PACE’s source in the title: ***P***acing Operator Learning to ***A***ccurate Optical Field Simulation for ***C***omplicated Photonic D***e***vices.\n\nIt has two meanings: we **pace** to a new SOTA in AI for photonic device simulation; our learning flow decouples the hard task into two progressively tasks in a **pacing** manner.\n\n**Q.7 Ood testing**\n\nFirstly, we’d like to stress that our generalization experiment is done on a specific wavelength range (1.53-1.565 μm), which is C-band (an interested range for device design). We train with sampled wavelengths in C-band and test its ood generation on unseen frequencies. It is a vital test to prove the usefulness of PACE in helping device design within an interested wavelength range.\n\nAs you suggested, we tested the accuracy outside the C-band (See reb-Fig.6). PACE shows good accuracy on neighboring wavelengths while holding a 10-15% error at a further range. This is expected since wave propagation is sensitive to frequency. It can be mitigated by incorporating sampled wavelengths into training.\n\n**Q.8 Speed-up experiments**\n\nIn speed-up evaluation on various device scales, we only test the speedup by varying device geometry sizes; therefore, no training is performed. \n\nFor resolution invariance, our cross-axis factorized neural operator (Fig.2) uses the same components, the Fourier integral operator and linear layers, with factorized FNO and FNO. Hence, we hold the same resolution-invariant property theoretically. But, as stated in [4], no model can really deliver a real zero-shot since finer scales will contain new physics interactions (re-training is preferred).\n\n**Q.9 Seemingly inconsistent colorbars A.9**\n\nIn A.9, each column consistently uses the same test case to evaluate various methods **using the same color bar**. Across different test cases between columns, the optical field may yield distinct statistics, e.g., total energy, leading to different color bars. We put a wrong figure for Dil-ResNet in Fig.10, which may be the reason to confuse you(updated in reb-Fig.2). \n\n**Q.10 Error in log space**\n\nWe show the absolute error of one device in log space in reb-fig.7, which highlights small errors to help understand the error pattern. As we can see, the error increases from left to right due to challenge 3 in Sec.3.1 (non-uniform learning complexity). Besides, error residual occurs especially when scattering or local resonance happens (challenge 1: complex light-matter interaction).\n\n**Q.11 3D cases**\n\nWe leave the investigation of our operator for 3D FDFD optical simulation for future study. Collecting data in a 3D case would require much more time and computation resources. However, we believe our operator will scale well in 3D cases as a type of factorized neural operator kernel like FactFNO.\n\n**Q.12 Writing issues**\n\nThank you so much for carefully reading and for your valuable suggestions. Based on your feedback, we will fix typos and improve the figures. \n\nWe will consistently use \"complicated\" for device and \"complex\" for light-matter interaction.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "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": "UVtgl39pno", "reviewer": "Reviewer_weV6", "summary": "This paper proposes a neural operator (PACE) with improved accuracy, efficiency, and speed for electromagnetic field simulation. PACE is a network block in the surrogate machine learning model.", "questions": "1. How does this paper reflect “complicated” photonic devices, especially when compared with NeurOLight? I found that the datasets of both papers are custom-generated with random configurations and have a similar data volume. I expect the authors to provide more detailed explanations about the differences, particularly regarding the use of the term “complicated”, to avoid misleading readers into interpreting this paper as a simple incremental study.\n\n2. This paper mentions that the NeurOLight method fails for “real-world” photonic devices, but I did not see any evaluation experiments of the proposed method on real-world datasets in this paper.\n\n3. Is the speed improvement compared against CPU algorithms (line 308)? If the proposed method uses a GPU while the baseline Angler algorithm uses only a CPU, I do not think this is a fair comparison, and the corresponding 577x time boost could be misleading.\n\n4. Please avoid using exaggerated language, for example: “boost the prediction fidelity to an unprecedented level” (line 12), and “extremely challenging cases” (line 57).\n\n5. Unexplained descriptions, for example: “inspired by human learning” (line 57). There is no explanation provided as to why this method is inspired by human learning.\n\n6. Some of the paper copy-paste from the NeurOLight paper, for example, equations 1 and 2. Corresponding references and explanations should be added.\n\n7. Figure 4 is difficult to understand, and the caption is too brief.\n\n8. In Table 1, the phrase “Improvement over all baselines (geo-means)” is rarely used and may not be considered as a good evaluation in research.", "limitations": "Limitations are discussed in this paper but some points may be missed. For example, this paper works on only simulated and customized datasets.", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 2, "contribution": 2, "strengths": "This paper conducted extensive experimental evaluations for the proposed method. Also, the accuracy and efficiency of machine learning surrogate models for physical simulation are of great importance, therefore I think even though a little improvement is valuable.", "weaknesses": "This paper, including the writing and the experiments, is very similar to the main comparison NeurOLight, which has code available online. Therefore I think it is important to discuss if it is only an incremental work or in-depth research.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you so much for your time in reviewing our paper and providing valuable feedback on our work! Please see our responses below:\n\n**Q.1 ***Incremental work*** over NeurOLight with ***little improvement***?**\n\nWe would like to emphasize that our work is not simply incremental over NeurOLight but offers **novel machine learning contributions** and a **substantial accuracy improvement** over NeurOLight.\n\nFirst, our work aims to tackle the complicated photonic device simulation problem where NeurOLight shows unsatisfactory accuracy.\n* Please refer to Common Response Q. 2 for a detailed explanation of real-world and complicated devices.\n* TL; DR. Our work successfully established a new strong SOTA accuracy on complicated photonic devices with >39% lower error compared to the strongest NeurOLight.\n\nSecond, we have clear and unique machine-learning contribution:\n* **Novel neural operator design**: We propose a parameter-efficient cross-axis factorized operator design with multiple key components to unlock superior accuracy by deeply analyzing the challenges in complicated photonic device simulation in Sec. 3.1.\n* **Novel cascaded learning flow**: We divide the simulation task into two progressively latent tasks, reducing learning complexity and successfully showing highly competitive accuracy on hard cases.\n\n**Q.2 Real-world dataset**\n\nPlease refer to Common Response Q.2 for a detailed description of real-world.\n\nIn summary, the device we use has a real-world geometry, permittivity distribution, light source, and wavelength range. We randomly generate the device configurations to obtain sufficient data to learn the underlying rules in the photonic Maxwell PDE. This randomness also ensures thorough testing of our model’s effectiveness in predicting the optical field, not just fitting some specific examples.\n\nOur claim is that NeurOLight fails for real-world **complicated** photonic devices with complex light-matter interactions, such as scattering and resonance, as reported in their original paper. On the dataset where NeurOLight fails (Etched MMI 3x3), PACE achieves 39% lower error compared to NeurOLight.\n\n**Q.3 Speed-up comparison**\n\nWe focus on 2-D device-level finite difference frequency domain (FDFD) simulation. The numerical solver uses a CPU-based direct sparse linear solver (not the iterative solver), which is faster than the GPU version, as acknowledged in SPINS-B [1]. In 3D cases, the numerical solver uses an iterative solver where the GPU version is faster than the CPU version.\nNote that SPINS-B derives its conclusion using the slower scipy. However, in this work, we compare with the highly optimized Intel Pardiso solver on a 20-core CPU (NeurOlight only compares with slow scipy). We already provide a strong baseline and ensure a fair comparison when evaluating speedup.\n\nWe agree that 577x is somehow misleading. We will separately discuss the speedup over scipy and pardiso. We will change our wording to “In terms of runtime, our PACE model demonstrates 154-577x and 11.8-12x simulation speedup over numerical solver using scipy or highly-optimized pardiso solver, respectively”\n\n**Q.4 Unexplained human learning**\n\nPlease check our common response Q.3 to see how we are inspired by human learning in cognition.\n\n**Q.5 Similar equation and notation to NeurOLight**\n\nFor the Maxwell PDE description in Preliminary 2.2, as it is a common Maxwell equation, we honor the original notation. For the problem setup section, as we stated and cited in Line 104, we follow NeurOLight in formulating the optical device simulation as an operator learning problem. Therefore, we intentionally use the same notation and symbols to avoid confusion and ensure clarity. \n\nWe will adjust the wording and add more specific citations to prevent any potential misunderstanding.\n\n**Q.6 Why use Geo-means**\n\nTo compare different benchmarks, we first calculate the average improvement over all baseline methods on the same benchmark, then use the geometric mean (geo-mean) to assess the effectiveness of PACE across different benchmarks. The geo-mean fairly evaluates results across benchmarks, avoiding domination by large improvements that can skew the arithmetic mean. For example, for values 0.1, 0.5, and 0.9, the geo-mean is ~0.35, while the arithmetic mean is 0.5.\n\nWe will add this justification in the revised manuscript.\n\n**Q.7 Writing and unclear Fig. 4**\n\nThank you so much for pointing out places that can further improve our paper’s quality.\nDue to we cannot directly upload the revised version now, we will carefully modify the language in the final version to avoid exaggeration.\n\nFig. 4's caption will be updated to, “The proposed cascaded learning flow with two stages. The first stage learns an initial and rough solution, followed by the second stage to revise it further. A cross-stage distillation path is used to transfer the learned knowledge from the first stage to the second stage.” We will also modify Fig.4 to make it easier to understand.\n\n**Q.8 Limitation for customized dataset**\n\nGiven that no complete and open-sourced dataset is available for the photonic Maxwell PDE problem, we follow NeurOLight to generate a dataset with real-world device geometry by randomly varying device configurations, which is a common way to generate PDE benchmarks. Our devices reflect the characteristics of real-world devices and are sufficiently complex with complex light-matter interaction inside. \n\nMoreover, aware of the lack of an open-source dataset, we open-source the dataset to facilitate the community.\n\n\n**Reference**\n\n[1] SPINS-B (as we cannot provide the direct link in rebuttal, you can check the documentation of SPINS-B and read the claim in the faq section)", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "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:
We sincerely appreciate the constructive feedback provided by all the reviewers. We are truly inspired by their acknowledgment of our paper’s contribution and strong accuracy compared to various baselines as a new state-of-the-art in AI for optic simulation.

We appreciate the chance to address common comments in this response.

**Q.1 Relevance to NeurIPS audience: (Reviewer ZkVJ
 & TfAG)**

Our work tackles the challenging photonic device simulation problem with a novel **neural operator design** and a new **learning flow** and establishes a **strong SOTA** on AI for optics simulation compared to previous SOTA (NeurOLight @ NeurIPS’22). 

It aligns well with the interest of the NeurIPS:
* Photonic analog computing has gained much momentum, promising a next-generation efficient AI computing paradigm. A line of works on optics for AI [1-5] and AI for optics[6] appeared at NeurIPS or other ML conferences. Moreover, NeurIPS held a “Machine Learning with New Compute Paradigms” workshop last year and will do so this year. Our work can serve as an ultra-fast, high-fidelity surrogate model to accelerate the exploration of new customized devices and the design of scalable next-generation photonic computing systems, closing the loop of optics-AI synergy.

We have intellectual contributions to the AI for PDE community:
* A novel parameter-efficient cross-axis factorized neural operator.
* A new learning flow that solves hard PDE problems from coarse to fine.
* Our discrete permittivity challenges are similar to those in multi-scale PDE [7-8]. Since they don’t open-source their dataset, we compare the most recent [8] with PACE on our benchmark(see 5.3 Discussion.4), where PACE shows a significantly lower N-MSE error (10.59% vs 17.4%)
* We open-source the dataset and plan to include more devices and examples. This will contribute a new PDE problem to the AI for the PDE community, and the new PDE problem will have a huge impact on photonic design.

We believe our work will inspire more "AI for optics" works, and our exploration shares insights for "parameter-efficient ML PDE surrogate model" and "how to enhance neural operator for challenging PDE cases."

**Q.2 Clarification of ***real-world*** and ***complicated*** photonic devices (Reviewer ZkVJ
 & weV6)**

Let’s first recall the problem formulation: find a mapping from the observation space  Ω, 𝜖, 𝑤, 𝐽 to the underlying optical field 𝑈. 

We claim to be **real-world** since
* The used photonic device really exists[9-10], meaning the geometry (Ω), and permittivity(𝜖) are real-world.
* The input light source 𝐽 is real, and we test our model with various randomly generated 𝐽 combinations.
* The wavelength (𝑤) range we focus on is C-band (1.53-1.565), a commonly interested range for device design.
* Moreover, permittivity distribution 𝜖 is discrete and contrastive considering the real manufacturing limits. However, many PDE benchmarks generate coefficients that are slowly/smoothly transited.

Our devices are **complicated** since (details at Sec 3.1)
* Complex light-matter interaction, e.g., scattering and resonance.
* Sensitivity to local minor structures
* Rich frequency information in the optical field

And NeurOLight fails to give decent accuracy on these **complicated** devices, as they reported.

**Q.3: unclear explanation of human learning (Reviewer TfAG & weV6)**

Thanks for pointing out that our term “human learning” is unclear. 
We added the explanation here and will include it in the revised version.

Existing ML for PDE solving work typically learns a model in a one-shot way by directly learning the underlying relationship from input-output pairs. Unlike AI systems, humans don’t learn new and difficult tasks in a one-shot manner; instead, they learn skills progressively, starting with easier tasks and gradually moving to harder ones. For example, instead of directly learning how to solve equations, students first learn basic operations, such as addition and multiplication, and then move on to solving complex equations.

Inspired by this human learning process, unlike previous work that directly learns a one-stage model, we propose to divide the challenging optical field prediction problem into two sequential latent tasks. The first learns to generate a rough initial solution, followed by the second to refine details.

[1] Li, Yingjie, et al. "RubikONNs: Multi-task Learning with Rubik's Diffractive Optical Neural Networks." IJCAI 2023.

[2] Gu, Jiaqi, et al. "L2ight: Enabling on-chip learning for optical neural networks via efficient in-situ subspace optimization." NeurIPS 2021.

[3] Gu, Jiaqi, et al. "Efficient on-chip learning for optical neural networks through power-aware sparse zeroth-order optimization." AAAI 2021.

[4] Ohana, Ruben, et al. "Photonic differential privacy with direct feedback alignment." NeurIPS, 2020.

[5] Gupta, Sidharth, et al. "Don't take it lightly: Phasing optical random projections with unknown operators." NeurIPS 2019.

[6] Gu, Jiaqi, et al. "Neurolight: A physics-agnostic neural operator enabling parametric photonic device simulation." NeurIPS 2022.

[7] Liu, Xinliang et.al. "Mitigating spectral bias for the multiscale operator learning with hierarchical attention." arXiv preprint arXiv:2210.10890 (2022).

[8] Bo, Xu, et al. “Dilated convolution neural operator for multiscale partial differential equations.” arXiv preprintarXiv:2408.00775v1 (2024).

[9] Mohammad, Tahersima, et.al. “Deep neural network inverse design of integrated photonic power splitters.” Sci. Rep., 2019.

[10] Sanaz Zarei et al. “Integrated photonic neural network based on silicon metalines”. Optics Express 2020.

Meta-review: The paper proposes a new neural operator setup for solving Maxwell equations in 2D. All reviewers recommend acceptance, so it is difficult to argue; I will add, however, that some of the claims are quite misleading. There is a speedup compared to the classical solver, but nothing is said about the speedup for the comparable accuracy; thus, this has to be correctly phrased in the final version. It is ok to beat n.o. baseline, but it is not ok to claim being faster but for a different accuracy.
Decision: Accept (poster)