SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Samuel Goldman
Computational and Systems Biology MIT
samlg@mit.edu
John Bradshaw
Chemical Engineering MIT
jbrad@mit.edu
Jiayi Xin
Statistics and Actuarial Science
Pokfulam, Hong Kong
xinjiayi@connect.hku.hk
Connor W. Coley
ccoley@mit.edu
Computational predictions of mass spectra from molecules have enabled the discovery of clinically relevant metabolites. However, such predictive tools are still limited as they occupy one of two extremes, either operating (a) by fragmenting molecules combinatorially with overly rigid constraints on potential rearrangements and poor time complexity or (b) by decoding lossy and nonphysical discretized spectra vectors. In this work, we use a new intermediate strategy for predicting mass spectra from molecules by treating mass spectra as sets of molecular formulae, which are themselves multiset of atoms. After first encoding an input molecular graph, we decode a set of molecular subformulae, each of which specify a predicted peak in the mass spectrum, the intensities of which are predicted by a second model. Our key insight is to overcome the combinatorial possibilities for molecular subformulae by decoding the formula set using a prefix tree structure, atom-type by atom-type, representing a general method for ordered multiset decoding. We show promising empirical results on mass spectra prediction tasks.
As the primary tool to discover unknown small molecule structures from biological samples, tandem mass spectrometry (MS/MS) experiments have enabled the identification of numerous important molecules implicated in health and disease [4, 38, 51]. Tandem mass spectrometers are capable of isolating, fragmenting, and measuring the resulting fragment masses of small molecules from a sample, producing a signature (a mass spectrum) for each detected molecule (Figure 1 top).
Figure 1: Tandem mass spectrometers measure fragmentation patterns of molecules, resulting in characteristic peaks that are indicative of their structure. SCARF simulates these fragmentation patterns in silico.
Computationally predicting mass spectra from molecules in silico (Figure 1, bottom) is thus a longstanding and important challenge. Not only does this assist practitioners in better understanding the fragmentation process, but it also enables the identification of molecules from newly observed spectra by comparing an observed spectrum to virtual spectra generated from a database of candidate molecules. While a large library of empirical mass spectra could theoretically serve the same purpose, the size of such libraries is limited by the slow and expensive process of acquiring pure chemical standards and measuring their spectra, motivating computational prediction.
We argue that there are three core, interrelated desiderata for a forward molecule-to-spectrum simulation model, or "spectrum predictor". An ideal spectrum predictor should be (i) accurate, being able to predict the exact set of fragment masses and intensities with a precision comparable to experimental measurements; (ii) physically inspired, to avoid making physically nonsensical ("invalid") suggestions and to provide interpretations of the chemical species responsible for each peak for the benefit of human expert chemists; and (iii) fast, such that it is computationally inexpensive to predict spectra for many (e.g., millions) hypothetical molecules.
Unfortunately, many existing spectrum predictors do not meet these criteria. Methods to date have tended to follow one of two approaches: (a) physically motivated fragmentation approaches or (b) molecule-to-vector (or "binned") approaches (Figure 2A-B). Fragmentation approaches (e.g., [2, 17, 40, 52]; Figure 2A) take an input molecule and suggest bonds that may break, creating fragments that are scored by ML algorithms or curated rulesets. While interpretable, these methods are often slow and restrictive; certain mass spectrum peaks are generated by complex chemical rearrangements within the collision cell that cannot be approximated by bond breaking alone. That is, the bonds in observed fragments are not a subset of those in the original molecule [9, 11]. On the other hand, binned prediction approaches (e.g., [50, 55, 58]; Figure 2B) are less physically grounded, using neural networks to directly learn a mapping from molecules to vectors representing discretized versions of the spectra. These methods, while fast, lack interpretability and due to discretization have a mass precision lower than that of most modern spectrometers, limiting their accuracy.
We propose to address the shortcomings of previous work by predicting mass spectra from molecules at the level of molecular formulae (e.g.,  $\mathrm{C}_x\mathrm{N}_y\mathrm{O}_z\mathrm{H}_w\dots$ ) and introduce a new method, Subformulae Classification for Autoregressively Reconstructing Fragmentations (SCARF) to do so. Because the molecular formula for each input molecule is known, each subformula in the predicted set of peaks is constrained to contain a subset of the atoms in the original formula. Our primary contributions are:
- posing mass spectrum prediction as a two step process: first generating the set of molecular formulae for the fragments, then associating these formulae with intensities;
- overcoming the combinatorial subformula option space by learning to generate formula prefix trees;
- demonstrating the empirical benefit of SCARF in predicting experimental mass spectra quickly and accurately using two separate datasets, providing a benchmark for future work.
We provide a short introduction of tandem mass spectrometry suitable for a general machine learning audience, detail previous approaches to modeling this process as they relate to our proposed approach SCARF, and explain how such tools can be utilized to discover molecules from new spectra. We refer interested readers to [25] for further details on the physical process of mass spectrometry.
2.1 Tandem mass spectrometry
Tandem mass spectrometers (MS/MS) measure fragmentation patterns of molecules in a multi-stage process. The input to the process is a solution containing a precursor molecule,  $\mathcal{M} \in \mathcal{X}$ , associated with a molecular formula,  $\mathcal{F}$ , defining the counts of each element present; for instance  $\mathcal{F} = \mathrm{C}_{16}\mathrm{O}_4\mathrm{H}_{12}$  for the precursor molecule shown in Figure 1. The precursor molecule is first ionized (i.e., made charged), often by bonding or associating with an adduct (e.g., a proton,  $\mathrm{H^{+}}$ ) present in the solution. The charged product is then measured by a mass analyzer (MS1), where its mass-to-charge ratio  $(m / z)$  is measured.
Figure 2: Overview of various approaches to spectrum prediction. A. Fragmentation prediction approaches use heuristics and scoring rules to break down the molecule into fragments and their associated intensities. B. Binned prediction approaches discretize the possible mass-to-charge values and predict intensities for each possible bin. C. Formula prediction approaches predict spectra as sets of molecular formulae and intensities. Our model SCARF utilizes a two stage approach, first by predicting the product formulae present (constrained by the precursor formula), which defines the x-axis locations of the peaks, before secondly assigning intensities to these formulae (defining the peaks' y-axis values).
This precursor ion is then filtered into a collision cell. Here, through interactions with an inert gas, the precursor ion is broken down into a set of one or more product ions, each of which is associated with a new chemical formula; for example, one might be  $f^1 = \mathrm{C}_7\mathrm{OH}_4$  for the process shown in Figure 1. Finally, this set of product ions is measured by a second mass analyzer (MS2), along with the set of their intensities,  $y^{i}\in \mathbb{R}^{+}$  (i.e., their relative frequencies over several repetitions of this process), creating for each ion what is referred to as a peak. The collection of all peaks makes up a molecule's mass spectrum, and is commonly represented as a plot of intensities versus  $m / z$  (Figure 1, right).
2.2 Predicting mass spectra from molecules (spectrum predictors)
Fragmentation prediction. A complex but physically grounded strategy is to model the bond breakage processes occurring in the collision cell (Figure 2A). Examples include MetFrag [52], MAGMa [40], and CFM-ID [2], which recursively fragment molecules (either bond or atom removals) to generate fragment predictions. These methods combine expert rules and local scoring methods to enumerate molecular fragmentation trees to predict spectra. CFM-ID [2] learns subsequent fragmentation transition probabilities between fragments with an expectation maximization algorithm to determine intensities at each fragment. Rule-based methods and full tree enumeration reduce the flexibility of these approaches, and along with the inherent ambiguity in the fragmentation process, limit this strategy's overall accuracy and speed.
Binned prediction. An increasingly popular and straightforward approach to spectra prediction is to map molecules to discretized 1D mass spectra from either molecular fingerprint [50] or graph inputs [55, 58] (Figure 2B). Specifically, these methods divide the  $m/z$  axis into fixed-width "bins" and predict an aggregate statistic of the peaks found in each bin (such as their maximum or summed intensity). While more flexible and end-to-end than fragmentation-based approaches, these methods do not impose the same physical constraints or shared information across fragments, making them less interpretable and susceptible to making invalid predictions. Further, discretizing the input spectrum inherently restricts the precision of such models compared to exact-mass predictions.
Formula prediction. We introduce the strategy of predicting spectra at the level of molecular formulae, an intermediate between binned and fragmentation prediction (Figure 2C). Simultaneous to our work, two groups have separately explored formula prediction strategies [34, 59]. However, to generate plausible subformulae candidates, they either generate a fixed vocabulary of formulae [34] or restrict their model to molecules under 48 atoms for exhaustive enumeration [59], which is smaller than many compounds of interest. We overcome the combinatorial problem of formula generation using prefix trees, allowing our method to scale and eliminating the need for large, fixed vocabularies.
2.3 Mass spectrum libraries
One important use of spectrum predictors is in building large in silico libraries of molecule spectra to augment the small size of existing, experimentally derived databases (on the order of  $10^{4}$ ) which are expensive to curate. These spectra libraries are then leveraged downstream in different ways, for example for training molecular property predictors directly from mass spectra [46]. Another common application of spectra libraries is to infer an unknown molecule's structure from a newly observed spectra – a particularly hard problem, with only  $13\%$  of spectra measured from clinical samples identifiable using current elucidation tools [6]. In this problem, spectra libraries are used as part of a process called retrieval: The newly observed spectra is compared with the existing spectra in the library using a fixed or learned spectral distance function, such as cosine distance [5, 24], and the molecules associated with the closest spectra are returned as possible matches. In practice, the retrieval process is constrained to choosing among isomers (i.e., molecules with the same molecular formula, and therefore molecular weight, but with different bond configurations) due to the high resolution of modern mass spectrometers (i.e., absolute errors on the order of  $10^{-4}$  to  $10^{-3}m/z$  for MS1 measurements) [13, 33, 53].
Given the varied use cases of spectra libraries, we focus on evaluating spectrum predictors in terms of both (a) their prediction accuracies (\$4.2), using metrics such as "cosine similarity", and (b) their use in generating virtual spectral libraries to assist with retrieval (\$4.3).
Here, we describe our model, SCARF, for predicting mass spectra from precursor molecules via first predicting subformulae of the precursor molecule, referred to as product formulae. Building upon the notation introduced in the previous section, we continue to denote precursor molecules  $\square$  as  $\mathcal{M} \in \mathcal{X}$ , and their associated formula vector as  $\mathcal{F} \in \mathbb{N}_0^e$ , defining at each position,  $j \in \{1, \dots, e\}$ , the count of each possible chemical element present,  $\mathcal{F}_j$  (with zero indicating none of that chemical element is present). Likewise, we define the set of  $n$  product formulae as  $\{\pmb{f}^i\}_{i=1}^n$ , and associate with each an intensity,  $y^i$ . Note that the mass  $\boxed{2}$  corresponding to a given formula (and, as such, the x-axis location of the peak on a mass spectrum) is determined deterministically from the counts of each elements present.
At a high level, SCARF generates mass spectra through the composition of two learned functions:
$$
\left\{\left(\boldsymbol {f} ^ {i}, y ^ {i}\right) \right\} _ {i = 1} ^ {n} = g _ {\theta} ^ {\text {W e a v e}} \left(g _ {\theta} ^ {\text {T h r e a d}} (\mathcal {M}), \mathcal {M}\right), \tag {1}
$$
first mapping from the original molecule to a set of product formulae,  $g_{\theta}^{\mathrm{Thread}}: \mathcal{M} \mapsto \{\pmb{f}^i\}_{i=1}^n$ , and then mapping from this set of formulae (and the original molecule) to the respective intensities,  $g_{\theta}^{\mathrm{Weave}}: (\{\pmb{f}^i\}_{i=1}^n, \mathcal{M}) \mapsto \{(f^i, y^i)\}_{i=1}^n$ . The particularities of both functions are described in detail below. The specific architectures and hyperparameters used are deferred to the appendix; model code can be found at https://github.com/samgoldman97/ms-pred
3.1 SCARF-Thread : Generating product formulae via generating prefix trees
SCARF-Thread is tasked to learn a mapping to the set of product formulae,  $\{\pmb{f}^i\}_{i=1}^n$ , given the original molecule. Naively, one might try to define this model autoregressively, predicting the set formula by formula, chemical element by chemical element. However, such an approach soon runs into a number of problems as (i) the predictions are not invariant to set and ordering permutations; (ii) the time complexity of prediction would scale poorly, being proportional to both the number of elements and number of product formulae (i.e.,  $\mathcal{O}(e \times n)$ ); and (iii) the predictions would likely contain duplicates.
<sup>1</sup>We model and discuss uncharged molecules and formulae, despite mass spectrometry measuring the masses of adduct ions. In practice, we reduce all molecules to uncharged candidates by simply shifting all the spectra weights by the  $m/z$  of their respective adducts, which we assume to be equal to the (known) adduct of the parent molecule.
2We assume singly charged adducts (as is common practice, [2]), such that masses and mass-to-charge ratios are interchangeable.
Figure 3: Illustration of the SCARF-Thread architecture. A. The formulae of the product fragments can be represented using a prefix tree. SCARF-Thread predicts this tree for new molecules at test time. It does so by expanding each node at a given depth in parallel, treating the counts of subsequent elements as dependent only on the counts of elements predicted so far (i.e., the prefix) and the original molecular structure. B. The SCARF-Thread predictive task at the  $\mathrm{C_7}$  node from the prefix tree diagram shown in A. Here the network takes as input (i) an embedding of the overall molecule; (ii) a vector representing the counts of each element in the prefix so far (counts yet to be predicted are represented using a special token), (iii) the difference of the counts predicted so far from the precursor molecule, and (iv) a one-hot representation of the element for which the counts are currently being predicted. The network predicts which counts are valid next nodes in the prefix tree (where counts that are greater than those in the original precursor molecular formula are automatically masked out as invalid). See also Alg. A.1.
We therefore take a different approach using the insight that the set of all product formulae can be compactly represented as a prefix tree (Figure 3A). In this tree, edges at a given depth represent valid counts of a particular chemical element, which are often identical across multiple product formulae (shown in the circles). By following each path from the root node to the different leaf nodes, we can reconstruct each product formula (as the orange dashed path does for a single product formula).
We thus propose SCARF-Thread as an autoregressive generator to define a probability distribution over such a prefix tree (Alg. A.1). We assume that each product formula is a subset of the precursor formula, meaning that the precursor formula sets an upper bound on the maximum number of each element. At each node in the tree (corresponding to a prefix  $f_{<j}^{\prime}$ ), we pose the prediction of the set of child nodes (corresponding to the set of valid counts of the subsequent element) as a multi-label binary classification problem (Figure 3B). Concretely, we use a neural network module for this task, giving it as input a context vector representing the node being expanded:
$$
\boldsymbol {c} ^ {\prime} = \left[ \operatorname {g n n} (\mathcal {M}), \operatorname {c o u n t s} \left(\boldsymbol {f} _ {<   j} ^ {\prime}\right), \operatorname {c o u n t s} \left(\boldsymbol {\mathcal {F}} - \boldsymbol {f} _ {<   j} ^ {\prime}\right), \text {o n e - h o t} (j) \right], \tag {2}
$$
where  $\mathtt{gnn}(\cdot)$  specifies a neural encoding of the molecular graph (SA.5.2),  $\mathtt{counts}(\cdot)$  specifies a count-based encoding of the associated prefix (SA.5.3), and one-hot() specifies a one-hot encoding of the node's depth (or equivalently, which element the predicted count is for). In our experiments, we use a fixed ordering of the chemical elements (SA.2), but optimizing or even learning the tree construction order could be carried out [45].
Formulae as differences. Following Wei et al. [50], we find it helpful to not only parameterize product formulae in terms of their element counts, but also in terms of the elements that they have lost,
3While it is possible for fragments to fuse together, potentially taking the count of a chemical element over the number in the original precursor formula, we postpone the extension to modeling such rare events to future work.
i.e., their difference from the precursor formula. On the input side, this is already covered by including in the context vector a count-based embedding of the prefix formula minus the product formula  $(\mathrm{counts}(\mathcal{F} - \boldsymbol{f}_{< j}^{'}))$ . However, on the output side this is achieved by combining the probabilities of a "forward" and a "difference" network:
$$
p \left(f _ {j} ^ {\prime} = a \mid \boldsymbol {f} _ {<   j} ^ {\prime}, \mathcal {M}\right) = \boldsymbol {\alpha} _ {a} \boldsymbol {\sigma} \left(\mathrm {M L P} ^ {F} \left(\boldsymbol {c} ^ {\prime}\right)\right) _ {a} + \left(\mathbf {1} - \boldsymbol {\alpha}\right) _ {a} \boldsymbol {\sigma} \left(\mathrm {M L P} ^ {D} \left(\boldsymbol {c} ^ {\prime}\right)\right) _ {\mathcal {F} _ {j} - a}, \tag {3}
$$
where  $\mathsf{MLP}^F (\cdot)$  and  $\mathsf{MLP}^D (\cdot)$  specify multi-layer perceptrons (MLPs) for predicting the probability of observing a count of  $a$  and a loss of  $\mathcal{F}_j - a$  atoms respectively;  $\alpha$  is a variable (output from a third, unshown network) deciding how to weight these predictions; and  $\sigma (\cdot)$  is the element-wise sigmoid function.
3.2 SCARF-Weave: Predicting intensities given product formulae
Figure 4: The SCARF-Weave network, which takes in the product formulae (e.g., predicted by SCARF-Thread) and predicts their intensities. We use a Set Transformer architecture [29], such that our model takes in the details of the other product formulae present when predicting intensities.
Given the product formulae outputs from SCARF-Thread, SCARF-Weave predicts corresponding intensities at each formula. This is a set-to-set problem, well suited for any equivariant set2set architecture [56, §3.1]. In our experiments, we use a Set Transformer [29, 44], which enables the model to consider all the formulae present in the mass spectrum (and their possible interactions) when predicting final intensities.
We choose to represent formula in the set similarly to the context vectors used in SCARF-Thread. For each input, we concatenate a vector embedding of the initial molecular graph with count-based embeddings of the product formula and its difference from the precursor formula (Figure 4). We again defer the particularities of the embedding functions to the Appendix (§A.5).
3.3 Training and inference
Provided with a dataset of molecules and formula-labeled mass spectra, we could train the two components of SCARF separately. However, in practice we find it beneficial to first train SCARF-Thread and then train SCARF-Weave on its outputs so that the distribution the latter model sees is the same at training and prediction time. SCARF-Weave is trained using a cosine loss (A.5.5), as this most closely resembles the "retrieval" setting (4.3).
SCARF-Thread is trained using the binary cross entropy losses associated with the multi-label classification tasks at each non-leaf node in the prefix tree. We use teacher forcing, i.e., we train on each level of the tree in parallel by conditioning on the ground-truth set of prefixes at each stage. In our experiments, when generating the set of product formulae from this model we always pick the top 300. Empirically, we find that this provides better performance than picking a variable number based on a likelihood threshold.
We evaluate SCARF on spectra prediction (§ 4.2) and molecule identification in a retrieval task (§ 4.3).
We train and validate SCARF on two libraries: a gold standard commercial tandem mass spectrometry dataset, NIST20 [35], as well as a more heterogeneous public dataset, NPLIB1, extracted from the
Table 1: Model coverage (higher better) of true peak formulae as determined by MAGMa at various max formula cutoffs for the NIST20 and NPLIB1 datasets. Best result for each column is in bold. Results are computed for a single test set; all re-trained models (i.e., Autoregressive and SCARF variants) are averaged across three random seeds.
Dataset	NIST20	NPLIB1
Coverage @	10	30	300	1000	10	30	300	1000
Random	0.009	0.026	0.232	0.532	0.004	0.014	0.126	0.336
Frequency	0.173	0.275	0.659	0.830	0.090	0.151	0.466	0.688
CFM-ID	0.197	0.282	-	-	0.170	0.267	-	-
Autoregressive	0.204	0.262	0.309	0.317	0.072	0.082	0.095	0.099
SCARF-Thread-D	0.248	0.425	0.839	0.941	0.158	0.284	0.681	0.856
SCARF-Thread-F	0.249	0.476	0.855	0.943	0.155	0.306	0.708	0.859
SCARF-Thread	0.308	0.552	0.907	0.968	0.164	0.309	0.724	0.879
GNPS database [48] by Dührkop et al. [14] and subsequently processed by Goldman et al. [19]. We prepare both datasets by extracting and preprocessing spectra, as well as filtering to compounds that (a) are under 1,500 Da (i.e., typically under 100 heavy atoms), (b) only contain predefined elements, and (c) are only charged with common positive-mode adduct types (\$4.1).
Overall, NIST20 contains 35,129 total spectra with 24,403 unique structures, and 12,975 unique molecular formulae; NPLIB1 contains 10,709 spectra, 8,553 unique structures and 5,433 unique molecular formulae. Both datasets are evaluated using a structure-disjoint  $90\% / 10\%$  train/test split with  $10\%$  of training data held out for validation, such that all compounds in the test set are not seen in the train and validation sets.