AlphaFold 3 を最大70倍高速化する「AlphaFast」:GPUでMSAを高速生成する新パイプライン

論文タイトル

AlphaFast: High-throughput AlphaFold 3 via GPU-accelerated MSA construction

出典

AlphaFast: High-throughput AlphaFold 3 via GPU-accelerated MSA construction
AlphaFold 3 (AF3) enables accurate biomolecular modeling but is limited by slow, CPU-bound multiple sequence alignment (MSA) generation. We introduce AlphaFast,...

要旨

AlphaFold 3 の計算ボトルネックである MSA(多重配列アラインメント)生成を GPU によって高速化するパイプライン「AlphaFast」を提案した論文です。

解説など

MSA(Multiple Sequence Alignment:多重配列アラインメント)生成は構造予測における計算コストの大きなボトルネックです。MSAを作るには、

  • UniRef90
  • MGnify
  • Small BFD
  • UniProt

などの巨大データベースを検索する必要があり、AF3ではこの処理が CPUベースの JackHMMER に依存しています。

本研究では、このボトルネックを解消するために AlphaFast というパイプラインが開発されました。

主なアイデアはシンプルで、CPUベースのMSA検索をGPUに置き換えるというものです。

具体的には、

  • JackHMMER → MMseqs2-GPU に置き換え
  • AlphaFold 3 の折りたたみモデル自体は変更しない

という「ドロップイン型」の設計になっています。

AlphaFastでは

  1. すべての配列をまとめてバッチ化
  2. GPUで一括検索
  3. 検索結果を後処理してMSAを生成

という処理の流れになります。

さらに以下の最適化も行われています。

① バッチ検索

複数配列をまとめて検索することでGPUの並列性を最大限活用します。

② 非同期処理

あるデータベースの検索中に、前の検索結果の後処理を同時に実行します。

③ 2段階パイプライン

JAX初期化とMSA生成のVRAM競合を防ぐため、

  • Phase1:MSA生成
  • Phase2:構造予測

の2段階処理になっています。

これにより、4×H200 GPUでは約71.2高速化になります。

本手法は、完全にローカル環境で高速化できる手法であり、またMSAの生成と構造予測は独立モジュールとして機能するので他の構造予測モデルにも活用することが可能です。

プログラムコード

GitHub - RomeroLab/alphafast: AlphaFast: ultra-high-throughput AlphaFold3 inference with MMSeqs2-GPU
AlphaFast: ultra-high-throughput AlphaFold3 inference with MMSeqs2-GPU - RomeroLab/alphafast