[论文解读] Are Protein Language Models Compute Optimal?

Yaiza Serrano, Álvaro Ciudad+1 more|arXiv (Cornell University)|Jun 11, 2024

Machine Learning in Bioinformatics被引用 4

一句话总结

本研究通过将自然语言处理中的缩放定律适配至蛋白质语言模型（pLMs），探究了计算最优的pLMs，发现无论模型大小或数据规模如何，性能均会在固定损失值处达到平台期。采用单次训练，一个3500万参数的模型实现了11.43的困惑度，与更大的模型如ESM-2（150亿参数）和xTrimoPGLM（1000亿参数）相当，表明当训练最优时，较小的模型在计算效率上更具优势。

ABSTRACT

While protein language models (pLMs) have transformed biological research, the scaling laws governing their improvement remain underexplored. By adapting methodologies from NLP scaling laws, we investigated the optimal ratio between model parameters and training tokens within a fixed compute budget. Our study reveals that pLM sizes scale sublinearly with compute budget, showing diminishing returns in performance as model size increases, and we identify a performance plateau in training loss comparable to the one found in relevant works in the field. Our findings suggest that widely-used pLMs might not be compute-optimal, indicating that larger models could achieve convergence more efficiently. Training a 35M model on a reduced token set, we attained perplexity results comparable to larger models like ESM-2 (15B) and xTrimoPGLM (100B) with a single dataset pass. This work paves the way towards more compute-efficient pLMs, democratizing their training and practical application in computational biology.

研究动机与目标

• 探究在固定FLOP预算下，广泛使用的蛋白质语言模型（pLMs）是否为计算最优。
• 确定在固定计算预算下，最小化训练损失时，模型参数与训练词元数量之间的最优比例。
• 评估当前的pLMs（如ESM-2和xTrimoPGLM）是否遵循计算最优的缩放定律，或是否表现出收益递减现象。
• 使用困惑度作为关键指标，评估单次训练的微型模型（3500万参数）与更大模型的性能表现。
• 推导在不同计算水平下，最小化损失的最优模型尺寸与数据集大小的缩放定律。

提出的方法

• 将Hoffmann等人（2022）和Kaplan等人（2020）提出的NLP缩放定律适配至pLMs，用于建模在固定FLOP预算下，模型大小（N）、数据集大小（D）与训练损失之间的关系。
• 在UniRef50的随机采样子集上训练一系列pLMs（500万至6.5亿参数），通过调整词元数量（从65亿到208亿）来模拟不同的计算分配。
• 使用训练损失曲线的样条插值方法，将FLOPs映射到每个计算水平下的最优模型大小与词元数量，从而精确识别出最高效的配置。
• 拟合联合缩放定律：$ L(N,D) = ig[ (N_c/N)^{\alpha_N/\alpha_D} + (D_c/D) \big]^{\alpha_D} $，以建模损失随模型与数据规模的变化。
• 通过损失平台期的多项式拟合，识别出固定计算预算下的最优模型大小（例如，$10^{17}$ FLOPs时为506亿参数，$10^{18}$ FLOPs时为1570亿参数），预测了在不同计算水平下的最优规模。
• 使用困惑度评估模型在保留集上的性能，将单次训练的3500万参数模型与更大的模型（如ESM-2，150亿参数；xTrimoPGLM，1000亿参数）进行比较。

Figure 1: FLOPs vs loss with token-based learning rate decay.

实验结果

研究问题

• RQ1 在固定计算预算下，蛋白质语言模型的模型参数与训练词元数量之间应保持何种最优比例？
• RQ2 广泛使用的pLMs（如ESM-2和xTrimoPGLM）是否为计算最优？还是表现出收益递减现象？
• RQ3 当采用最优计算分配与早停策略时，小型pLM能否实现与大型模型相当的性能？
• RQ4 在pLM预训练中，随着计算预算的增加，最优模型规模如何变化？
• RQ5 在相同数据上进行多轮训练是否能显著提升性能，超过单次训练中观察到的损失平台期？

主要发现

• 在UniRef50的2000万条序列子集上单次训练的3500万参数模型，实现了11.43的困惑度，其计算效率优于更大的模型如ESM-2（150亿参数）和xTrimoPGLM（1000亿参数）。
• 研究发现，无论模型大小或数据规模如何，损失平台期均保持一致，表明性能增益在达到平台期后不再随模型或数据规模增加而提升，暗示当前pLM训练存在根本性限制。
• 在10^{17} FLOPs的计算预算下，最优模型大小约为506亿参数，而在10^{18} FLOPs下增至1570亿参数，表明最优模型大小随计算预算呈次线性增长。
• 更大的模型因样本效率更高，可在更少FLOPs下达到损失平台期，表明尽管参数量更高，但模型容量的提升反而提高了计算效率。
• 最优模型大小随计算预算呈次线性增长，最优词元数量也呈次线性增长，表明随着模型规模扩大，性能增益的边际效益递减。
• 结果表明，当前的pLMs（如ESM-2和xTrimoPGLM）并非计算最优，因为它们需要显著更多的FLOPs才能达到与更小、更优优化的模型相当的性能平台期。

Figure 2: FLOPs vs loss without learning rate decay.