CoChat
Sign Up
Cited by 120SARS-CoV-2 and COVID-19 ResearchGenomics and Phylogenetic StudiesMachine Learning in Bioinformatics

The DOI matches its Crossref record; 63 of the 63 references it cites were found in a registry or an index.

GenSLMs: Genome-scale language models reveal SARS-CoV-2 evolutionary dynamics

Open access
Zvyagin, Maxim, Brace, Alexander, Hippe, Kyle, Deng, Yuntian, Zhang, Bin, Bohorquez, Cindy Orozco, Clyde, Austin, Kale, Bharat, Perez-Rivera, Danilo, Ma, Heng, Mann, Carla M., Irvin, Michael, Ozgulbas, Defne G., Vassilieva, Natalia, Pauloski, James Gregory, Ward, Logan, Hayot-Sasson, Valerie, Emani, Murali, Foreman, Sam, Xie, Zhen, Lin, Diangen, Shukla, Maulik, Nie, Weili, Romero, Josh, Dallago, Christian, Vahdat, Arash, Xiao, Chaowei, Gibbs, Thomas, Foster, Ian, Davis, James J., Papka, Michael E., Brettin, Thomas, Stevens, Rick, Anandkumar, Anima, Vishwanath, Venkatram, Ramanathan, Arvind·2023·The International Journal of High Performance Computing Applications
DOI 10.1177/10943420231201154PMC PMC9709791

Published version of a preprint

Open this paper in CoChat — free accountOpen PDF

Abstract

We seek to transform how new and emergent variants of pandemic-causing viruses, specifically SARS-CoV-2, are identified and classified. By adapting large language models (LLMs) for genomic data, we build genome-scale language models (GenSLMs) which can learn the evolutionary landscape of SARS-CoV-2 genomes. By pre-training on over 110 million prokaryotic gene sequences and fine-tuning a SARS-CoV-2-specific model on 1.5 million genomes, we show that GenSLMs can accurately and rapidly identify variants of concern. Thus, to our knowledge, GenSLMs represents one of the first whole-genome scale foundation models which can generalize to other prediction tasks. We demonstrate scaling of GenSLMs on GPU-based supercomputers and AI-hardware accelerators utilizing 1.63 Zettaflops in training runs with a sustained performance of 121 PFLOPS in mixed precision and peak of 850 PFLOPS. We present initial scientific insights from examining GenSLMs in tracking evolutionary dynamics of SARS-CoV-2, paving the path to realizing this on large biological data.

Overview

The claims and statistics of this paper haven’t been checked yet.

Run the check — free, no account needed. It reads the open-access text and checks the sources it cites.

A few free checks a day · papers already verified open instantly · results are public

Integrity signals

Registry

Crossref · 10.1177/10943420231201154

Cites

63 of 63 references verified0 citation errors · 0 not checkable

Statistics

Figures

