PACE: Pacing Operator Learning to Accurate Optical Field Simulation for Complicated Photonic Devices

Thank you so much for your time in reviewing our paper and providing valuable feedback on our work! Please see our responses below:

Q.1 Incremental work over NeurOLight with little improvement?

We 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.

First, our work aims to tackle the complicated photonic device simulation problem where NeurOLight shows unsatisfactory accuracy.

• Please refer to Common Response Q. 2 for a detailed explanation of real-world and complicated devices.
• TL; DR. Our work successfully established a new strong SOTA accuracy on complicated photonic devices with >39% lower error compared to the strongest NeurOLight.

Second, we have clear and unique machine-learning contribution:

• 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.
• 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.

Q.2 Real-world dataset

Please refer to Common Response Q.2 for a detailed description of real-world.

In 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.

Our 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.

Q.3 Speed-up comparison

We 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. Note 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.

We 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”

Q.4 Unexplained human learning

Please check our common response Q.3 to see how we are inspired by human learning in cognition.

Q.5 Similar equation and notation to NeurOLight

For 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.

We will adjust the wording and add more specific citations to prevent any potential misunderstanding.

Q.6 Why use Geo-means

To 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.

We will add this justification in the revised manuscript.

Q.7 Writing and unclear Fig. 4

Thank you so much for pointing out places that can further improve our paper’s quality. Due to we cannot directly upload the revised version now, we will carefully modify the language in the final version to avoid exaggeration.

Fig. 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.

Q.8 Limitation for customized dataset

Given 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.

Moreover, aware of the lack of an open-source dataset, we open-source the dataset to facilitate the community.

Reference

[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)

Thank you so much for your appreciation of our contributions and for providing valuable feedback on our work! Please see our responses below:

Q.1 Use “coarse to fine” instead of “divide and conquer”

Thank 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.

Q.2 Testing errors of traditional numerical solver

Since 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.

In 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.

Q.3 Include run-time in Tab. 1

We 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.

Our PACE shows excellent accuracy, superior parameter efficiency, and better run-time compared to the previous SOTA NeurOLight.

We’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].

Moreover, 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.

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.

Q.4 Incomplete ablation study on pre-normalization and double skip:

We 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:

As 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.

Q.5 Visualization of feature map before/after non-linear activation in our explicitly designed high-frequency projection path

We 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.

Q.6 Writing

Thank 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.

We 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.

[1] Tran, Alasdair, et al. "Factorized fourier neural operators."ICLR 2023.

Dear Reviewer 9WdN,

We sincerely appreciate your valuable comments and suggestions for our paper.

We have carefully addressed all your comments in our response above. Regarding your questions,

• We have added an experiment for the ablation study to test the effectiveness of pre-normalization and double skip. The results confirm that pre-normalization and double skip are not the main sources of improvement. Our proposed new neural operator block with many essential model design considerations, such as cross-axis factorized operator and high-frequency projection path, is the key to unlocking the improvement.
• We have provided extra visualization of the feature map before/after nonlinear activation in Rebuttal Figure 1. It validates our model design choice by clearly showing that the high-frequency feature is generated after the nonlinear.
• Runtime is provided as you suggested, and the tab. 1 will be updated in the revised version. (please check our response Q3)

As the discussion period is drawing to a close in 1 day, we would be very grateful if you could take a moment to review our responses. If our replies have satisfactorily addressed your concerns, we would greatly appreciate it if you could acknowledge this and consider raising the rating.

If you have any remaining questions or concerns, please do not hesitate to reach out to us. We look forward to your feedback.

Dear Reviewer 9WdN,

Thank you once again for your review. As the deadline for the author-reviewer discussion approaches, we noticed that we haven't received any comment from you.

We have addressed all your questions with additional experiments and clarifications:

• We have added an experiment for the ablation study to test the effectiveness of pre-normalization and double skip. The results confirm that pre-normalization and double skip are not the main sources of improvement. Our proposed new neural operator block with our carefully designed model components, such as cross-axis factorized operator and high-frequency projection path, is the key to unlocking the improvement.
• We have provided extra visualization of the feature map before/after nonlinear activation in Rebuttal Figure 1. It validates our model design choice by clearly showing that the high-frequency feature is generated after the nonlinear.
• We provide the runtime (please check our response Q3) and the tab. 1 will be updated in the revised version.

As the discussion period is coming to an end, we would greatly appreciate any additional feedback you might have. If our responses have clarified your understanding of our paper, we sincerely hope you might consider raising the rating.

Thank you again for your effort in reviewing our paper.

Best regards, Authors of Paper 472

Thank you so much for your valuable feedback on our work! Please see our responses below:

Q.1 Narrow audience and other benchmarks

We'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)

The 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.

We 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.

Q.2 Used errors:

We 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).

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.

Why normalized MSE, not regular MSE? We 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.

As 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).

Q.3 Speedup over numerical methods with less accuracy:

We 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).

As 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.

Q.4 Unclear human learning

Please see common response Q.3.

Q.5 Spectrum of the predicted field

The 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.

Q.6 Etymology of PACE:

Sorry for not annotating PACE’s source in the title: Pacing Operator Learning to Accurate Optical Field Simulation for Complicated Photonic Devices.

It 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.

Q.7 Ood testing

Firstly, 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.

As 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.

Q.8 Speed-up experiments

In speed-up evaluation on various device scales, we only test the speedup by varying device geometry sizes; therefore, no training is performed.

For 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).

