Pre-trained Maldi Transformers improve MALDI-TOF MS-based prediction

For the last decade, matrix-assisted laser desportion/ionization time-of-flight mass spectrometry (MALDI-TOF MS) has been the reference method for species identification in clinical microbiology. Ham- pered by a historical lack of open data, machine learning research towards models specifically adapted to MALDI-TOF MS remains in its infancy. Given the growing complexity of available datasets (such as large-scale antimicrobial resistance prediction), a need for models that (1) are specifically designed for MALDI-TOF MS data, and (2) have high representational capacity, presents itself. Here, we introduce Maldi Transformer , an adaptation of the state-of-the-art transformer architecture to the MALDI-TOF mass spectral domain. We propose the first self-supervised pre-training technique adapted to mass spectra. The technique is based on shuffling peaks across spectra, and pre-training the transformer as a peak discriminator . Extensive benchmarks confirm the efficacy of this novel de- sign. The final result is a model exhibiting state-of-the-art (or competitive) performance on downstream prediction tasks. In addition, we show that Maldi Transformer’s identification of noisy spectra may be leveraged towards higher predictive performance. All code supporting this study is distributed on PyPI and is packaged under: https://github.com/ gdewael/maldi-nn

MALDI-TOF MS ⋅ Neural networks ⋅ Transformers 1. Introduction Matrix-assisted laser desportion/ionization time-of- flight mass spectrometry (MALDI-TOF MS) is a pro- teomic technique commonly used to identify micro- bial species (Dauwalder et al., 2023). First introduced to clinical microbiology at the end of 2010, its routine use is characterized by low-cost, speed, and reliability (Weis et al., 2020a). MALDI-TOF MS generates spec- tra containing peaks signifying mostly ribosomal pro- teins (Seng et al., 2009). As such, the spectra can serve as fingerprints indicative of species identity (Bizzini et al., 2011). For bacterial species identification, clinical diag- nostic labs will typically use the solutions provided by MALDI-TOF MS manufacturers. These solutions are built on large, proprietary, in-house databases (Van Belkum et al.,2012). The models used in such so- lutions presumably rely on querying certain marker peaks to a large database (Florio et al., 2018). This strategy works reasonably well for identification of most species, but some strains remain problematic to identify this way (Cao et al., 2018; Vrioni et al., 2018). In addition, the peak-matching approach does not suit more-difficult prediction tasks such as strain typing (Hettick et al., 2006), antimicrobial resistance predic- tion (Weis et al., 2022), and virulence factor detection (Gagnaire et al., 2012). In these cases, researchers turn to machine learning in order to possibly mine more-intricate patterns from the spectra. Historically, machine learning for MALDI-TOF MS has been hampered by a lack of large open data. Because of this, the nascent field has not often progressed beyond off-the-shelf learning techniques (Weis et al., 2020a). Only a handful of examples exist of more advanced machine learning methods specif- ically adapted to a MALDI-TOF-based task (Vervier et al., 2015; Weis et al., 2020b; Mortier et al., 2021; De Waele et al., 2023). As such, the design of machine learning algorithms specifically for MALDI-TOF mass spectra has not been studied yet in sufficient detail. Typically, MALDI-TOF MS-based machine learning

either discretize the (𝑚/𝑧)-axis, or pre- select a number of peak locations as features based on training set characteristics. Discretizing the (𝑚/𝑧)- axis results in fixed-length spectral representations, where every feature consists of a bin. If bins are cho- .CC-BY-NC-ND 4.0 International licenseavailable under a (which was not certified by peer review) is the author/funder, who has granted bioRxiv a license to display the preprint in perpetuity. It is made The copyright holder for this preprintthis version posted January 22, 2024. ; https://doi.org/10.1101/2024.01.18.576189doi: bioRxiv preprint Pre-trained Maldi Transformers improve MALDI-TOF MS-based prediction sen too large, resolution is lost and multiple peaks may be lumped together . On the other hand, small bins result in a higher-dimensional representation where most bins contain no peaks, unnecessarily adding model complexity (Weis et al., 2020b; Mortier et al., 2021). Other works select a fixed set of (𝑚/𝑧) locations based on training set characteristics (Tran et al., 2021). Features for every spectrum are then derived by encoding the peak height (or binary peak presence) for every selected 𝑚/𝑧 location. The disad- vantage here is that spectra may contain important peaks not included in the fixed set of selected 𝑚/𝑧 locations. Hence, the model is unable to cope with patterns differing too much from general training set characteristics (e.g. rare peaks). The information in MALDI-TOF mass spectra is de- fined by their peaks (many associated with riboso- mal subunits) (Ryzhov and Fenselau, 2001). As such, we argue that, ideally, a machine learning model for MALDI-TOF mass spectra operates directly on sets of peaks as inputs. This, however , constitutes a non- trivial machine learning setup. A kernel-based tech- nique that processes spectra in this way was previ- ously proposed by Weis et al. (2020b). However , their

