FGeneBERT: function-driven pre-trained gene language model for metagenomics

Abstract

Metagenomic data, comprising mixed multi-species genomes, are prevalent in diverse environments like oceans and soils, significantly impacting human health and ecological functions. However, current research relies on K-mer, which limits the capture of structurally and functionally relevant gene contexts. Moreover, these approaches struggle with encoding biologically meaningful genes and fail to address the one-to-many and many-to-one relationships inherent in metagenomic data. To overcome these challenges, we introduce FGeneBERT, a novel metagenomic pre-trained model that employs a protein-based gene representation as a context-aware and structure-relevant tokenizer. FGeneBERT incorporates masked gene modeling to enhance the understanding of inter-gene contextual relationships and triplet enhanced metagenomic contrastive learning to elucidate gene sequence-function relationships. Pre-trained on over 100 million metagenomic sequences, FGeneBERT demonstrates superior performance on metagenomic datasets at four levels, spanning gene, functional, bacterial, and environmental levels and ranging from 1 to 213 k input sequences. Case studies of ATP synthase and gene operons highlight FGeneBERT's capability for functional recognition and its biological relevance in metagenomic research.

Keywords: DNA; metagenomics; pre-trained language model; transformer.

Figure 1

Motivaion. Two types of complex relationships between gene sequences and functions in metagenomics. One-to-Many problem means that the same gene may display different functions based on the genomic context; for example, ATP synthase works differently in plants, heterotrophic bacteria, and humans. Many-to-One problem shows that multiple genes may perform the same function; for instance, different genes from different bacteria, e.g., Cpf1, Cas1, etc., produce the same resistance function within the immune system CRISPR.

Figure 2

Overview of FGeneBERT. A metagenomic sequence is converted into ordered protein-based gene representations via a Context-Aware Tokenizer. Next, we pre-train a Gene Encoder with , 15% of these tokens are masked to predict labels . Meanwhile, we introduce to distinguish gene sequences. The data augmentation and negative sampling modules generate positive samples and negative samples , respectively. Finally, after fine-tuning, FGeneBERT can handle various downstream tasks.

Figure 3

Framework of Context-Aware Tokenizer. Gene sequences extracted from metagenomic sequence are translated into amino acid sequences and encoded into ESM-2 representations as . These representations are then concatenated sequentially every representation to form gene groups .

Figure 4

Visualization of attention. The attention value is from the first head of the last (19th) attention layer. Darker shading indicates higher attention weight.

Figure 5

Ablation studies of our proposed modules on four downstream tasks.

Figure 6

T-SNE visualization of different embeddings for ATP synthases. Each dot denotes a sequence and is grouped according to different functions.

Figure 7

Comparative analysis on tokenization efficiency: time(s) vs. memory (MB). Each point denotes a specific dataset, with the size indicating its scale.

Figure 8

Sensitivity w.r.t. hyper-parameters of CARD dataset on AMR gene family.