BioNeMo Framework: a modular, high-performance library for AI model development in drug discovery

编码生物学和化学的人工智能模型正在为高通量和高质量的计算机辅助药物开发开辟新的途径。然而，这些模型的训练越来越依赖于计算规模，最近的蛋白质语言模型（pLM）需要在数百个图形处理单元（GPU）上进行训练。我们引入了BioNeMo框架，以促进在数百个GPU上训练计算生物学和化学的人工智能模型。其模块化设计允许将各个组件（如数据加载器）集成到现有的工作流程中，并对社区贡献开放。我们通过一些用例详细介绍了BioNeMo框架的技术特性，例如pLM的预训练和微调。在256个NVIDIA A100 GPU上，BioNeMo框架可以在4.2天内训练一个基于BERT的、拥有三十亿参数的pLM，该模型可以处理超过一万亿个标记。BioNeMo框架是开源的，供所有人免费使用。

该论文试图解决在生物学和化学领域训练大规模AI模型时面临的计算资源需求问题，尤其是蛋白质语言模型(pLM)的训练，这通常需要大量的图形处理单元(GPU)。这是一个在高通量和高质量的虚拟药物开发中日益突出的问题。

论文的关键思路是引入BioNeMo框架，一个模块化设计的工具，旨在支持跨数百个GPU的生物和化学AI模型的高效训练。与现有的方法相比，BioNeMo框架不仅提高了训练效率，还通过其开放性和模块化设计促进了社区贡献和现有工作流程的集成。

论文展示了BioNeMo框架在实际应用中的性能，例如在256个NVIDIA A100 GPU上，能够在4.2天内完成对一个包含30亿参数的BERT基础pLM的训练，处理超过一万亿个令牌。此外，BioNeMo框架是开源且免费的，为研究人员提供了强大的工具。论文还详细描述了框架的技术特点，并通过具体的用例说明了其灵活性和可扩展性。

近期在这个领域的一些相关研究包括：1. 'ProteinMPNN: A Fast and Accurate Conditional Protein Design Model'，该研究提出了一种用于条件蛋白质设计的快速准确模型。2. 'ESM-2: A Pre-trained Protein Sequence Model'，该研究介绍了ESM-2，一种预训练的蛋白质序列模型。3. 'AlphaFold-Multimer: Improved Protein Complex Prediction Using Evolutionary and Geometric Constraints'，该研究通过进化和几何约束改进了蛋白质复合体的预测。