Figure 1: Overview of GenSLM models for predictive modeling of SARS-CoV-2 evolution. The inputs to GenSLM are nucleotide sequences, encoded at the codon level (every three nucleotide represents a codon; hence the 20 natural amino acid language is described by 64 codons). These inputs are successively fed into transformer blocks (referred to as layers (Li)), which ultimately results in learning a semantic embedding z→ space from which one may obtain the probability of any given sequence token p(X→i∣X→[N]\{i}), where N represents the sequence length and 𝑖 represents a particular position in the entire genome.
Figure 1: Overview of GenSLM models for predictive modeling of SARS-CoV-2 evolution. The inputs to GenSLM are nucleotide sequences, encoded at the codon level (every three nucleotide represents a codon; hence the 20 natural amino acid language is described by 64 codons). These inputs are successively fed into transformer blocks (referred to as layers (Li)), which ultimately results in learning a semantic embedding z→ space from which one may obtain the probability of any given sequence token p(X→i∣X→[N]\{i}), where N represents the sequence length and 𝑖 represents a particular position in the entire genome.CC BY-NC-ND · PubMed Central
Figure 2: GenSLMs learned latent space describes biologically meaningful properties for SARS-CoV-2 genomes. (A) The embeddings from GenSLMs are visualized with t-distributed stochastic neighbor embedding (t-SNE) and each gene sequence is represented as a dot in the 2D plot. We paint each sequence by its variant ID – although we have more than 515 PANGO (Rambaut et al., 2020) lineages represented in the data, we only show those with WHO designated labels. (B) The latent space can also be painted with the MAFFT-determined alignment score (Yamada et al., 2016) with respect to an Omicron genome; clustering in the distance measures is clearly visible. Visualizing the sequence log-likelihood (blue bar) and the cross-protein attention (orange lines) from (C) Delta and (D) Omicron SARS-CoV-2 strains highlights how different the co-evolutionary patterns are in these lineages. It is interesting to note that while the Spike protein from Delta strain shows coupling to nsp3, nsp5, and other proteins, these couplings are not observed in the Omicron strain.
Figure 2: GenSLMs learned latent space describes biologically meaningful properties for SARS-CoV-2 genomes. (A) The embeddings from GenSLMs are visualized with t-distributed stochastic neighbor embedding (t-SNE) and each gene sequence is represented as a dot in the 2D plot. We paint each sequence by its variant ID – although we have more than 515 PANGO (Rambaut et al., 2020) lineages represented in the data, we only show those with WHO designated labels. (B) The latent space can also be painted with the MAFFT-determined alignment score (Yamada et al., 2016) with respect to an Omicron genome; clustering in the distance measures is clearly visible. Visualizing the sequence log-likelihood (blue bar) and the cross-protein attention (orange lines) from (C) Delta and (D) Omicron SARS-CoV-2 strains highlights how different the co-evolutionary patterns are in these lineages. It is interesting to note that while the Spike protein from Delta strain shows coupling to nsp3, nsp5, and other proteins, these couplings are not observed in the Omicron strain.CC BY-NC-ND · PubMed Central
Figure 3: Illustration of diffusion-based hierarchical modeling. To predict a codon (such as TAA), we use both the previous codons within the context window (we use size 3 shown in green for illustration) and the high-level representations z.
Figure 3: Illustration of diffusion-based hierarchical modeling. To predict a codon (such as TAA), we use both the previous codons within the context window (we use size 3 shown in green for illustration) and the high-level representations z.CC BY-NC-ND · PubMed Central
Figure 4: Diffusion-based hierarchical modeling of SARS-CoV-2 genomes results in generation of sequences that captures the correct context of various open reading frames (ORFs). (A) Comparison of statistics measured on generated sequences and on real data for the ORFs. Diffusion-based hierarchical LM has a global high-level plan whereas the baseline can only take into account the previous 1023 codons. (B) Generated sequences (light blue) from the model overlaid on the phylogenetic tree demonstrate that these sequences are similar to observed strains.
Figure 4: Diffusion-based hierarchical modeling of SARS-CoV-2 genomes results in generation of sequences that captures the correct context of various open reading frames (ORFs). (A) Comparison of statistics measured on generated sequences and on real data for the ORFs. Diffusion-based hierarchical LM has a global high-level plan whereas the baseline can only take into account the previous 1023 codons. (B) Generated sequences (light blue) from the model overlaid on the phylogenetic tree demonstrate that these sequences are similar to observed strains.CC BY-NC-ND · PubMed Central
Figure 5: Conceptual overview of our workflow. A “Thinker” orchestrates data flow between two applications, namely the sequence generator and the Bayesian optimization to drive the generated sequences towards a target property using reward-guided beam search, where μ represents the mixing constant used to balance the reward function against the log likelihood of generating the next token.
Figure 5: Conceptual overview of our workflow. A “Thinker” orchestrates data flow between two applications, namely the sequence generator and the Bayesian optimization to drive the generated sequences towards a target property using reward-guided beam search, where μ represents the mixing constant used to balance the reward function against the log likelihood of generating the next token.CC BY-NC-ND · PubMed Central
Figure 6: (A) Scaling results on Polaris and Selene systems for MSL=2048; (B) Scaling behavior of DDP vs. DeepSpeed runs on Selene (C) Scaling results on Polaris and Selene systems for MSL=10240;
Figure 6: (A) Scaling results on Polaris and Selene systems for MSL=2048; (B) Scaling behavior of DDP vs. DeepSpeed runs on Selene (C) Scaling results on Polaris and Selene systems for MSL=10240;CC BY-NC-ND · PubMed Central
Figure 7: Workflow utilization measured by the number of active workers (applications actively serving requests) as a function of workflow runtime measured on 224 nodes of Polaris (896 A100 GPUs). The warm-able application design realizes 97% utilization, enabling 1.9X more sequences to be generated compared to a cold start baseline.
Figure 7: Workflow utilization measured by the number of active workers (applications actively serving requests) as a function of workflow runtime measured on 224 nodes of Polaris (896 A100 GPUs). The warm-able application design realizes 97% utilization, enabling 1.9X more sequences to be generated compared to a cold start baseline.CC BY-NC-ND · PubMed Central

Verification

Claims the paper makes

The claims and statistics of this paper haven’t been checked yet.

Run the check — free, no account needed. It reads the open-access text and checks the sources it cites.

A few free checks a day · papers already verified open instantly · results are public

Claims other papers make about this one

Each is checked against this paper’s own text: does it say what the citing paper says it does?

Self-check

Percentages, sums and reported differences re-computed from the paper’s own numbers.

The self-check needs the paper’s full text.

Citations

Cites

Show all 57 references

