[논문 리뷰] Learning the Language of Protein Structure

Benoit Gaujac, Jérémie Donà외 4명|arXiv (Cornell University)|2024. 05. 24.

Biomedical Text Mining and Ontologies인용 수 5

한 줄 요약

벡터 양자화(autoencoder)를 도입하여 단백질 구조를 토큰 유사 코드로 이산화하고, 설계된 단백질 골격의 시퀀스 모델 기반 생성을 가능하게 한다. 학습된 코드북에 대해 학습된 GPT 모델로 경쟁력 있는 de novo 구조 생성을 시연한다.

ABSTRACT

Representation learning and \emph{de novo} generation of proteins are pivotal computational biology tasks. Whilst natural language processing (NLP) techniques have proven highly effective for protein sequence modelling, structure modelling presents a complex challenge, primarily due to its continuous and three-dimensional nature. Motivated by this discrepancy, we introduce an approach using a vector-quantized autoencoder that effectively tokenizes protein structures into discrete representations. This method transforms the continuous, complex space of protein structures into a manageable, discrete format with a codebook ranging from 4096 to 64000 tokens, achieving high-fidelity reconstructions with backbone root mean square deviations (RMSD) of approximately 1-5 Å. To demonstrate the efficacy of our learned representations, we show that a simple GPT model trained on our codebooks can generate novel, diverse, and designable protein structures. Our approach not only provides representations of protein structure, but also mitigates the challenges of disparate modal representations and sets a foundation for seamless, multi-modal integration, enhancing the capabilities of computational methods in protein design.

연구 동기 및 목표

• 3D 단백질 구조를 이산적 연속 표현으로 변환하여 시퀀스 모델링을 가능하게 한다.
• 단백질 구조를 위한 유한 코드북을 가진 이산적이고 잠재적으로 저차원 류공간(latent space)을 학습한다.
• 확장 가능한 토큰 어휘를 제공하면서 재구성 정확도(RMSD ~1-5 Å)를 높게 달성한다.
• 학습된 코드북으로 GPT 모델을 학습시켜 새로운 설계 가능 단백질 구조를 생성하는 다운스트림 유용성을 시연한다.

제안 방법

• N-Cα-C-O 백본 좌표를 3층 MPNN 기반 그래프 인코더로 인코딩하여 다운샘플링된 잔기에 대한 임베딩을 생성한다.
• Finite Scalar Quantization(FSQ)을 사용하여 잠재 임베딩을 이산화하고 코드북(최대 64k 코드)을 형성한다.
• 양자화된 잠재 코드에서 3D 백본 좌표로 복원하는 구조 모듈은 AlphaFold의 프레이밍에 영감을 받아 Frame Align Point Error(FAPE) 손실로 최적화한다.
• 약 310k 개의 PDB 엔트리에 대해 증강과 유사한 다운샘플링 및 회전 불변 특성으로 프레임 맞춤 및 견고한 재구성을 보장하며 학습한다.
• 재구성은 RMSD와 TM-score로 평가하고, 토큰화된 구조에 대해 디코더 전용 GPT를 학습시켜 생성 능력을 시연한다.
• 공개 코드를 제공하고 생성 품질을 특화된 확산 모델(FrameDiff, RFDiffusion)과 비교한다.

실험 결과

연구 질문

• RQ1 단백질 구조를 대재앙적 정보 손실 없이 유한한 토큰 어휘로 효과적으로 이산화할 수 있을까?
• RQ2 이산 잠재 공간이 표준 시퀀스 모델(예: GPT)을 활용한 de novo 단백질 구조 생성에 효과적으로 작동할 수 있을까?
• RQ3 벡터 양자화 autoencoder에서 코드북 크기와 다운샘플링을 변화시킬 때 재구성 정확도와 설계 가능성 간의 트레이드오프는 무엇인가?
• RQ4 생성된 구조가 설계 가능성, 참신성, 다양성 측면에서 참조 데이터와 어떻게 비교되는가?

주요 결과

• 다운샘플링 없이(codebook 64k, r=1) 테스트 구조에서 RMSD ~1.59 Å 및 TM-score 0.95를 달성했다.
• 양자화가 코드북을 붕괴시키지 않는다; 더 큰 코드북은 재구성을 개선하면서도 비슷한 정확도를 유지한다.
• 다운샘플링과 더 작은 코드북은 재구성 오차를 증가시키지만 평균적으로 TM-score가 0.5 이상으로 유지되어 전형적 전체 폴드가 유지됨을 시사한다.
• 토큰화된 구조에 대해 학습된 간단한 디코더 전용 GPT는 확산 기반대비 자체 일관성 지표가 경쟁력 있는 새로운 설계 가능 백본을 생성할 수 있다.
• 생성된 구조는 비자명한 보조 2차 구조 요소를 보이며 설계 가능성 평가를 위한 ESM-예측 구조와 잘 정렬된다.
• FrameDiff 및 RFDiffusion과 비교했을 때 GPT 기반 접근은 설계 가능성 및 참신함/다양성 프로필에서 경쟁력을 보이며 샘플링 전략에 따라 트레이드오프가 발생한다.