Neural Graph Matching Improves Retrieval Augmented Generation in Molecular Machine Learning

We thank the reviewer for recognizing the state-of-the-art accuracy of MARASON and the organization of this paper. We believe there are several misunderstandings and we clarify them as follows.

The paper does not describe a generation process and does not utilize RAG.

The problem MARASON aims to tackle is generating MS/MS from a given molecular structure. On the technical side, our first-stage model is an auto-regressive generator that predicts one bond-breaking event at each step.

The paper does not clarify which specific type of mass spectrometry it supports.

We will emphasize that MARASON is developed with ESI-MS/MS in the main text. As discussed in section 4.4.1, “We trained our models on the NIST 2020 dataset with 530,640 high-energy collision-induced dissociation (HCD) spectra and 25,541 unique molecular structures.” We believe we made it clear that our experiment setting contains MS/MS spectra (more specifically, Orbitrap spectra) with an instrument type label “HCD” in the NIST database.

Please report separately the results on the FT-HCD dataset and the FT-CID dataset.

To clarify, as discussed in section 4.4.1, we only use spectra labeled as “HCD” instrument type in NIST. We do not find any entries with instrument type labels CID or FT-CID in NIST. We believe we are working with an open-source training and testing framework that supports most baselines and all methods are trained and tested on the same dataset, ensuring fair comparison.

Questions on model architecture details and GNN layers.

We would like to clarify that MARASON adopts the same backbone architecture as ICEBERG, including both the GNN and transformer layers.

Regarding the GNN blocks shown in Figure 2: the first and third GNNs operate on molecular graphs and are structurally identical to their counterparts in ICEBERG. The second GNN, discussed in Section 3.3.2 under “DAG Hierarchical Embedding Learning,” serves as the neural graph matching module. While all GNN blocks share the same architecture, our strategy involves using separate weights depending on the module's function—i.e., whether it is used for intensity prediction or for graph matching.

We will include more detailed discussions and highlight such differences in Figure 2 in future revisions.

I recommend citing additional references for the correlation between spectrum and structural similarities.

Thank you for the suggestion, we plan to include the following references on molecular networking. Please also let us know if there are any specific references in your mind.

• Aron AT et al. Nature protocols. 2020
• Wang M et al. Nature biotechnology. 2016

MoMS-Net and 2DMolMS are missing.

We do not find any open-source implementation for MoMS-Net and the dataset statistics in Table A1 are different from ours, making it challenging to include MoMS-Net for comparison. We do not find any references to 2DMolMS—if you meant 3DMolMS, it is the first entry in Figure 3 and Table 1.

How to prevent data leakage is unclear.

Thanks for the suggestion. One possible concern of running RAG on MS/MS simulation is that the retrieved spectra might be considered as a source of data leakage that simplifies the problem. Therefore, in numerical evaluations, we restrict retrieving spectra only from training data; in real-world cases, there is no doubt about using any reference spectra available for MARASON. We will elaborate in future revisions.

Why retrieve three reference spectra with different collision energies?

Since the reference spectra may not have the exact collision energy of interest, MARASON tries to interpolate from three collision energies. We did not try including more references because the more reference spectra we have, the more computational overhead there will be.

Does this framework explicitly account for experimental variables?

Our framework handles ionization mode by different adduct types. We incorporate collision energy as an important instrument variable. For others, we assume they are the same across Orbitrap instruments. If any important factors are missing, please let us know and we would be more than glad to address them in future work.

Testing on NIST23 for results with a larger dataset

We purchased the NIST23 license, but the raw data file (.SDF) we used for NIST20 is not available at least in the NIST23 distribution we purchased. We will certainly explore NIST23 after figuring out how to extract training data from it.

Within the timeframe of rebuttal, we try to address this concern from another direction by showing that MARASON achieves state-of-the-art retrieval accuracy on MassSpecGym (please refer to the table in response to reviewer VQAR).

Recommend to elaborate a bit more on Supplementary Material.

As we cannot edit the supplementary materials at this time, we will try to include more details in future revisions. If there are any specific points that feel unclear to you, please let us know.

We truly appreciate your recognition of our contribution to the mass spectrometry field and our technical novelty of introducing neural graph matching. We conduct statistical tests and perform preliminary MassSpecGym experiments following your suggestions. We will work actively to complete the new results in future revisions.

Conducting statistical tests for the results described in Section 4.2.1 will make claims stronger.

Thank you for the suggestion. We first generate results with 3 random seeds for the scaffold split and perform a t-test on ICEBERG (w/ collision energy) and MARASON. The P-values are 0.005 and 0.002 for random split and scaffold split, respectively, both above the 95% confidence interval of statistical significance.

It would also be interesting to see examples of predicted spectra compared to the original spectra and the spectra of the reference compound. For reference, similar qualitative results were presented in earlier works. Such plots would demonstrate how similar the predicted and reference spectra are.

