Graph Neural Networks for Interferometer Simulations

We thank the reviewer for their detailed review. We agree with the reviewer that a more detailed ablation study and more complete comparisons of results would improve the paper; they shall be added before the end of the discussion period.

In response to the reviewer’s questions and comments:

1. We will add the ablation study. We would like to clarify that some of the more advanced architectures were in fact tried. In particular, we did experiment with Graph Transformers, but saw no significant advantage; this will be included in the ablation study.

2. Edge features are used; in particular, the separation between optics and their indices of refraction are encoded as edge features in our graph.

3. Figure 3 demonstrates the success of the model at the power prediction task. We will add discussion of this to the paper.

4. Regarding the question on FINESSE: In order to compute the power at each node, FINESSE does need to compute other quantities (in particular the amplitudes of each individual Hermite Gauss mode, and the beam parameters of each cavity). However, we still feel that this is a fair comparison, for two reasons. First, we consider the power to be a valuable metric for measuring the quality of a particular interferometer design. The fact that FINESSE must compute other quantities in order to get to that power is precisely where the computational cost of the FINESSE simulation arises from, and why we believe that this model serves a useful function. Secondly, the intensity prediction model predicts the spatial distribution, and thus indirectly predicts the modal decomposition, which is also what FINESSE computes under the hood. Seeing as the intensity prediction model achieves similar speedups to FINESSE/SIS, we feel that this comparison is still fair.

5. Regarding the direct applicability of this model to gravitational wave interferometry, we will add more discussion of this in the revision. However, as a short synopsis: This model is promising in that it does capture many of the important physical phenomena present in the interferometer. Namely, it is able to predict spatial intensity profiles, and thus modal decompositions of the EM-field. This is enough to be useful in accelerating optimizations of basic design parameters like cavity lengths and mirror radii of curvature. The interferometer topologies considered here range from simple (Fabry-Perot) to close to full interferometer complexity (half-ALIGO). Therefore, we feel that the most important limitation of this work is not in the complexity of the interferometer topologies considered, but in the fact that we only model a subset of the physics (power and spatial intensity distributions). To increase the applicability of these models, we would ideally need to account for more of the physical effects that are present, such as higher order mode scattering due to thermal lensing and surface defects in the mirrors. However, we will reiterate that the goal of this work was not to model all of the complexity of the LIGO interferometer, but instead to demonstrate that GNNs are capable of capturing the complex dependencies of the field propagation in interferometers, and thus are a promising candidate for simulation acceleration.

We thank the reviewer again for the comments. We hope the revisions we made have satisfactorily addressed them.

1. We would like to clarify that for the intensity prediction model, the model described as “MLP model” still included GAT layers. We meant that an MLP was used to learn the radial intensity distribution instead of a KAN. This has been clarified in the text.
2. We have added a much more detailed and full presentation of results in Table 2.
3. We have added more discussion on the next steps required to make the model more useable in interferometer design.

1. Due to limited time and compute resources, we had to choose which experiments to prioritize. The intensity prediction model takes significantly longer to train, and so we chose to prioritize the power prediction task, as this is the more important metric for interferometer design. However, we can prepare more complete results here for the camera-ready, should the paper be accepted. In particular, we can do a similar comparison in the style of Table 2 for the intensity prediction task.

2. The authors are unclear on what the reviewer means by "parameters". The ground truth data is generated in the following way: A "base" interferometer design is randomly chosen, and then a random walk in parameter space is performed starting at the base design. This is repeated for a large number of base designs. If the reviewer could clarify which parameters they are referring to, we would be happy to clarify what values were used.

3. We agree that the GraphTransformer is roughly on par with the GAT+MLP combo. However, we chose not to evaluate it much further, again due to the time constraint, and because the inference time using GraphTransformer was higher than with GAT, without performance gain. As the core objective of our model is to enable fast simulation, we felt that it was lower priority to attempt to explore models that provide marginal improvements in accuracy at the cost of higher computational cost. A similar logic applies to the GAT+KAN model. KAN is roughly an order of magnitude slower at inference time than the GAT+MLP combo, and additionally performs slightly worse on two of the three tested topologies. We would be happy to add exact figures for the computational efficiency of each tested model.

4. For the experiments added in the revision, we tested with only a single seed, yes. This seed was fixed across all experiments.

