Diffusion : 논문리뷰 : HIERARCHICAL PROTEIN BACKBONE GENERATIONWITH LATENT AND STRUCTURE DIFFUSION

쉬운 설명

LSD는 단백질 구조를 한 번에 다 만들려고 하지 않고, 먼저 **LDM으로 단백질의 대략적인 '뼈대 설계도'(contact map)**를 그린 다음, 이 설계도를 바탕으로 SDM이 '실제 원자 위치'(atomic coordinates)를 정교하게 조립하는 2단계 방식을 사용합니다. 이렇게 단계를 나누니 복잡하고 다양한 단백질 데이터(AFDB)에서도 잘 학습하고, 특히 '설계도'를 그리는 단계에서 "더 안정적일 것 같은(low PAE) 설계도"나 "더 흥미로운(high LRC) 설계도"가 되도록 **방향을 유도(guidance)**할 수 있게 됩니다.

 

 

LSD (Latent and Structure Diffusion) 논문 학습 노트

Purpose of the Paper

• 기존 연구 한계 극복:
  • Structure Diffusion Models (SDMs)는 큰 dataset (예: AFDB)으로 확장 시 low quality structure로 인한 bias 문제 및 최적화/일반화 어려움 발생 (equivariance, physical constraints).
  • 기존 Latent Diffusion Models (LDMs)은 protein backbone generation에서 SDM 대비 성능 부족.
• 새로운 접근 방식 제시:
  • 단백질 생성을 계층적 (hierarchical) 으로 분해: Coarse-grained (contact map) 생성 후 Fine-grained (atomic coordinates) 생성.
  • Contact map을 중간 표현 (intermediate representation) 으로 활용하여 LDM과 SDM을 결합.
  • 이를 통해 대규모 dataset (AFDB) 학습을 용이하게 하고, latent space에서의 guidance를 통해 생성물의 속성 (designability, diversity, novelty) 제어 가능성 제시.

Key Contributions & Novelty

• Contribution:
  • LSD Framework: LDM (contact map 생성)과 SDM (atomic coordinate 생성)을 결합한 새로운 hierarchical generative model 제안.
  • Contact Map Generation: LDM을 사용하여 단백질의 coarse topology를 나타내는 contact map을 생성하는 단계 도입.
  • High-Level Guidance: Protein backbone generation 최초로 LDM latent space에서 PAE (Predicted Alignment Error) 및 LRC (Long Range Contacts) guidance 적용, 생성 단백질의 designability, diversity, novelty 제어 시연.
  • AFDB Scaling: 다양성이 높지만 quality가 낮은 구조가 포함된 AFDB dataset에서 성공적으로 학습 및 생성 능력 입증. FrameFlow 단독 사용 시 발생하는 mode collapse 문제 완화.
  • SSD Metric: 생성된 sample의 secondary structure 분포가 training data 분포와 얼마나 유사한지 측정하는 새로운 metric (Secondary Structure Distance) 제안.
• Novelty:
  • LDM-SDM Hierarchical Combination for Proteins: 단백질 구조 생성에 LDM (coarse topology)과 SDM (fine geometry)을 계층적으로 결합한 접근 방식 자체의 독창성.
  • Latent Space Guidance for Properties: 단백질의 latent space에서 직접적으로 PAE, LRC 같은 high-level property를 제어하기 위한 classifier-free guidance 메커니즘 적용.
  • Handling Diverse Datasets: Contact map 중간 단계를 통해 SDM 단독 사용보다 AFDB와 같은 크고 다양한 dataset을 더 효과적으로 학습하고 일반화할 수 있음을 보임.

Experimental Highlights

• Datasets: AlphaFold DataBase (AFDB) - Foldseek clustered, pLDDT > 80, length < 128 필터링 (282,936개 구조).
• Metrics: Designability (Des), Diversity (Div), Novelty (Nov), Designable Pairwise TM-score (DPT), H/S/C percentage, Secondary Structure Distance (SSD - 제안 metric).
• Baselines: FrameFlow (PDB/AFDB), GENIE2, ProteinSGM, ESM3, LatentDiff.
• Key Results:
  • LSD는 기존 LDM baseline인 LatentDiff 대비 모든 metric에서 월등한 성능 달성 (Table 2).
  • LSD는 AFDB dataset에서 학습 시 mode collapse (대부분 alpha-helix만 생성)를 보인 FrameFlow AFDB와 달리 다양한 secondary structure 생성 (Fig 2).
  • Guidance 효과 검증:
    • PAE guidance (LSD_PAE)는 Designability를 크게 향상시킴 (최대 94%, Table 10). W_PAE 증가 시 AF2 mean PAE 감소 확인 (Fig 11).
    • LRC guidance (LSD_LRC)는 Diversity, Novelty, strand 비율을 증가시킴 (최고 Nov 0.61, Table 10).
    • Joint guidance (LSD_J)는 Diversity를 극대화하며 다른 metric 간의 균형을 맞춤 (최고 Div 296, Table 10).
  • SOTA SDM인 GENIE2와 비교 시, LSD는 특정 metric (Div, Nov, SSD)에서 경쟁력 있는 성능을 보이며 생성 속도가 훨씬 빠름 (L=100 기준 0.3분 vs 2.3분). 단, GENIE2는 단일 설정으로 전반적인 SOTA 성능 달성.
  • 생성된 contact map (ĉ_ψ(z⁽⁰⁾))과 최종 구조의 contact map (c(x)) 간의 높은 일치도 확인 (ROCAUC 0.99, PRAUC 0.92), SDM이 contact map 정보를 효과적으로 활용함을 시사 (Fig 3).

Limitations and Future Work

• Limitations:
  • Single Setting Performance: SOTA SDM (GENIE2)과 달리, 단일 guidance 설정으로 모든 metric에서 동시에 최적 성능을 달성하기 어려움 (특정 metric 최적화를 위한 guidance tuning 필요).
  • Unconditional Generation: 현재 unconditional backbone generation에 국한됨.
  • Architectural Constraints: 계산 효율성을 위해 의도적으로 제외했지만, triangle update layer 같은 더 복잡한 architecture component를 사용하면 성능이 더 향상될 수 있음.