Thank you so much for the suggestion, unfortunately, we do not have the option to include visualization results with Openreview. We will add the visualization of reference spectra, predicted spectra, and ground-truth spectra in the appendix in future revisions.

To enhance this paper, an evaluation of a second dataset is recommended.

We truly appreciate your suggestion. We have run our method on the recently developed MassSpecGym benchmark [1], which is a publicly accessible library with spectra collected from MoNA, MassBank, and GNPS. We share an initial result of its performance here:

MARASON outperforms FraGNNet, the state-of-the-art on MassSpecGym, in terms of retrieval accuracy. We get the aforementioned result within the tight rebuttal time frame, and we will keep working on this benchmark with a comprehensive study in future revisions.

[1] Bushuiev et al. MassSpecGym: A benchmark for the discovery and identification of molecules. NeurIPS 2024

Section 3.2.2 says: "All reference intensities at the same collision energy are processed by a set transformer, followed by an average pooling layer that merges intensity embeddings per fragment from three collision energies." I understand that there are three different energies for the reference compound, yet here processing intensities at the same collision energy are described. Could you elaborate on this process in more detail? What vectors serve as input to the set transformer?

You are right about how we process reference spectra at each collision energy. Our purpose of using three reference spectra is that the target collision energy may not have a close match energy in the reference database and multiple reference spectra could be considered together to interpolate. We first concatenate spectral peaks with their corresponding collision energies and the target collision energy. After that, we feed the concatenated spectral vectors into a set transformer and eventually a linear layer, whereby our design of the linear layer is to “move” the peak intensities to the target energy level. Finally, we use average pooling to collect information from the same peaks at different energy levels to generate a reference spectrum embedding that is a learned interpolation from three energies.

How is this model trained? Are there any new components in the loss function? To train neural graph matching, do you use any additional training steps, or is it trained alongside intensity prediction as part of the vector in Equation 7?

The first-stage model that generates fragments is trained in the same way as ICEBERG-Generate. The second-stage model that predicts intensities is trained end-to-end, where the cosine loss between predicted spectra and ground truth spectra is the only supervision. We are able to make such a design choice because being fully differentiable is one of the major advantages of neural graph matching.

Thank you for agreeing with the novelty and technical solidity of our paper. We provide more random seeds, new MassSpecGym results, and MassFormer baseline with collision energy as per your comments. Please find our reply to your questions as follows.

Why does the scaffold split experiment have only 1 random seed?

In the ICEBERG paper, there is only one random seed with scaffold split. Following your suggestion, we run MARASON on 3 random seeds on scaffold split with the following updated results:

From the results, MARASON’s variance to random seeds is less significant compared to the accuracy improvement (for handy reference, ICEBERG w/ collision energy has a cosine similarity of 0.711) and the trend is consistent with either random or scaffold split. We will update the results in future revisions.

Request to evaluate on MassSpecGym

Although the rebuttal time frame was quite short, we retrained our model and were able to get the following retrieval accuracies on MassSpecGym. We demonstrate the superiority of MARASON on MassSpecGym for candidates with the same formula

We will perform more benchmarking and update results in future revisions.

Questions on model design choices and the "neutral loss" embedding.

We would like to clarify that the model architecture was adopted from ICEBERG and we hypothesized that the GNN architecture of ICEBERG learns the structural information essential for graph matching. To elaborate on encoding the “neutral loss”, since $\mathcal{M}, \mathcal{F}_i, \mathcal{M} - \mathcal{F}_i$ represent precursor, ionized fragment, and neutral loss, respectively, subtracting the embedding of $\mathcal{F}_i$ from $\mathcal{M}$ represents the embedding of the neutral loss. We will elaborate on these details in the revised manuscript.

Why are both the forward and reverse graphs needed in DAG embedding?

When embedding the fragmentation DAG, the forward graph handles parent-to-children message passing and the reverse graph handles children-to-parent. We believe the update functions of both directions should be distinct to avoid degeneration into an undirected graph. Empirically, this strategy was more powerful than using bi-directional GNNs during our model development.

How different collisions are incorporated?

We treat spectra with different collision energies as distinct spectra. The collision energy value becomes another input dimension that is concatenated to the GNN input. In NIST retrieval experiments, we compute spectra at all collision energies, compare each one with its corresponding real spectrum, and compute the average cosine similarity over all recorded collision energies. The preliminary MassSpecGym experiments follow the official benchmark setting, which we will elaborate upon in the revised version.

Are the data splits identical to FraGNNet?

As discussed in the FraGNNet paper (Appendix I), their NIST benchmark follows MassFormer and ICEBERG to ensure fair comparison, which should align with our benchmark.

It's also possible to improve other models with RAG.

