[论文解读] Training Compute-Optimal Protein Language Models

Xingyi Cheng, Bo Chen+4 more|arXiv (Cornell University)|Nov 4, 2024

Machine Learning in Bioinformatics被引用 4

一句话总结

本文通过为蛋白质序列数据量身定制的缩放法则，提出了在固定计算预算下蛋白质语言模型的计算最优训练策略。基于1940亿 token 的数据集并训练超过300个模型，研究发现：在固定计算预算下，数据量和模型规模的次线性缩放可实现更优性能，且通过估计的有效转移 token 数（Effectively Transferred Tokens）量化了因果语言模型（CLM）预训练向掩码语言模型（MLM）任务的有效迁移，使性能优于 ESM-2 和 PROGEN2 等基线模型，且在相同计算约束下表现更优。

ABSTRACT

We explore optimally training protein language models, an area of significant interest in biological research where guidance on best practices is limited. Most models are trained with extensive compute resources until performance gains plateau, focusing primarily on increasing model sizes rather than optimizing the efficient compute frontier that balances performance and compute budgets. Our investigation is grounded in a massive dataset consisting of 939 million protein sequences. We trained over 300 models ranging from 3.5 million to 10.7 billion parameters on 5 to 200 billion unique tokens, to investigate the relations between model sizes, training token numbers, and objectives. First, we observed the effect of diminishing returns for the Causal Language Model (CLM) and that of overfitting for the Masked Language Model~(MLM) when repeating the commonly used Uniref database. To address this, we included metagenomic protein sequences in the training set to increase the diversity and avoid the plateau or overfitting effects. Second, we obtained the scaling laws of CLM and MLM on Transformer, tailored to the specific characteristics of protein sequence data. Third, we observe a transfer scaling phenomenon from CLM to MLM, further demonstrating the effectiveness of transfer through scaling behaviors based on estimated Effectively Transferred Tokens. Finally, to validate our scaling laws, we compare the large-scale versions of ESM-2 and PROGEN2 on downstream tasks, encompassing evaluations of protein generation as well as structure- and function-related tasks, all within less or equivalent pre-training compute budgets.

研究动机与目标

• 在固定计算预算下，识别蛋白质语言模型训练的最优计算资源分配策略。
• 通过引入宏基因组序列扩展训练数据，缓解 CLM 和 MLM 中的过拟合与性能饱和问题。
• 为因果语言模型（CLM）和掩码语言模型（MLM）目标，推导适用于蛋白质序列数据的特定缩放法则。
• 通过迁移学习验证 CLM 预训练在下游 MLM 微调中的有效性。
• 基于推导出的缩放法则，重新优化现有模型设置（如 ESM-2、PROGEN2），在相同计算预算下提升性能。

提出的方法

• 通过整合 Uniref 与宏基因组序列，构建了包含 1940 亿 token 和 9.39 亿条序列的大型数据集（UniMeta200B），以增强多样性并减少过拟合。
• 在 50 亿至 2000 亿条唯一 token 的训练数据范围内，于 A100 GPU 上训练了 300 多个参数规模从 350 万到 107 亿的模型，以在受控计算预算下研究缩放行为。
• 对模型损失与性能在模型规模和训练 token 数之间的关系拟合幂律缩放函数，估算出 CLM 的缩放指数为 0.57，MLM 为 0.77。
• 提出“有效转移 token 数”（Dₜ）概念，用于量化从 CLM 到 MLM 的迁移效率，从而实现两项目标间计算资源的最优分配。
• 基于推导出的缩放法则，重新分配 ESM-2（3B）和 PROGEN2-xlarge 的模型规模与训练 token 数，训练出 72 亿和 107 亿参数的新模型。
• 在 A100 GPU 上采用 FlashAttention、bfloat16 精度及优化的学习率调度（余弦衰减，2.5% 学习率预热），结合数据并行策略，显著加速训练。

实验结果

研究问题

• RQ1 蛋白质语言建模中，CLM 与 MLM 的缩放法则有何不同？其各自的幂律指数为何？
• RQ2 通过在 Uniref 之外引入宏基因组序列以提升训练数据多样性，能否有效缓解 CLM 和 MLM 中的过拟合与性能饱和？
• RQ3 在固定计算预算下，CLM 预训练在多大程度上可有效迁移至 MLM 任务以提升性能？
• RQ4 当 CLM 和 MLM 预训练均需进行时，计算资源在两者之间的最优分配策略是什么？
• RQ5 基于缩放法则重新分配模型规模与训练 token 数，是否能在性能上超越 ESM-2 和 PROGEN2 等既有模型？

主要发现

• MLM 的缩放法则具有更高的指数（0.77），而 CLM 为 0.57，表明在模型规模增长方面，MLM 更能从计算资源增加中获益。
• 计算量提升 10 倍时，MLM 实现模型规模增长 6 倍、数据量增加 70%；而 CLM 仅实现模型规模增长 4 倍、数据量增加 3 倍。
• CLM 预训练可有效迁移至 MLM，其迁移效率通过“有效转移 token 数”（Dₜ）量化，从而实现两项目标间计算资源的最优配置。
• 在 ESM-2（3B）和 PROGEN2-xlarge 的计算预算下重新分配资源，分别训练出 72 亿和 107 亿参数的新模型，其在下游任务中表现优于原始模型。
• 在新策略下训练的 107 亿参数模型，在相同计算预算下，于蛋白质生成与结构/功能预测任务中达到当前最优性能。
• 本研究证实，通过宏基因组序列提升数据多样性，对避免 CLM 和 MLM 目标中的过拟合与性能饱和至关重要。