• Future Work:
  • Architecture/Model Improvement: Triangle update layer 도입, SDM 부분을 FrameFlow 대신 GENIE2 등 더 성능 좋은 모델로 교체, Autoguidance/Stable Diffusion 3 등 최신 LDM 기법 적용.
  • Task Extension: All-atom generation, conditional generation (binder design, motif scaffolding 등).
  • Scalability: 더 큰 단백질 및 복합체 (complex) 생성으로 확장.
  • Latent Space Enrichment: Side-chain 좌표 및 sequence 정보를 latent space에 포함하여 더 풍부한 생성 및 design 가능성 탐색.

Overall Summary

LSD는 단백질 backbone 생성을 위한 새로운 hierarchical generative model으로, LDM을 이용한 contact map 생성 (coarse)과 SDM을 이용한 atomic coordinate 생성 (fine) 단계를 결합한다. 이 접근법은 AFDB와 같은 크고 다양한 dataset으로의 확장을 용이하게 하며, latent space에서 PAE 및 LRC guidance를 통해 생성물의 designability, diversity, novelty를 제어할 수 있는 새로운 가능성을 열었다. 비록 특정 metric 최적화를 위해 guidance 조절이 필요하지만, 기존 LDM 성능을 크게 능가하고 SOTA SDM과 경쟁력 있는 결과를 보여주며, 제어 가능하고 확장성 있는 단백질 생성 연구에 중요한 기여를 한다.

 

 

ABSTRACT

우리는 coarse와 fine-grained details를 분리하는 hierarchical protein backbone generative model을 제안합니다.

LSD라고 불리는 우리의 approach는 두 개의 stages로 구성됩니다: contact map으로 decoded되는 latents를 sampling하고, 그 다음 contact map에 conditioned된 atomic coordinates를 sampling합니다.

LSD는 large datasets로 scaling하면서 바람직한 properties를 향해 protein generation을 제어하는 새로운 방법을 가능하게 합니다.

특히, AlphaFold DataBase (AFDB)는 다양한 structure topologies 때문에 매력적이지만, 낮은 designability로 어려움을 겪습니다.

우리는 AFDB에서 LSD를 train하고, AlphaFold2 Predicted Alignment Error와 long range contacts를 향한 latent diffusion guidance가 generated samples에서 designability, diversity, novelty를 명시적으로 균형 맞출 수 있음을 보여줍니다.

우리의 results는 structure diffusion models와 competitive하며, prior latent diffusion models보다 outperforms합니다.

 

 

이 이미지는 Figure 1: Overview로, 논문에서 제안하는 LSD (Latent and Structure Diffusion) 모델의 전체적인 작동 방식을 보여주는 그림입니다.

이 모델은 단백질 구조 generation 과정을 크게 두 단계(Coarse와 Fine)로 나눕니다.

1. Coarse 단계 (개략적인 단계):
   • 이 단계에서는 단백질의 개략적인 representation인 contact map을 생성합니다. Contact map은 단백질 잔기(residue)들이 서로 얼마나 가깝게 있는지를 나타내는 2D 지도와 같습니다.
   • 과정:
     • 임의의 Noise에서 시작합니다.
     • Latent diffusion 모델을 사용하여 Noise로부터 Latents (잠재 변수 Z)를 생성합니다. Latents는 contact map의 정보를 압축적으로 담고 있는 중간 representation입니다.
     • 생성된 Latents (Z)를 Decoder에 입력하여 최종적으로 contact map을 만들어냅니다.
   • 특징: 이미지 아래 캡션에서 설명하듯이, 이 Coarse 단계에서 guidance 기법을 적용하여 생성되는 contact map이 특정 high-level properties (예: 낮은 PAE, 높은 LRC 등 원하는 특성)를 갖도록 유도할 수 있습니다.
2. Fine 단계 (세부적인 단계):
   • 이 단계에서는 최종적인 3차원 원자 Structure (단백질 구조)를 생성합니다.
   • 과정:
     • Coarse 단계에서 생성된 contact map을 조건(condition)으로 입력받습니다. 즉, contact map이 제시하는 골격을 따라가도록 합니다.
     • 별도의 Noise에서 시작합니다.
     • Structure diffusion 모델 (텍스트에서 언급된 FrameFlow 같은 모델일 수 있음)을 사용하여, 입력받은 contact map 조건과 Noise로부터 최종적인 원자 수준의 3D Structure를 생성합니다.

요약:

결론적으로, 이 모델은 복잡한 3D 단백질 구조를 한 번에 직접 생성하는 대신,

1. 먼저 Latent diffusion을 이용해 제어 가능한 coarse representation인 contact map을 만들고 (Coarse 단계),
2. 그 다음, 이 contact map 정보를 바탕으로 Structure diffusion을 이용해 세부적인 fine 구조를 생성하는 (Fine 단계) hierarchical (계층적) 방식을 사용합니다.

이를 통해 large datasets에 대한 scalability를 높이고, 생성 과정에서 원하는 properties를 더 효과적으로 제어하려는 목표를 가지고 있습니다.

 

 

 

1 INTRODUCTION

protein backbone generation을 위한 diffusion models 전반의 과제는 large datasets로 scaling하는 것이었습니다: 이상적으로는 향상된 diversity와 generalization의 이점을 얻어야 하지만, 경험적으로 이는 low quality protein structures에서 비롯된 원치 않는 biases를 초래합니다. 본 연구에서는 원하는 properties를 제어할 수 있는 능력을 갖추고 AlphaFold DataBase (AFDB)로 scales하는 diffusion model을 개발하는 것을 목표로 합니다. 이전 approaches는 atomic coordinates에 대해 structure-based diffusion models (SDMs)를 사용하지만, 이는 bond lengths 및 angles와 같은 equivariance 및 physical constraints를 존중하는 데 어려움을 야기합니다. 불행히도, 이는 deep learning models의 optimization과 generalization을 저해할 수 있습니다 – structure prediction, conformer generation, material design에 대한 최근 연구들에서는 equivariance와 physical constraints를 제거함으로써 개선된 results를 발견했습니다.