We absolutely agree with that; we select ICEBERG as the base model because it is the current open-source state-of-the-art. Another conclusion is that neural graph matching is vital for the success of RAG—just concatenating the reference spectrum to the neural network even harms the test performance, as shown in Table 2. Ultimately, we have demonstrated a RAG approach that can be readily adapted to other modeling tasks in the molecular spectroscopy domain.

Is it possible to improve other baselines with collision energy?

Yes. We applied the same design to MassFormer and validated changes to retrieval accuracies on NIST’20:

It shows a marginal improvement but still does not outperform MARASON.

Why did the authors choose 0.1 Da? Is it possible to demonstrate with more fine-grained binning?

Since we want to compare all methods under the same metric, every model’s output should be transformed into the same mass resolution. 0.1 Da is the result of balancing between the mass resolution of HCD/Orbitrap spectra and the feasibility of implementing binned-prediction baselines. Empirically, experimentalists also find rounding to 0.1 Da resolution is sufficient in practice. MARASON (and ICEBERG) are adaptable with any mass resolution and can be run with higher resolution; however, selecting a finer-grained resolution will make binned-prediction methods unfeasible as a 0.01 Da resolution with a maximum mass of 1500 Da will require a 150,000-dim output in their neural networks.

Typo on page 3

Thanks for pointing this out. It has been fixed.

Thank you for recognizing our state-of-the-art performance and our technical soundness. We update with experiments on the recently developed MassSpecGym benchmark, evaluation with non-similar reference structures, and elaboration on the ablation study to address your concerns. We are more than happy to clarify any further questions.

The model's training and evaluation were exclusively conducted using the NIST dataset, with no validation performed on alternative publicly accessible mass spectral libraries (e.g., MoNA).

We truly appreciate your suggestion. We have run our method on the recently developed MassSpecGym benchmark [1], which is a publicly accessible library with spectra collected from MoNA, MassBank, and GNPS. We share an initial result of its performance here:

MARASON outperforms FraGNNet, the state-of-the-art on MassSpecGym, in terms of retrieval accuracy. We get the aforementioned result within the tight rebuttal time frame, and we will keep working on this benchmark with a comprehensive study in future revisions.

A bit more clarification on why we focused on NIST: as discussed under “software and data” on page 9, NIST is the only database where all spectra have collision energy annotations. In most open-source MS/MS libraries, collision energy labels are at least partially missing, making it challenging to develop a RAG model with them. It is still a reasonable assumption to have access to collision energy values for prospective use because it is an experimental variable set on MS/MS instruments.

[1] Bushuiev et al. MassSpecGym: A benchmark for the discovery and identification of molecules. NeurIPS 2024

Under circumstances where the mass spectral database lacks structural analogs, could the RAG strategy potentially lose efficacy due to failed similarity-based retrieval?

Thank you for bringing up this important point. MARASON benefits from RAG with Tanimoto similarity > 0.3 (90.7% of the testing set), under low Tanimoto similarity (< 0.3) regimes, the performance of MARASON is still close to the non-RAG baseline.

We group test instances based on the Tanimoto similarity between the retrieved structure and the target structure. We report cosine similarities (higher is better) as follows

All results are from a random split on NIST20. MARASON (non-RAG) is from the first entry in Table 2. This study also highlights a strategy for further performance improvements: use standard MARASON when the retrieved structure has Tanimoto similarity > 0.3 and use the non-RAG version otherwise. There is a trend of decreased performance for non-RAG MARASON when the Tanimoto similarity drops, because a lower Tanimoto similarity means there are fewer similar structures in the training set, making it a more challenging out-of-distribution instance. We will incorporate new results and discussions in future revisions.

Does the observed enhancement primarily stem from additional information introduced by spectral retrieval or from the structural alignment capability enabled by neural graph matching?

Thank you for the thought-provoking question. We are happy to elaborate based on our ablation study presented in Table 2.

Simply retrieving the reference spectrum and concatenating it with the input to the neural network—a naive RAG-style approach—does not improve performance for MS/MS generation. In fact, it slightly degrades it: the resulting cosine similarity is 0.737, representing a 0.3% decrease compared to the non-RAG baseline.

Our findings suggest that effective fragment-level matching is crucial for realizing the benefits of retrieval. For instance, applying a simple Hungarian matching algorithm raises the cosine similarity to 0.746, which is a 0.9% improvement over the baseline. Further, introducing a learnable neural graph matching module with carefully designed architecture yields a cosine similarity of 0.757, amounting to a 2.4% relative improvement over the non-RAG model.

While gains in retrieval accuracy are even more substantial—partly because the cosine similarity metric is more saturated—the top-1 retrieval accuracies indicate that there remains considerable room for improvement.

In summary, the structural alignment capability enabled by the neural graph matching module is the primary driver of MARASON's performance improvement.