[논문 리뷰] Diffusion Language Models Are Versatile Protein Learners

Xinyou Wang, Zaixiang Zheng외 4명|arXiv (Cornell University)|2024. 02. 28.

Natural Language Processing Techniques인용 수 16

한 줄 요약

DPLM은 진화 규모의 서열에 대해 사전 학습된 이산 확산 기반의 단백질 언어 모델로, 새로운 서열을 생성하고 다운스트림 예측 작업을 위한 강력한 표현 학습기로서 작동하며 부분 서열, 교차 모달, 분류기 가이드를 포함한 다목적 조건화 옵션을 제공합니다.

ABSTRACT

This paper introduces diffusion protein language model (DPLM), a versatile protein language model that demonstrates strong generative and predictive capabilities for protein sequences. We first pre-train scalable DPLMs from evolutionary-scale protein sequences within a generative self-supervised discrete diffusion probabilistic framework, which generalizes language modeling for proteins in a principled way. After pre-training, DPLM exhibits the ability to generate structurally plausible, novel, and diverse protein sequences for unconditional generation. We further demonstrate the proposed diffusion generative pre-training makes DPLM possess a better understanding of proteins, making it a superior representation learner, which can be fine-tuned for various predictive tasks, comparing favorably to ESM2 (Lin et al., 2022). Moreover, DPLM can be tailored for various needs, which showcases its prowess of conditional generation in several ways: (1) conditioning on partial peptide sequences, e.g., generating scaffolds for functional motifs with high success rate; (2) incorporating other modalities as conditioner, e.g., structure-conditioned generation for inverse folding; and (3) steering sequence generation towards desired properties, e.g., satisfying specified secondary structures, through a plug-and-play classifier guidance. Code is released at \url{https://github.com/bytedance/dplm}.

연구 동기 및 목표

• 강력한 생성 능력과 예측 능력을 결합한 다목적 단백질 LM의 필요성을 제시한다.
• 단백질 서열의 생성과 이해를 통합하기 위한 이산 확산 사전 학습을 제안한다.
• DPLM이 구조적으로 그럴듯하고 새로운 단백질 서열을 생성하며 다운스트림 작업에 대해 우수한 표현을 제공하는 능력을 입증한다.
• 부분 서열 조건화, 교차 모달 조건화, 컨트롤 가능한 생성용 플러그 앤 플레이 가이던스 등 다양한 조건화 모드를 보여준다.

제안 방법

• 단백질 서열에 대해 언어 모델링의 일반화로서 원칙적 프레이밍을 제공하는 이산 확산 확률 프레임워크를 채택한다.
• 노이즈 스케줄 아래의 전방 확산에서 Cat(x^(t)|x^(t-1))와 흡수 상태 [X]를 정의해 마스킹을 흉내 낸다.
• 역방향 노이즈 제거 목적을 재매개변수화하여 masked-LM 및 자동회귀 LM의 특수 사례(식(4))로 축약된다.
• UniRef50(~45M 서열, ~14B 토큰)에서 사전 학습하며 모델 규모를 최대 3B 파라미터까지 확장하고, 마스킹 LM 사전 학습-그다음 확산 목적의 두 단계 전략을 따른다.
• 완전히 노이즈가 낀 시작점에서의 반복적 제거를 통해 생성 가능성을 제공하며, 마스케이트-프리딕트 샘플링과 유사하다.
• 시퀀스 조건화, 어댑터를 통한 교차 모달 조건화, 그리고 제어 가능한 생성을 위한 이산 분류기 가이던싱과 같은 유연한 조건화를 도입한다.

실험 결과

연구 질문

• RQ1 단백질 서열의 생성과 이해 모두에서 탁월한 통일된 모델을 이산 확산 사전 학습이 제공할 수 있는가?
• RQ2 확산 기반 사전 학습 후 DPLM의 다운스트림 작업 성능이 ESM2 같은 확립된 단백질 LM과 비교해 어떤가?
• RQ3 부분 서열, 교차 모달, 분류기-guided 등의 어떤 조건화 메커니즘이 실용적이고 제어 가능한 단백질 서열 설계를 가능하게 하는가?
• RQ4 확산 기반 사전 학습이 길이에 걸쳐 구조적으로 그럴듯하고 새로운 단백질 서열을 생성하는가?
• RQ5 DPLM이 강력한 표현 학습자로서 기능하면서도 높은 품질의 조건부 생성을 가능하게 하는가?

주요 결과

• DPLM은 길이 전반에 걸쳐 구조적으로 그럴듯하고 새로운 다채로운 단백질 서열을 생성하며, 접힘 가능도(pLDDT 점수)가 높은 성능을 모델 규모 확장과 함께 보인다.
• DPLM은 다운스트림 예측 작업에 대해 ESM2와 비교하여 우수한 표현을 제공하며, 일부 설정에서 구조-인식 모델에 근접한다.
• 더 큰 DPLM 모델일수록 무조건적 생성 및 다운스트림 작업에서 더 나은 성능을 보이며, 단백질 LM에 대한 스케일링 법칙을 시사한다.
• 모티프 골격화(Motif scaffolding), 교차 모달 조건화(예: 구조 조건화 생성) 및 구조와 같은 특성을 조절하기 위한 플러그 앤 플레이 분류기 가이던스를 통해 조건부 생성을 지원한다.
• 이산 확산은 단백질 서열 생성 및 표현 학습에서 Masked-LM 및 AR-LM보다 더 효과적임이 입증되며, 두 단계 학습 전략이 생성 품질을 향상시킨다.
• 모티프 골격화 실험에서 DPLM이 기준선보다 더 높은 성공률과 모티프 보존을 달성했고, 구조 인식 조건화가 추가 이득을 제공한다.