[論文レビュー] Training Compute-Optimal Protein Language Models

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

Machine Learning in Bioinformatics被引用数 4

ひとこと要約

本稿は、タンパク質配列データに特化したスケーリング法則を確立することで、計算量最適化されたタンパク質言語モデルのトレーニング戦略を提案する。1940億トークンのデータセットを用い、300を超えるモデルをトレーニングすることで、固定計算量予算下でのデータ量およびモデルサイズの非線形的スケーリングが優れた性能をもたらすことを示した。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における過学習と性能の飽和を是正すること。
• Causal Language Model (CLM)およびMasked Language Model (MLM)の両目的に対して、タンパク質配列データに特化したスケーリング法則を導出すること。
• 転移学習を用いたCLM事前学習の有効性を、下流のMLMファインチューニングにおいて検証すること。
• 導出されたスケーリング法を用いて、既存のモデル設定（例：ESM-2、PROGEN2）を再最適化し、同じ計算量予算内で性能を向上させること。

提案手法

• Unirefとメタゲノム配列を組み合わせることで、9億3900万配列、1940億トークンのデータセット（UniMeta200B）を構築し、多様性を向上させ、過学習を低減した。
• 50億〜2000億個の固有トークンを用い、3.5M〜10.7Bパラメータのモデルを300個以上、制御された計算量予算下でトレーニングし、スケーリング特性を調査した。
• モデルサイズおよびトレーニングトークン数の両方において、モデル損失と性能にパワー則スケーリング関数をフィットさせ、CLM（0.57）およびMLM（0.77）のスケーリング指数を推定した。
• CLMからMLMへの転送効率を定量化するための「効果的転送トークン数（Dₜ）」の概念を導入し、両目的間での計算量の最適割り当てを可能にした。
• 導出されたスケーリング法を用いて、ESM-2（3B）およびPROGEN2-xlargeのモデルサイズとトレーニングトークンを再割り当てし、7.2Bおよび10.7Bパラメータの新しいモデルをトレーニングした。
• A100 GPUを用い、データ並列化とともにFlashAttention、bfloat16、最適化された学習率スケジュール（コサイン減衰、2.5％ウォームアップ）を適用し、トレーニングを高速化した。

実験結果

リサーチクエスチョン

• RQ1 タンパク質言語モデリングにおけるCLMおよびMLMのスケーリング法はどのように異なり、それぞれのパワー則指数は何か？
• RQ2 Unirefを越えるトレーニングデータの多様性の向上により、CLMおよびMLMにおける過学習と性能の飽和は緩和可能か？
• RQ3 固定計算量予算下で、CLM事前学習がMLM性能の向上に効果的に転送できる範囲はどの程度か？
• RQ4 CLMおよびMLMの両方の事前学習が必要な場合、計算量を両目的間で最適に割り当てる方法は何か？
• RQ5 スケーリング法に基づいたモデルサイズおよびトレーニングトークンの再割り当てにより、ESM-2 や PROGEN2 などの既存モデルと比較して、下流タスクの性能が向上するか？

主な発見

• MLMのスケーリング法はCLM（0.57）よりも高い指数（0.77）を示しており、モデルサイズの増大に伴い、MLMが計算量の増加に対してより大きな恩恵を受けることが示された。
• 計算量が10倍に増加すると、MLMではモデルサイズが6倍に、データ量が70％増加する一方、CLMではモデルサイズが4倍に、データ量が3倍に増加した。
• CLM事前学習はMLMへの効果的な転送を可能にし、効果的転送トークン数（Dₜ）によってその転送効率が定量化可能であり、両目的間での計算量の最適割り当てが可能になった。
• ESM-2（3B）およびPROGEN2-xlargeの計算量予算下で再割り当てを行った結果、それぞれ7.2Bおよび10.7Bパラメータのモデルが得られ、元のモデルを下流タスクで上回った。
• 新しい戦略に従ってトレーニングされた10.7Bパラメータのモデルは、同じ計算量予算内で、タンパク質生成および構造／機能予測タスクで最先端の性能を達成した。
• 本研究は、メタゲノム配列によるデータ多様性が、CLMおよびMLM両目的における過学習および性能の飽和を回避するために極めて重要であることを確認した。