단백질 : 논문리뷰 : ProSST: Protein Language Modeling with Quantized Structure and Disentangled Attention

쉬운 설명:

이 논문(ProSST)은 단백질을 이해할 때 아미노산 서열(sequence)만 보는 것의 한계를 넘어서, 3D 구조(structure) 정보까지 함께 활용하는 똑똑한 방법을 제안합니다.

1. 구조를 '쉬운 딱지'로 만들기 (Structure Quantization): 단백질 각 부분(residue) 주변의 복잡한 3D 모양을 GVP라는 도구로 분석해서, 미리 정해둔 몇 가지 '모양 유형' (structure codebook, K개) 중 가장 비슷한 것의 번호(token)를 붙여줍니다. 마치 복잡한 레고 조각에 '모서리 블록', '평면 블록' 같은 간단한 딱지를 붙이는 것과 같습니다. Foldseek보다 주변 블록까지 더 자세히 보고 딱지를 붙입니다.
2. 서열과 '모양 딱지' 관계 따로 보기 (Disentangled Attention): 단백질 서열과 위에서 붙인 '모양 딱지'를 분석할 때, 그냥 한꺼번에 섞어서 보는 대신, "서열-서열 관계", "서열-모양 딱지 관계", "모양 딱지-서열 관계" 등을 따로따로 주목해서 봅니다. 이렇게 하면 서열과 구조가 서로 어떻게 영향을 주는지 더 명확하게 파악할 수 있습니다.

이렇게 구조 정보를 더 똑똑하게 표현하고 서열과의 관계를 명시적으로 분석함으로써, ProSST는 단백질의 돌연변이 효과 예측이나 다양한 기능 예측에서 기존 모델들보다 더 좋은 성능을 보여줍니다.

 

ProSST 학습 노트

Purpose of the Paper:

• 기존 Protein Language Models (PLMs) 대부분이 단백질의 3D structure 정보를 통합하지 못하고 sequence 정보에만 의존하는 한계를 극복하고자 함.
• 기존 structure-aware PLM (e.g., Foldseek 기반)의 한계점 개선:
  • Foldseek: 개별 residue 주변의 매우 제한된 정보 (이전/다음 residue)만 사용하므로 local structure 표현력 부족.
  • 기존 Transformer: Sequence와 structure token 간의 관계를 명시적으로 모델링하는 능력 부족.
• 핵심 목표: 단백질 sequence와 structure 정보를 효과적으로 통합하고, 이 둘 간의 관계를 명시적으로 학습하는 새로운 PLM 개발.

Key Contributions & Novelty:

• Structure Quantization Module:
  • Contribution: 3D 단백질 구조를 discrete token sequence로 변환.
    • 각 residue의 local structure (해당 residue + 주변 40개 neighbor residues)를 pre-trained GVP (Geometric Vector Perceptron) encoder로 encoding.
    • C.A.T.H dataset의 local structure embedding에 대해 pre-trained k-means clustering 모델을 사용하여 K개의 centroid (structure codebook) 생성.
    • 각 residue의 local structure embedding을 가장 가까운 centroid index (structure token)로 quantize.
  • Novelty: Foldseek 대비 더 풍부한 local structural context (neighborhood 정보)를 GVP encoder를 통해 학습하고, 이를 기반으로 structure vocabulary (K개의 discrete token)를 생성하여 구조를 표현.
• Sequence-Structure Disentangled Attention:
  • Contribution: Standard self-attention을 대체하여 sequence token과 structure token 간의 상호작용을 명시적으로 모델링.
  • Attention 계산을 5가지 component로 분해: residue-to-residue (R2R), residue-to-structure (R2S), residue-to-position (R2P), structure-to-residue (S2R), position-to-residue (P2R). (Equation 4 참고)
  • Novelty: Sequence 정보와 structure 정보의 관계를 분리된 attention 가중치로 학습하여, 두 정보 modality 간의 복잡한 상호 의존성을 보다 효과적으로 포착.
• Large-scale Pre-training:
  • Contribution: 18.8 million 개의 단백질 구조 데이터셋 (AlphaFold DB 기반)에서 Masked Language Model (MLM) objective로 pre-train 수행. MLM loss 계산 시, corrupted sequence 정보와 uncorrupted structure token sequence를 모두 context로 활용 (Equation 6).
  • Novelty: 대규모 구조 데이터를 활용하고, pre-training 단계에서부터 sequence와 structure 정보를 함께 고려하여 contextual representation 학습.

Experimental Highlights:

• Zero-Shot Mutation Effect Prediction:
  • ProteinGYM benchmark에서 SOTA 달성 (Spearman ps = 0.504), 기존 sequence-based 및 structure-aware 모델 (SaProt 등) 성능 상회 (Table 1).
  • 특히 Stability 예측 성능이 우수하며, 이는 structure 정보 통합의 이점을 잘 보여줌 (Appendix A.4, Table A6).
• Supervised Fine-Tuning Tasks:
  • DeepLoc, Metal Ion Binding, Thermostability, GO annotation prediction 등 다양한 downstream task에서 SOTA 또는 경쟁력 있는 성능 달성 (Table 2).
  • 특히 DeepLoc, Metal Ion Binding에서 가장 우수한 성능을 보였고, GO prediction에서는 더 큰 모델인 ESM-GearNet과 비슷하거나 약간 낮은 성능을 보임.
• Ablation Studies:
  • Quantized Structure: Structure codebook size (K)가 증가할수록 성능 향상 (K=2048에서 최적). 제안된 GVP 기반 quantization 방식이 Foldseek, DSSP 방식보다 우수 (Table 3).
  • Disentangled Attention: 제안된 5가지 attention component 모두 성능에 기여하며, 특히 R2S (Residue-to-Structure)가 중요. Standard self-attention 대비 우수한 성능 확인. Positional Encoding (PE) 제거 시 성능 크게 하락 (Table 4, Figure 3b).
  • Disentangled attention은 올바른 structure token이 주어졌을 때 효과를 발휘하며, attention 메커니즘 자체만으로는 성능 향상이 어려움을 확인 (Appendix A.5, Table A8).

Limitations and Future Work:

