ProkBERT Family: Genomic Language Models for Microbiome Applications

Abstract

Machine learning offers transformative capabilities in microbiology and microbiome analysis, deciphering intricate microbial interactions, predicting functionalities, and unveiling novel patterns in vast datasets. This enriches our comprehension of microbial ecosystems and their influence on health and disease. However, the integration of machine learning in these fields contends with issues like the scarcity of labeled datasets, the immense volume and complexity of microbial data, and the subtle interactions within microbial communities. Addressing these challenges, we introduce the ProkBERT model family. Built on transfer learning and self-supervised methodologies, ProkBERT models capitalize on the abundant available data, demonstrating adaptability across diverse scenarios. The models learned representations align with established biological understanding, shedding light on phylogenetic relationships. With the novel Local Context-Aware (LCA) tokenization, the ProkBERT family overcomes the context size limitations of traditional transformer models without sacrificing performance or the information rich local context. In bioinformatics tasks like promoter prediction and phage identification, ProkBERT models excel. For promoter predictions, the best performing model achieved an MCC of 0.74 for E. coli and 0.62 in mixed-species contexts. In phage identification, they all consistently outperformed tools like VirSorter2 and DeepVirFinder, registering an MCC of 0.85. Compact yet powerful, the ProkBERT models are efficient, generalizable, and swift. They cater to both supervised and unsupervised tasks, providing an accessible tool for the community. The models are available on GitHub and HuggingFace.