is difficult to scale to larger sample sizes. Con- sequently, in this work, we propose a deep learning- based solution operating on sets of peaks: a trans- former for MALDI-TOF data. Transformers elegantly operate on input sets via their permutation-invariant self-attention operations (Vaswani et al., 2017). Self- attention can be interpreted as message passing on a complete directed graph (i.e. peaks are nodes, and all nodes are connected to each other). Following its rise to dominance in natural language processing, transformers are increasingly adopted towards biological data modalities (Clauwaert and Waegeman, 2020; Jumper et al., 2021; Avsec et al., 2021; Elnaggar et al., 2021). Most closely related to the MALDI-TOF domain, transformers have also been adapted to operate on tandem mass spectra for de novo peptide sequencing (Yilmaz et al., 2022). This ever-increasing adoption of transformers across sub- fields of machine learning speaks to their general- ity. Because they can be viewed as operating on a complete digraph, they place no inductive bias on the learned patterns between input tokens. This property lends transformers supreme representational capa- bilities, but is also the reason why they are typically described as data-hungry. Consequently, transform- ers are usually mentioned in the same breath as self- supervised learning (SSL) (Liu et al., 2021). SSL is the paradigm within deep learning wherein a supervised learning task is designed for data that has not explic- itly been labeled (Balestriero et al., 2023). This task is (usually) not a useful prediction problem in itself, but rather serves to pre-train a large model. In doing so, greater downstream performance on tasks of in- terest can be obtained. SSL is, therefore, most useful in scenarios where labeled data is limited, such as the MALDI-TOF MS domain (Weis et al., 2020a). In the present study, we introduce Maldi Trans- former , a deep learning architecture for processing MALDI-TOF mass spectra. The inputs to Maldi Trans- former consist of spectra as sets of peaks. To fully take advantage of the transformer architecture’s represen- tational capacity, we propose a novel self-supervised pre-training strategy. The strategy relies on discrimi- nating real peaks from noisy ones introduced to the spectrum via shuffling peaks across samples. We pre-train Maldi Transformer on the large open DRI- AMS database (Weis et al.,2020a). Maldi Transformer obtains strong performance on downstream bench- marks, demonstrating the power of the approach. 2. Methods 2.1. Data To train and test models, we use three large (> 10 000 spectra) MALDI-TOF datasets, two of which are avail- able to the public domain. First, the recently published DRIAMS database (Weis et al., 2022), is used to both pre-train Maldi Transformer and to fine-tune it on antimicrobial re- sistance (AMR) prediction (binary classification using a dual-branch recommender system). DRIAMS con- tains a total of 250 070 spectra, originating from four hospitals in Switzerland. For AMR prediction, the same data splits are used as in earlier work (De Waele et al., 2023). Briefly, DRIAMS-A spectra from be- fore 2018 are split to the training fraction, whereas DRIAMS-A spectra measured during 2018 are evenly split between validation and test set. For pre-training, the same splits are used, but all spectra from DRIAMS- B, -C, and -D are additionally added to the training set. Apart from AMR measurements, DRIAMS con- tains species labels for many spectra. These labels are derived from the species identification pipelines included with the MALDI-TOF MS machines. After processing (see AppendixA), the pre-training DRIAMS dataset spans 469 species. The second dataset consists of the public Robert Koch-Institute (RKI) database (Lasch et al., 2023). The final used dataset is a private historical database con- taining more than 2400 taxonomic reference strains, 2 .CC-BY-NC-ND 4.0 International licenseavailable under a (which was not certified by peer review) is the author/funder, who has granted bioRxiv a license to display the preprint in perpetuity. It is made The copyright holder for this preprintthis version posted January 22, 2024. ; https://doi.org/10.1101/2024.01.18.576189doi: bioRxiv preprint Pre-trained Maldi Transformers improve MALDI-TOF MS-based prediction Table 1 | Data sources used, along with their use and sizes (in number of spectra). Dataset Used for Train-Val-Test Size DRIAMS Pre-training 207 172 - 21 440 - 21 443𝑎,𝑏 DRIAMS-A Fine-tuning on AMR prediction 28 331 - 4 994 - 4 999𝑐 RKI Fine-tuning on species identification 8 442 - 1 350 - 1 263 LM-UGent Fine-tuning on species identification 88 267 - 8 710 - 8 700 𝑎 Contains all spectra in DRIAMS with at least 200 detected peaks. 𝑏 Of which 97 783, 14 055, and 14 183 have species labels in train, val and test splits, respectively. 𝑐 Numbers reflect spectra. In total, 409 395, 76 431, and 76 133 AMR labels across 65 drugs are associated with those splits. cultured and analyzed by the Laboratory of Microbi- ology at Ghent University. In this manuscript, both datasets will be abbreviated as RKI and LM-UGent, respectively. The RKI dataset contains MALDI-TOF mass spectra from highly pathogenic bacteria, cov- ering similar species as in DRIAMS. The LM-UGent dataset, on the other hand, includes a broader taxo- nomic range. Both datasets are used for fine-tuning on species identification (multi-class classification). As in Mortier et al. (2021), in order to create a chal- lenging training-validation-test split, spectra are split in such a way that there is no overlap in terms of strains. As a consequence, the models are tested whether they can identify unseen strains of (seen) species. The following rules are used in data splitting: all spectra for a species are assigned to the training set if that species only has one strain in the dataset. If the species has more than one strain, strains are split such that 80% of strains of that species are assigned to the training set, and the other 20% evenly split be- tween validation and test set (with a floor value of at least one strain being assigned to either validation or test set). The total number of species in the RKI train- ing set spans 270, of which 106 and 108 are present in the validation and test set, respectively. For the LM- UGent dataset, these numbers are 1088, 200, and 202, for the training, validation and test set, respectively. For more details on the LM-UGent dataset, the reader is referred to Mortier et al. (2021). Table 1 lists a sum- mary of the sizes of all used data. All MALDI-TOF mass spectra are preprocessed us- ing standard practices (Gibb and Strimmer,2012). Fol- lowing Weis et al. (2022), all spectra undergo the fol- lowing steps: (1) square-root transformation of the in- tensities, (2) smoothing using a Savitzky-Golay filter with half-window size of 10, (3) baseline correction using 20 iterations of the SNIP algorithm, (4) trimming to the 2000-20000 Da range, (5) intensity calibration so that the total intensity sums to 1. For Maldi Trans- former , peaks are then detected on the preprocessed spectrum using the persistence transformation, in- troduced by Weis et al. (2020b). While the original publication proposes this algorithm to nullify other preprocessing steps, we find that prior preprocess- ing steps help peak detection (see Appendix B). As in Weis et al. (2020b), we select the highest 200 peaks for every spectrum as inputs 1. Maldi Transformer is compared against baseline methods, which require a fixed-length input. For these models, instead of run- ning peak detection as a last preprocessing step, spec- tra are instead binned to a 6000-dimensional vector by summing together intensities in intervals of 3 Da. 2.2. Maldi Transformer Model To introduce Maldi Transformer , let us denote a spectrum as a set of 𝑚 peaks S = {( (𝑚/𝑧)𝑖 , 𝐼𝑖)}𝑚 𝑖=1, with each peak characterized by its (𝑚/𝑧) value and (preprocessed) intensity 𝐼. An an- notated dataset D then consists of a set of 𝑛 samples D = {(S𝑖, 𝑦 spec 𝑖 )}𝑛 𝑖=1, with 𝑦spec the species label. Maldi Transformer requires a single input repre- sentation 𝒙 ∈ R𝑑 for each peak ( (𝑚/𝑧)𝑖 , 𝐼𝑖). To achieve this, intensities𝐼 are linearly transformed to a 𝑑-dimensional space. Similarly, (𝑚/𝑧)values are em- bedded to sinusoidal positional encodings (Vaswani et al., 2017). For any (𝑚/𝑧) value, the positional en- coding (𝑃𝐸) at feature index 𝑗 is 𝑃𝐸 ((𝑚/𝑧), 𝑗) = ⎧⎪⎪ ⎪⎨⎪⎪ ⎪⎩ 𝑠𝑖𝑛 ( (𝑚/𝑧) 10⋅10000 𝑗/𝑑 ), if 𝑗 is even 𝑐𝑜𝑠 ( (𝑚/𝑧) 10⋅10000(𝑗−1)/𝑑 ), if 𝑗 is odd Note that this formulation is identical to the one de- scribed in Vaswani et al. (2017), apart from the factor 10 division of the(𝑚/𝑧)positions. This division is per- formed to bring the numerical range of (𝑚/𝑧)values (2000 Da - 20000 Da) closer to the numerical range of positional indices for which this equation was origi- nally designed2. The sinusoidal (𝑚/𝑧) embedding and linear inten- sity 𝐼 embedding are summed to a single input rep- resentation per peak 𝒙𝑖∈{1,...,𝑚}. A trainable [CLS] vector is prepended to the input for spectrum-level 1Note that Maldi Transformer can, in principle, deal with vari- able number of peaks per spectrum. 2Language transformers are often trained with maximum se- quence lengths between 512 and 2048. (Beltagy et al., 2020) 3 .CC-BY-NC-ND 4.0 International licenseavailable under a (which was not certified by peer review) is the author/funder, who has granted bioRxiv a license to display the preprint in perpetuity. It is made The copyright holder for this preprintthis version posted January 22, 2024. ; https://doi.org/10.1101/2024.01.18.576189doi: bioRxiv preprint Pre-trained Maldi Transformers improve MALDI-TOF MS-based prediction prediction, following common practice (Devlin et al., 2018; Dosovitskiy et al., 2020). The final input to the encoder-only transformer is, hence, 𝑿 ∈ R201×𝑑, with 201 the number of peaks plus the [CLS]token, and 𝑑 the hidden dimensionality of the model. The encoder-only transformer processes 𝑿 to a spectrum- level output embedding 𝒑[CLS] and output peak-level embeddings 𝒑𝑖∈{1,...,𝑚}. The design of the transformer encoder blocks follows current state-of-the-art prac- tices (Appendix E Figure 6) (Narang et al., 2021). An overview of the model is visualized in Figure 1. Pre-training task design To boost the performance of Maldi Transformer on supervised tasks with lim- ited labeled data, a novel self-supervised pre-training strategy is designed. We propose to pre-train Maldi Transformers as peak discriminators. That is, in a batch of spectra, some peaks are randomly sampled to use for training. Following Devlin et al. (2018), we sample 15% of the peaks. Half of those sampled peaks are shuffled among all spectra, while the others are kept as part of their original spectrum. Using the shuf- fled peaks as negative "noise" peaks in a spectrum, a discriminative model is trained to distinguish the noise peaks from the sampled original ones using the cross-entropy loss. More formally, for a spectrum S, let us denote its version with shuffled peaks as Sshuff. Further , let 𝒊train = {𝑖1, ..., 𝑖𝑠} denote the indices of its sampled peaks. Maldi Transformer processes Sshuff to output peak-level embeddings 𝒑𝑖∈{1,...,𝑚}, which are then pro- cessed to predictions ˆ𝑦 peak 𝑖∈{1,...,𝑚} = 𝜎 (𝒘 ⊺ 𝑝 𝒑𝑖∈{1,...,𝑚}), with 𝒘𝑝 a learnable weight vector . The peak discrim- ination loss is, then, given by Lpeaks = E ( ∑ 𝒊train − I (S𝑖 𝑗 = S shuff 𝑖 𝑗 )log ( ˆ𝑦 peak 𝑖 𝑗 ) − I (S𝑖 𝑗 ≠ S shuff 𝑖 𝑗 )log (1 − ˆ𝑦 peak 𝑖 𝑗 ) ) with I the indicator function. Complementary to the peak discrimination strat- egy, the spectrum embedding 𝒑[CLS] is sent to a multi-class linear output head to predict the mi- crobial species identity 𝑦spec of the original spec- trum3. The final pre-training loss function is the sum of peak discrimination binary cross-entropy and species identification multi-class cross-entropy: L = Lpeaks + 𝜆 Lspec, with 𝜆 ∼ Bern(0.01). The species identification loss Lspec is only randomly ap- plied in 1% of the training steps. This is performed Table 2 | Pre-training configurations for different Maldi Transformer sizes. Species clf 𝜆 refers to the probability that the species identification loss is ap- plied per training step. Model Hyperpar ameter S M L XL # params 1.65M 3.27M 6.92M 14.84M # layers 4 6 8 10 hidden dim 𝑑 160 184 232 304 # heads 8 8 8 8 Learning rate 5e-4 5e-4 5e-4 3e-4 Pre-training steps 500 000 500 000 500 000 400 000 Species clf 𝜆 Bern(0.01) Bern(0.01) Bern(0.01) Bern(0.005) because, empirically, overfitting of the species classi- fication task is observed when applied at every train- ing step (see Appendix E Figure 7). A visual represen- tation of the entire pre-training approach is shown in Figure 1. A more-thorough description of the whole pre-training strategy can be found in Appendix C Al- gorithm 1. Our novel peak discrimination pre-training strat- egy is proposed due to conceptual difficulties with porting established pre-training approaches to the MALDI-TOF MS domain, a point further elaborated on in Appendix C. We benchmark this novel pre-training strategy against alternatives in §3.2. Model configurations We train Maldi Transformer in four different sizes: S, M, L, and XL. Model sizes are chosen so that the total weight numbers roughly cor- respond to the ones in De Waele et al. (2023). Table 2 lists the size of all models, along with some hyperpa- rameter settings. The Adam optimizer is used to pre- train all models in BFloat16 mixed precision (Kingma and Ba, 2014). Gradients are clipped to a norm of 1. A batch size of 1024 is applied for all models. A linear learning rate warm-up is applied over the first 2500 steps, after which the learning rate remains constant (Table 2). During training, a dropout of 0.2 is applied in the GLU feedforward and over the attention matrix. 2.3. Downstream tasks As mentioned in §2.1, Maldi Transformer’s perfor- mance is validated on three downstream supervised 3While this is not a purely self-supervised training objective, we justify its use with two reasons. First, in language, pre-training with a sentence-level [CLS]task boosts downstream performance (De- vlin et al., 2018). Second, due to the integrated nature of MS man- ufacturers’ species identification pipelines, bacterial species labels are relatively easy to come by. As such, no manual labeling is nec- essary to obtain these labels. Additionally, it has to be noted that, as not all spectra in DRIAMS carry a species label, the species clas- sification loss is only calculated for labeled spectra. 4 .CC-BY-NC-ND 4.0 International licenseavailable under a (which was not certified by peer review) is the author/funder, who has granted bioRxiv a license to display the preprint in perpetuity. It is made The copyright holder for this preprintthis version posted January 22, 2024. ; https://doi.org/10.1101/2024.01.18.576189doi: bioRxiv preprint Pre-trained Maldi Transformers improve MALDI-TOF MS-based prediction Figure 1 | Left: Representation of MALDI-TOF mass spectra. A raw spectrum is preprocessed, after which a topological persistence transformation is performed to detect peaks, resulting in a sparse representation of the spectrum (Weis et al., 2020b). Middle: Maldi Transformer . A peak is characterized by an intensity𝐼 and (𝑚/𝑧) value. Both are embedded to higher dimensional space by a linear layer and sinusoidal embedding, respec- tively, and are then summed. A[CLS] token is prepended and the resulting vectors𝒙 are sent through multiple Transformer encoder layers. The resulting output vectors𝒑𝑖∈{1,...,𝑚} can be used as peak embeddings. Addition- ally, 𝒑[CLS]can be used as a summary embedding of the whole spectrum. Right: Proposed peak discrimination pre-training strategy. In a mini-batch of spectra, 15% of peaks are randomly sampled, half of which are shuffled among all spectra in the batch. The resulting spectra are encoded with Maldi Transformer . The resulting peak embeddings 𝒑𝑖∈{1,...,𝑚} are sent through a linear output head trying to distinguish original peaks (blue) from shuffled "noise" peaks (orange). Spectrum embeddings𝒑[CLS]are sent through a separate linear head to predict species identity (green). tasks. For all three tasks, the pre-trained model is plugged in at initialization and all weights are fine- tuned (i.e. no weight freezing). A task-specific lin- ear output head𝑾 𝑜𝑢𝑡 projecting the spectrum embed- ding 𝒑[CLS]to the desired output space is trained from scratch. For AMR prediction, Maldi Transformer is used as a spectrum embedder in a dual-branch neu- ral network recommender system. To compare its performance against previous results, recommenders are trained with the four best-scoring drug embed- ders in De Waele et al. (2023). To benchmark Maldi Transformer in terms of species identification, the RKI and LM-UGent datasets are used. Species identifi- cation is compared to MLP baselines, Logistic Regres- sion, Random Forest, and k-nearest neighbors (k-NN) models. Details on the exact training setups for each downstream tasks are found in Appendix D. The LM-UGent dataset covers a broader species di- versity than the clinical species found in the pre- training set. As such, this dataset can be considered out-of-distribution for the pre-trained Maldi Trans- former . For this reason, a domain adaptation step on the pre-trained Maldi Transformer is performed be- fore supervised fine-tuning. The domain adaptation step consists of pre-training Maldi Transformer for 20 000 additional steps in the same fashion, but now us- ing the LM-UGent dataset, instead of DRIAMS4. 3. Results 3.1. Maldi Transformer improves perfor- mance on downstream tasks Pre-training curves for Maldi Transformer are shown in Appendix E Figure 8. After pre-training, Maldi Transformer is fine-tuned w.r .t. a downstream task. 4By performing domain adaptation, in contrast to the other su- pervised tasks, the task-specific output head 𝑾 𝑜𝑢𝑡 does not need to be initialized from scratch anymore, but can be copied from the pre-trained model, as it already has the right dimensions. 5 .CC-BY-NC-ND 4.0 International licenseavailable under a (which was not certified by peer review) is the author/funder, who has granted bioRxiv a license to display the preprint in perpetuity. It is made The copyright holder for this preprintthis version posted January 22, 2024. ; https://doi.org/10.1101/2024.01.18.576189doi: bioRxiv preprint Pre-trained Maldi Transformers improve MALDI-TOF MS-based prediction Maldi Transformer’s downstream performance is compared to preprocessed and binned baselines. For AMR prediction, it is compared to MLPs of similar sizes. For species identification, it is additionally com- pared to non-neural network baselines. Figure 2 shows the experimental results for all downstream tasks. For AMR prediction, models are evaluated in terms of micro ROC-AUC and instance- wise ROC-AUC 5. For species identification, they are evaluated in terms of species- and genus-level accu- racy. In general, Maldi Transformer obtains supe- rior performance in comparison to other tested meth- ods. For AMR prediction, Maldi Transformer consis- tently outperforms all MLP models in terms of mi- cro ROC-AUC. In terms of the instance-wise ROC-AUC, Maldi Transformer is sometimes outperformed by MLP models, but the best-scoring model overall is still one using Maldi Transformer (i.e. the large model paired with a Morgan fingerprint drug embedder). For species identification on the RKI dataset, Maldi Transformer does not always provide a clear advan- tage. On species-level, it is consistently outperformed by both Logistic Regression and similar-sized MLPs. For genus-level accuracy, however , Maldi Transform- ers surpass their MLP counterparts. While a k-NN model provides competitive performance, the best overall model in terms of genus-level accuracy is the large Maldi Transformer . On the larger and more-difficult LM-UGent dataset,