• Limitations:
  • Local structure construction 및 GVP encoding 과정의 높은 계산 비용 (Computational cost).
  • Sequence 데이터 대비 부족한 structure 데이터 양에 대한 의존성.
  • 예측된 구조의 품질 (e.g., AlphaFold pLDDT)이 낮은 disordered protein region에서는 성능이 저하될 수 있음 (Appendix A.7).
• Future Work:
  • Structure quantization process의 효율성 개선 (Speed up).
  • 더 큰 structure dataset을 활용한 pre-training.
  • 모델 파라미터 확장 (Expanding parameters).
  • Sequence-only dataset 입력 처리를 위한 Masked Structure Training (MST) 방식 개선 등 (Appendix A.6).

Overall Summary:

ProSST는 기존 PLM의 structure 정보 부재 문제를 해결하기 위해, GVP encoder와 k-means clustering 기반의 novel structure quantization 방식과 sequence-structure disentangled attention 메커니즘을 제안한 Transformer 기반 PLM이다. 이 두 핵심 기술을 통해 단백질의 sequence와 local structure 정보를 효과적으로 통합하고 상호 관계를 명시적으로 학습한다. 대규모 구조 데이터셋 pre-training을 통해 zero-shot mutation prediction 및 다양한 downstream task에서 SOTA 성능을 달성하며, 정교한 구조 정보 통합의 중요성을 입증했다. 이는 향후 단백질 기능 및 특성 예측 연구에 기여할 수 있는 발전된 방법론을 제시한다.

 

 

 

 

Abstract

Protein language models (PLMs)은 다양한 단백질 기능 prediction task에서 놀라운 능력을 보여주었습니다.

하지만 단백질 기능은 구조와 복잡하게 연결되어 있지만, 기존의 대부분 PLMs는 단백질 구조 정보를 통합하지 않습니다.

이 문제를 해결하기 위해, 우리는 단백질 sequences와 구조를 모두 원활하게 통합하는 Transformer-based protein language model인 ProSST를 소개합니다.

ProSST는 structure quantization module과 disentangled attention을 갖춘 Transformer architecture를 통합합니다.

structure quantization module은 먼저 단백질 구조를 잔기 수준의 지역 구조로 직렬화한 다음 이를 밀집된 vector space로 embeds하여 3D 단백질 구조를 이산적인 tokens의 sequence로 변환합니다.

이 벡터들은 pre-trained된 clustering model에 의해 이산적인 structure tokens으로 quantized됩니다.

이 tokens는 효과적인 단백질 structure representation 역할을 합니다.

또한, ProSST는 sequence-structure disentangled attention을 통해 단백질 잔기 token sequences와 structure token sequences 간의 관계를 명시적으로 학습합니다.

우리는 masked language model objective를 사용하여 수백만 개의 단백질 구조에 대해 ProSST를 pre-train하여 단백질의 포괄적인 문맥적 representations를 학습할 수 있도록 합니다.

제안된 ProSST를 평가하기 위해, 우리는 zero-shot 돌연변이 효과 prediction과 여러 supervised downstream tasks에 대한 광범위한 실험을 수행했으며, 여기서 ProSST는 모든 baselines 중에서 state-of-the-art performance를 달성합니다.

우리의 code와 pre-trained models는 공개적으로 사용 가능합니다.

 

 

 

ProSST Abstract 정리노트 (AI 연구자용)

• Problem: 기존 Protein Language Models (PLMs)는 단백질 기능 예측에 중요함에도 불구하고 3D 구조 정보를 통합하지 못하는 한계가 있음.
• Proposed Model: ProSST - 단백질 sequence와 structure 정보를 함께 통합하는 새로운 Transformer-based protein language model.
• Key Innovation 1: Structure Representation (structure quantization):
  • 3D 단백질 구조를 residue-level local structure로 직렬화(serialize).
  • 이를 dense vector space로 embedding.
  • Pre-trained clustering model을 사용하여 이 vector들을 discrete structure tokens으로 quantization함. (연속적인 3D 정보를 Transformer가 처리하기 용이한 이산적 token sequence로 변환하는 핵심 아이디어).
• Key Innovation 2: Integration Mechanism (sequence-structure disentangled attention):
  • Transformer architecture 내부에 sequence-structure disentangled attention을 도입.
  • 이를 통해 protein residue token sequences와 structure token sequences 간의 관계를 명시적으로 learning함.
• Training & Evaluation:
  • 수백만 개의 단백질 구조 데이터셋에서 masked language model objective를 사용하여 pre-train됨.
  • Zero-shot mutation effect prediction 및 다수 supervised downstream tasks에서 state-of-the-art performance 달성, 제안된 구조 통합 방식의 효과 입증.
• Contribution: 3D 구조 정보를 효과적으로 이산화(quantize)하고 이를 sequence 정보와 함께 representation learning에 통합하는 새로운 architecture 및 방법론 제시. Code 및 pre-trained models 공개.

---

쉬운 설명 :

단백질은 아미노산 서열(sequence)이라는 설계도와 그 설계도로 만들어진 실제 3차원 구조(structure)를 가지고 있어요. 단백질의 기능은 이 3차원 구조와 아주 밀접한 관련이 있죠. 그런데 기존의 단백질을 이해하려는 AI 모델(Protein Language Models, PLMs)들은 주로 설계도(sequence)만 읽고, 실제 구조는 잘 보지 못하는 경우가 많았어요.

이 논문에서는 ProSST라는 새로운 AI 모델을 만들었어요. 이 모델의 가장 큰 특징은 단백질의 설계도(sequence) 정보와 실제 3차원 구조(structure) 정보를 함께 사용한다는 점이에요.

ProSST는 복잡한 3차원 구조 정보를 AI가 이해하기 쉬운 형태로 바꾸는 특별한 방법(structure quantization module)을 사용해요. 3D 구조를 작은 부분들로 나누어 특징을 뽑아낸 다음, 이것들을 미리 학습된 분류기(pre-trained clustering model)를 이용해 몇 가지 종류의 '구조 딱지'(discrete structure tokens)로 만들어요. 마치 복잡한 모양을 몇 개의 간단한 키워드로 요약하는 것과 비슷해요.

그리고 나서, 설계도 정보(sequence tokens)와 이 구조 딱지 정보(structure tokens)가 서로 어떤 관련이 있는지를 AI가 잘 배울 수 있도록 특별히 설계된 주의 집중 메커니즘(sequence-structure disentangled attention)을 사용해요.