우리는 sampled된 contact maps에 generation 과정을 conditioning함으로써 SDMs를 개선할 수 있다고 가정합니다. contact map은 각 residues 쌍이 서로 짧은 distance 내에 있는지 여부를 나타내는 2D binary matrix로 정의됩니다. contact map은 단백질의 fold topology를 설명하기에 충분하며, coordinates는 biomolecular conformations를 포착하고 단백질의 function을 설명할 수 있습니다. 이전 연구들은 generative models이 의미 있는 representations을 학습하는 데 실패하며 alignment 또는 high-level conditioning으로부터 이점을 얻을 수 있음을 보여주었습니다. 우리는 후자의 직관에 따라 contact maps를 generate하기 위해 latent diffusion models (LDMs)를 사용하고, 이어서 contact map에 conditioned된 atomic coordinates를 generate하기 위해 SDMs를 사용하는 것을 제안합니다. 우리의 방법인 LSD (Latent and Structure Diffusion)는 large datasets에 대한 training을 가능하게 하고, 그 다음 원하는 properties를 위해 guiding할 수 있도록 합니다.

우리의 주요 technical contribution은 contact map generation을 위한 LDMs 개발입니다. 우리는 protein structures를 L × K 차원의 latent space로 encode합니다. 여기서 L은 단백질 길이이고 K는 latent dimensionality입니다. 그 다음 decoder는 주어진 latents로부터 contact map을 reconstructs합니다. 다음으로, 우리는 학습된 latent space에서 sample하기 위해 수정된 Diffusion Transformer (DiT)를 train하고, FrameFlow는 contact maps와 latents에 conditioned된 atomic coordinates를 generate하도록 trained됩니다. 마지막으로, 우리는 개선된 designability와 diversity를 위해 더 낮은 평균 Predicted Alignment Error (PAE)와 증가된 Long Range Contacts (LRC)를 향하도록 latent generation을 guiding함으로써 generated structures에 대한 high-level control을 시연합니다.

우리는 FrameFlow만으로는 학습하기 어렵다는 것을 보여주는 AlphaFold Database (AFDB)에서 training함으로써 protein backbone generation에 대해 LSD를 evaluate합니다. 우리의 results는 LSD가 더 diverse한 fold topologies로의 generalization을 개선함을 보여주며, 이는 LDMs와 SDMs를 결합하는 것이 AFDB와 같은 large datasets로 scaling하는 데 유익할 수 있음을 시사합니다. 우리는 LSD를 PAE 또는 LRC guidance와 함께 사용할 때 개선된 results를 달성하며, 이는 LDMs를 사용한 high-level protein generation control의 novel capability를 보여줍니다. LSD는 AFDB에서 SDMs와 competitive하며 protein backbone generation을 위한 기존 LDMs보다 outperforms합니다. 우리의 contributions는 다음과 같이 요약됩니다:

1. 우리는 contact map generation을 위해 LDMs를 사용하고 atomic coordinate generation을 위해 SDMs를 사용하는 novel hierarchical protein generative model인 LSD를 개발합니다.
2. 우리는 protein backbone generation을 위한 high-level guidance의 첫 번째 사례를 시연합니다. PAE guidance는 designability를 향상시키는 반면 LRC guidance는 diversity와 novelty를 향상시킵니다.
3. AFDB에서 trained되었을 때, LSD는 AFDB에 대한 state-of-the-art SDMs와 competitive하며, protein backbone generation을 위해 유일하게 공개적으로 사용 가능한 LDMs보다 outperforms합니다.

 

연구자를 위한 정리 노트: 1 INTRODUCTION 섹션

LSD: Hierarchical Diffusion for Scalable & Controllable Protein Backbone Generation (Intro Summary)

• Problem: 기존 protein backbone generation diffusion models (특히 atomic coordinates 기반 SDMs)는 AFDB 같은 large datasets로의 scaling 및 generalization에 어려움을 겪음. 이는 low quality 데이터로 인한 biases와 equivariance 등 physical constraints 처리의 복잡성에서 기인함. (오히려 constraint 제거가 성능 향상에 도움이 된다는 최근 연구 동향 언급)
• Proposed Model (LSD): Coarse (topology) 생성과 Fine (structure) 생성을 분리하는 새로운 hierarchical generative model 제안.
  • Stage 1 (Coarse): Latent Diffusion Model (LDM) (modified DiT 기반)을 사용하여 contact map 생성. 단백질의 fold topology를 latent space에서 학습하고 decode함.
  • Stage 2 (Fine): Structure Diffusion Model (SDM) (FrameFlow 사용)을 사용하여 Stage 1의 contact map에 conditioned된 atomic coordinates 생성.
• Key Innovations:
  1. Hierarchical (LDM → SDM) Architecture: Large datasets (AFDB) scaling 및 generalization 향상 목표. Coarse (contact map) representation을 중간 단계로 사용하여 노이즈/저품질 데이터에 대한 강건성 확보 및 fine-grained 생성의 conditioning 강화.
  2. High-Level Guidance via Latent Diffusion: Stage 1 (latent generation)에서 guidance를 적용하여 최종 생성물의 properties 제어.
     • PAE guidance → Designability 향상.
     • LRC guidance → Diversity & Novelty 향상.
     • 논문 주장: Protein backbone generation 분야에서 이러한 high-level guidance의 첫 시연.
• Main Claims (Intro):
  • LSD는 AFDB에서 generalization 성능 개선.
  • 제안된 guidance 메커니즘은 designability/diversity 제어에 효과적.
  • AFDB에서 state-of-the-art SDMs와 competitive한 성능.
  • 기존 공개 LDMs 대비 outperform.

---

쉬운 설명 : 1 INTRODUCTION 섹션

이 연구는 AI를 사용해 새로운 단백질 구조(모양)를 효과적으로 만드는 방법에 대한 것입니다.

기존 방법의 어려움: 이전 AI 모델들은 단백질의 3D 구조를 한 번에 만들려고 했습니다. 하지만 알파폴드 데이터베이스(AFDB)처럼 아주 방대한 양의 실제 단백질 데이터를 학습시키려고 하면, 데이터 안에 품질이 낮은 정보도 섞여있어서 AI가 원치 않는 방향으로 학습하거나, 원자 하나하나의 물리적 규칙(결합 길이, 각도 등)을 정확히 따르도록 만들기가 매우 복잡하고 오히려 성능을 떨어뜨리기도 했습니다.

이 논문의 새로운 아이디어 (LSD 모델): 복잡한 단백질 생성 과정을 두 단계로 나누어 접근합니다.

