[논문 리뷰] xTrimoPGLM: Unified 100B-Scale Pre-trained Transformer for Deciphering the Language of Protein

Bo Chen, Xingyi Cheng외 13명|arXiv (Cornell University)|2024. 01. 11.

Machine Learning in Bioinformatics인용 수 9

한 줄 요약

단일 통합 단백질 언어 모델로 이해와 생성을 공동 학습하며, 100B 매개변수와 1T 토큰으로 확장되어 18개 단백질 이해 벤치마크에서 강력한 결과를 달성하고 PLM 기반 3D 구조 예측 및 제어 가능한 서열 생성을 가능하게 한다.

ABSTRACT

Protein language models have shown remarkable success in learning biological information from protein sequences. However, most existing models are limited by either autoencoding or autoregressive pre-training objectives, which makes them struggle to handle protein understanding and generation tasks concurrently. We propose a unified protein language model, xTrimoPGLM, to address these two types of tasks simultaneously through an innovative pre-training framework. Our key technical contribution is an exploration of the compatibility and the potential for joint optimization of the two types of objectives, which has led to a strategy for training xTrimoPGLM at an unprecedented scale of 100 billion parameters and 1 trillion training tokens. Our extensive experiments reveal that 1) xTrimoPGLM significantly outperforms other advanced baselines in 18 protein understanding benchmarks across four categories. The model also facilitates an atomic-resolution view of protein structures, leading to an advanced 3D structural prediction model that surpasses existing language model-based tools. 2) xTrimoPGLM not only can generate de novo protein sequences following the principles of natural ones, but also can perform programmable generation after supervised fine-tuning (SFT) on curated sequences. These results highlight the substantial capability and versatility of xTrimoPGLM in understanding and generating protein sequences, contributing to the evolving landscape of foundation models in protein science.

연구 동기 및 목표

• 단백질을 위한 자동인코딩(autoencoding)과 자기회귀(auto-regressive) 목표를 결합한 통합 프레임워크의 필요성을 제시한다.
• 100B 매개변수와 1T 훈련 토큰으로 통합 단백질 언어 모델의 규모를 확장한다.
• 모델이 단백질 이해 벤치마크를 개선하고 고급 구조 예측 및 생성을 가능하게 한다는 것을 보여준다.
• 단일 서열 구조 예측 및 제어 가능한 서열 생성을 위한 PLM 기반의 더 빠른 경로를 보여준다.
• 대형 단백질 파운데이션 모델의 배포에 대한 한계점 및 실용적 고려사항을 논의한다.

제안 방법

• 양방향 어텐션과 자기회귀 목표를 갖춘 백본으로 일반 언어 모델(GLM)을 채택한다.
• 이해도를 높이기 위해 양방향 접두부 영역에 MLM 목표를 도입한다.
• 두 단계로 사전 학습: 먼저 ~400B 토큰에서 MLM을 학습; 그런 다음 약 600B 토큰에서 20%/80%의 비율로 통합 MLM+GLM을 학습한다.
• xTrimoPGLM-100B(100B 매개변수)를 ~940M 개의 고유 시퀀스(~200B 잔기)에서 96대의 NVIDIA DGX 시스템의 A100 GPU를 사용하여 학습한다.
• 단일 서열 구조 예측을 위한 PLM 표현에 접합하는 접힘 모듈을 도입해 xTrimoPGLM-Fold(xT-Fold)을 개발하고, 4비트 양자화 및 FlashAttention를 이용한다.
• 출력을 대상 특성에 정렬하기 위해 감독 학습 미세조정(SFT)과 강화 학습 자기훈련(ReST)을 통해 단백질 서열 생성을 가능하게 한다.

실험 결과

연구 질문

• RQ1 단일화된 사전 학습 목표가 단백질 이해와 생성 작업을 공동으로 지원할 수 있는가?
• RQ2 100B 매개변수와 1T 토큰으로 확장하는 것이 단백질 이해 벤치마크에서의 성능에 어떤 영향을 미치는가?
• RQ3 PLM 기반 접근법이 MSA 기반 방법과 견줄 수 있는 단일 서열 구조 예측(xT-Fold)을 제공할 수 있는가?
• RQ4 SFT와 ReST를 사용한 프로그래머블 생성 및 제어 가능한 단백질 합성의 가능성은 무엇인가?

주요 결과

• xTrimoPGLM-100B는 4개 범주에 걸쳐 18개 단백질 이해 작업 중 15개에서 SOTA 벤치마크를 상회한다.
• 모델은 OOD 단백질 세트의 두 곳에서 다른 비교모델(예: ESM2-15B, ProGen2-xlarge)보다 낮은 퍼플렉시티를 달성한다.
• xT-Fold는 TM-스코어 0.86(CAMEO) 및 0.70(CASP15)을 달성하여 일부 PLM 기반 경쟁자들을 능가하고 MSAs 보강 방법에 근접하다.
• 생성된 단백질은 다양한 구조를 보이고 예측 신뢰도(중위 pLDDT ~85.4)가 높으며 PDB 엔트리와의 서열 유사도는 낮아 새로운 폴드를 탐색하고 있음을 시사한다.
• SFT와 ReST는 생성된 서열을 원하는 특성에 맞춰 제어 가능하게 정렬할 수 있게 하며, 동일한 프로토콜에서 종종 ProGen2 및 ProtGPT2를 능가한다.
• xTrimoPGLM 프레임워크는 관찰 가능한 확장 경향을 보이며: 더 큰 모델일수록 성능이 더 좋아지는 경향이 있으며, 복잡한 작업에서 현저한 이득을 보인다.