이렇게 설계도와 구조 정보를 함께 학습한 ProSST는, 단백질의 기능(예: 작은 변화가 단백질 기능에 미치는 영향 예측 - zero-shot mutation effect prediction)을 예측하는 여러 시험(downstream tasks)에서 기존 모델들보다 훨씬 좋은 성능(state-of-the-art performance)을 보여주었어요. 즉, 단백질을 더 깊이 이해하려면 설계도뿐만 아니라 실제 구조 정보까지 통합하는 것이 효과적이라는 것을 보여준 연구예요.

 

 

1 Introduction

단백질의 기능을 predicting하는 것은 생명 과학에서 가장 중요한 영역 중 하나입니다. 최근 수십 년 동안 단백질 sequence databases는 기하급수적으로 성장했으며, 이로 인해 data-driven 방식으로 large-scale models를 사용하여 단백질 sequences의 근본적인 representations를 학습하는 것이 가능해졌습니다. Natural language processing의 pre-trained language models에서 영감을 받아 많은 pre-trained Protein Language Models (PLMs)가 등장했습니다. 뛰어난 단백질 representation 능력 덕분에 이들은 단백질 관련 tasks에서 생물정보학의 기본적인 도구가 되었습니다.

단백질의 기능은 그 구조에 의해 결정됩니다. 하지만 대부분의 PLMs는 주로 단백질 sequences를 modeling하는 데 중점을 두어 구조 정보의 중요성을 간과하며, 이러한 현상의 한 가지 중요한 이유는 구조 data의 부족입니다. 다행히 AlphaFold나 RoseTTAFold와 같은 몇몇 뛰어난 연구들이 제안되었으며, 이들은 단백질 구조를 정확하게 predict할 수 있습니다. 이러한 연구들은 단백질 구조 dataset을 수백만 개로 크게 확장하고 large-scale structure-aware PLMs의 pre-training을 가능하게 합니다. 그 이후 주요 challenge는 단백질 구조 정보를 PLMs에 효과적으로 통합하는 방법입니다. 구체적으로, 기존의 structure-aware PLMs는 먼저 Foldseek를 사용하여 단백질 구조를 이산적인 structure tokens로 변환한 다음, 이 구조 data를 Transformer architecture에 통합합니다. 하지만 여러 tasks에서 유망한 performance를 달성했음에도 불구하고 이 접근 방식은 여전히 두 가지 주요 문제에 직면해 있습니다. 첫째, Foldseek는 단백질 내 잔기의 구조를 이전 잔기와 다음 잔기의 features만을 고려하여 encode합니다. 이 representation은 불충분하며, 단백질 기능에 중요한 촉매 부위나 결합 포켓과 같은 단백질의 국소 구조에서의 미묘한 차이를 간과할 수 있습니다. 둘째, 일반적인 Transformer architecture는 단백질 sequences와 structure token sequences 간의 관계를 명시적으로 model하는 능력이 부족하여 구조적 단서를 효과적으로 활용하기 어렵게 만듭니다.

본 논문에서는 structure-aware pre-trained protein language model인 ProSST (Protein Sequence-Structure Transformer)를 개발합니다. 구체적으로 ProSST는 주로 두 개의 modules로 구성됩니다: structure quantization module과 sequence-structure disentangled attention을 갖춘 Transformer입니다. structure quantization module은 GVP (Geometric Vector Perceptron) encoder에 기반하며, 이는 잔기 구조를 그 국소 구조 내의 이웃과 함께 encode하고 encoding vector를 discrete tokens로 quantize할 수 있습니다. 개별 잔기만 고려하는 Foldseek와 비교할 때, 이 encoder는 잔기의 미세 환경으로부터 더 많은 정보를 고려할 수 있습니다. sequence-structure disentangled attention module은 Transformer model의 self-attention module을 대체합니다. 이는 Transformer model이 단백질 sequence tokens와 structure tokens 간의 관계를 명시적으로 model하도록 하여 단백질 sequences와 구조의 더 복잡한 features를 포착할 수 있게 합니다. ProSST가 contextual representation을 포괄적으로 학습할 수 있도록, 우리는 1,880만 개의 단백질 구조를 포함하는 대규모 dataset에서 Masked Language Modeling (MLM) objective를 사용하여 우리 model을 pre-train합니다. 요약하자면, 우리의 주요 contributions는 다음과 같습니다:

• 우리는 단백질 구조를 discrete tokens의 sequence로 변환할 수 있는 단백질 구조 quantizer를 제안합니다. 이러한 token sequences는 단백질 내 잔기의 국소 구조 정보를 효과적으로 represent합니다.
• 우리는 단백질 구조와 잔기 간의 관계를 명시적으로 학습하여 structural token sequences와 아미노산 sequences의 더 효율적인 통합을 촉진하는 disentangled attention mechanism을 제안합니다.

제안된 ProSST를 평가하기 위해, 우리는 zero-shot 돌연변이 효과 prediction과 다수의 supervised downstream tasks에 대한 광범위한 실험을 수행했으며, 제안된 model은 모든 baselines 중에서 state-of-the-art 결과를 달성합니다. 이 외에도, 우리는 ProSST의 각 설계의 효과를 입증하기 위해 상세한 ablations를 제공합니다.

 

ProSST Introduction 정리노트 (AI 연구자용)

• 배경 및 동기: 단백질 기능 예측은 중요합니다. PLMs는 sequence 데이터를 잘 사용하지만, 단백질 기능은 3D 구조에 크게 의존합니다. AlphaFold 등 정확한 구조 예측기의 등장으로 대규모 구조 datasets이 생성되었고, 이는 structure-aware PLMs 개발을 가능하게 했습니다.
• 기존 Structure-Aware PLMs의 한계:
  • 구조 토큰화를 위해 Foldseek 같은 방법에 의존: Foldseek는 주로 인접 잔기의 features만 사용하므로, 중요한 국소 구조 세부 정보(예: 촉매 부위)를 놓치는 불충분한 representation을 생성할 수 있습니다.
  • 단순한 Transformer architecture: 표준 self-attention은 protein sequences와 structure token sequences 간의 관계를 명시적으로 model하지 못해 효과적인 통합을 저해합니다.