1. 먼저, 대략적인 설계도 그리기 (Coarse 단계): 단백질의 전체적인 형태, 즉 어떤 부분들이 서로 가깝게 붙어있는지를 나타내는 **contact map(접촉 지도)**이라는 간단한 2D 설계도를 먼저 그립니다. 이때 Latent Diffusion이라는 AI 기술을 사용합니다. 이 단계의 핵심은, "이런 특성을 가진 단백질을 만들고 싶다" (예: 실제로 만들기 쉬운 구조(designability 향상), 혹은 아주 새로운 구조(diversity/novelty 향상))는 목표를 AI에게 알려주어(guidance) 설계도 생성 방향을 유도할 수 있다는 점입니다.
2. 다음, 상세 모델링 (Fine 단계): 1단계에서 만든 설계도(contact map)를 바탕으로, Structure Diffusion이라는 다른 AI 기술을 사용해 실제 원자들의 정확한 3D 위치를 결정하여 최종 단백질 구조를 완성합니다.

왜 이 방법이 더 나은가?

• 복잡한 문제를 두 단계로 나누어 AI가 더 쉽게 학습하고 생성할 수 있도록 합니다.
• 방대한 데이터(AFDB)를 더 잘 활용하고, 데이터 내의 노이즈에 덜 민감하게 만들 수 있습니다 (일단 중요한 골격 정보인 contact map에 집중하므로).
• 특히 첫 단계에서 guidance를 통해 최종 단백질의 중요한 특성들을 조절할 수 있는 새로운 제어 능력을 제공합니다.

결론: 이 LSD라는 새로운 2단계 방식은 큰 데이터셋에서도 잘 작동하며, 특히 guidance를 통해 원하는 특성을 가진 단백질을 만들도록 제어하는 능력을 보여줍니다. 성능 면에서도 기존의 좋은 모델들과 견줄 만합니다.

 

---

2 BACKGROUND: LATENT DIFFUSION MODELS

Latent Diffusion Models (LDMs)는 두 가지 components를 사용합니다: data를 latent space에 embed하기 위한 autoencoder와, latent space로부터 samples를 generate하여 decoder를 통해 다시 data로 decoded하는 diffusion model입니다. 이 섹션에서는 두 components에 대한 background를 제공합니다.

Notation. 괄호가 있는 Superscript는 time을 나타내고 subscripts는 matrix 또는 vector의 index를 나타내는 데 사용됩니다. 즉, ( x^{(t)}_i )는 diffusion process에서 index i와 time t에서의 data입니다. Scalars와 probability distributions는 time에 대해 subscripts를 사용합니다.

2.1 AUTOENCODER: MAPPING DATA TO LATENT SPACE

Autoencoders는 data ( x \sim p_{\text{data}} )를 latent variable z로 maps하는 encoder ( p_{\phi}(z|x) )와 z를 다시 x로 maps하는 decoder ( p_{\psi} (x|z) )로 구성됩니다. encoder와 decoder는 각각 weights ϕ와 ψ를 가진 neural networks로 parameterized됩니다. 우리는 variational upper bound를 minimizing하여 trained되는 Variational AutoEncoder (VAE)를 사용합니다: [ E_{p_{\text{data}}(x)} \left[ E_{p_{\phi}(z|x)} \left[- \log p_{\psi}(x|z) \right] \right] \underset{\text{Reconstruction}}{\underbrace{+ \lambda KL \left[ p_{\phi}(z|x) || \mathcal{N}(0, I) \right]}} \underset{\text{Regularization}}{\underbrace{}} ] 여기서 λ는 Kullback-Leibler divergence (KL)에 대한 regularization weight입니다. 추가적인 regularization terms가 활용되기도 했습니다. 우리는 향후 연구에서 더 많은 regularizations를 탐색할 계획입니다.

2.2 DIFFUSION MODEL OVER LATENT SPACE

autoencoder가 trained되면, latent diffusion process의 target distribution으로 ( p^{(0)} = p_{\phi} )를 정의합니다. 다음으로, diffusion model은 identity I를 가진 Gaussian noise ( z^{(1)} \sim \mathcal{N}(0, I) )로부터의 samples를 latents ( z^{(0)} \sim p_0 )로 map하는 것을 학습하여 data를 generate하도록 trained되며, 이는 다시 data ( p_{\psi} (x|z^{(0)}) )로 decoded됩니다. time variable ( t \in [0, 1] )은 time-dependent process에 기반하여 noise와 latents 간의 mapping을 제어합니다: [ z(t) = \alpha(t) z^{(0)} + \sigma(t) z^{(1)}. \quad (2) ] ( \alpha(t) )와 ( \sigma(t) )에 대한 많은 선택지가 제안되었지만, 우리는 flow models에서 널리 사용되는 간단한 linear interpolation을 사용합니다: ( \alpha(t) = 1 - t ), ( \sigma(t) = t ). 식 (2)의 conditional distribution과 score는 다음과 같이 analytically computed될 수 있습니다: [ \nabla_{z(t)} \log q_t(z(t) | z^{(0)}) = \frac{z(t) - \alpha(t) z^{(0)}}{\sigma(t)^2} \quad \text{where} \quad q_t(z(t) | z^{(0)}) = \mathcal{N}(z(t); \alpha(t) z^{(0)}, \sigma(t)^2 I). \quad (3) ] 우리는 ( q_{t=1}(\cdot|z^{(0)}) = p_1(\cdot) )와 ( q_{t=0}(\cdot|z^{(0)}) = p_0(\cdot) )를 요구합니다. 이는 score matching objective를 통해 marginal score ( \nabla_{z(t)} \log p_t(z(t)) )를 근사하는 학습을 가능하게 합니다. 동등하게, 우리는 denoising autoencoder objective를 최적화합니다: [ E_{p_0} \left[ z^{(0)} | z^{(t)} \right] \approx \arg \min_{\hat{z}{\theta}} E{q_t(z^{(t)} | z^{(0)})} \left[ \frac{1}{\sigma(t)^2} | \hat{z}{\theta}(z^{(t)}, t) - z^{(0)} |_2^2 \right] U(t; 0, 1) p_0(z^{(0)}) \quad (4) ] 여기서 U는 uniform distribution이고 ( 1/\sigma(t)^2 )는 t에 걸쳐 동등한 loss weighting을 장려하기 위한 weighting term입니다. ( \hat{z}{\theta} )는 true latents를 predict하도록 trained된 weights θ를 가진 neural network입니다. 식 (3)을 사용하여, marginal score는 ( s_{\theta} (z^{(t)}, t) = \frac{z^{(t)} - \alpha(t)\hat{z}_{\theta} (z^{(t)}, t)}{\sigma(t)^2} )로 approximated될 수 있습니다. score가 learned되면, ( z^{(1)} )에서 시작하여 ( z^{(0)} )로 향하는 reverse SDE를 integrating하여 samples ( z^{(0)} )를 얻을 수 있습니다: [ dz^{(t)} = \left[ a(t)z^{(t)} - b(t)^2 s(z^{(t)}; t) \right] dt + \gamma \cdot b(t) dw^{(t)} \quad (5) ] 여기서 ( w^{(t)} )는 Wiener process이고, ( a(t) = \frac{\partial}{\partial t} \log \alpha(t) )이며, ( b(t)^2 = 2\sigma(t)\left( \frac{\partial\sigma(t)}{\partial t} - a(t)\sigma(t) \right) )입니다. γ는 noise의 variance를 제어하기 위한 scale이며 기본적으로 γ = 1로 설정됩니다. 이전 연구들에서는 sample quality를 개선하기 위해 γ < 1로 설정하는 것이 효과적임을 발견했습니다. 본 연구에서는 모든 samples에 대해 Euler-Maruyama integrator를 사용합니다. 우리는 식 (5)가 식 (2)의 interpolant에 대한 reverse SDE임을 App. A에서 derivation합니다.

 

 

