Intel Gaudi 2上でのタンパク質言語モデルProtSTの高速化

ProtSTモデル概要

ProtSTは、テキストプロンプトに基づくタンパク質設計のためのマルチモーダル言語モデルで、MILAとIntel LabsがICML 2023で最初に公開しました。このモデルは、アミノ酸配列の細胞内局在予測などのタスク向けに設計されており、配列が核、細胞膜、細胞質、またはミトコンドリアのどれに属するかを識別します。

  • Zero-shot Performance: ProtST-ESM-1b バリアントは、細胞内局在予測において、最先端の few-shot クラスifier を上回る性能を即座に示します。
  • Accessibility: このモデルと関連データセットは再構築され、mila-intel 組織を通じて Hugging Face Hub で共有されています。

Intel Gaudi 2上での推論性能

ProtST-SubcellularLocalization データセット（長さが79から1,999のアミノ酸配列2,772件）で実施した推論テストにより、Intel Gaudi 2 が NVIDIA A100 80GB PCIe に比べて大幅な性能向上をもたらすことが示されました。

モデルをフル bfloat16 精度でバッチサイズ1で実行した場合、結果は以下の通りです。

  • Accuracy: NVIDIA A100 と Intel Gaudi 2 はどちらも 0.44 の同一精度を達成しました。
  • Speed: Intel Gaudi 2 は A100 に比べて 1.76 倍速い推論速度を実現しました。

ファインチューニングとスケーリング機能

下流タスクで ProtST をファインチューニングすることでモデリング精度が向上します。二値ロケーションタスク（タンパク質が膜結合か可溶かを判定）では、ProtST-ESM1b-for-sequential-classification モデルを bfloat16 精度で ProtST-BinaryLocalization データセットを用いてファインチューニングしました。

訓練精度と速度

ファインチューニング結果は元の研究論文の結果とほぼ一致し、さまざまなハードウェア構成で約92.5% の精度を達成しました。

ハードウェア比較

単一の Gaudi 2 アクセラレータと単一の NVIDIA A100 を比較すると：

  • Fine-tuning Speed: 単一の Gaudi 2 は単一の A100 より 2.92 倍速いです。
  • Scalability: Gaudi 2 上の分散トレーニングは、アクセラレータ数を 4 台または 8 台に増やすと、ほぼ線形にスケールします。

実装リソース

開発者は以下のツールを使用して Intel Gaudi 2 上で ProtST をデプロイできます：

  • Optimum for Intel Gaudi: transformers ベースのスクリプトを Gaudi 2 ハードウェアへ移植する作業を簡素化するオープンソースライブラリです。
  • Intel Gaudi Documentation: ハードウェア実装のための公式技術ガイドです。

Sources