Title: PACE: Pacing Operator Learning to Accurate Optical Field Simulation for Complicated Photonic Devices
Venue: NeurIPS 2024
Accepted: yes
Keywords: AI for Scienece

Optical simulation

Neural opeartor

AI for PDE
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).


Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis 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.\n\nSTRENGTHS:\nThis 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.\n\nWEAKNESSES:\nThis 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.\n\nQUESTIONS:\n1. 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.\n\nLIMITATIONS:\nLimitations are discussed in this paper but some points may be missed. For example, this paper works on only simulated and customized datasets.", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2024-07-03"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe 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.\n\nSTRENGTHS:\nStrengths:\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.\n\nWEAKNESSES:\n1. 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.\n\nQUESTIONS:\n1. 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?\n\nLIMITATIONS:\nThe 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.", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 4, "DATE": "2024-07-13"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis 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.\n\nSTRENGTHS:\n1. 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.\n\nWEAKNESSES:\n1. 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.\n\nQUESTIONS:\n1. 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.\n\nLIMITATIONS:\nNA", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 4, "DATE": "2024-07-13"}

comments: 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.

STRENGTHS:
Efficiently and accurately simulating photonic devices is an important problem with broad applications.

THe 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". 

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

QUESTIONS:
Are the "real-world" devices actually manufactured and measured or merely simulated? I found section 5.1 unclear.

LIMITATIONS:
Yes
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 3
DATE: 2024-07-14