This Large Language Model Predicts COVID Variants

A group of academic and commercial researchers from institutions like Argonne National Laboratory, NVIDIA, and the University of Chicago among others, trained an LLM to predict new and emergent variants of pandemic-causing viruses, especially SARS-CoV-2, the one behind COVID-19. 

The research team, also the finalist for the Gordon Bell special prize for high performance computing-based Covid-19 research, was able to achieve this milestone by using Genome-scale Language Models (GenSLMs). Unlike the protein language models (PLMs) which are trained on nucleic acids (DNA/RNA) or protein, GenSLMs are trained on genome-scale data to identify mutations at nucleotide scale. 

The team was able to demonstrate the scaling of GenSLMs by developing AI model on GPU-based supercomputers, like the NVIDIA A100 Tensor Core GPU, and using AI-hardware accelerators, achieving a performance of over 1.54 zettaflops in training runs, making way for the largest biological language models created to date.

The current version of the model, trained on 2.5 billion parameters, took over a month to complete with around 4000 GPUs required for the same. In all, the team spent four months working on the project before releasing the paper and the code to the public.   

The code is available on GitHub and can be accessed here. 

The shift from protein-level to gene-level data is considered to be an important breakthrough in AI-based biological research as it opens up applications related to protein annotation workflows, metagenome reconstruction, protein engineering, and biological pathway design. 

LLMs are generally trained on human languages, where a dozen letters can be arranged in a permutation of lakhs of words, however nucleotides consist of four letters—A, T, G and C in DNA, or A, U, G AND C in RNA—arranged in different sequences as genes. To solve the challenge of breaking down the large number of nucleotides (3 billion in humans, 30,000 in coronaviruses) making up a genome into meaningful units, NVIDIA researchers came up with a hierarchical diffusion method that will enable LLMs to treat long strings of nucleotides as sentences. 

The new model, trained on nucleotide sequences, consist of much larger vocabulary compared to the PLMs, and is able to pick finer details, representing a much larger repertoire of biological properties, and a generation of output that has reasonable fidelity to the intrinsic organisation of the SARS-CoV-2 sequences.

Anima Anandkumar, senior director of AI research at NVIDIA, said, “We developed a diffusion model that operates at a higher level of detail that allows us to generate realistic variants and capture better statistics.”   

The model was trained on more than 110 million gene sequences using open-source data collected from the Bacterial and Viral Bioinformatics Resource Center. Later, it was fine-tuned with 1.5 million high-quality genome sequences for the COVID virus. Since the model was trained on a large dataset, it was able to generalise to other sequences—distinguish between them, and predict potential mutations of the COVID genome that could further the research on different variants of the virus. 

Further, the GenSLM model could also be integrated with protein structure models like AlphaFold and OpenFold, leading researchers to simulate viral structure and understand the role of genetic mutations in a virus' ability to impact its host.