Research works printed but not linked: 14

  1. Bradley T, 2012, Gpu Performance Analysis and Optimisation
  2. Deepspeed, 2022, Flops Profiler - Deepspeed
  3. Hall S, 2021, Training Giant Neural Networks Using Weight Streaming on Cerebras Wafer-Scale Systems
  4. Hugging face, 2022, Transformers: State-Of-The-Art Machine Learning For Pytorch, Tensorflow, And Jax
  5. Kenton JDMWC, 2019, Proceedings of NAACL-HLT, 4171
  6. Moritz P, Nishihara R, Wang S, et al. (2018) Ray: a distributed framework for emerging AI applications. In 13th USENIX Symposium on Operating Systems Design and Implementation (OSDI 18). Carlsbad, CA, 2018, pp. 561–577. https://www.usenix.org/conference/osdi18/presentation/moritz
  7. Nvidia, 2022, Automated Mixed Precision
  8. NVIDIA, 2022, Nvidia Nsight Systems
  9. Page AJ, 2016, Microbial Genomics, 5, 436
  10. Pytorch, 2022, Pytorch Lightning
  11. Radford A, 2018, Improving Language Understanding by Generative Pre-training
  12. Top500, 2022, June 2022 – Top500
  13. 10.1101/2022.08.08.22278553 · DOI
  14. Wang RE, Durmus E, Goodman N, et al. (2022) Language modeling via stochastic processes. In International Conference on Learning Representations. 2022

Cited by

Sorted by citations

See all 120

Related

Metadata

Authors
Zvyagin, M., Brace, A., Hippe, K., Deng, Y., Zhang, B., Bohorquez, C. O., Clyde, A., Kale, B., Perez-Rivera, D., Ma, H., Mann, C. M., Irvin, M., Ozgulbas, D. G., Vassilieva, N., Pauloski, J. G., Ward, L., Hayot-Sasson, V., Emani, M., Foreman, S., Xie, Z., Lin, D., Shukla, M., Nie, W., Romero, J., Dallago, C., Vahdat, A., Xiao, C., Gibbs, T., Foster, I., Davis, J. J., Papka, M. E., Brettin, T., Stevens, R., Anandkumar, A., Vishwanath, V., & Ramanathan, A.
Affiliations
  • Argonne National Laboratory (US)
  • Harvard University (US)
  • Cerebras Systems (United States) (US)
  • Northern Illinois University (US)
  • Department of Biochemistry, University of Illinois-Urbana Champaign, Champaign, IL, USA
  • University of Chicago (US)
  • Data Science and Learning Division, Argonne National Laboratory, Lemont, IL, USA
  • Nvidia (United States) (US)
  • NVIDIA Inc., Santa Clara, CA, USA
  • University of Illinois Urbana-Champaign (US)
Corresponding author
  • Vishwanath, Venkatram
  • Ramanathan, Arvind
ORCID
Venue
The International Journal of High Performance Computing Applications
Publisher
SAGE Publications
Published
2023-10-27
Volume
37
Issue
6
Pages
683-705
ISSN
  • 0890-2720 (ISSN-L)
  • 1094-3420
  • 1741-2846
DOI
10.1177/10943420231201154
PMC
PMC9709791
OpenAlex
W4387994442
Cited by
120
References
57
Citations by year
2026: 30 · 2025: 40 · 2024: 35 · 2023: 14
FWCI
24.63
Type
journal article
Language
en
Open access
Licence
Locations
  • The International Journal of High Performance Computing Applications (publishedVersion)
  • bioRxiv (Cold Spring Harbor Laboratory) (acceptedVersion, cc-by-nc-nd)
  • PubMed (publishedVersion)
  • OSTI OAI (U.S. Department of Energy Office of Scientific and Technical Information) (submittedVersion)
  • PubMed Central (submittedVersion, cc-by-nc-nd)
  • Espace ÉTS (ETS) (submittedVersion)
Funders
Exascale Computing Project
Topics
  • Health Sciences › Medicine › Infectious Diseases › SARS-CoV-2 and COVID-19 Research
  • Life Sciences › Biochemistry, Genetics and Molecular Biology › Molecular Biology › Genomics and Phylogenetic Studies
  • Life Sciences › Biochemistry, Genetics and Molecular Biology › Molecular Biology › Machine Learning in Bioinformatics
Keywords
SARS-CoV-2, COVID-19 variants of concern, variant identification, viral evolution, foundation models, mixed-precision training

About these checks

Registry, retraction, Cites and statistics results are computed, not judged, with one exception: a printed reference that no index lists is read once by a model before it counts as a citation error. Claims are extracted by a model and kept only when their quote appears verbatim in the text; each cited source is then checked, with a model’s help, for whether it says what the claim relies on. Those verdicts can be wrong — read the claim’s quote and the cited source before relying on one.

Statistics and arithmetic are recomputed, not judged.

Confidence bands are shown only for fields calibrated on our annotated test set. None is calibrated yet, so no band is shown.

What counts as a claim ›How we measure confidence ›How we check this ›

Are you an author, or do you see a problem with this page? Request a review