• 제안 해결책: ProSST (Protein Sequence-Structure Transformer): 새로운 structure-aware pre-trained protein language model입니다.
• 핵심 기술 혁신:
  1. Structure Quantization Module: GVP (Geometric Vector Perceptron) encoder를 사용합니다. Foldseek와 달리, 이 encoder는 잔기의 이웃/미세환경을 고려하여 더 풍부한 국소 구조 encoding을 수행한 후 discrete structure tokens으로 quantization 합니다.
  2. Sequence-Structure Disentangled Attention: 표준 self-attention module을 대체합니다. 이 mechanism은 sequence tokens와 structure tokens 간의 상호작용 및 관계를 명시적으로 model하도록 설계되었습니다.
• 학습 (Training): 1,880만 개의 단백질 구조 dataset에서 Masked Language Modeling (MLM) objective를 사용하여 pre-trained되었습니다.
• 주요 기여 (Claimed Contributions):
  • 더 정보력 있는 이산적(discrete) 구조 representations를 생성하는 새로운 구조 quantizer.
  • 명시적인 sequence-structure 관계 학습 및 더 나은 통합을 위한 disentangled attention mechanism.
• 검증 (Validation): zero-shot 돌연변이 효과 prediction 및 supervised downstream tasks에서 state-of-the-art 결과를 달성했으며, ablation 연구로 뒷받침됩니다.

 

2 Related Work

2.1 Protein Representation Models

입력 modality에 따라, 단백질 representation models는 세 가지 카테고리로 나눌 수 있습니다: sequence-based models, structure-based models, 그리고 structure-sequence hybrid models.

Sequence-based models. Sequence-based models는 단백질을 잔기 tokens의 sequence로 취급하며, 광범위한 sequence datasets에서 unsupervised pre-training을 위해 Transformer model을 사용합니다. Pre-training objective에 따라, 현재 models는 BERT-based models, GPT-based models, 그리고 span-mask based models로 더 나눌 수 있습니다. 구체적으로, ESM-series models, ProteinBert, TAPE를 포함한 BERT-style models는 training phase에서 masked tokens를 복구하는 것을 목표로 합니다. Tranception, ProGen2, ProtGPT2와 같은 GPT-style models는 auto-regressive 방식으로 token sequences를 점진적으로 generate합니다. 마지막으로, span-mask를 training objective로 사용하는 models에는 Ankh, ProtT5, xTrimo가 있습니다.

Structure-based models. 단백질 구조는 단백질 기능성에서 지배적인 역할을 합니다. 따라서, 구조 정보를 활용하는 models는 일반적으로 더 정확한 predictions를 얻습니다. 최근에는 CNN-based models와 GNN-based models를 포함하여 단백질 구조 representation 학습에 다양한 기술이 적용되었으며, 특히 GNN-based models는 node 또는 edge attributes를 통해 단백질 특이적 features를 통합하는 데 상당한 다재다능함을 보여주었습니다. 더욱이, 최근 단백질 folding models의 발전은 structure-based models가 광범위한 단백질 구조 datasets에 접근할 수 있게 했습니다. 이는 단백질 구조 단서를 활용하는 PLMs 개발에 대한 관심 증가로 이어졌습니다.

Structure-sequence hybrid models. 단백질의 sequence와 structure 정보를 모두 통합하는 Hybrid models는 더 효과적인 단백질 representations를 제공합니다. 예를 들어, LM-GVP model은 ProtBERT-BFD embeddings를 GVP model의 input features로 사용하며, ESM-GearNet은 ESM-1b representations를 GearNet과 통합하는 다양한 방법을 연구합니다. 유사하게, 최근 ProtSSN model은 ESM-2 embeddings를 EGNN model의 input으로 활용하여 주목할 만한 발전을 이루었습니다. ESM-IF1과 MIF-ST는 모두 inverse folding을 목표로 하며, 구조를 활용하여 해당 단백질 잔기를 predict하는 반면, ProstT5는 잔기 sequences와 그들의 structure token sequences 간의 transformation을 pre-training objective로 사용합니다. SaProt는 foldseek에 의해 생성된 structure tokens를 사용하여 structure-aware vocabulary를 구축합니다. SaProt와 ProstT5 모두 그들의 pre-training datasets를 위해 대규모 구조 databases를 광범위하게 활용합니다. ProSST 역시 hybrid structure-sequence model입니다. 이전 연구와 비교하여, ProSST는 구조 단서를 활용하기 위해 진보된 structure quantization 방법과 더 나은 attention formulation을 개발합니다.

2.2 Protein Structure Quantization

단백질 구조를 represent하는 가장 직관적인 방법은 좌표, 이면각, 거리 맵과 같은 continuous features를 사용하는 것입니다. 하지만, pre-training에서 이러한 continuous features를 직접 사용하는 것은 overfitting을 초래할 수 있습니다. 이 문제는 training set(model predictions에서 파생됨)과 test set(wet-lab 실험으로 측정됨) 간의 구조 representations 불일치에서 발생합니다. 이 격차를 해소하기 위한 다리로서, structure quantization은 몇몇 연구에서 조사되었습니다. 이러한 방법들은 이산적인(discrete) 2차 구조를 생성하는 방식에 따라 두 그룹으로 나눌 수 있는데, DSSP와 같은 물리적 계산(physical computing) 기반의 방법과, Foldseek와 같은 deep learning 기반의 방법이 있으며, 후자는 structure-aware PLMs에 성공적으로 적용되었습니다. ProSST의 structure quantization module 역시 learning-based 접근 방식에 의존하지만 Foldseek보다 더 상세한 잔기 구조 representation을 제공합니다.

 

Related Work (Section 2) 정리노트 (AI 연구자용)

• Protein Representation Models 분류: 연구 동향은 크게 3가지로 나뉨: sequence-based models, structure-based models, structure-sequence hybrid models. ProSST는 hybrid model에 속함.
  • Sequence-based models: (ESM, ProtT5 등) sequence 정보만 사용하며, pre-training objective (BERT-style, GPT-style, span-mask)에 따라 다양함. 구조 정보 부재가 본질적 한계.
  • Structure-based models: (GNN-based 등) 구조 정보 활용. 단백질 folding models (AlphaFold 등) 발전으로 대규모 구조 datasets 접근이 용이해져 PLMs 개발 관심 증가.
  • Structure-sequence hybrid models: Sequence와 structure 정보 통합 시도.
    • 기존 접근법: sequence model embeddings를 structure model (GVP, EGNN 등)의 input features로 사용 (LM-GVP, ProtSSN); representation 통합 (ESM-GearNet); inverse folding (ESM-IF1); Foldseek 기반 structure tokens 활용 (ProstT5, SaProt - pre-training objective 또는 vocabulary).
    • ProSST의 차별점: 이전 연구 대비, (1) 더 진보된 structure quantization 방법 (Foldseek 대비 GVP encoder 활용) 및 (2) 구조 정보 활용을 위한 더 나은 attention formulation (sequence-structure disentangled attention)을 개발함.
