[論文レビュー] Are Protein Language Models Compute Optimal?

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

Machine Learning in Bioinformatics被引用数 4

ひとこと要約

本研究では、NLPスケーリング法則をpLMに適応させることで、計算量最適なタンパク質言語モデル（pLM）を調査し、モデルサイズやデータスケールに関係なく、固定された損失で性能が飽和することを明らかにした。単一パス学習を用いることで、3500万パラメータのモデルが11.43のパープレキシティを達成した。これは、150億パラメータのESM-2 や1000億パラメータのxTrimoPGLM と同等の性能であり、最適に学習された小さなモデルが、計算効率が優れていることを示している。

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.

研究の動機と目的

• 広く使われているタンパク質言語モデル（pLM）が、固定されたFLOP予算下で計算量最適であるかどうかを調査すること。
• 固定された計算予算内での訓練損失を最小化するための、モデルパラメータとトレーニングトークン数の最適な比率を特定すること。
• ESM-2 や xTrimoPGLM などの現在のpLMが、計算量最適なスケーリング法則に従っているのか、あるいは収益逓減を示しているのかを評価すること。
• パープレキシティを主な指標として用い、1パスで学習された小さなモデル（3500万パラメータ）の性能を、より大きなモデルと比較すること。
• 異なる計算レベルにおける損失を最小化するための最適なモデルサイズとデータセットサイズのスケーリング法則を導出すること。

提案手法

• Hoffmannら（2022）およびKaplanら（2020）のNLPスケーリング法則を適応し、固定FLOP予算下でのモデルサイズ（N）、データセットサイズ（D）、および訓練損失の関係をモデル化した。
• UniRef50のランダムに抽出されたサブセットを用い、500万から6億5000万パラメータのpLMを、65億〜208億トークンの範囲で変動させたデータで学習した。
• FLOPsと最適なモデルサイズおよびトークン数の対応関係を正確に特定できるよう、訓練損失曲線のスプライン補間を用いた。
• 結合スケーリング法則：$ L(N,D) = ig[ (N_c/N)^{\alpha_N/\alpha_D} + (D_c/D) \big]^{\alpha_D} $ をフィッティングし、損失をモデルサイズとデータサイズの関数としてモデル化した。
• 損失の飽和領域を多項式フィッティングにより特定し、固定計算予算（例：$10^{17}$ FLOPs）における最適モデルサイズを同定した。$10^{17}$ FLOPsでは約506億パラメータ、$10^{18}$ FLOPsでは1570億パラメータが最適と予測された。
• ホールドアウトデータセット上でパープレキシティを用いてモデル性能を評価し、1パスで学習された3500万パラメータのモデルを、ESM-2（150億）やxTrimoPGLM（1000億）といった大規模モデルと比較した。

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

実験結果

リサーチクエスチョン

• RQ1 固定計算予算下でのタンパク質言語モデルにおける、モデルパラメータとトレーニングトークン数の最適な比率は何か？
• RQ2 ESM-2 や xTrimoPGLM といった広く使われているpLMは、計算量最適であるのか、それとも収益逓減を示しているのか？
• RQ3 最適な計算割り当てと早期停止を用いた場合、小さなpLMが大規模モデルと同等の性能を達成できるか？
• RQ4 pLMの事前学習において、計算予算が増加するに従い、最適なモデルサイズはどのように変化するか？
• RQ5 同じデータを複数回通す複数パス学習は、1パス学習で観察された損失の飽和を超えて、性能を著しく向上させるのか？

主な発見

• UniRef50の2000万シーケンスのサブセットを1パスで学習した3500万パラメータのモデルが、パープレキシティ11.43を達成した。これは、150億パラメータのESM-2 や1000億パラメータのxTrimoPGLM よりも計算効率に優れていることを示している。
• モデルサイズに関係なく、一貫した損失の飽和が観察された。これは、モデルサイズやデータスケールにかかわらず、性能向上が限界に達することを示しており、現在のpLM学習における根本的な限界があることを示唆している。
• 計算予算が$10^{17}$ FLOPsの場合、最適なモデルサイズは約506億パラメータであり、$10^{18}$ FLOPsでは1570億パラメータに増加する。計算量に伴い最適モデルサイズが非線形的に増加するが、これは計算量に伴うサブ線形的スケーリングである。
• 大規模モデルは、より高いサンプル効率を示すため、より少ないFLOPsで損失の飽和に到達する。これは、パラメータ数が多くても、計算効率が向上することを示しており、容量の増加が計算効率の向上に寄与している。
• 最適モデルサイズは計算予算に伴いサブ線形的に増加し、最適トークン数についても同様にサブ線形的に増加する。これは、モデルサイズが増加するに従い、性能向上の収益が逓減することを示している。
• 結果から、ESM-2 や xTrimoPGLM といった現在のpLMは計算量最適ではないと示唆される。これらのモデルは、より小さな最適化済みモデルが達成する性能の飽和に到達するのに、はるかに多くのFLOPsを要している。

Figure 2: FLOPs vs loss without learning rate decay.