Q.9 Seemingly inconsistent colorbars A.9

In 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).

Q.10 Error in log space

We 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).

Q.11 3D cases

We 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.

Q.12 Writing issues

Thank you so much for carefully reading and for your valuable suggestions. Based on your feedback, we will fix typos and improve the figures.

We will consistently use "complicated" for device and "complex" for light-matter interaction.

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

[2] Bo Xu and Lei Zhang. “Dilated convolution neural operator for multiscale partial differential equations.” arXiv preprintarXiv:2408.00775v1 (2024).

[3] Timurdogan, Erman, et al. "APSUNY process design kit (PDKv3. 0): O, C and L band silicon photonics component libraries on 300mm wafers." Optical Fiber Communication Conference. Optica Publishing Group, 2019.

[4] Github of Latent-Spectral-Models

Dear Reviewer TfAG,

We sincerely appreciate your valuable comments and time. You provided a very high-quality discussion and made important suggestions to further improve our paper writing and figures.

We have carefully addressed all your comments in our response above, with extra experiments as you suggested. We would like to point out that our paper aligns with NeurIPS’s interests and offers novel ML contributions to the AI for PDE community. We built a new and stronger SOTA compared to the previous SOTA (NeurOLight @ NeurIPS’22).

As the discussion period is drawing to a close in 1 day, we would be very grateful if you could take a moment to review our responses. If our replies have satisfactorily addressed your concerns, we greatly appreciate it if you could acknowledge this in the discussion thread and consider raising the rating.

If you have any remaining questions or concerns, please do not hesitate to contact us. We look forward to your feedback.

Dear Reviewer TfAG,

Thank you once again for your review. As the deadline for the author-reviewer discussion approaches, we noticed that we haven't received any comment from you.

We have addressed all your questions with additional experiments and clarifications.

• Regarding your main concern about our scope, we would like to point out that our paper aligns with NeurIPS’s interests and offers novel ML contributions to AI for the PDE community (See global response Q.1), with a new parameter-efficient cross-axis factorized neural operator and a new learning flow for solving challenging PDE cases.
• We would point out that no one model can rule all PDE cases. Our work shows important insights into how to enhance neural operators for challenging PDE cases. Moreover, we compare with one SOTA work for the challenging multi-scale PDE problem where previous neural operators also fail (similar to our challenge in that the PDE coefficient is highly contrasted). Because they didn't release the dataset, we compared our work with theirs on our benchmark. Our work shows significantly better accuracy than theirs, proving it can provide important insight into the multi-scale PDE problem.

As the discussion period is ending, we would greatly appreciate any additional feedback you might have. If our responses have clarified your understanding of our paper, we sincerely hope you might consider raising the rating.

Thank you again for your effort in reviewing our paper.

Best regards, Authors of Paper 472

Dear authors, Thanks for the detailed point-by-point clarifications to my review. I went through the rebuttal and the attached pdf and believe that it addresses most of the points I raised.

The following questions have been answered sufficiently: Q1, Q2, Q4, Q5, Q6, Q7, Q8, Q9, Q10, Q11, Q12.

Please add the plots and explanations for Q5, Q7, Q9, Q10 in the camera-ready version.

Your explanation for relevance to NeurIPS (and energy-efficient ML hardware in general) should also be incorporated in the papers as it can serve as a strong motivation for general ML readers. Also add the origin of PACE acronym somewhere early in the paper.

I am not totally convinced by the speedup claims as the comparisons are not one-to-one wrt accuracy and computational time/hardware, but the authors made a good-faith effort to address this. Further exploration would be outside the scope of this paper at this point. For more information on assessing speed-up claims for ML-PDE models, please refer to this paper[1], specifically eq 1 on page 11.

I also agree with the other reviewer that coarse-to-fine is a better description than divide-and-conquer. You can also frame it as an iterative model. In this context, using the term human learning feels a little bit hyperbolic.

Finally, please proof-read the paper for typos that are mentioned in Q12 and also for those that might have been missed.

I am hoping that the mentioned changes are incorporated in the camera-ready version of the paper, and raising the score assuming this.

[1] McGreivy and Hakim, Weak baselines and reporting biases lead to overoptimism in machine learning for fluid-related partial differential equations. https://arxiv.org/abs/2407.07218

Thank you so much for your valuable feedback on our work! Please see our responses below:

Q.1 Somewhat narrow audience

We 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.

Q.2 Confusing description in speedup

Thank 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”.

Q.3 Unclear description of “real-world” devices

The 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)

As 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.

Q.4 Typos

Thank 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.

Thank you so much for reviewing our paper and this response. We look forward to further discussions with you if you have any other questions!

Dear Reviewer ZkVJ,

We sincerely appreciate your valuable time in reviewing our paper.

We have carefully addressed all your comments in our response above. We would like to point out that our paper aligns with NeurIPS’s interests and offers novel ML contributions to AI for the PDE community. We built a new and stronger SOTA compared to the previous SOTA (NeurOLight @ NeurIPS’22). Moreover, we follow your suggestion to clarify the unclear description, and those parts will be merged into the revised version. Thank you again for your suggestion, which is crucial to improve our paper writing.

As the discussion period draws to a close in 1 day, we would be very grateful if you could take a moment to review our responses. If our replies have satisfactorily addressed your concerns, we would greatly appreciate it if you could acknowledge this in the discussion thread.

If you have any questions or concerns, please do not hesitate to contact us. We look forward to your feedback.