• Protein Structure Quantization:
  • 필요성: Continuous features (좌표 등) 직접 pre-training 시, training set (predicted 구조)과 test set (실험 구조) 간 representation 불일치로 overfitting 발생 가능. Quantization이 이 간극 해소.
  • 기존 방법: 물리 계산 기반 (DSSP 등) 및 deep learning 기반 (Foldseek). 특히 Foldseek는 이전 structure-aware PLMs (SaProt, ProstT5)에서 활용됨.
  • ProSST의 접근: Learning-based 방식에 속하지만, Foldseek보다 더 상세한(detailed) 잔기 구조 representation 제공을 목표로 함 (GVP encoder 활용).

---

쉬운 설명 :

AI가 단백질을 이해하는 방법들에 대한 이전 연구들을 소개하는 부분이에요.

1. AI가 단백질을 배우는 방법들

• 설계도만 읽기 (sequence-based): 단백질의 아미노산 서열(설계도)만 보고 배우는 AI 모델들이에요. (ESM, ProtT5 같은 모델들) 하지만 실제 모양 정보를 놓친다는 단점이 있죠.
• 모양만 보기 (structure-based): 단백질의 3D 구조(모양) 정보만 보고 배우는 모델들이에요. (GNN 기반 모델 등) 요즘 AlphaFold 같이 모양 예측을 잘하는 AI 덕분에 이 분야 연구가 활발해졌어요.
• 설계도와 모양 함께 보기 (structure-sequence hybrid): 설계도와 모양 정보를 둘 다 활용하는 모델들이에요. ProSST가 바로 여기에 속해요. 다른 모델들은 설계도 정보를 모양 정보 이해에 도움 주거나(LM-GVP), 모양 정보로 설계도를 예측하거나(inverse folding), 'Foldseek'라는 도구로 만든 단순한 모양 코드를 활용(ProstT5, SaProt)하는 등 다양한 시도를 했어요.

2. 단백질 모양을 AI에게 어떻게 알려줄까? (Protein Structure Quantization)

• 단백질의 3D 모양은 복잡한 좌표 값 등으로 표현되는데, 이걸 AI 학습(pre-training)에 그대로 쓰면 문제가 생길 수 있어요. AI가 학습 때 본 (예측된) 모양과 실제 실험에서 얻는 모양이 미세하게 달라서, AI가 학습 데이터에만 너무 익숙해지는 overfitting이 발생할 수 있거든요.
• 그래서 연구자들은 복잡한 3D 모양 정보를 간단한 '모양 코드' 또는 '딱지'(discrete tokens)로 변환하는 방법(structure quantization)을 사용해요. 마치 복잡한 자동차 모양을 "세단", "SUV" 같은 간단한 분류로 나누는 것과 비슷해요.
• 이전 연구들(ProstT5, SaProt)은 'Foldseek'라는 도구가 만든 모양 코드를 많이 썼어요.

ProSST는 뭐가 다른가?

이 섹션에서는 ProSST가 기존 연구들, 특히 Foldseek 기반의 hybrid 모델들과 비교해서 두 가지 점에서 더 발전했다고 주장해요.

1. 더 똑똑한 모양 코드 생성: ProSST는 Foldseek 대신 더 자세한 주변 환경 정보까지 고려하는 방법(GVP encoder 기반 structure quantization)으로 모양 코드를 만들어서, 단백질의 국소적인 모양을 더 잘 표현할 수 있다고 해요.
2. 더 나은 설계도-모양 연결: ProSST는 설계도 정보와 이 새로운 모양 코드 정보를 효과적으로 연결하고 그 관계를 잘 배우기 위해 특별히 설계된 attention 방식(sequence-structure disentangled attention)을 사용한다고 해요.

 

 

3 Method

이 섹션에서는 ProSST의 architecture를 소개합니다. ProSST는 주로 두 개의 modules를 포함합니다: structure quantization (섹션 3.1) module과 sequence-structure disentangled attention을 갖춘 transformer-based model (섹션 3.2)입니다.

3.1 Structure Quantization Module

structure quantization module은 잔기의 국소 구조(local structure)를 discrete token으로 변환하는 것을 목표로 합니다. 먼저, 국소 구조는 pre-trained된 structure encoder를 사용하여 dense vector로 encode됩니다. 그 후, pre-trained된 k-means clustering model이 encode된 vector를 기반으로 국소 구조에 카테고리 레이블(category label)을 할당합니다. 마지막으로, 카테고리 레이블이 잔기에 structure token으로 할당됩니다. structure quantization의 pipeline은 그림 1에 나와 있습니다.

Structure representation. 우리는 단백질 구조를 두 가지 구별되는 수준으로 분류합니다: 단백질 구조(protein structure)와 국소 구조(local structure). 단백질 구조는 모든 잔기를 포함하는 단백질의 완전한 architecture를 나타냅니다. 국소 구조는 특정 개별 잔기에 초점을 맞춥니다. 특정 잔기를 중심으로 하고 3차원 공간에서 그 주변의 가장 가까운 40개 잔기를 함께 포함하여 잔기의 국소 환경(local environment)을 설명합니다. 단백질 구조와 비교할 때, 국소 구조는 더 세분화된(finer granularity) 단위이며, 이는 잔기 구조를 더 정확하게 설명할 수 있게 합니다. 따라서 L개의 잔기를 포함하는 단백질은 하나의 단백질 구조와 L개의 국소 구조를 가집니다. 구조의 수준이 다르지만, 우리는 이를 represent하기 위해 graphs를 사용할 수 있습니다. 공식적으로, 우리는 구조를 graph G = (V, E)를 사용하여 represent하며, 여기서 V와 E는 각각 잔기 수준의 nodes와 edges를 나타냅니다. 주어진 node v ∈ V 에 대해, 이는 잔기 자체의 잔기 유형 정보 없이 잔기의 구조 정보만을 포함합니다. 이는 structure encoder가 오직 구조 단서에만 집중하도록 보장합니다. edge 집합 E = {e_{ij}}는 v_j가 v_i의 상위 40개 최근접 이웃 중 하나인 모든 i, j를 포함하며, 이는 Cα 원자 간의 거리에 의해 결정됩니다.