연구자를 위한 정리 노트: 2 BACKGROUND: LATENT DIFFUSION MODELS 섹션

LDM Background Summary (Section 2 - Specifics for This Paper)

• Framework: 표준적인 Latent Diffusion Model (LDM) 설정을 따르며, 두 가지 주요 components 활용:
  1. Autoencoder: Variational AutoEncoder (VAE) (Vahdat et al., 2021 참고)를 사용하여 data (문맥상 단백질 구조 관련 정보)를 latent space ( z )로 매핑하고 다시 복원. 표준 ELBO (Reconstruction + KL Regularization)를 통해 trained. 추가적인 regularization에 대한 향후 연구 계획 언급됨.
  2. Diffusion Model (on Latent Space ( z )): Encoder에 의해 생성된 latents의 분포 ( p_{\phi} )를 학습.
     • Forward Process: Latents ( z^{(0)} )와 noise ( z^{(1)} ) 사이의 선형 보간(linear interpolation) 스케줄 채택: ( z(t) = (1-t) z^{(0)} + t z^{(1)} ). 이는 flow models에서 대중화된 방식임을 언급.
     • Training: ( 1/\sigma(t)^2 = 1/t^2 ) loss weighting을 사용하는 denoising autoencoder objective ( ( \hat{z}_{\theta} )를 통해 ( z^{(t)} )로부터 ( z^{(0)} ) 예측)를 최적화.
     • Sampling: 선택된 선형 보간 스케줄에 해당하는 reverse SDE (App. A에서 유도)를 Euler-Maruyama integrator를 사용하여 적분함으로써 sampling 수행. Sample quality 향상을 위해 noise scale γ (< 1) 튜닝 가능성 언급.
• 핵심 요약: 표준 LDM 구성 요소를 사용하지만, VAE 공식, diffusion 보간 스케줄(선형, flow-model 스타일), denoising 타겟(( z^{(0)} )), sampling integrator (Euler-Maruyama) 등 구체적인 선택 사항들을 명시함.

---

쉬운 설명 : 2 BACKGROUND: LATENT DIFFUSION MODELS 섹션

이 섹션은 이 논문에서 사용하는 핵심 AI 기술 중 하나인 Latent Diffusion Models (LDMs)가 어떻게 작동하는지 설명합니다. LDM은 새로운 데이터를 생성하는 (예: 이미지나 이 논문에서는 단백질 관련 정보) 강력한 방법입니다.

LDM은 크게 두 부분으로 구성됩니다:

1. Autoencoder (오토인코더): 데이터를 압축하고 복원하는 역할
   • Encoder (압축기): 원본 데이터(예: 단백질 구조 정보)를 받아서 훨씬 작은 크기의 핵심 정보, 즉 **latent 코드(잠재 코드)**로 압축합니다. 데이터의 본질적인 특징만 남기는 과정과 같습니다. (이 논문에서는 VAE라는 특정 종류의 오토인코더를 사용합니다).
   • Decoder (복원기): 압축된 latent 코드를 다시 원래 데이터 형태로 복원합니다.
   • 학습: Autoencoder는 원본 데이터를 잘 압축했다가 다시 최대한 원본과 비슷하게 복원하도록 학습시킵니다.
2. Diffusion Model (디퓨전 모델): 새로운 '압축된 코드'를 생성하는 역할
   • 목표: Autoencoder가 만들어내는 실제 latent 코드와 아주 유사한, 새로운 latent 코드를 만들어내는 방법을 학습하는 것입니다.
   • 학습 방법 (잡음 추가 과정 - Forward Process): 실제 latent 코드에 점진적으로 노이즈(무작위 잡음)를 추가합니다. 처음에는 조금, 나중에는 많이 추가해서 결국 완전한 노이즈 상태로 만듭니다. (수식 ( z(t) = \alpha(t) z^{(0)} + \sigma(t) z^{(1)} )는 시간 t에 따라 원본 코드 ( z^{(0)} )와 노이즈 ( z^{(1)} )를 섞는 과정을 보여줍니다.) AI는 각 단계에서 노이즈가 얼마나 섞였는지를 보고 원래 코드를 예측하도록 학습합니다.
   • 생성 방법 (잡음 제거 과정 - Reverse Process): 학습이 끝나면, 이제 반대로 완전한 노이즈에서 시작합니다. AI는 학습한 내용을 바탕으로 점차 노이즈를 제거해 나갑니다. 매 단계에서 '이 정도 노이즈를 제거하면 원본 코드에 더 가까워질 거야'라고 예측(이 예측값을 score라고 하며, ( \hat{z}_{\theta} )라는 신경망 모델이 이 역할을 함)하며 노이즈를 걷어냅니다. 이 과정을 계속 반복하면 최종적으로 그럴듯한 새로운 latent 코드가 생성됩니다. (수학적으로는 reverse SDE를 풀거나 Euler-Maruyama integrator 같은 방법을 사용해 이 단계를 진행합니다.)

