解析基因组的“语言”：戈登贝尔奖决赛选手使用大型语言模型来预测新冠病毒变异株

来自美国阿贡国家实验室、NVIDIA、芝加哥大学等组织机构的研究员开发了一个处理基因组规模数据的先进模型,并入围戈登贝尔COVID-19研究特别奖决赛。

这一戈登贝尔特别奖旨在表彰基于高性能计算的COVID-19研究。一位决赛入围选手教会了大型语言模型(LLMs)一种新的语言——基因序列,使这些模型能够提供基因组学、流行病学和蛋白质工程方面的洞察。

这项开创性的成果发表于10月,是由来自美国阿贡国家实验室、NVIDIA、芝加哥大学等组织机构的二十多名学术和商业研究员合作完成。

该研究团队训练了一个LLM来追踪基因突变,并预测需要关注的SARS-CoV-2(导致COVID-19的病毒)变异株。虽然迄今为止大多数应用于生物学的LLM都是在小分子或蛋白质的数据集上训练的,但这一项目是在原始核苷酸序列(DNA和RNA的最小单位)上训练的首批模型之一。

负责带领该项目的阿贡国家实验室计算生物学家Arvind Ramanathan表示:“我们假设从蛋白质水平到基因水平的数据有助于我们构建出更易于理解新冠病毒变异株的模型。通过训练模型去追踪整个基因组及其进化过程中的所有变化,我们不仅能够更好地预测COVID,还能预测已掌握足够基因组数据的任何疾病。”

戈登贝尔奖被誉为HPC领域的诺贝尔奖。今年的戈登贝尔奖将在本周的SC22上由美国计算机协会颁发。该协会代表着全球约10万名计算领域的专家,自2020年开始向使用HPC推进COVID-19研究的杰出研究员颁发特别奖。

在一种只有四个字母的语言上训练大型语言模型

长期以来,LLM一直在接受人类语言的训练,这些语言通常由几十个字母组成,可以排列组合成数万个单词,并连接成长句和段落。而生物学语言只有四个代表核苷酸的字母,即DNA中的A、T、G和C,或RNA中的A、U、G和C。这些字母按不同顺序排列成基因。

虽然较少的字母看似会降低AI学习的难度,但实际上生物学语言模型要复杂得多。这是因为人类的基因组由超过30亿个核苷酸组成,而冠状病毒的基因组由大约3万个核苷酸组成,因此很难将基因组分解成不同、有意义的单位。

Ramanathan表示: “在理解基因组这一 ‘生命代码’的过程中,我们所面对的一个主要挑战是基因组中的庞大测序信息。核苷酸序列的意义可能会受另一序列的影响,以人类的文本做类比,这种影响的范围不仅仅是文本中的下一句话或下一段话,而是相当于一本书中的整个章节。”

参与该项目协作的NVIDIA研究员设计了一种分层扩散方法,使LLM能够将约1500个核苷酸的长字符串当作句子来处理。

论文共同作者、NVIDIA AI研究高级总监、加州理工学院计算+数学科学系布伦讲席教授Anima Anandkumar表示:“标准语言模型难以生成连贯的长序列,也难以学习不同变异株的基本分布。我们开发了一个在更高细节水平上运作的扩散模型,该模型使我们能够生成现实中的变异株,并采集到更完善的统计数据。”

预测需要关注的新冠病毒变异株

该团队首先使用细菌和病毒生物信息学资源中心的开源数据,对来自原核生物(像细菌一样的单细胞生物)超过1.1亿个基因序列进行了LLM预训练,然后使用150万个高质量的新冠病毒基因组序列,对该模型进行微调。

研究员还通过在更广泛的数据集上进行预训练,确保其模型能够在未来的项目中推广到其他预测任务,使其成为首批具备此能力的全基因组规模的模型之一。

在对COVID数据进行了微调后,LLM就能够区分病毒变异株的基因组序列。它还能够生成自己的核苷酸序列,预测COVID基因组的潜在突变,这可以帮助科学家预测未来需要关注的变异株。

在长达一年时间内积累的SARS-CoV-2基因组数据的训练下,该模型可以推断出各种病毒株之间的区别。左边的每个点对应一个已测序的SARS-CoV-2病毒株,并按变异株颜色编码。右图放大了该病毒的一个特定毒株,它捕捉到了该毒株特有的病毒蛋白进化耦合关系。图片由美国阿贡国家实验室的Bharat Kale、Max Zvyagin和Michael E. Papka提供。

Ramanathan表示: “大多数研究员一直在追踪新冠病毒突刺蛋白的突变,尤其是与人类细胞结合的域。但病毒基因组中还有其他蛋白质也会经历频繁的突变,所以了解这些蛋白质十分重要。”

论文中提到,该模型还可以与AlphaFold、OpenFold等常见的蛋白质结构预测模型整合,帮助研究员模拟病毒结构,研究基因突变如何影响病毒感染其宿主的能力。OpenFold是NVIDIA BioNeMo LLM服务中包含的预训练语言模型之一。NVIDIA BioNeMo LLM服务面向的是致力于将LLM应用于数字生物学和化学应用的开发者。

利用GPU加速超级计算机大幅加快AI训练速度

该团队在由NVIDIA A100 Tensor Core GPU驱动的超级计算机上开发AI模型,包括阿贡国家实验室的Polaris、美国能源部的Perlmutter以及NVIDIA的Selene系统。通过扩展到这些强大的系统,他们在训练中实现了超过1500 exaflops的性能,创建了迄今为止最大的生物语言模型。

Ramanathan 表示: “我们如今处理的模型有多达250亿个参数,预计这一数量未来还会大幅增加。模型的尺寸、基因序列的长度、以及所需的训练数据量,都意味着我们的确需要搭载数千颗GPU的超级计算机来完成复杂的计算。”

研究员估计,训练一个具有25亿参数的模型版本,需要约4000个GPU耗时一个多月。该团队已经在研究用于生物学的LLM,在公布论文和代码之前,他们在这个项目上已耗时约四个月。GitHub页面上有供其他研究员在Polaris和Perlmutter上运行该模型的说明。

NVIDIA BioNeMo框架可在NVIDIA NGC中心上的GPU优化软件中抢先体验。该框架将帮助研究员在多个GPU上扩展大型生物分子语言模型。作为NVIDIA Clara Discovery药物研发工具集的一部分,该框架将支持化学、蛋白质、DNA和RNA数据格式。

在SC22上莅临NVIDIA展台并观看以下特别演讲回放:

头图是研究员的LLM所测序的新冠病毒株。每个点都按新冠病毒变异株进行了颜色编码。该图由美国阿贡国家实验室的Bharat Kale、Max Zvyagin和Michael E. Papka提供。

关于 NVIDIA

自1993年成立以来,NVIDIA(NASDAQ: NVDA)一直是加速计算领域的先驱。NVIDIA 1999年发明的GPU驱动了PC游戏市场的增长,并重新定义了现代计算机图形,开启了现代AI时代,推动了元宇宙的创建。NVIDIA现在是一家全栈计算公司,其数据中心规模的产品正在重塑整个行业。更多信息,请访问https://nvidianews.nvidia.com/。