Structure encoder. 위에서 언급한 구조 정의에 기초하여, 우리는 geometric vector perceptrons (GVP)를 (국소) structure encoder로 사용합니다. 특히, GVP는 구조 feature extraction 함수 \pi_{\theta}(G) \in R^{l \times d}로 represent될 수 있으며, 여기서 l은 nodes의 수, d는 embedding dimension, θ는 trainable parameters입니다. 우리는 GVP를 position-wise multi-layer perceptron (MLP)를 포함하는 decoder와 통합하여 auto-encoder model을 형성합니다. 전체 model은 de-noising pre-training objective를 사용하여 trained됩니다. 이 과정에서, 우리는 3D Gaussian noise로 Cα 좌표를 교란시키고(perturb) (Figure 1A), RF-Diffusion에 따라 회전 행렬의 manifold 상에서 브라운 운동(Brownian motion)을 사용합니다. 그 다음 model은 구조를 원래의 노이즈 없는(noise-free) 상태로 복구하는 task를 수행합니다. C.A.T.H dataset에서 trained된 후 (부록 A.2 참조), 우리는 decoder를 제외하고 encoder의 mean pooled output만을 구조의 final representation으로 활용합니다. structure encoder는 단백질 구조에 대해 trained되었지만, 국소 구조를 효과적으로 encode할 수 있습니다. 따라서, 단백질 구조의 graph G에 대한 encoding은 다음과 같습니다: r = \frac{1}{l} \sum_{i=1}^{l} \pi_{\theta}(g_i), 여기서 g_i는 graph G 내 i번째 잔기와 관련된 국소 구조의 graph를 represent하며, \pi_{\theta}(g_i) \in R^d는 i번째 node에 대한 encoder의 output입니다. 여기서 r \in R^d는 encoder의 mean pooled output이며 국소 구조의 vectorized representation입니다.

Local structure codebook. structure codebook은 단백질 구조를 represent하는 dense vectors를 discrete tokens로 quantizes합니다 (Figure 1B). 이를 구축하기 위해, 우리는 C.A.T.H dataset의 모든 잔기의 국소 구조를 continuous latent space로 embed하기 위해 structure encoder를 사용합니다 (부록 A.2 참조). 그 다음 우리는 이 latent space 내에서 K개의 centroids를 식별하기 위해 k-means algorithm을 적용하며, 이를 {e_i}_{i=1}^{K}로 표기합니다. 이 centroids는 그림 1B에 표시된 바와 같이 structure codebook을 구성합니다. 임의의 국소 구조 embedding에 대해, 이는 codebook 내 가장 가까운 vector e_j에 의해 quantized되고, j가 structure token 역할을 합니다. 본 논문에서, clustering number K는 structure vocabulary size라고도 합니다.

Protein serialization and quantization. 일반적으로, 단백질 sequence 내 위치 i에 있는 잔기에 대해, 우리는 먼저 그 국소 구조만을 기반으로 graph g_i를 구축한 다음, structure encoder를 사용하여 이를 continuous vector r_i로 embed합니다. 그 다음 우리는 codebook을 사용하여 이 vector에 structure token s_i \in \{1, 2, ..., K\}를 잔기의 structure token으로 할당합니다. 전반적으로, 전체 단백질 구조는 structure tokens의 sequence로 serialized되고 quantized될 수 있습니다 (Figure 1C).

3.2 Sequence-Structure Disentangled Attention

Residue Embedding DeBerta에서 영감을 받아, 우리는 잔기 sequences와 구조 sequences 및 상대적 위치(relative positions)의 attention을 결합하기 위해 확장된 형태의 disentangled attention을 사용합니다. 구체적으로, 단백질 sequence 내 위치 i에 있는 잔기는 세 가지 항목으로 represent될 수 있습니다: R_i는 그 잔기 token hidden state를 나타내고, S_i는 잔기 수준 국소 구조의 embedding을 represent하며, P_{i|j}는 위치 j의 token과의 relative position의 embedding입니다. 잔기 i와 잔기 j 사이의 cross attention A_{i,j} 계산은 다음과 같이 9개의 components로 분해될 수 있습니다: [ A_{i,j} = {R_i, S_i, P_{i|j}} \times {R_j, S_j, P_{j|i}}^\top = R_i R_j^\top + R_i S_j^\top + R_i P_{j|i}^\top + S_i R_j^\top + S_i S_j^\top + S_i P_{j|i}^\top + P_{i|j} R_j^\top + P_{i|j} S_j^\top + P_{i|j} P_{j|i}^\top. ]

방정식 1에서 공식화된 바와 같이, 잔기 쌍의 attention weight는 잔기 tokens, 구조 tokens, 그리고 relative positions를 포함하는 별도의 행렬들로 계산될 수 있습니다. 이 행렬들은 residue-to-residue, residue-to-structure, residue-to-position, structure-to-residue, structure-to-structure, structure-to-position, position-to-residue, position-to-structure, position-to-position과 같은 다양한 interactions에 활용됩니다. 우리 model은 잔기에 대한 contextual embeddings 학습에 집중하기 때문에, structure-to-structure (S_i S_j^\top), structure-to-position (S_i P_{j|i}^\top), position-to-structure (P_{j|i} S_j^\top), 그리고 position-to-position (P_{j|i} P_{j|i}^\top) 항들은 잔기에 대한 관련 정보를 제공하지 않으므로 크게 기여하지 않습니다. 결과적으로, 이 항들은 attention weight 계산의 우리 implementation에서 제거됩니다. 그림 2에 표시된 바와 같이, 우리의 sequence-structure disentangled attention mechanism은 5가지 유형의 attention을 포함합니다.