5. The graph transformer implementation used is the one implemented in PyTorch Geometric (https://pytorch-geometric.readthedocs.io/en/latest/generated/torch_geometric.nn.conv.TransformerConv.html). No positional encodings are used, and yes, we used the same 20 layer architecture as the GAT.

We thank the reviewer for the feedback provided. However, we would like to clarify a few points.

1. The work does for the most part use standard layers in our model. However:

a. We do modify some aspects of the training procedure to reflect the physics of our system. For example, we enforce energy conservation in the power prediction model by using a modified loss function, and we use the symmetries of the EM field to reduce the number of parameters of the intensity prediction model. While not direct improvements to GNN architecture itself, we feel that domain specific considerations like these make our work an example of techniques that can be used to enforce physical constraints in ML for science applications.

b. The goal of our work is not to provide architectural improvements to a network architecture, but rather to demonstrate an application of GNNs. While future work could certainly pursue larger architectural changes (for example, drawing from research in neural operator learning or multiphysics emulation), this work is a first foray into applying ML to interferometer simulation (and more broadly scientific instrument design).

c. We will also point out that the paper explores certain relatively new approaches in machine learning for science, like Kolmogorov Arnold networks, which were only introduced earlier this year, and are as of now still being studied to determine if and where they provide advantages over MLPs. Our work clearly demonstrates a use case in which KANs greatly outperform a similar MLP based architecture, and thus we feel that this paper may also be of interest to the broader ML community.

2. We agree with the reviewer’s criticism that the comparison to other model architectures is insufficient. We note that this is the first work applying ML to interferometer simulations, and thus we cannot compare it to other work done for this specific problem. However, we will add more detailed studies comparing our model performance to graph based and non-graph based ML model architectures, in order to help clarify the design decisions that we made.

3. We understand that the certain aspects of the content or the figures may have been unclear. If the reviewer has specific questions or critiques, the authors would be happy to hear them and to clarify them both here and in the paper itself. On the specific point of instrument design, we discuss this in the introduction and the conclusion, as this is the context in which this work was conceived. However, the specific challenges of instrument design more broadly are less relevant to this work than the specifics of interferometer physics that we discuss.

We appreciate the reviewer’s time in providing this review and thank the reviewer for the useful feedback.

Regarding the reviewers questions: “What are the practical benefits of using this AI model… Have you observed that the AI model can provide better manufacturing design parameters?” Other works (Richardson et. al 2022, for example) identify the existence of better manufacturing parameters than the current LIGO design, but have been restricted to exploring a subset of the interferometer’s degrees of freedom, due to the computational cost of simulation. Thus, the main benefit of our AI model, as the reviewer states, is in the accelerated simulation times. During the design phase of GW detectors, there is a vast parameter space that must be searched to identify the design which will maximize the sensitivity. This design phase is currently ongoing for Cosmic Explorer, a future interferometer, and so tools to improve this optimization/exploration of design space are intended to assist in that effort. This will be clarified in the paper text.

“Can you provide more insights on the rationale of designing such a 20-layer GAT + 6 layer MLP or 1 layer KAN kind of model?” This was an empirical decision, based on experiments conducted with different architectural parameters. A more detailed ablation study will be added to the paper, along with a discussion of the architectural choices.

Regarding the comments on clarifications to the dataset: We will modify the figures and explanations to clarify how the dataset is mapped to a graph, as well as the requested table. For the reviewer’s ease, we will provide a short description here as well. The datasets each contain around 10,000 graphs. Each graph represents one interferometer simulation, where each node in the graph represents a component of the EM field at a particular optic (mirror, laser, etc.). The edges represent the properties of the space between those optics (separation distance, index of refraction, etc.). A full-ALIGO dataset is not provided, as while it would be the end goal of such a study to model a full-ALIGO like interferometer, the half-ALIGO model captures much of the important physics also present in the full-ALIGO case.

We thank the reviewer again for the comments. We hope the revisions we made have satisfactorily addressed them.

1. We add a summary of datasets in Table 1, along with more description of how the data was collected in the text of the paper (Section 4.1).
2. We add a table comparing the performance of the network with depth, to explain why we chose the architectural parameters that we did.

We thank the reviewer for the in-depth and constructive critique. At a high level, we agree that the paper would benefit from a more in-depth ablation study and better embedding into the existing literature. As doing these requires running further experiments and understanding the references that the reviewer provides, we will require some time to revise the paper to address these critiques. We will do so by the end of the discussion period.

Regarding the questions that the reviewer poses about the simulation speedups and fidelities, we provide the following clarifications:

We believe that part of the confusion is the dual contexts in which we use the word “fidelity.” When we refer to costly, high-fidelity optical simulations, we mean that in order for the optical simulations to produce physically accurate results, a large number of higher order beam modes must be included. Our GNN model operates at this high fidelity. When we say that the model predictions would still be useful at low fidelities, we mean that a somewhat inaccurate model is still useful as a heuristic for search space exploration. We have changed the terminology to say that the model predictions do not have to be very high accuracy.

The GNN is 815x faster than an SIS simulation of the Half-ALIGO interferometer topology. The GNN is 159x faster than a FINESSE simulation of the same topology. The FINESSE model is run on CPU, as FINESSE does not currently have GPU support. The SIS simulations are also run on CPUs, though SIS does support GPUs. We will include a comparison with SIS run on GPU as another baseline. All three models are run at equivalent “fidelities”.

We will add revisions regarding the depth of evaluation critique once the experiments conclude.