전체 과정 요약:

새로운 단백질 관련 데이터(여기서는 contact map)를 만들고 싶을 때,

1. Diffusion Model을 사용해 (노이즈 제거 과정을 통해) 새로운 latent 코드를 생성합니다.
2. 이 생성된 latent 코드를 미리 학습된 Autoencoder의 Decoder 부분에 넣습니다.
3. Decoder가 이 코드를 해석하여 최종적인 새로운 데이터(예: 단백질 contact map)를 만들어냅니다.

이렇게 하면 복잡한 원본 데이터를 직접 다루는 대신, 더 다루기 쉬운 latent 코드 공간에서 생성 작업을 수행하므로 더 효율적이고 안정적으로 새로운 데이터를 만들 수 있습니다.

 

 

 

3 METHOD: LATENT AND STRUCTURE DIFFUSION (LSD)

이 섹션에서는 hierarchical protein backbone generation을 위한 우리의 방법을 제시합니다. 이 방법은 세 가지 components로 구성됩니다: structure-to-contact autoencoder (Sec. 3.1), autoencoder latent space에서 latents를 sample하는 LDM, 그리고 sampled latents로부터 structures를 sample하는 SDM (Sec. 3.2)입니다. Training과 sampling에 대해서는 Sec. 3.3에서 논의합니다. 마지막으로, Sec. 3.4에서는 PAE 및 LRC guidance에 대해 논의합니다.

3.1 STRUCTURE-TO-CONTACT AUTOENCODER

단백질 backbone의 atomic coordinates를 ( x \in \mathbb{R}^{L \times 3 \times 3} )로 표기합니다. 여기서 L은 단백질의 길이(residues의 수)이며, 3은 각 residue의 Nitrogen, Carbonα, Carbon atom에 해당합니다. encoder로는 ProteinMPNN architecture를 사용하여 x를 latent ( p_{\phi}(x) = z \in \mathbb{R}^{L \times K} )로 embed합니다. 여기서 K는 latent dimensionality이며 ( K \ll L )입니다. 우리는 각 latent의 mean과 variance Gaussian parameters를 출력하도록 ProteinMPNN을 수정합니다.

우리는 latent space에서 x의 coarse representation을 학습하는 것을 목표로 합니다. 우리의 approach는 decoder가 contact map ( c \in {0, 1}^{L \times L} )를 predict하도록 train하는 것입니다. 여기서 ( x_i, x_j )의 Carbonα atoms 간의 distance가 8 Å 미만이면 ( c_{ij} = 1 )이고, 그렇지 않으면 ( c_{ij} = 0 )입니다. Decoder ( p_{\psi} )는 latents의 Kronecker product를 취하고 contact map probabilities를 predict합니다: 모든 i, j에 대해 ( p_{\psi}(z_i \otimes z_j) ). 우리는 ( p_{\psi} )를 ReLU activations를 가진 3개 layer의 multi-layer perceptron으로 parameterize합니다. Kronecker product ( z_i \otimes z_j \in \mathbb{R}^{K \times K} )는 ( z_i )와 ( z_j )의 entries 간 모든 가능한 products의 matrix입니다. Encoder와 decoder는 eq. (1)을 사용하여 trained되는데, reconstruction term은 ground truth contacts의 negative log-likelihood를 minimize하도록 수정됩니다: [ E_{p_{\phi}(z|x)} \left[ \frac{1}{|Z_0|} \sum_{(i,j) \in Z_0} -\log p_{\psi}(c_{ij}=0 | z_i \otimes z_j) + \frac{1}{|Z_1|} \sum_{(i,j) \in Z_1} -\log p_{\psi}(c_{ij}=1 | z_i \otimes z_j) \right] ] 여기서 ( Z_0 = {(i, j) : c_{ij} = 0} )와 ( Z_1 = {(i, j) : c_{ij} = 1} )은 각각 ( c_{ij} = 0 )과 ( c_{ij} = 1 )인 indices의 집합입니다. contacts의 sparsity 때문에, 우리는 각 class를 해당 propensity에 따라 별도로 weight합니다. 약식 표기를 위해, 전체 decoded contact probabilities를 ( \hat{c}{\psi}(z) \in [0, 1]^{L \times L} )로 지칭할 것이며, 여기서 ( \hat{c}{\psi}(z){ij} = p{\psi}(z_i \otimes z_j) )입니다.

3.2 LATENT AND STRUCTURE DIFFUSION MODELS

LDM은 latent score function ( s_{\theta}(z^{(t)}, t) )를 학습해야 합니다. 여기서 ( z^{(t)} )는 eq. (2)에서 정의된 encoded latents ( z^{(0)} = p_{\phi}(x) )의 noisy version입니다. Neural network architecture로는 computer vision에서 성공적으로 사용되었기 때문에 Diffusion Transformer (DiT)를 사용합니다. 우리는 각 residue latent ( z^{(t)}_i )를 token으로 취급하여 우리의 목적에 맞게 DiT를 adapt합니다. absolute positional encodings 대신 Rotary Positional Encodings (RoPE)를 사용합니다.

SDM은 FrameFlow의 수정된 버전으로, denoised atomic coordinates ( \hat{x}{\varphi}(x^{(t)}, \hat{c}{\psi}(z^{(0)}), z^{(0)}, t) )를 predict하도록 trained됩니다. 여기서 ( \varphi )는 FrameFlow neural network weights입니다. 우리는 ( z^{(0)} )와 ( \hat{c}_{\psi}(z^{(0)}) )를 FrameFlow에 제공되는 초기 1D 및 2D features 세트에 concatenating하여 SDM을 condition합니다. FrameFlow는 SE(3) 상에서의 flow matching을 사용하며, 이는 diffusion models의 probabilistic Ordinary Differential Equation (ODE) perspective와 동일합니다. 우리는 FrameFlow의 training 및 sampling procedure를 따르되, 우리의 latent conditioning을 추가하고, self-conditioning과 rotation annealing은 사용하지 않습니다. 우리는 open sourced code의 용이성 때문에 다른 SDMs보다 FrameFlow를 선택했지만, 우리의 framework에서는 다른 어떤 SDM도 사용할 수 있습니다. Architecture choice로 인해, 우리의 LDM은 SE(3) invariant인 반면 SDM은 SE(3) equivariant입니다.