다음 부분에서는 sequence-structure disentangled attention의 작동을 설명하기 위해 single-head attention을 예로 사용합니다. 먼저, i번째 잔기에서 j번째 잔기까지의 relative position을 \delta(i, j)로 정의합니다: [ \delta ( i, j ) = \begin{cases} 0 & \text{if } i - j \leq - L_{max} \ 2 L_{max} - 1 & \text{if } i - j \geq L_{max} \ i - j + L_{max} & \text{otherwise} \end{cases} ] 여기서, L_{max}는 우리가 고려하는 최대 상대 거리(maximum relative distance)를 represent하며, implementation에서는 1024로 설정됩니다. 표준 self-attention operation과 유사하게, 구조, 잔기, 상대 위치에 대한 query, key와 잔기에 대한 value의 computation은 다음과 같습니다: [ Q^r = RW^{q_r} \quad K^r = RW^{k_r} \quad V^r = RW^{v_r} \ Q^s = SW^{q_s} \quad K^s = SW^{k_s} \ Q^p = PW^{q_p} \quad K^p = SW^{k_p} ] 그리고 잔기 i에서 잔기 j까지의 attention score \hat{A}_{i,j}는 다음과 같이 계산될 수 있습니다: [ \hat{A}{i,j} = Q_i^r K_j^{r\top} + Q_i^r K_j^{s\top} + Q_i^r K{\delta(i,j)}^{p\top} + K_j^r Q_i^{s\top} + K_j^r Q_{\delta(j,i)}^{p\top} ] 여기서 Q^r_i는 행렬 Q^r의 i번째 행을 represent하고, K^r_j는 K^r의 j번째 행을 나타냅니다. Q^s_i와 K^s_j는 각각 Q^s와 K^s의 i번째 및 j번째 행입니다. 항 K^p_{\delta(i,j)}는 상대 거리 \delta(i, j)에 의해 index된 K^p 내의 행을 가리키고, Q^p_{\delta(j,i)}는 상대 거리 \delta(j, i)에 의해 index된 Q^p 내의 행을 가리킵니다. attention scores를 normalize하기 위해, scaling factor 1/\sqrt{5d}가 \hat{A}에 적용됩니다. 이 scaling은 model training의 안정성을 보장하는 데 중요하며, 특히 large-scale language models를 다룰 때 그렇습니다. 모든 \hat{A}_{ij}는 attention matrix를 형성하며, final output 잔기 hidden state는 R_o입니다: [ R_o = \text{softmax}(\frac{\hat{A}}{\sqrt{5d}})V^r, ] 이는 다음 layer의 hidden state에 대한 input으로 사용됩니다.

3.3 Pre-Training Objective

ProSST는 structure-conditioned masked language modeling으로 pre-trained됩니다. 이 접근 방식에서, 각 input sequence x는 잔기의 일부를 특별한 mask token 또는 다른 잔기로 대체하여 noised됩니다. ProSST의 objective는 손상된(corrupted) sequence와 그 structure token sequence s를 context로 활용하여 input sequence에서 noise 처리된 원래 tokens를 predict하는 것입니다: [ \mathcal{L}{MLM} = E{x \sim X} E_{M} \sum_{i \in M} -\log p(x_i | x_{/M}, s). ] 우리는 noise 처리 및 back-propagation을 위한 loss 계산을 위해 집합 M에서 무작위로 15%의 indices를 선택합니다. 선택된 각 index i에서, 잔기를 mask token으로 대체할 확률이 80%, 임의의 잔기 token으로 교체할 확률이 10%이며, 나머지 잔기는 변경되지 않습니다. training objective는 부분적으로 noise 처리된 sequence x_{/M}과 noise 처리되지 않은 structure tokens를 contextual cues로 삼아, 각 noise 처리된 잔기 x_i에 대한 음의 로그 우도(negative log-likelihood)를 최소화하는 것입니다. 따라서, noise 처리된 tokens를 정확하게 predict하기 위해, 이 objective는 model이 잔기 간의 dependencies뿐만 아니라 잔기와 구조 간의 관계도 학습할 수 있도록 합니다. pre-training dataset 및 hyper-parameter configuration의 세부 정보는 부록 A.2에서 찾을 수 있습니다.

Method (Section 3) 정리노트 (AI 연구자용)

• 전체 Architecture: 두 가지 주요 component로 구성됨: Structure Quantization Module (3.1) 및 Sequence-Structure Disentangled Attention을 사용하는 Transformer model (3.2).
• Structure Quantization Module (3.1):
  • 목표: 잔기의 국소 구조(local structure)를 discrete token으로 변환.
  • Pipeline: 국소 구조 → (pre-trained) Structure Encoder → dense vector → (pre-trained) k-means → 카테고리 레이블 → structure token.
  • 국소 구조 정의: 특정 잔기 + Cα 거리 기준 최근접 40개 이웃 잔기. Residue node(구조 정보만 포함)와 top-40 이웃 edge로 구성된 graph G=(V,E)로 표현.
  • Structure Encoder:
    • GVP (Geometric Vector Perceptron) 사용.
    • De-noising pre-training objective (좌표 교란 후 복구)를 사용하는 auto-encoder의 일부로 trained됨 (C.A.T.H dataset 사용).
    • 학습 후 encoder 부분만 사용. 국소 구조 graph g_i에 대한 encoder의 mean pooled output π_θ(g_i)가 dense vector r_i가 됨.
  • Local Structure Codebook:
    • C.A.T.H dataset의 모든 국소 구조를 pre-trained structure encoder로 embedding하여 생성.
    • 이 embeddings에 k-means algorithm을 적용하여 K개의 centroids {e_i} 식별. 이 centroids가 codebook임. (K: structure vocabulary size).
  • Quantization 프로세스: 새로운 국소 구조 g_i의 embedding r_i에 대해 codebook 내 가장 가까운 centroid e_j를 찾고, index j를 structure token s_i로 할당.
  • Serialization: 전체 단백질 구조는 structure tokens의 sequence s = s_1, ..., s_L로 변환됨.
