[Paper Review] FGBERT: Function-Driven Pre-trained Gene Language Model for Metagenomics

Chenrui Duan, Zelin Zang+7 more|arXiv (Cornell University)|Feb 24, 2024

Machine Learning in Bioinformatics5 citations

TL;DR

FGBERT introduces a function-aware metagenomic pre-trained model with MGM and TEM-CL, achieving state-of-the-art results across eight downstream tasks.

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 FGBERT, a novel metagenomic pre-trained model that employs a protein-based gene representation as a context-aware and structure-relevant tokenizer. FGBERT incorporates Masked Gene Modeling (MGM) to enhance the understanding of inter-gene contextual relationships and Triplet Enhanced Metagenomic Contrastive Learning (TMC) to elucidate gene sequence-function relationships. Pre-trained on over 100 million metagenomic sequences, FGBERT demonstrates superior performance on metagenomic datasets at four levels, spanning gene, functional, bacterial, and environmental levels and ranging from 1k to 213k input sequences. Case studies of ATP Synthase and Gene Operons highlight FGBERT's capability for functional recognition and its biological relevance in metagenomic research.

Motivation & Objective

• Motivate improved representations for multi-species metagenomic data beyond K-mer-based approaches.
• Develop context-aware, protein-based gene tokenization to capture gene context and structure.
• Propose Masked Gene Modeling (MGM) for inter-/intra-gene contextual learning.
• Propose Triple Enhanced Metagenomic Contrastive Learning (TEM-CL) to model gene sequence–function relationships.
• Demonstrate end-to-end improvements on multiple downstream functional and pathogenicity tasks.

Proposed method

• Use a protein-based gene representation as a context-aware tokenizer based on ESM-2 to encode each gene sequence with an orientation-aware 1281D embedding.
• Introduce Masked Gene Modeling (MGM) to pre-train a Gene Encoder by masking 15% of gene tokens and predicting them, combining Feature Reconstruction Loss (FRL) and Probability Prediction Loss (PPL).
• Integrate Triplet Enhanced Metagenomic Contrastive Learning (TEM-CL) with data augmentation and hard negative sampling to cluster genes by function using a SupCon-Hard loss variant.
• Employ a data augmentation pipeline that generates positive gene samples via mutations and computes negatives via cluster-aware sampling based on EC-function groups.
• Jointly optimize L_Total = L_MGM + lambda * L_Tri to learn co-representations of metagenomic gene sequences and functions.

Experimental results

Research questions

• RQ1 How can metagenomic sequences be tokenized in a biologically meaningful way beyond fixed K-mers?
• RQ2 Can context-aware, function-relevant representations improve downstream metagenomic tasks across taxonomy, function, and ecology?
• RQ3 Do MGM and TEM-CL jointly improve modeling of inter-gene contexts (One-to-Many) and gene–function mappings (Many-to-One)?
• RQ4 Is a protein-based tokenizer beneficial for long metagenomic sequences in terms of performance and efficiency?

Key findings

• FGBERT achieves state-of-the-art macro F1 and weighted F1 scores across eight downstream tasks compared to a wide baseline set.
• On operon prediction (E-K12), FGBERT attains 61.8% M.F1 and 65.4% W.F1, outperforming alternatives.
• On CARD-AMR categories (A, D, R), FGBERT achieves 78.6% M.F1 and 90.1% W.F1 (CARD-A), 57.4% M.F1 and 85.2% W.F1 (CARD-D), 69.4% M.F1 and 91.4% W.F1 (CARD-R).
• On VFDB, ENZYME, PATRIC, and NCycDB, FGBERT attains 75.7% M.F1 / 90.2% W.F1 (VFDB), 99.1% M.F1 / 98.8% W.F1 (ENZYME), 99.3% M.F1 / 99.0% W.F1 (PATRIC), 99.5% M.F1 / 99.2% W.F1 (NCycDB).
• Ablation shows MGM and TEM-CL both contribute to performance, with MGM providing a larger impact across tasks.