引言
蛋白質語言模型 (Protein Language Models, PLM) 已成為蛋白質結構與功能預測及設計的有力工具。在 2023 年國際機器學習會議 (ICML) 上，MILA 和英特爾實驗室聯合釋出了 ProtST 模型，該模型是個可基於文字提示設計蛋白質的多模態模型。此後，ProtST 在研究界廣受好評，不到一年的時間就積累了 40 多次引用，彰顯了該工作的影響力。
PLM 最常見的任務之一是預測氨基酸序列的亞細胞位置。此時，使用者輸入一個氨基酸序列給模型，模型會輸出一個標籤，以指示該序列所處的亞細胞位置。論文表明，ProtST-ESM-1b 的零樣本亞細胞定位效能優於最先進的少樣本分類器 (如下圖)。
為了使 ProtST 更民主化，英特爾和 MILA 對模型進行了重寫，以使大家可以透過 Hugging Face Hub 來使用模型。大家可於 此處 下載模型及資料集。
本文將展示如何使用英特爾 Gaudi 2 加速卡及 optimum-habana 開源庫高效執行 ProtST 推理和微調。英特爾 Gaudi 2 是英特爾設計的第二代 AI 加速卡。感興趣的讀者可參閱我們 之前的博文，以深入瞭解該加速卡以及如何透過 英特爾開發者雲 使用它。得益於 optimum-habana，僅需少量的程式碼更改，使用者即可將基於 transformers 的程式碼移植至 Gaudi 2。
對 ProtST 進行推理
常見的亞細胞位置包括細胞核、細胞膜、細胞質、線粒體等，你可從 此資料集 中獲取全面詳細的位置介紹。
我們使用 ProtST-SubcellularLocalization 資料集的測試子集來比較 ProtST 在英偉達 A100 80GB PCIe 和 Gaudi 2 兩種加速卡上的推理效能。該測試集包含 2772 個氨基酸序列，序列長度範圍為 79 至 1999。
你可以使用 此指令碼 重現我們的實驗，我們以 bfloat16 精度和 batch size 1 執行模型。在英偉達 A100 和英特爾 Gaudi 2 上，我們獲得了相同的準確率 (0.44)，但 Gaudi 2 的推理速度比 A100 快 1.76 倍。單張 A100 和單張 Gaudi 2 的執行時間如下圖所示。
微調 ProtST
針對下游任務對 ProtST 模型進行微調是提高模型準確性的簡單且公認的方法。在本實驗中，我們專門研究了針對二元定位任務的微調，其是亞細胞定位的簡單版，任務用二元標籤指示蛋白質是膜結合的還是可溶的。
你可使用 此指令碼 重現我們的實驗。其中，我們在 ProtST-BinaryLocalization 資料集上以 bfloat16 精度微調 ProtST-ESM1b-for-sequential-classification。下表展示了不同硬體配置下測試子集的模型準確率，可以發現它們均與論文中釋出的準確率 (~92.5%) 相當。
下圖顯示了微調所用的時間。可以看到，單張 Gaudi 2 比單張 A100 快 2.92 倍。該圖還表明，在 4 張或 8 張 Gaudi 2 加速卡上使用分散式訓練可以實現近線性擴充套件。
總結
本文，我們展示瞭如何基於 optimum-habana 輕鬆在 Gaudi 2 上部署 ProtST 推理和微調。此外，我們的結果還表明，與 A100 相比，Gaudi 2 在這些任務上的效能頗具競爭力: 推理速度提高了 1.76 倍，微調速度提高了 2.92 倍。
如你你想在英特爾 Gaudi 2 加速卡上開始一段模型之旅，以下資源可助你一臂之力:
- optimum-habana 程式碼庫
- 英特爾 Gaudi 文件
感謝垂閱！我們期待看到英特爾 Gaudi 2 加速的 ProtST 能助你創新。
英文原文: https://hf.co/blog/intel-protein-language-model-protst
原文作者: Julien Simon，Jiqing Feng，Santiago Miret，Xinyu Yuan，Yi Wang，Matrix Yao，Minghao Xu，Ke Ding
譯者: Matrix Yao (姚偉峰)，英特爾深度學習工程師，工作方向為 transformer-family 模型在各模態資料上的應用及大規模模型的訓練推理。