• Sequence-Structure Disentangled Attention (3.2):
  • 영감: DeBERTa의 disentangled attention 기반.
  • 위치 i별 입력 Representation:
    • R_i: 잔기 token hidden state.
    • S_i: 잔기 수준 국소 structure token s_i의 embedding.
    • P_{i|j}: i와 j 간 relative position의 embedding.
  • 전체 Attention 분해: 위치 i, j 간 cross attention A_{i,j}는 9개 term으로 분해됨 (Eq. 1).
  • 제거된 Terms: 잔기의 contextual embeddings 학습에 집중하기 위해, 잔기 hidden state와 직접 관련 없는 term (S-S, S-P, P-S, P-P)은 implementation에서 제외됨.
  • 유지된 Attention Terms (5가지 유형): R-R, R-S, R-P, S-R, P-R (Eq. 4 Â_{i,j} 계산식에 반영됨).
  • 계산 (single-head 예시):
    • Residue(r), Structure(s), Position(p)에 대한 Query(Q), Key(K), Value(V) 행렬 계산 (Eq. 3 - 원문 K^p = SW^{k_p} 유지).
    • 5개 interaction term 합으로 attention score Â_{i,j} 계산 (Eq. 4).
    • Scaling (1/√{5d}) 및 softmax 적용하여 attention weight 계산.
    • Residue 값 V^r의 가중합으로 final output hidden state R_o 계산 (Eq. 5). R_o는 다음 layer의 input이 됨.
    • Relative position δ(i, j)는 L_{max}=1024로 계산 (Eq. 2).
• Pre-Training Objective (3.3):
  • 방법: Structure-conditioned masked language modeling.
  • 프로세스: 입력 잔기 sequence의 15%를 noise 처리 (80% mask, 10% random token, 10% 유지).
  • 목표: 손상된 잔기 sequence x_{/M}와 손상되지 않은 structure token sequence s를 context로 사용하여 원본 token x_i를 predict (Eq. 6: L_{MLM} = ... -log p(x_i | x_{/M}, s)).
  • 효과: Model이 잔기 간 dependencies 뿐만 아니라, context로 주어지는 s를 통해 잔기와 구조 간 관계를 학습하도록 강제함.

---

쉬운 설명 :

ProSST가 실제로 어떻게 작동하는지 설명하는 부분이에요. 크게 두 부분으로 나눌 수 있어요.

1. 단백질의 부분적인 모양 이해하기 (Structure Quantization Module)

• 목표: 단백질 각 부분의 복잡한 3D 모양을 간단한 '모양 코드'(discrete token)로 바꾸는 거예요.
• 1단계 (부분 모양 파악): 단백질의 각 부위(아미노산 잔기)마다 자기 자신과 주변에서 가장 가까운 40개 친구들을 함께 살펴서 그 부분의 '동네 환경'(국소 구조)이 어떤지 파악해요.
• 2단계 (모양 특징 배우기 - Structure Encoder): 미리 훈련된 똑똑한 AI(GVP encoder)를 사용해요. 이 AI는 약간 망가진 모양 정보를 보고 원래 모양을 복원하는 훈련(de-noising objective)을 통해, 다양한 부분적인 모양들의 특징을 잘 뽑아내도록 학습되었어요.
• 3단계 (모양 코드 사전 만들기 - Local Structure Codebook): 학습된 모양 특징 추출 AI(GVP encoder)로 데이터베이스에 있는 수많은 단백질의 부분 모양들을 분석해요. 비슷한 모양 특징을 가진 것끼리 그룹으로 묶는 k-means 방법을 써서 대표 모양 K개를 뽑아요. 이 K개의 대표 모양 설명이 바로 '모양 코드 사전'(codebook)이 돼요.
• 4단계 (모양 코드 부여 - Quantization): 이제 새로운 단백질의 어떤 부분 모양을 보면, 모양 특징 추출 AI로 분석한 다음, 만들어둔 모양 코드 사전에서 가장 비슷한 대표 모양 코드를 찾아요. 그 대표 코드의 번호(1~K 중 하나)가 이 부분의 최종 '모양 코드'(structure token)가 되는 거죠.
• 결과: 이렇게 하면 단백질 전체 모양이 s_1, s_2, s_3... 같은 간단한 모양 코드들의 나열(sequence)로 표현돼요.

2. 설계도와 모양 코드를 똑똑하게 연결하기 (Sequence-Structure Disentangled Attention)

• 목표: 단백질 설계도(아미노산 sequence, R) 정보와 위에서 만든 모양 코드(structure token, S) 정보, 그리고 각 부분의 위치(P) 정보를 AI 두뇌(Transformer)가 잘 엮어서 이해하도록 돕는 거예요.
• 특별한 집중 방식: 그냥 R, S, P 정보를 마구 섞는 대신, ProSST는 특별한 집중 방식(disentangled attention)을 써서 다음과 같은 관계들을 따로따로 살펴봐요.
  • 설계도의 한 부분과 다른 부분의 관계 (R-R)
  • 설계도와 다른 위치의 모양 코드 간의 관계 (R-S)
  • 설계도와 위치 간의 관계 (R-P)
  • 모양 코드가 설계도에 미치는 영향 (S-R)
  • 위치가 설계도에 미치는 영향 (P-R)
• (왜 다른 관계는 뺐을까?) 모양 코드끼리(S-S) 또는 위치와 모양 코드(P-S)의 관계는 최종적으로 설계도 정보를 업데이트하는 데 직접적으로 중요하지 않다고 판단해서 여기서는 제외했어요.
• 계산 방식: 이 5가지 종류의 관계 점수를 각각 계산한 다음 합쳐서, AI가 단백질의 각 부분을 이해할 때 다른 부분들의 설계도 정보와 모양 코드 정보를 얼마나 중요하게 참고할지 결정해요.

3. ProSST는 어떻게 학습할까? (Pre-Training Objective)

• 학습 게임: '빈칸 채우기' 게임(Masked Language Modeling)을 해요.
• 게임 방법: 단백질 설계도(sequence)를 가져와서 약 15%의 아미노산을 가려요 (대부분은 MASK 기호로 바꿈).
• 미션: AI(ProSST)에게 가려진 원래 아미노산이 무엇이었는지 맞추게 해요.
• 핵심 힌트: 정답을 맞출 때, AI는 가려진 설계도 정보뿐만 아니라, 가려지지 않은 완전한 모양 코드 순서(structure token sequence)도 함께 볼 수 있어요!
• 학습 효과: 빈칸을 잘 채우려면 AI는 설계도 단어들 사이의 관계뿐만 아니라, 설계도와 모양 코드 사이의 관계도 반드시 배워야 해요. 즉, 설계도와 모양 정보를 연결해서 이해하도록 훈련되는 거죠.