[論文レビュー] RITA: a Study on Scaling Up Generative Protein Sequence Models

Daniel Hesslow, Niccoló Zanichelli+3 more|arXiv (Cornell University)|May 11, 2022

Machine Learning in Bioinformatics被引用数 59

ひとこと要約

RITAは最大1.2Bパラメータの自己回帰型タンパク質配列モデルを280M以上のUniRef-100配列で訓練し、スケーリングが次アミノ酸予測、適合性予測、酵素機能予測のパフォーマンスにどのように影響するかを体系的に研究する。

ABSTRACT

In this work we introduce RITA: a suite of autoregressive generative models for protein sequences, with up to 1.2 billion parameters, trained on over 280 million protein sequences belonging to the UniRef-100 database. Such generative models hold the promise of greatly accelerating protein design. We conduct the first systematic study of how capabilities evolve with model size for autoregressive transformers in the protein domain: we evaluate RITA models in next amino acid prediction, zero-shot fitness, and enzyme function prediction, showing benefits from increased scale. We release the RITA models openly, to the benefit of the research community.

研究の動機と目的

• タンパク質設計を加速させるための大規模自己回帰型タンパク質配列モデルの設計を動機づける。
• モデルサイズの増加に伴う性能を分析することにより、タンパク質配列モデリングのスケーリング法則を確立する。
• RITAを下流タスクで評価する：次アミノ酸予測、ゼロショット適合性（ProteinGym）、および酵素機能予測。
• オープンな研究とベンチマーキングを可能にするために、コミュニティへモデルを公開する。

提案手法

• 150Bアミノ酸の事前学習データを用い、SmallからXLargeまでの4サイズのデコーダーのみのトランスフォーマーを最大1.2Bパラメータで訓練する。
• AliBiを用いたアブレーションの後、RoPE回転位置埋め込みを使用；言語モデリング損失を低減するためRoPEを選択。
• 主にUniRef-100で事前学習し、補助的にMetaclustとMGnifyデータを用い、事前クラスタリングは行わない。
• Megatron-Deepspeedを用いて150Bアミノ酸を対象に、100台以上のGPUでAdam最適化手法、バッチサイズ512、コンテキストサイズ1024で訓練。
• UniRef-100、MGnify、Metaclustの保持外タンパク質ファミリーでパープレキシティを評価；ProtGPT2とProtXLNetと比較。
• 下流タスクを評価：i) ProteinGym置換による変異効果予測、ii) SwissProt注釈を用いた酵素機能予測、iii) 目標タンパク質ファミリを生成するためのプロンプト調整。

実験結果

リサーチクエスチョン

• RQ1 自己回帰トランスフォーマーにおけるタンパク質配列モデルの能力は、モデルサイズの増加とともにどのようにスケールするか？
• RQ2 大規模なRITAモデルは、ベースラインと比較して次アミノ酸予測、適合性予測、酵素機能予測を改善するか？
• RQ3 保持外のタンパク質ファミリに対する生成品質を、プロンプト調整でさらに向上させることができるか？
• RQ4 計算資源制約の下でのタンパク質配列モデリングのスケーリング法則の特徴（べき指数など）は何か？

主な発見

• すべてのテストデータセットでモデルサイズとともにパープレキシティが改善される；RITA-XLはUniRef-100、Metaclust、MGnify、およびPfam保持アウトで最高のパープレキシティを達成。
• 変異効果（ProteinGym）予測では、RITAモデルの性能はサイズとともに向上し、LargeおよびXLargeバリアントではESM-1vを超えることがある。
• 酵素機能予測では、より大きいモデルほどトップK精度が向上し、XLargeが最良の結果を示す（トップ1からトップ10の値は本研究に示されている）。
• プロンプト調整は保持アウトのタンパク質ファミリ（PF03272）に対してベースモデルと比較してパープレキシティを大幅に低減し、制御可能な生成を示す。
• 下流タスク全体で、RITAモデルはタスク固有のベースラインに近づくか上回り、NLPのスケーリング法則に類似した計算駆動型のスケーリング挙動を示すが、語彙が小さいことと訓練ダイナミクスの違いにより顕著な逸脱がある。