are again more convincing. The best- performing species-level model is the medium-sized Maldi Transformer , convincingly beating other mod- els with 84% accuracy. The same results are obtained on genus-level accuracy, where Maldi Transformer consistently beats other models. We hypothesize that, on small MALDI-TOF datasets with relatively-simple prediction tasks (e.g. the RKI dataset), simple models are sufficient, but Maldi Transformer remains competitive. For comparatively-complex tasks, such as AMR pre- diction with dual-branch neural networks, and species identification across >1000 species, Maldi Transformer consistently delivers state-of-the-art performance. 5An explanation of the instance-wise ROC-AUC is given in De Waele et al. (2023). 3.2. Maldi Transformer pre-training ab- lation To further validate the efficacy of the proposed Maldi Transformer , its performance is compared against two alternative realizations of transformers on MALDI-TOF MS data (i.e. using two alternative pre-training strategies). The first one takes inspira- tion from the masked language model (MLM) BERT (Devlin et al., 2018). Here, intensity values of peaks are randomly masked out, and a transformer is pre- trained to predict the original intensities back using the mean squared error loss. In the second, a discrim- inative model much like our final proposed strategy is trained. The difference in this model is that neg- ative peaks are sampled from some estimated distri- bution of peaks, instead of generating negative peaks by shuffling. Both alternative strategies are described in greater detail in Appendix C. In Figure 3.2A, it is observed that both alternative pre-training techniques are outperformed by our fi- nal proposed strategy. Naively porting the MLM strat- egy to peak intensity regression underdelivers by a wide margin. We hypothesize that the intensity of a peak is of lesser importance compared to whether that peak is present or not, hence making a regres- sion model learn superfluous information. A more- biologically relevant training objective would be to make the model learn over peak co-occurrence. A negative peak sampling strategy paired with a binary discrimination objective would achieve such a task. It can be seen from Figure 3.2A that this task, how- ever close in concept to the final Maldi Transformer , still consistently underperforms. A possible explana- tion could be that the estimation of negative peaks re- lies on some estimated overall distribution of peaks. Any inaccuracies in this distribution result in "unre- alistic" negative peaks, which are easier to recognize by the model. Thus, instead of truly reasoning over peak co-occurrence, the model takes a shortcut and learns the inaccuracies of the underlying peak sam- pler . Our final peak shuffling strategy does not have this drawback, and is, therefore, ideally suited for the MALDI-TOF MS domain, also resulting in superior performances. Figure 3.2B shows ablation results when leaving different parts out of the final pre-training strategy. Maldi Transformer is compared to models where one of the two loss components are left out: either the peak discrimination loss (Lpeaks), or the species iden- tification loss (L spec). It is also compared to a trans- former model without pre-training whatsoever . It can be seen that the latter strategy (i.e. training a separate 6 .CC-BY-NC-ND 4.0 International licenseavailable under a (which was not certified by peer review) is the author/funder, who has granted bioRxiv a license to display the preprint in perpetuity. It is made The copyright holder for this preprintthis version posted January 22, 2024. ; https://doi.org/10.1101/2024.01.18.576189doi: bioRxiv preprint Pre-trained Maldi Transformers improve MALDI-TOF MS-based prediction A B C One-hot embedding Morgan Fingerprint Image 2-D CNN SMILES 1-D CNN One-hot embedding Morgan Fingerprint Image 2-D CNN SMILES 1-D CNN S M L XL S M L XL S M L XL S M L XL S M L XL S M L XL S M L XL S M L XL S M L XL S M L XL Neural net sizeNeural net size Baselines Baselines S M L XL S M L XL Neural net sizeNeural net size Baselines Baselines Micro ROC-AUC Instance-wise ROC-AUC Species-level accuracy Species-level accuracy Genus-level accuracy Genus-level accuracy Logistic Regression k-Nearest Neighbors Random Forest MLP Maldi Transformer Figure 2 | Barplots of all main experimental results. Error-bars indicate standard deviation over three inde- pendent runs. Neural network model sizes range from S to XL. Details per model size for Maldi Transformer and MLP models are listed in Table 2 and Appendix D Table 3, respectively. A: AMR prediction results on DRI- AMS. Performances are shown for models using four different drug embedders in a dual-branch recommender system. B: Species identification results on the RKI dataset. C: Species identification results on the LM-UGent dataset. Species-level Accuracy Species-level Accuracy Species-level Accuracy Species-level Accuracy Genus-level Accuracy Genus-level Accuracy Genus-level Accuracy Genus-level Accuracy Intensity MLM + MSE Negative peak sampler Maldi Transformer Maldi Transformer Maldi Transformer No Pre-training Maldi Transformer Only Maldi Transformer Only A B Micro ROC-AUC Micro ROC-AUC Instance-wise ROC-AUC Instance-wise ROC-AUC DRIAMS - AMR RKI LM-UGent DRIAMS - AMR RKI LM-UGent Figure 3 | Barplots of all ablation results. Only results for Medium-sized transformers are shown. The final model as discussed in §2.2 is labeled as Maldi Transformer . Error-bars indicate standard deviation over three independent runs. For DRIAMS AMR prediction, only results using a Morgan Fingerprint drug embedder are shown. For the LM-UGent dataset, results are shown without the domain adaptation step, as this step was also not performed for ablation models. See AppendixE Figure 9 for the effects of this step.A: Testing different ways of pre-training a transformer on MALDI-TOF MS data. B: Leaving out different parts of the final pre-training strategy. transformer from scratch for every supervised task), delivers subpar performance. In addition, only pre- training on species identification does not help much. The biggest gain is made from the peak discrimination task. Yet, overall, the effects are additive, showing that each part contributes to the final performance. 3.3. Maldi Transformer model analysis Due to its pre-training design, Maldi Transformer out- puts the probability of a peak belonging to its spec- trum for every peak. More formally, let us denote the output probability of a peak 𝑖 belonging to its spec- trum S as Pr(𝑦 peak 𝑖 = 1∣S). This quantity can be 7 .CC-BY-NC-ND 4.0 International licenseavailable under a (which was not certified by peer review) is the author/funder, who has granted bioRxiv a license to display the preprint in perpetuity. It is made The copyright holder for this preprintthis version posted January 22, 2024. ; https://doi.org/10.1101/2024.01.18.576189doi: bioRxiv preprint Pre-trained Maldi Transformers improve MALDI-TOF MS-based prediction obtained through the pre-trained Maldi Transformer: Pr(𝑦 peak 𝑖 = 1∣S) = ˆ𝑦 peak 𝑝 = 𝜎(𝒘 ⊺ 𝑝 ⋅ 𝒑𝑝). During pre- training, as peaks are shuffled between spectra, one expects to encounter both positives and negative true labels. In this section, we examine how peak output probabilities Pr(𝑦 peak 𝑖 = 1∣S) may be used during in- ference (i.e. without shuffling). The interpretation of probabilities requires a model to be calibrated. This property is further examined in AppendixE Figure 10.

in the following paragraphs are derived from the DRIAMS pre-training test set. Maldi Transformer denoises spectra In Figure 4A, a randomly selected Staphylococcus epidermidis 6 spectrum is visualized, each peak colored according to its output probability Pr(𝑦 peak 𝑖 = 1∣S). Most peaks are (correctly) assigned a high probability of belong- ing to their spectrum. The peaks with a low probabil- ity could be mistaken by the model, or could repre- sent "true" noise in the spectrum. Such noise may still be present in the final spectrum due to, for example (1) noisy readouts from the spectrometer , or (2) short- comings in preprocessing. In order to validate this hypothesis, it makes sense to look at patterns across multiple spectra. In Figure 4E, all S. epidermidis spec- tra are visualized together , each peak as a single dot. Black dots represent peaks that are predicted to be noise with high probability (Pr(𝑦 peak 𝑖 = 1∣S) ≤ 0.05). In the magnified parts of the plot, it can be seen that blue dots cluster together , meaning thatS. epidermidis spectra often have peaks in the same places. Black dots mainly fall outside or on the edges of those clus- ters, signifying that those peaks are rightly picked up by the model as noise. Consequently, Maldi Trans- formers can serve a broader purpose as spectrum de- noisers, in addition to their supervised learning capa- bilities. Noisy peaks are indicative of lower downstream performance In order to better grasp the effect of noisy peaks on spectra, peak predictions are exam- ined across all spectra in the DRIAMS test set. Fig- ure 4B shows the empirical cumulative distribution of peak output probabilities. Only approximately 5% of all peaks are assigned a probability of belonging to their respective spectrum Pr(𝑦 peak 𝑖 = 1∣S) of lower than 50%. Conversely, half of the peaks are predicted with a probability of98.5% or greater . This shows that while noisy peaks do exist, they typically make up a small percentage of the overall input spectrum. 6Staphylococcus epidermidis is the most occurring species in the DRIAMS pre-training test set. Noisy peaks may not uniformly occur across the dataset. Some spectra may have more noisy peaks than others. A good spectrum quality statistic can, hence, be the fraction of peaks that are confidently predicted as "true", e.g. with a probability greater than 95%. Figure 4C shows the distribution of spec- tra in function of this statistic. A slight left tail in the distribution signifies that some spectra have a large amount of noisy peaks. Aproximately 10% of the spec- tra have more than half of their peaks predicted with a probability smaller than 95%. Figure 4D shows that this statistic is also indicative of predictive perfor- mance. Spectra with more confidently "belonging" peaks have a higher species-level accuracy (on DRI- AMS test set labels). Species-level accuracy ramps up from 90% (or lower) for spectra with a lot of noisy peaks, up to nearly-perfect accuracy for spectra with (almost) all "predicted true" peaks 7. The fraction of "predicted true" peaks also correlates with prediction certainty (see Appendix E Figure 11). These results showcase Maldi Transformer’s ability to not only im- prove performance, but also provide further insights into the data. 4. Discussion As with many subfields of machine learning, size, quality, and diversity of data constitute a huge bottle- neck. In this study, an apparent ceiling in represen- tational capacity has been obtained given the avail- able public data. This is evidenced by the fact that the XL variant of Maldi Transformer typically does not give an advantage in downstream tasks over the M or L model. Its size (∼15M weights), however , is still small by self-supervised transformer standards. We hypothesize that the representational capacity of larger transformers is simply not necessary for the relatively-simple datasets that are available in the MALDI-TOF MS domain. Because of this, we argue that more efforts to collect and publish data may be the most important factor for MALDI-TOF-based ma- chine learning research to continue to flourish. As spectral data is routinely generated within hundreds, if not thousands, of hospitals, we envision collabora- tion efforts with healthcare (as inWeis et al.(2022)) to play a crucial role in this regard. Deep learning models have been hailed as feature extractors. For this class of models, it is typically ar- 7Note that DRIAMS species labels are produced from the MS manufacturers’ software and models. It can be expected that it is relatively easy for a model to reproduce the labels (predictions) from another model. Hence, the nearly-perfect species-level accu- racy on DRIAMS is not unexpected. 8 .CC-BY-NC-ND 4.0 International licenseavailable under a (which was not certified by peer review) is the author/funder, who has granted bioRxiv a license to display the preprint in perpetuity. It is made The copyright holder for this preprintthis version posted January 22, 2024. ; https://doi.org/10.1101/2024.01.18.576189doi: bioRxiv preprint Pre-trained Maldi Transformers improve MALDI-TOF MS-based prediction Peak Intensity Peak (m/z) All S. epidermidis spectra visualized together C D E Cumulative fraction of peaks B Peak Intensity Peak (m/z) Randomly selected S. epidermidis spectrum A Species-level accuracy Fraction of peaks Fraction of spectra Fraction of peaks Figure 4 | Maldi Transformer analysis on DRIAMS pre-training test data. Note that no shuffling of peaks across spectra is performed to generate these plots. Each spectrum contains their original detected peaks. A: A ran- domly selected spectrum, each peak colored by the model’s output probability of that peak belonging to that spectrum or not Pr(𝑦 peak 𝑖 = 1∣S). B: Empirical cumulative distribution of those output probabilities over all spectra in the DRIAMS test set. C: Histogram of fraction of peaks confidently (> 0.95) predicted as "true" peaks per spectrum. D: Species-level accuracy in function of fraction of peaks confidently (>0.95) predicted as "true" peaks per spectrum. E: All S. epidermidis spectra visualized together , every peak as a separate dot. Peaks confi- dently predicted as "noise" (≤ 0.05) are shown in black. Zoomed in subplots show that "noise" peaks originate outside clusters of usual peak locations. gued that it matters less how features are presented to the model, as they can compose the relevant fea- tures themselves in their hidden representations. We would caution against this perspective, and state that how inputs are presented to a deep learning model could have far-reaching impacts on what the model can learn. Taking language as an example, character- level language models traditionally underperform the same models trained on (sub)word-level. In this work, MALDI-TOF mass spectra are presented to the model by their peaks. Because of this, any preprocessing steps and algorithms to determine peaks play cru- cial roles in the final performance of the model. For this reason, experimentation with peak detection al- gorithms is a promising future research direction. While this work is not the first to fit a transformer model on mass spectral data (Yilmaz et al., 2022), it is the first to propose a self-supervised learning strat- egy adapted to this data modality. It is expected that the combination of shuffling peaks to obtain nega- tive examples, paired with a peak discriminator , could be useful in other mass spectral domains — namely, the types of mass spectra where the (co-)occurrence of peaks constitute the most-critical biological signal. For this reason, we hope that our ablations (§3.2) and

in Appendix C serve as useful guidelines for adapting this work to other mass spectral data modalities, generated with MS instruments using al- ternative ionisation and separation methods. As the pre-trained model can be interpreted as learning peak co-occurences, this property is exam- ined in §3.3. There, it is shown that Maldi Trans- former can be used to detect noisy peaks. The pre- dicted absence of noisy peaks is found to correlate with higher predictive performance. These insights go beyond those offered by off-the-shelf machine learning techniques. Paired with its state-of-the-art (or competitive) performance results, it can be con- cluded that Maldi Transformer enhances what biolog- ical practitioners get out of their MALDI-TOF MS data.

This work was supported by Research Foundation - Flanders (FWO) [PhD Fellowship fundamental re- search grant 1153024N to G.D.W.]. W.W. also received funding from the Flemish Government under the "Onderzoeksprogramma Artificiële Intelligentie (AI) Vlaanderen" Programme 9 .CC-BY-NC-ND 4.0 International licenseavailable under a (which was not certified by peer review) is the author/funder, who has granted bioRxiv a license to display the preprint in perpetuity. It is made The copyright holder for this preprintthis version posted January 22, 2024. ; https://doi.org/10.1101/2024.01.18.576189doi: bioRxiv preprint Pre-trained Maldi Transformers improve MALDI-TOF MS-based prediction

Olivier Dauwalder , Tiphaine Cecchini, Jean Philippe Rasigade, and François Vandenesch. Matrix assisted laser desorption ionisation/time of flight (maldi/tof) mass spectrometry is not done revolutionizing clin- ical microbiology diagnostic. Clinical Microbiology and Infection, 29(2):127–129, 2023. Caroline V Weis, Catherine R Jutzeler , and Karsten Borgwardt. Machine learning for microbial identi- fication and antimicrobial susceptibility testing on maldi-tof mass spectra: a systematic review. Clin- ical Microbiology and Infection, 26(10):1310–1317, 2020a. Piseth Seng, Michel Drancourt, Frédérique Gouriet, Bernard La Scola, Pierre-Edouard Fournier , Jean Marc Rolain, and Didier Raoult. Ongoing revolution in bacteriology: routine identification of bacteria by matrix-assisted laser desorption ion- ization time-of-flight mass spectrometry. Clinical infectious diseases, 49(4):543–551, 2009. Alain Bizzini, Katia Jaton, Daniel Romo, Jacques Bille, Guy Prod’hom, and Gilbert Greub. Matrix- assisted laser desorption ionization–time of flight mass spectrometry as an alternative to 16s rrna gene sequencing for identification of difficult-to- identify bacterial strains. Journal of clinical micro- biology, 49(2):693–696, 2011. Alex Van Belkum, Martin Welker , Marcel Erhard, and Sonia Chatellier . Biomedical mass spectrometry in today’s and tomorrow’s clinical microbiology lab- oratories. Journal of clinical microbiology, 50(5): 1513–1517, 2012. Walter Florio, Arianna Tavanti, Simona Barnini, Emilia Ghelardi, and Antonella Lupetti. Recent ad- vances and ongoing challenges in the diagnosis of microbial infections by maldi-tof mass spectrome- try. Frontiers in microbiology, 9:1097, 2018. Yan Cao, Lei Wang, Ping Ma, Wenting Fan, Bing Gu, and Shaoqing Ju. Accuracy of matrix-assisted laser desorption ionization–time of flight mass spectrom- etry for identification of mycobacteria: a systematic review and meta-analysis. Scientific reports, 8(1):1– 9, 2018. Georgia Vrioni, Constantinos Tsiamis, George Oikono- midis, Kalliopi Theodoridou, Violeta Kapsimali, and Athanasios Tsakris. Maldi-tof mass spectrometry technology for detecting biomarkers of antimicro- bial resistance: current achievements and future perspectives. Annals of translational medicine, 6 (12), 2018. Justin M Hettick, Michael L Kashon, James E Slaven, Yan Ma, Janet P Simpson, Paul D Siegel, Gerald N Mazurek, and David N Weissman. Discrimination of intact mycobacteria at the strain level: a combined maldi-tof ms and biostatistical analysis.Proteomics, 6(24):6416–6425, 2006. Caroline Weis, Aline Cuénod, Bastian Rieck, Olivier Dubuis, Susanne Graf, Claudia Lang, Michael Oberle, Maximilian Brackmann, Kirstine K Sø- gaard, Michael Osthoff, et al. Direct antimicrobial resistance prediction from clinical maldi-tof mass spectra using machine learning. Nature Medicine, 28(1):164–174, 2022. Julie Gagnaire, Olivier Dauwalder , Sandrine Boisset, David Khau, Anne-Marie Freydiere, Florence Ader , Michèle Bes, Gerard Lina, Anne Tristan, Marie- Elisabeth Reverdy, et al. Detection of staphylo- coccus aureus delta-toxin production by whole-cell maldi-tof mass spectrometry. PloS one, 7(7):e40660, 2012. Gaetan De Waele, Gerben Menschaert, and Willem Waegeman. An antimicrobial drug recommender system using maldi-tof ms and dual-branch neural networks. bioRxiv, pages 2023–09, 2023. Thomas Mortier , Anneleen D Wieme, Peter Van- damme, and Willem Waegeman. Bacterial species identification using maldi-tof mass spectrometry and machine learning techniques: A large-scale benchmarking study. Computational and Structural Biotechnology Journal, 19:6157–6168, 2021. Caroline Weis, Max Horn, Bastian Rieck, Aline Cuénod, Adrian Egli, and Karsten Borgwardt. Topo- logical and kernel-based microbial phenotype pre- diction from maldi-tof mass spectra.Bioinformatics, 36(Supplement_1):i30–i38, 2020b. Kévin Vervier , Pierre Mahé, Jean-Baptiste Veyrieras, and Jean-Philippe Vert. Benchmark of structured machine learning methods for microbial identifica- tion from mass-spectrometry data. arXiv preprint arXiv:1506.07251, 2015. Nam K Tran, Taylor Howard, Ryan Walsh, John Pep- per , Julia Loegering, Brett Phinney, Michelle R Salemi, and Hooman H Rashidi. Novel application of automated machine learning with maldi-tof-ms for rapid high-throughput screening of covid-19: A proof of concept. Scientific reports, 11(1):8219, 2021. Victor Ryzhov and Catherine Fenselau. Characteriza- 10 .CC-BY-NC-ND 4.0 International licenseavailable under a (which was not certified by peer review) is the author/funder, who has granted bioRxiv a license to display the preprint in perpetuity. It is made The copyright holder for this preprintthis version posted January 22, 2024. ; https://doi.org/10.1101/2024.01.18.576189doi: bioRxiv preprint Pre-trained Maldi Transformers improve MALDI-TOF MS-based prediction tion of the protein subset desorbed by maldi from whole bacterial cells. Analytical chemistry, 73(4): 746–750, 2001. Dzmitry Bahdanau, Kyunghyun Cho, and Yoshua Bengio. Neural machine translation by jointly learning to align and translate. arXiv preprint arXiv:1409.0473, 2014. Ashish Vaswani, Noam Shazeer , Niki Parmar , Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Łukasz Kaiser , and Illia Polosukhin. Attention is all you need. Advances in neural information processing systems, 30, 2017. Xiao Liu, Fanjin Zhang, Zhenyu Hou, Li Mian, Zhaoyu Wang, Jing Zhang, and Jie Tang. Self-supervised learning: Generative or contrastive. IEEE transac- tions on knowledge and data engineering, 35(1):857– 876, 2021. Randall Balestriero, Mark Ibrahim, Vlad Sobal, Ari Morcos, Shashank Shekhar , Tom Goldstein, Florian Bordes, Adrien Bardes, Gregoire Mialon, Yuandong Tian, et al. A cookbook of self-supervised learning. arXiv preprint arXiv:2304.12210, 2023. John Jumper , Richard Evans, Alexander Pritzel, Tim Green, Michael Figurnov, Olaf Ronneberger , Kathryn Tunyasuvunakool, Russ Bates, Augustin Žídek, Anna Potapenko, et al. Highly accurate pro- tein structure prediction with alphafold. Nature, 596(7873):583–589, 2021. Ahmed Elnaggar , Michael Heinzinger , Christian Dal- lago, Ghalia Rehawi, Yu Wang, Llion Jones, Tom Gibbs, Tamas Feher , Christoph Angerer , Martin Steinegger , et al. Prottrans: Toward understanding the language of life through self-supervised learn- ing. IEEE transactions on pattern analysis and ma- chine intelligence, 44(10):7112–7127, 2021. Roshan M Rao, Jason Liu, Robert Verkuil, Joshua Meier , John Canny, Pieter Abbeel, Tom Sercu, and Alexander Rives. Msa transformer . InInternational Conference on Machine Learning, pages 8844–8856. PMLR, 2021. Žiga Avsec, Vikram Agarwal, Daniel Visentin, Joseph R Ledsam, Agnieszka Grabska-Barwinska, Kyle R Tay- lor , Yannis Assael, John Jumper , Pushmeet Kohli, and David R Kelley. Effective gene expression pre- diction from sequence by integrating long-range in- teractions. Nature methods, 18(10):1196–1203, 2021. Jim Clauwaert and Willem Waegeman. Novel trans- former networks for improved sequence labeling in genomics. IEEE/ACM Transactions on Compu- tational Biology and Bioinformatics, 19(1):97–106, 2020. Jörg Franke, Frederic Runge, and Frank Hutter . Proba- bilistic transformer: Modelling ambiguities and dis- tributions for rna folding and molecule design. Ad- vances in Neural Information Processing Systems, 35:26856–26873, 2022. Melih Yilmaz, William Fondrie, Wout Bittremieux, Se- woong Oh, and William S Noble. De novo mass spec- trometry peptide sequencing with a transformer model. In International Conference on Machine Learning, pages 25514–25522. PMLR, 2022. Peter Lasch, Maren Stämmler , and Andy Schneider . Version 4 (20230306) of the maldi-tof mass spec- trometry database for identification and classifi- cation of highly pathogenic microorganisms from the robert koch-institute (rki), 2023. URL https: //doi.org/10.5281/zenodo.7702375. Sebastian Gibb and Korbinian Strimmer . Maldiquant: a versatile r package for the analysis of mass spec- trometry data. Bioinformatics, 28(17):2270–2271, 2012. Iz Beltagy, Matthew E Peters, and Arman Cohan. Long- former: The long-document transformer . arXiv preprint arXiv:2004.05150, 2020. Jacob Devlin, Ming-Wei Chang, Kenton Lee, and Kristina Toutanova. Bert: Pre-training of deep bidi- rectional transformers for language understanding. arXiv preprint arXiv:1810.04805, 2018. Alexey Dosovitskiy, Lucas Beyer , Alexander Kolesnikov, Dirk Weissenborn, Xiaohua Zhai, Thomas Unterthiner , Mostafa Dehghani, Matthias Minderer , Georg Heigold, Sylvain Gelly, et al. An image is worth 16x16 words: Transformers for image recognition at scale. arXiv preprint arXiv:2010.11929, 2020. Sharan Narang, Hyung Won Chung, Yi Tay, William Fedus, Thibault Fevry, Michael Matena, Karishma Malkan, Noah Fiedel, Noam Shazeer , Zhenzhong Lan, et al. Do transformer modifications transfer across implementations and applications? arXiv preprint arXiv:2102.11972, 2021. Jimmy Lei Ba, Jamie Ryan Kiros, and Geoffrey E Hinton. Layer normalization. arXiv preprint arXiv:1607.06450, 2016. Noam Shazeer . Glu variants improve transformer . arXiv preprint arXiv:2002.05202, 2020. 11 .CC-BY-NC-ND 4.0 International licenseavailable under a (which was not certified by peer review) is the author/funder, who has granted bioRxiv a license to display the preprint in perpetuity. It is made The copyright holder for this preprintthis version posted January 22, 2024. ; https://doi.org/10.1101/2024.01.18.576189doi: bioRxiv preprint Pre-trained Maldi Transformers improve MALDI-TOF MS-based prediction Dan Hendrycks and Kevin Gimpel. Gaussian error lin- ear units (gelus). arXiv preprint arXiv:1606.08415, 2016. Diederik P Kingma and Jimmy Ba. Adam: A

for stochastic optimization. arXiv preprint arXiv:1412.6980, 2014. Alec Radford, Karthik Narasimhan, Tim Salimans, Ilya Sutskever , et al. Improving language under- standing by generative pre-training. 2018. Tom Brown, Benjamin Mann, Nick Ryder , Melanie Subbiah, Jared D Kaplan, Prafulla Dhariwal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, et al. Language models are few-shot learn- ers. Advances in neural information processing sys- tems, 33:1877–1901, 2020. A. DRIAMS species processing DRIAMS contains species labels for many spectra. These species labels are derived from the species identification pipelines included with the MALDI-TOF MS machines. As such, some preprocessing steps are taken to make the labels as presentable as possible to an ML model. Firstly, labels in DRIAMS containing the string"not reliable identification" are integrally deleted. Second, many spectra are labeled as "MIX!species", indicating that the spectrum potentially contains an impure mixture of species. As such, species labels for these spectra are also not used. Additionally, species labels that occur fewer than five times in the train- ing set are removed. Finally, species labels occurring only in the validation or test set, but not in the training set, are similarly removed. After processing, DRIAMS contains 469 different species labels. Note that the previous steps involve removing of labels, not spec- tra themselves. Corresponding spectra are still kept in the dataset, albeit without a species label. B. On preprocessing and persis- tence transformation While Weis et al.(2020b) argue that persistence trans- formation nullifies the need for the parameter-heavy chain of preprocessing steps, we advocate for the op- posite. To support this stance, a simple exploratory visualization is made. In Figure 5, all Bacillus anthracis8 spectra in the RKI training set are put through two preprocessing chains. The first only does (1) trimming to the 2000-20000 Da range, (2) intensity calibration so that the total inten- sity sums to 1, and (3) persistence transformation, and keeps the 200 highest peaks. The second preprocess- ing chain performs all those steps preceded by vari- ance stabilization, smoothing, and background re- moval, as in §2.1. Then, across all preprocessed spec- tra, the occurrence of detected peaks in bins of 3 Da is counted and plotted in Figure 5. The resulting profile can be considered a summary profile of all detected peaks in Bacillus anthracis spectra. The detected peaks with extra preprocessing result in a cleaner profile, with peaks more concentrated in regions that clearly correspond to biologically-informative signal (especially notable when inspecting the left tail of the spectra). In contrast, without prior preprocess- ing, detected peaks are seemingly more-randomly dis- tributed, hinting that noise in the spectrum negatively affects peak detection. Figure 5 | Visualization of summary profiles of Bacil- lus anthracis RKI training spectra obtained when do- ing peak detection with (R) or without (L) preprocess- ing first. C. Design of a self-supervised pre- training task A large part of this study concerns the design of a self-supervised pre-training task for MALDI-TOF MS data. This is motivated by the scarcity of data in this domain. Through self-supervised learning, a greater representational capacity can be obtained, maximally utilizing patterns in the data. The following para- graphs chronicle thoughts and experiments w.r .t. the design of the SSL task. MALDI-TOF MS spectra can be represented as sets of peaks. As each peak is characterized by an inten- sity and an (𝑚/𝑧) value, a natural parallel is drawn with language data. Instead of a sequence of words, 8Bacillus anthracis is chosen because it is the most occurring species in the RKI training dataset. 12 .CC-BY-NC-ND 4.0 International licenseavailable under a (which was not certified by peer review) is the author/funder, who has granted bioRxiv a license to display the preprint in perpetuity. It is made The copyright holder for this preprintthis version posted January 22, 2024. ; https://doi.org/10.1101/2024.01.18.576189doi: bioRxiv preprint Pre-trained Maldi Transformers improve MALDI-TOF MS-based prediction a sequence of peaks is processed. A crucial differ- ence is that, for language transformers, positional in- dices for each text "token" are set to the integer range numbering 0 to 𝑛 − 1, with 𝑛 the number of tokens in the input sequence. With MALDI-TOF MS spec- tra, positional indices (i.e. (𝑚/𝑧) values) are irregu- larly spaced and real-valued. Additionally, whereas words have a categorical identity in text, MALDI-TOF MS intensity is also real-valued. These two factors have to be taken into consideration when porting self- supervised learning techniques from one modality to the mass spectral domain. In language transformer pre-training, perhaps the two most-classically quoted techniques are masked language modeling (MLM) (Devlin et al.,2018) and au- toregressive modeling (AR) (Radford et al., 2018). The GPT series of models best exemplifies the success of the latter category (Brown et al., 2020). For MALDI- TOF MS, however , it is unclear how to autoregres- sively model a set of peaks. Autoregressive models im- ply a certain ordering of data, whereas this perspec- tive is ill-fitting for the set-valued peaks input. For example, how does one train to generate the "next" peak given the previous peaks, if it is unclear how to define what the "next" peak is? For example, does one choose to order peaks by their height, or by their (𝑚/𝑧) value? If the model predicts the second-next peak instead of the next, is that necessarily wrong? If not, how to efficienctly construct the loss to take this into account? Instead of answering the issues with autoregres- sively modeling a set-valued input, we may consider the alternative strategy: MLM, introduced by BERT (Devlin et al., 2018). For example, instead of masking out words, intensities can be masked, and a model can be trained to recover the intensities with the use of the mean-square error loss. This brings us to first ab- lated pre-training technique in §3.2. The training for this strategy is performed identical to the final strat- egy outlined in §2.2. The only difference being that the peak discrimination loss Lpeaks is exchanged for the mean-square error loss on masked peaks. Peaks to train on are similarly selected with 15% probabil- ity. Of those 15%, 80% are assigned masked intensities and 20% are left unchanged. After fitting a regression MLM on intensities with limited success, a logical next step is to design a self- supervised classification task. The following para- graphs describe the exact procedure resulting in the negative peak sampler model in §3.2. Taking inspiration from contrastive learning (Liu et al., 2021), a per-peak classification task can ask a model whether a peak belongs to the rest of the spectrum or not. Just as in contrastive learning, this strategy requires to sample negative peaks to de- liver negative samples. One way to generate nega- tive peaks is to randomly generate them from some estimated distribution of peaks. Here, we estimate the distribution of peaks in two steps. First, all peak locations in the training dataset are collected: 𝑃train = { (𝑚 𝑧 )𝑗 »» » » » » (𝑚 𝑧 )𝑗 ∈ 𝑆𝑖} 𝑛 𝑖=1 , and a probability mass function over discrete bins is calculated: Pr(𝑏𝑖) = » » » » » » {𝑝 𝑗∣𝑝 𝑗 ∈ 𝑃train ∧ 𝑝 𝑗 ∈ 𝑏𝑖}» » » » » » » » » » »» {𝑝 𝑗∣𝑝 𝑗 ∈ 𝑃train}»» » » »» (1) with bins chosen by 1 Da intervals: 𝑏𝑖 ∈ ]𝑖, 𝑖 + 1], for 𝑖 ∈ {2000, 2001, ..., 19 999}. In other words, a probability mass function over 1 Da bins is cre- ated by counting how many times peaks fall into each bin in the training data. Next, within each bin, quantiles of the found intensities therein are calculated: 𝑄𝑏𝑖 = {𝑞0.00, 𝑞 0.01, ..., 𝑞 0.99, 𝑞 1.00}. To sample a negative peak, an (𝑚/𝑧) value is first sampled by uniformly sampling a location within a sampled bin: (𝑚/𝑧) ∼ Pr(𝑏𝑖) + Unif[0,1]. After , an intensity is drawn by uniformly sam- pling within a uniformly sampled interquan- tile range: 𝐼 ∼ Unif[𝑞𝑗, 𝑞𝑗+1 ∣ 𝑞 ∈ 𝑄𝑏𝑖 ], with 𝑗 ∼ Unif{0.00, 0.01, ...0.99}. The training for this strategy is performed identical to the final strategy outlined in §2.2. The only difference being that, in- stead of shuffling peaks around to generate negatives, negative peaks are now sampled. In every training step, 15% of the peaks are selected for training, half of which are exchanged for sampled negative ones. As discussed in §3.2, generating negative peaks is not ideal for forcing the model to reason over peak co- occurrences. The model is allowed to learn any mis- representation in negatives as a shortcut. A way to circumvent the issues with generating negative peaks, is to use real ones. One way to present real peaks as negatives, is to take them from other spectra. This is where we land on the final self-supervised strategy that we test, and ultimately propose in our main text (see §2.2, Figure 1, and Algorithm 1). D. Downstream tasks and base- lines For all three downstream tasks, the pre-trained model is plugged in at initialization and all weights are fine- 13 .CC-BY-NC-ND 4.0 International licenseavailable under a (which was not certified by peer review) is the author/funder, who has granted bioRxiv a license to display the preprint in perpetuity. It is made The copyright holder for this preprintthis version posted January 22, 2024. ; https://doi.org/10.1101/2024.01.18.576189doi: bioRxiv preprint Pre-trained Maldi Transformers improve MALDI-TOF MS-based prediction Algorithm 1: Maldi Transformer peak dis- crimination pre-training. Input: 𝑿 ∈ R𝑏×𝑚×𝑑, 𝒚spec ∈ R𝑏 # 𝑏 = batch size, 𝑚 = number of peaks, 𝑑 = hidden dim # 𝑿 = encoded batch of spectra, 𝒚spec = species labels # 𝒚spec contains int labels and NaN if unlabeled 1 Let 𝑼 ∈ {0,1}𝑏×𝑚, with 𝑈𝑖 𝑗 ∼ Bern(0.15)# Sample 15% of peaks 2 Let 𝑷 ∈ {0,1}𝑏×𝑚, with 𝑃𝑖 𝑗 ∼ Bern(0.50) # Sample 50% of peaks 3 𝒑 = [(𝑖, 𝑗) ∶ 𝑈𝑖 𝑗 = 1 ∧ 𝑃𝑖 𝑗 = 1] # Positive label indices 4 𝒏 = [(𝑖, 𝑗) ∶ 𝑈𝑖 𝑗 = 1 ∧ 𝑃𝑖 𝑗 = 0] # Negative label indices 5 𝑿 shuff = 𝑿[permute(𝒏)]# Shuffle neg peaks 6 𝒑𝑖∈{1,...,𝑚}, 𝒑[CLS] = MaldiTransformer(𝑿 shuff) 7 ˆ𝒀 peak = 𝜎(𝒘⊺ 𝑝 ⋅ 𝒑𝑖∈{1,...,𝑚}) # Projection of peak embeddings 8 ˆ𝒀 spec = Softmax(𝑾 ⊺ 𝑠 ⋅ 𝒑[CLS]) # Projection of spec embed 9 Lpeaks = E (∑(𝑖, 𝑗) 𝑈𝑖 𝑗 =1 −I (𝑋𝑖 𝑗 = 𝑋shuff 𝑖 𝑗 )log (ˆ𝑌 peak 𝑖 𝑗 ) − I (𝑋𝑖 𝑗 ≠ 𝑋shuff 𝑖 𝑗 ) − log (1 − ˆ𝑌 peak 𝑖 𝑗 ) )# Binary CE on peaks 10 Lspec = E (∑𝑏 𝑖 − log ˆ𝑌 spec 𝑖,𝑐 ) # with 𝑐 = 𝑦 spec 𝑖 the class index 11 L = Lpeaks + 𝜆 ⋅ Lspec, with 𝜆 ∼ Bern(0.01) Output: L # Final loss tuned (i.e. no weight freezing). A task-specific lin- ear output head𝑾 𝑜𝑢𝑡 projecting the spectrum embed- ding 𝒑[CLS]to the desired output space is trained from scratch. All downstream models are similarly trained with the Adam optimizer with a batch size of 128 and dropout of 0.2. A linear learning rate warm-up over the first 250 steps is applied, after which the rate is kept constant. For AMR prediction, training of Maldi Transformer- based recommenders is performed identical to the MLP-based baselines in De Waele et al. (2023). Briefly explained, for every combination of spectrum em- bedder (four sizes: S, M, L, and XL) and drug em- bedder (four types), four different learning rates ({1e-5,5e-5, 1e-4, 5e-4})are tested. For all these dif- ferent combinations, three models are trained (using different random seeds for model initialization and batching of data). For every spectrum and drug em- bedder combination, only results from the best learn- ing rate are presented; that is, the learning rate re- sulting in the best average validation micro ROC-AUC for that combination. The validation set is checked ev- ery tenth of an epoch. Models are trained for a max- imum of 50 epochs, and their training is halted early when validation micro ROC-AUC hasn’t improved for 10 validation set checks. The checkpoint of the best performing model (in terms of validation micro ROC- AUC) is used as the final model. The baselines for the AMR prediction task are the models described in De Waele et al. (2023), which describes the recom- mender model structure in greater detail. The pre-training of Maldi Transformer for species identification is performed in a similar way. The differences consist of: (1) 𝑾 𝑜𝑢𝑡 returning 270, or 1088, for the RKI and LM-UGent dataset, respec- tively, (2) five different learning rates are tested: ({1e-5,5e-5, 1e-4, 5e-4, 1e-3}), and (3) validation species-level accuracy is tracked to halt training early (for a maximum of 250 epochs, and it is only checked once per epoch). As species identification is a multi-class classification task, models are then optimized using a softmax operation, combined with the cross-entropy loss. Species identification is compared to MLP baselines, Logistic Regression, Ran- dom Forest, and k-nearest neighbors (k-NN) models. All of these baselines are trained on preprocessed and binned spectra (see §2.1). The S, M, L, and XL MLPs are identical in construction to the spectrum embedders in De Waele et al. (2023), but with 𝑛- dimensional outputs instead of 64, with 𝑛 the number of classes (see Table 3). MLP baselines are trained using the same strategy as Maldi Transformers. That is, for all model sizes, five different learning rates ({1e-5,5e-5, 1e-4, 5e-4, 1e-3}) are trained in triplicates. Model results from the best learning rate (in terms of validation species-level accuracy) are presented. Model training halts before its maximum of 250 epochs if validation species-level accuracy hasn’t increased in 10 epochs, and the model with the best validation accuracy is saved. Table 3 | All tested model sizes for the MLP base- line Hidden sizes represent the evolution of the hid- den state dimensionality as it goes through the model, with every hyphen defining one fully connected layer . 𝑛 represents the number of output nodes. 𝑛 equals 64, 270, and 1088 for DRIAMS AMR prediction, species identification on RKI, and LM-UGent, respectively. Size # Weights Hidden sizes S 1.58M 6000-256-128- 𝑛 M 3.25M 6000-512-256-128- 𝑛 L 6.85M 6000-1024-512-256-128- 𝑛 XL 15.09M 6000-2048-1024-512-256-128- 𝑛 For non-neural baseline classifiers (Random For- est, Logistic Regression, and k-NN), a grid-search is performed to find optimal hyperparameters. Given the non-stochastic nature of their implementations, only one model is trained after tuning and, hence, only one test performance is reported. The param- eter grid for Random Forest consists of {max_depth = [25,50,75,100], min_samples_split = [2,5,10], max_features = [10,25,50,100]}. All random forests are trained with 200 trees. For Logistic Regres- sion: {standardscaling = [True,False], L2_norm = [1e3,1e2,10,1,0.1,1e-2,1e-3]}. And for k-NN: 14 .CC-BY-NC-ND 4.0 International licenseavailable under a (which was not certified by peer review) is the author/funder, who has granted bioRxiv a license to display the preprint in perpetuity. It is made The copyright holder for this preprintthis version posted January 22, 2024. ; https://doi.org/10.1101/2024.01.18.576189doi: bioRxiv preprint Pre-trained Maldi Transformers improve MALDI-TOF MS-based prediction {standardscaling = [True,False], n_neighbors = [1,2,3,4,5,6,7,9,10,25]. E. Figures and tables supporting the methods and results sec- tions Figure 6 | Utilized transformer encoder block. The network uses pre-LayerNorms and GeLU gated lin- ear units (GLU) in the feedforward (FF) networks (Ba et al., 2016; Shazeer, 2020). Peak discrimination loss Species classification loss Training step Training step Figure 7 | Pre-training dynamics using different strategies for combining the peak discrimination loss Lpeaks and the species identification loss Lspec. Loss on the validation set is shown. The figure is shown only for the first 200 000 pre-training steps using a medium-sized Maldi Transformer to illustrate. It is observed that if Lspec is applied at every step (using a weight of 0.01 or 1.00), this component starts over- fitting well before the Lpeaks component is converged (training is performed for 500 000 steps total). By not applying Lspec at every training step, the Adam opti- mizer momentum is dominated by Lpeaks. As Lspec is the "easier" task and more prone to overfitting, this regime benefits the final model. Species classification loss Peak discrimination loss Training step Training step Figure 8 | Pre-training validation loss curves for all Maldi Transformer model sizes. Species-level accuracy Genus-level accuracy Maldi Transformer Maldi Transformer (No self-supervised fine-tuning on LM-UGent) Figure 9 | LM-UGent performance improves when, prior to supervised fine-tuning, the pre-trained model is first fine-tuned using the self-supervised training task. In the main text, we refer to this step as the do- main adaptation step. 15 .CC-BY-NC-ND 4.0 International licenseavailable under a (which was not certified by peer review) is the author/funder, who has granted bioRxiv a license to display the preprint in perpetuity. It is made The copyright holder for this preprintthis version posted January 22, 2024. ; https://doi.org/10.1101/2024.01.18.576189doi: bioRxiv preprint Pre-trained Maldi Transformers improve MALDI-TOF MS-based prediction Number of obs. in bin Figure 10 | Calibration curve for pre-training peak discrimination. On the DRIAMS test-set, the usual peak shuffling and peak discrimination is performed (see §2.2). Predictions are then sorted and split into equal frequency bins (on the x-axis). Within those bins, the average predicted value is plotted against the actual fraction of positive true labels in that bin. A calibrated model requires predictions to be inter- pretable in a frequentist manner , i.e. a sample with an output probability of 80% is expected to be posi- tive 80% of the times. Hence, the aforementioned plot is expected to follow the diagonal. The plot shows that the model is reasonably-well calibrated, with slight overconfidence. Fraction of peaks Max species prediction logit Figure 11 | Max species prediction logit𝑾 ⊺ 𝑠 ⋅ 𝒑[CLS] in function of number of confidently "belonging" peaks for a spectrum for DRIAMS test set spectra. A higher max logit for the species prediction task corresponds to a higher max output probability post-softmax, and, hence, higher confidence. It is displayed that the model is more confident in its prediction for spectra with a higher number of confidently belonging peaks. 16 .CC-BY-NC-ND 4.0 International licenseavailable under a (which was not certified by peer review) is the author/funder, who has granted bioRxiv a license to display the preprint in perpetuity. It is made The copyright holder for this preprintthis version posted January 22, 2024. ; https://doi.org/10.1101/2024.01.18.576189doi: bioRxiv preprint