METAGENE-1: Metagenomic Foundation Model for Pandemic Monitoring

🔼 Figure 2 illustrates the process of creating the metagenomic dataset used to train the METAGENE-1 model. Starting with wastewater samples (left), which naturally contain DNA and RNA from diverse organisms (center, e.g., bacteria, viruses, human, etc.), high-throughput sequencing is applied. This generates millions of paired-end reads (right), each several hundred base pairs long. The combined length of these sequences in the final dataset exceeds 1.5 trillion base pairs.

read the caption

Figure 2: Overview of the metagenomic data collection and sequencing pipeline for model pretraining. The process begins with the collection of wastewater (left), which contains genomic fragments from a diverse collection (e.g., tens of thousands) of constituent organisms (center). These samples are processed via high-throughput metagenomic sequencing to produce millions of paired-end reads (right), each consisting of hundreds of base pairs. The complete dataset comprises over 1.5 trillion base pairs of metagenomic sequences used for model pretraining.

🔼 This figure is a visualization of the taxonomic composition of the metagenomic dataset used to pretrain the METAGENE-1 model. The dataset was analyzed using Kraken 2 software, which identifies the taxonomic origin of short sequences (reads) by comparing them to a database of known organisms. The resulting taxonomic assignments are displayed using the Krona software for visualization as a treemap. The visualization shows the relative abundance of different taxonomic groups present in the dataset, including bacteria, viruses, eukaryotes (mainly Homo sapiens), and unclassified sequences. The visualization highlights the high diversity and complexity of the metagenomic data, which is typical of wastewater samples and comprises a range of organisms, making it suitable for training a general-purpose foundation model.

read the caption

Figure 3: Metagenomic composition of the METAGENE-1 pretraining dataset, estimated via Kraken 2 (Wood et al., 2019) sequence classification, and visualized via Krona (Ondov et al., 2011). See Figure 7 for a more-detailed view.

🔼 The figure shows a graph of the z-loss during the pretraining process of the METAGENE-1 model. The z-loss is a metric used to monitor the stability of the training process, and its behavior can indicate potential issues such as instability or divergence. A stable training process will typically show a smooth, gradually decreasing z-loss curve. Significant fluctuations or spikes in the z-loss can signal instability, which may require intervention to prevent the training from failing. The plot allows for visualization of the training stability of the model over time.

read the caption

Figure 4: We show z𝑧zitalic_z-loss during pretraining, which aids and gives an indicator of stability.

🔼 This figure displays the training and validation loss curves for the METAGENE-1 model during its pretraining phase. The left panel shows the training loss, which is the loss calculated on the training dataset during each iteration of training. The right panel shows the validation loss, which is the loss calculated on a separate held-out portion of the metagenomic dataset to assess the model’s generalization ability. The plots provide insights into the model’s learning progress and help identify potential issues such as overfitting (where the training loss decreases while validation loss increases). The slight oscillations visible in the training curve result from the pseudo-random data shuffling employed for training efficiency.

read the caption

Figure 5: METAGENE-1 loss curves during pretraining. We show training loss (left), and validation loss on a held out metagenomic sample (right).

🔼 Figure 6 shows the distribution of length-normalized cross-entropy loss values generated by METAGENE-1 across four different datasets: metagenomics, random sequences, human genome sequences, and mouse genome sequences. The length normalization accounts for differences in sequence length between these datasets, making the comparisons more meaningful. The figure is a histogram showing the frequency distribution of these loss values for each dataset. This visualization helps illustrate the model’s ability to distinguish between in-distribution (metagenomics) and out-of-distribution (OOD) data, which is relevant for anomaly detection tasks.

read the caption

Figure 6: Distribution of the length-normalized cross entropy loss across all datasets, given by METAGENE-1.

🔼 This table presents the results of the pathogen detection benchmark. Four different datasets (PATHOGEN-DETECT-1 through 4), each derived from separate wastewater sequencing deliveries and excluding data used for model pretraining, were used for evaluation. The model’s performance on each dataset is evaluated using the Matthews Correlation Coefficient (MCC), a measure commonly used for evaluating classification model performance. The header row displays the macro-averaged MCC across all four datasets, providing a single summary metric. Details on how these datasets were constructed and the specifics of the evaluation methodology are provided in Section 5.2 of the paper.

read the caption