3.3 MULTI-STAGE TRAINING AND SAMPLING

VAE training loss는 Sec. 3.1에 설명되어 있고, LDM 및 SDM losses는 Sec. 3.2에 설명되어 있습니다. 모든 models의 end-to-end training이 가능하지만, 우리는 이것이 unstable하고 optimize하기 어렵다는 것을 발견했습니다. 대신, 우리는 latent diffusion training 동안 autoencoder가 frozen되는 방식의 training procedure를 사용합니다. 이는 세 가지 stages를 포함합니다: (1) VAE 자체를 pre-training하기, (2) VAE와 SDM을 jointly training하기, (3) VAE weights를 freezing하고 LDM만 training하기. 우리의 experiments에서는 모든 stages에 걸쳐 동일한 optimizer와 learning rate를 사용합니다. App. B.1에서 training setup에 대한 자세한 내용을 제공합니다.

Sample하기 위해, 우리는 먼저 LDM으로부터 eq. (4)를 사용하여 latents ( z^{(0)} )를 generate하고 contact map ( \hat{c}{\psi}(z^{(0)}) )를 얻습니다. 그런 다음 ( z^{(0)} )와 ( \hat{c}{\psi}(z^{(0)}) ) 둘 다 SDM에 제공되어 SE(3) flow를 사용하여 ( \hat{c}_{\psi}(z^{(0)}) )에 conditioned된 atomic coordinates를 sample합니다. Fig. 1은 sampling process를 보여줍니다. 다음 섹션에서는 원하는 high-level properties를 향한 guidance를 사용한 sampling에 대해 설명합니다.

3.4 PAE AND LRC GUIDANCE

각 latent ( z^{(0)} )와 관련된 class label과 같은 property를 y라고 표기합시다. 각 noised latent ( z^{(t)} )로부터 property를 predict하는 classifier를 train한 다음, 이를 사용하여 LDM을 원하는 class label로 guide하는 방법이 제안되었습니다. 이는 Bayes rule을 사용하여 property conditioned score를 approximate함으로써 달성됩니다: [ \nabla_{z^{(t)}} \log p_t(z^{(t)}|y) = \nabla_{z^{(0)}} \log p_t(z^{(t)}) + \nabla_{z^{(t)}} \log p_t(y|z^{(t)}) \approx s_{\theta} (z^{(t)}; t) + s(z^{(t)}; t) \quad (5) ] 여기서 ( s(z^{(t)}; t) )는 ( \nabla_{z^{(0)}} \log p_0(y|z^{(t)}) )를 approximate하도록 parameterized됩니다. 그런 다음 우리는 eq. (5)를 score로 eq. (8)에 대입하여 ( p(z^{(0)}|y) )로부터 대략적으로 sample합니다: [ dz^{(t)} = \left[ a(t)z^{(t)} - b(t)^2 \left( s_{\theta} (z^{(t)}; t) + s(z^{(t)}; t) \right) \right] dt + \gamma \cdot b(t)dw^{(t)}. \quad (6) ] 우리는 PAE와 long range contacts를 향해 guiding하기 위한 ( s(z^{(t)}; t) )의 여러 옵션을 설명합니다.

Long range contact (LRC). Protein generative models은 종종 alpha-helical structures를 주로 선호하는 경향을 보이는데, 이는 protein datasets에서 alpha-helices의 prevalence 때문이며, 이는 생성된 fold topologies의 diversity를 제한할 수 있습니다. 이 bias를 해결하기 위해, 우리는 decoder의 contact map predictions ( p_{\psi} )를 활용하여 더 많은 LRCs를 향한 guidance를 사용합니다. LRC는 sequence distance가 12보다 큰 contact ( c_{ij} )로 정의됩니다. ( Z_{LR} = {(i, j) : |i - j| > 12} )를 sequence distance가 12보다 큰 pairwise indices의 집합이라고 합시다. ( y = {c_{ij} = 1 \forall(i, j) \in Z_{LR}} )를 LRC property로 사용하여, LRC guidance score를 다음과 같이 정의합니다: [ s_{\text{LRC}}(z^{(t)}; t) = e^{-r_{\text{LRC}} \cdot (1-t)} \cdot \nabla_{z^{(t)}} \left[ \frac{1}{|Z_{LR}|} \sum_{(i,j) \in Z_{LR}} \log p_{\psi} \left( c_{ij} = 1 | \hat{z}{\theta} (z^{(t)}, t)i \otimes \hat{z}{\theta} (z^{(t)}, t)j \right) \right] ] 여기서 ( r{\text{LRC}} \in \mathbb{R} )은 score coefficient의 decay를 제어하는 hyperparameter입니다. 더 많은 LRCs를 향한 guidance를 수행하기 위해, eq. (6)에서 ( s(z^{(t)}; t) ) 대신 ( s{\text{LRC}}(y, z^{(t)}; t) )를 대입합니다.

Predicted Alignment Error (PAE). 낮은 PAE는 종종 protein design success와 correlated되므로, 우리는 낮은 PAE를 가진 protein structures를 generating하는 데 관심이 있습니다. 우리의 주요 metric인 designability 또한 낮은 PAE와 correlates됩니다 (Fig. 10 참조). 우리는 각 latent noised latent ( z^{(t)} )에 해당하는 structure x의 global average PAE를 predict하도록 weights ( \vartheta )를 가진 neural network ( f_{\vartheta}(\hat{z}{\theta}(z^{(t)}, t), t) )를 train합니다. Data curation, architecture, training details는 App. B.2를 참조하십시오. PAE는 scalar value이므로, classifier guidance가 직접적으로 작동하지 않습니다. 대신, 우리는 ( p_t(y|z^{(t)}) )를 낮은 PAE predictions에 높은 probability를 할당하는 Boltzmann distribution으로 대체합니다: ( p^{\text{PAE}}t(z^{(t)}) \propto e^{-\omega{\text{PAE}} \cdot f{\vartheta}(\hat{z}{\theta}(z^{(t)},t),t)} ) (weight ( \omega{\text{PAE}} \in \mathbb{R} )). 우리는 PAE guidance score를 다음과 같이 정의합니다: [ s_{\text{PAE}}(z^{(t)}, t) = \nabla_{z^{(t)}} \log p^{\text{PAE}}t(z^{(t)}) = -\omega{\text{PAE}} \nabla_{z^{(t)}} f_{\vartheta}(\hat{z}{\theta}(z^{(t)}, t), t). ] 이것이 principled하지는 않지만, 우리는 ( s{\text{PAE}} )가 낮은 PAE를 향해 guiding하는 데 intuitive하며 실제 practice에서 잘 작동함을 발견했습니다. 낮은 PAE를 향한 guidance를 수행하기 위해, eq. (6)에서 s 대신 ( s_{\text{PAE}} )를 대입하고 experiments에서 다른 ( \omega_{\text{PAE}} ) 값들에 대해 sweep합니다.

Joint guidance. PAE를 낮추는 것은 개선된 designability와 correlated되지만 diversity와 novelty를 희생시킬 수 있습니다. 반면에, LRCs를 증가시키는 것은 더 diverse하고 novel한 protein structures를 낳지만 designability를 해칠 수 있습니다. 따라서 둘 다 결합하는 것이 합리적입니다. 우리는 두 properties를 모두 향해 guiding할 때 다음 score를 사용할 것입니다: [ s_J(z^{(t)}; t) = s_{\text{PAE}}(z^{(t)}; t) + s_{\text{LRC}}(z^{(t)}; t) ] 여기서 우리는 여전히 각 guidance의 strength를 제어하기 위한 두 개의 hyperparameters ( \omega_{\text{PAE}} )와 ( r_{\text{LRC}} )를 가집니다.

쉬운 설명 : 3 METHOD: LATENT AND STRUCTURE DIFFUSION (LSD) 섹션

이 섹션은 LSD 모델이 실제로 어떻게 단백질 구조를 만들어내는지 그 방법을 자세히 설명합니다. 크게 3가지 부분과 특별한 '조종' 기능(guidance)으로 구성되어 있어요.

1. 설계도(Contact Map)를 그리는 오토인코더 (Structure-to-Contact Autoencoder)

• 먼저, AI(Encoder, ProteinMPNN 기반)가 주어진 단백질의 3D 구조(원자 위치 x)를 보고, 그 구조의 핵심 특징만 담은 **압축 코드(latent (z))**를 만듭니다.
• 다른 AI(Decoder, MLP 기반)는 이 압축 코드 (z)를 입력받아, 단백질의 **설계도(contact map (c))**를 그리는 법을 배웁니다. 이 설계도는 단백질의 어떤 부분들이 서로 가까이 접혀 있는지를 보여주는 2D 지도 같은 거예요. (기술적으로는 압축 코드 조각들의 Kronecker product라는 정보를 이용해 그립니다.)
• 이 오토인코더는 3D 구조 → 압축 코드 → 설계도 과정을 잘 수행하도록 학습됩니다.

2. 새로운 설계도 코드 생성(LDM) 및 최종 구조 조립(SDM)

• LDM (Latent Diffusion Model): 이 AI(DiT 아키텍처 사용)는 앞서 학습한 '진짜' 압축 코드 (z)와 아주 유사한 새로운 압축 코드 (z^{(0)})를 생성하는 역할을 합니다. (이전 섹션에서 설명한 노이즈 제거 방식으로 만듭니다.)
• SDM (Structure Diffusion Model): 이 AI(FrameFlow 아키텍처 사용)는 LDM이 만든 새로운 압축 코드 (z^{(0)}) 와, 그 코드로부터 오토인코더가 그린 설계도 (\hat{c}_{\psi}(z^{(0)})), 이 두 가지 정보를 모두 활용하여 최종적인 **상세 3D 단백질 구조 (x)**를 조립합니다. 즉, 설계도(뼈대)와 추가적인 코드 정보(세부 특징)를 함께 보고 최종 모양을 정교하게 만드는 거죠.

3. 단계별 학습 및 생성 (Multi-stage Training and Sampling)

• 학습: 이 모델들을 한 번에 학습시키는 것이 어려워서, 3단계로 나누어 순서대로 학습시켰습니다. (1단계: 오토인코더 학습 → 2단계: 오토인코더와 SDM 함께 학습 → 3단계: 오토인코더는 고정하고 LDM만 학습)
• 생성: 새로운 단백질을 만들 때는, (1) LDM으로 새로운 압축 코드를 만들고, (2) 오토인코더로 설계도를 그리고, (3) SDM이 이 코드와 설계도를 바탕으로 최종 3D 구조를 만듭니다.

4. 생성 방향 조종하기 (PAE and LRC Guidance - 특별 기능!)

• 이것이 LSD의 중요한 특징입니다. LDM이 새로운 압축 코드 (z^{(0)})를 만드는 도중에 '이런 방향으로 만들면 좋겠어'라고 신호를 주어 결과물을 조종하는 기능입니다.
• LRC Guidance: 만들어지는 단백질이 너무 흔한 모양만 갖지 않고, **더 다양하고 새로운 구조(diversity/novelty 향상)**를 갖도록 유도합니다. LDM이 코드를 만들 때, '이 코드는 멀리 떨어진 부분끼리 연결(LRC)이 많은 구조를 만들 것 같아!' 라는 신호를 강화하는 방향으로 유도합니다.
• PAE Guidance: 만들어지는 단백질이 **더 안정적이고 실제로 합성하기 쉬운 구조(designability 향상)**가 되도록 유도합니다. LDM이 만드는 코드 중간 단계를 보고, 별도의 AI(f_vartheta)가 '이 코드는 최종 구조 안정성(PAE 점수)이 좋을 것 같아!' 라고 예측하면, 그 예측 점수가 좋아지는 방향으로 유도합니다.
• Joint Guidance: 위 두 가지 유도(LRC, PAE)를 동시에 적용하여 다양하면서도 안정적인 구조를 만들도록 시도합니다.

결론: LSD 모델은 단백질 구조를 (1) 설계도 코드 생성, (2) 설계도 그리기, (3) 설계도와 코드를 이용한 최종 구조 조립의 단계로 만들며, 특히 첫 단계에서 guidance를 통해 최종 결과물의 특성(다양성, 안정성 등)을 효과적으로 제어할 수 있는 방법을 제시합니다.