Table 2: Results on the Pathogen Detection benchmark. The metric used for all evaluations is MCC. The header row reports macro-averaged performance metrics. See Section 5.2 for details.

🔼 Table 3 presents the results of the Gene-MTEB benchmark, a newly introduced evaluation metric designed to assess the quality of genomic sequence embeddings generated by various foundation models. It evaluates the models’ zero-shot performance on eight classification and eight clustering tasks using data from the Human Microbiome Project and a held-out portion of the METAGENE-1 metagenomic dataset. The table compares the performance of METAGENE-1 against several other state-of-the-art genomic models, specifically DNABERT-2, DNABERT-S, and two variants of Nucleotide Transformer (NT-2.5b-Multi and NT-2.5b-1000g). The results highlight METAGENE-1’s performance across different genomic tasks.

read the caption

Table 3: Results on the Genomic Embedding (Gene-MTEB) benchmark. See Section 5.3 for details.

🔼 Table 4 presents the results of the Genome Understanding Evaluation (GUE) benchmark, comparing METAGENE-1’s performance against other state-of-the-art genomic foundation models. The GUE benchmark consists of 28 sequence-level classification tasks focused on various aspects of genomics. The primary evaluation metric is the Matthews Correlation Coefficient (MCC), except for the COVID-19 task which uses the F1 score. The table displays the MCC or F1 score for each model on each task and then shows the macro-averaged performance across all tasks. The final row shows the ‘Global Win %’, indicating the percentage of tasks where each model achieved the highest score. Results for models other than METAGENE-1 are taken from Zhou et al. (2023). This provides a comprehensive comparison of METAGENE-1’s performance across a broad range of genomic tasks.

read the caption

Table 4: Results on the Genome Understanding Evaluation (GUE) benchmark. Non-METAGENE-1 results are adapted from Zhou et al. (2023). The metric used for all evaluations is MCC, except for the COVID task, which uses F1 score. The header rows report macro-averaged performance metrics. The final row shows Global Win %, i.e., the percentage of tasks in which a given method achieves top score under the associated metric.

🔼 This table presents the out-of-distribution (OOD) detection performance of the METAGENE-1 model. It compares the model’s ability to distinguish metagenomic sequences from other data sources, including randomly generated sequences and sequences from human and mouse genomes. The evaluation metric used is F1 score, and the table shows the F1 scores and standard errors for each data source, along with the average sequence length for each data source. The results demonstrate METAGENE-1’s effectiveness in identifying metagenomic sequences as in-distribution data and other types of sequences as out-of-distribution data.

read the caption

Table 5: OOD detection performance between metagenomics sequences and other data sources.

🔼 This table details the hyperparameters used for fine-tuning various models on the pathogen detection task. It lists the model (DNABERT-2, DNABERT-S, Nucleotide Transformer variants, and METAGENE-1), the type of fine-tuning (full model or LoRA), the LoRA modules used (if applicable), LoRA rank, alpha, and dropout rate. Optimizer parameters (optimizer, momentum, learning rate, learning rate scheduler, warmup steps, weight decay, denominator epsilon), precision, batch size, and number of epochs are also provided. A note indicates that for DNABERT-S, the learning rate was halved due to observed oscillations in training loss.

read the caption

Table 6: Hyperparameter settings for the Pathogen Detection fine-tuning experiments. ΛΛ\Lambdaroman_Λ: for DNABERT-S, we halve the learning to 5e-5 as we observe clear oscillation behavior in the training loss.

🔼 Table 7 details the hyperparameters used during fine-tuning experiments for the Genome Understanding Evaluation (GUE) benchmark. It specifies that Low-Rank Adaptation (LoRA) was applied to either query-value or query-key-value-dense modules, with a rank of 8 and alpha of 16. A dropout rate of 0.1 was used. The AdamW optimizer was employed with a learning rate tuned across an equally-spaced grid from 1e-4 to 1e-3, along with other standard hyperparameter settings. All hyperparameters were selected based on their performance on validation sets.

read the caption

Table 7: Hyperparameter settings for the GUE fine-tuning experiments. ΛΛ\Lambdaroman_Λ: LoRA is applied to query-value or query-key-value-dense modules. ΩΩ\Omegaroman_Ω: learning rates are tuned over a equally-spaced grid of 1e-4, 2e-4, ⋯⋯\cdots⋯, 1e-3. All hyperparameters are selected according to performances on validation sets.