AI바라기의 인공지능
Diffusion : 논문 리뷰 : PROTCOMPOSER: COMPOSITIONAL PROTEINSTRUCTURE GENERATION WITH 3D ELLIPSOIDS 본문
Diffusion : 논문 리뷰 : PROTCOMPOSER: COMPOSITIONAL PROTEINSTRUCTURE GENERATION WITH 3D ELLIPSOIDS
AI바라기 2025. 4. 3. 20:36쉬운 설명:
ProtComposer는 단백질 구조를 디자인할 때, 레고 블록(ellipsoid)으로 대략적인 스케치를 먼저 하고, 그 스케치에 맞춰 세부 구조를 채워 넣는 방식과 유사합니다. 각 레고 블록은 특정 모양(helix/sheet)과 크기, 위치를 가지며, 이를 조절하여 원하는 단백질 구조를 만들거나 기존 구조를 수정할 수 있습니다. 특히, 흔하지 않은 레고 블록 조합(statistical model)을 사용함으로써 더 새롭고 다양한 단백질 구조 생성을 가능하게 합니다.
ProtComposer: Compositional Protein Structure Generation with 3D Ellipsoids 학습 노트
Purpose of the Paper:
- 기존 연구의 한계 극복: 기존 Machine Learning 기반 단백질 구조 생성 모델들은 주로 unconditional generation 또는 scaffold/binder inpainting에 국한되어, 생성될 단백질의 전반적인 spatial layout이나 고차원적 구조 (e.g., secondary structure 요소들의 배치)를 제어할 능력이 부재했음. 이는 생성되는 단백질의 다양성 부족 및 사용자 제어의 한계로 이어졌고, 특히 helix bundle과 같이 비교적 단순한 구조를 과도하게 생성하는 경향이 있었음 (low compositionality).
- 새로운 접근 방식 제시: Image generation 분야에서 bounding box나 blob representation을 이용한 제어 방식에 영감을 받아, 단백질 구조 생성을 위해 3D ellipsoid를 중간 수준(intermediate level)의 spatial layout "sketch"로 사용하는 새로운 제어 메커니즘을 제안함. 이를 통해 사용자가 substructure의 위치, 크기, 방향, 형태, semantic 정보(secondary structure type)를 제어하면서도 low-level 디테일은 모델이 생성하도록 함.
Key Contributions & Novelty:
- (Contribution) 3D Ellipsoid Representation for Layout Control: 단백질 substructure의 대략적인 위치, 형태 (covariance matrix), 크기 (residue count), semantic 정보 (secondary structure type: α-helix/β-sheet)를 포함하는 3D ellipsoid 세트를 protein layout specification 방식으로 제안함.
- (Novelty) Protein generation 분야에서 최초로 이러한 형태의 explicit하고 semantic 정보가 포함된 geometric layout control 방식을 도입하여, global conditioning과 fine-grained constraint 사이의 간극을 메움.
- (Contribution) Invariant Cross Attention (ICA) Mechanism: SE(3) equivariant flow model (Multiflow) 내에서 3D 공간의 ellipsoid 정보와 residue frame 정보를 효과적으로 통합하기 위한 새로운 attention 메커니즘을 개발함.
- (Novelty) Residue의 local frame을 활용하여 ellipsoid token으로부터 residue token으로 SE(3)-invariant 방식으로 정보를 전달하는 specific한 equivariant attention 구조를 제안함.
- (Contribution) Statistical Ellipsoid Layout Model: Unconditional generation 시, 다양하고 새로운 (novel) ellipsoid layout을 샘플링하기 위해 의도적으로 deep generative model 대신 간단한 통계적 모델 (Gaussian means, Wishart covariances + repulsion term)을 사용함.
- (Novelty) 학습 데이터 분포를 벗어나는 layout 생성을 보장하여 base model의 diversity/novelty 한계를 직접적으로 해결하고자, 복잡한 모델 대신 통계적 모델을 채택한 전략적 선택.
- (Contribution) Expanded Pareto Frontiers: Designability 대비 Novelty, Diversity, Helicity 측면에서 기존 방법들 (Multiflow annealing 조절, Chroma, RFDiffusion temperature 조절) 대비 현저히 확장된 Pareto frontier를 달성함.
- (Novelty) Layout conditioning이 inference-time 파라미터 튜닝보다 이러한 trade-off를 제어하는 데 더 효과적인 방법임을 실험적으로 입증함.
- (Contribution) Compositionality Metric: 생성된 단백질 구조의 복잡성 (e.g., 단일 helix bundle 편중 방지)을 정량화하기 위한 새로운 compositionality metric을 제안하고, ellipsoid conditioning이 이를 개선함을 보임.
- (Novelty) Designable하지만 구조적으로 단순한 단백질 생성 문제를 해결하기 위해 고안된 새로운 평가 지표.
- (Contribution) New Capabilities Enabled:
- 기존 단백질의 ellipsoid layout을 수정하여 구조를 편집 (structural editing)하는 기능.
- 사용자가 직접 디자인한 (hand-constructed) 복잡한 ellipsoid layout으로부터 단백질을 생성하는 기능.
- (Novelty) ML 기반 단백질 생성에서 layout 기반의 직접적 편집 및 복잡한 de novo 디자인 specification을 가능하게 한 최초의 기능들.
Experimental Highlights:
- Setup:
- Base Model: Pre-trained Multiflow (SE(3) flow matching model).
- Conditioning: Ellipsoids from PDB validation set (data ellipsoids) or statistical model (synthetic ellipsoids).
- Control: Classifier-free guidance (λ parameter).
- Metrics: Ellipsoid Consistency (Coverage, Accuracy, JSD 등), Designability (ProteinMPNN + ESMFold scRMSD), Diversity (Vendi score), Novelty (TM-Score vs PDB), Helicity, Compositionality.
- Baselines: Multiflow (varying annealing), Chroma (ShapeConditioner, varying temp), RFDiffusion (varying temp).
- Key Result 1 (Adherence & Control): 생성된 단백질 구조가 입력 ellipsoid layout에 매우 높은 수준으로 부합하며 (strong adherence), guidance strength(λ) 증가에 따라 adherence가 oracle 수준에 근접함 (Table 1, Fig 3). 이를 통해 정교한 구조 편집 (Fig 6) 및 hand-specified layout 기반 생성이 가능함을 입증. Chroma의 inference-time conditioning보다 우수한 adherence를 보임.
- Key Result 2 (Diversity/Novelty/Helicity): Synthetic ellipsoid conditioning을 통해 Designability vs. Diversity/Novelty/Helicity Pareto frontier를 크게 확장함 (Fig 4). Baseline 대비 동일 designability 수준에서 더 높은 diversity/novelty 및 PDB 분포에 가까운 helicity 비율을 달성.
- Key Result 3 (Compositionality): Baseline Multiflow 대비 더 복잡하고 다양한 구조(less prone to all-helix bundles)를 생성하며, PDB 통계에 더 가까운 compositionality를 보임 (Table 2, Fig 5).
Limitations and Future Work:
- (Limitation) 현재 ellipsoid annotation이 secondary structure (α/β)로 제한되어 있어, 기능적 정보나 더 세밀한 구조적 특징을 직접적으로 반영하지 못함.
- (Future Work) Ellipsoid annotation을 기능 명세 (functional specifications)나 다른 semantic 정보로 확장하여, 기능 중심의 단백질 디자인 (function-driven design)으로 나아갈 것을 제안함.
- (Limitation) Ellipsoid layout 생성을 위한 통계적 모델이 단순하여 자연 단백질 layout의 모든 복잡성을 포착하지 못할 수 있음.
- (Future Work) Ellipsoid layout을 위한 더 정교한 통계적 모델 또는 학습 기반 모델 개발 가능성 언급 (diversity와 biological plausibility 균형).
- (Limitation) 매우 새롭거나 복잡한 구조를 생성할 때 designability와의 trade-off가 여전히 존재함 (Pareto frontiers 및 Fig 7 결과에서 암시됨).
- (Future Work) Base generative model 또는 conditioning 메커니즘 개선을 통해 novel structure 생성 시 designability 저하를 최소화할 필요성.
Overall Summary:
ProtComposer는 3D ellipsoid layout (substructure 형태 및 secondary structure 포함)을 조건으로 SE(3) flow model(Multiflow)을 제어하는 새로운 단백질 구조 생성 프레임워크를 제안합니다. 새롭게 개발된 Invariant Cross Attention과 classifier-free guidance를 통해 높은 layout adherence를 달성하여, 정교한 구조 편집 및 사용자 지정 layout 기반 생성을 가능하게 합니다. 특히, 간단한 통계적 모델로부터 생성된 다양한 layout을 조건으로 사용할 경우, 기존 방법들보다 훨씬 향상된 novelty, diversity, compositionality를 보여주며 PDB 단백질 특성에 더 가까운 구조 생성을 가능하게 합니다.
ABSTRACT
우리는 substructure shapes와 semantics를 capturing하는 3D ellipsoids 세트를 통해 specified되는 spatial protein layouts에 conditioned된 protein structures를 generate하기 위해 ProtComposer를 개발합니다.
inference time에는, 직접 hand-constructed 하거나, 기존 proteins에서 extracted 하거나, statistical model로부터 가져온 ellipsoids에 condition하며, 각 option은 새로운 capabilities를 unlocking합니다.
ellipsoids를 Hand-specifying하는 것은 users가 protein substructures의 location, size, orientation, secondary structure, 그리고 approximate shape을 control할 수 있도록 enables합니다.
기존 proteins의 ellipsoids에 Conditioning하는 것은 substructure의 connectivity를 redesigning하거나 substructure properties를 editing하는 것을 enables합니다.
단순한 statistical model로부터 얻은 novel하고 diverse한 ellipsoid layouts에 conditioning함으로써, 우리는 designability, novelty, diversity 사이의 확장된 Pareto frontiers를 가진 protein generation을 improve합니다.
더 나아가, 이는 기존 generative models이 흔히 conceptually simple한 helix bundles를 oversample하는 것과 달리, PDB proteins와 matches하는 helix-fraction을 가진 designable proteins를 sampling하는 것을 enables합니다.
1 INTRODUCTION
Proteins는 광범위한 biological 및 chemical processes를 수행하는 복잡한 macromolecular machines입니다. rational protein design의 원대한 vision은 자연에서 발견되는 것과 유사한 complex하고 modular functions를 design할 수 있게 되는 것입니다. 여기서는 protein의 서로 다른 spatial parts가 coordinated fashion으로 작용하는 다른 properties를 가집니다. 그러나 현재 ML-based protein structure generation의 패러다임은 대체로 unconditional generation에 국한되거나, structure의 알려진 부분에 conditioned된 scaffolds 및 binders의 inpainting에 제한되어 있으며, generated protein의 higher-level spatial placement나 layout을 control할 능력이 없습니다. 이는 generated samples의 limited diversity와 control로 이어지며, protein generation을 image generation과 구별짓게 합니다. image generation에서는 이러한 수준의 control이 일반적이며 human users에게 새로운 capabilities를 제공합니다.
protein design toolbox의 이러한 격차를 해소하기 위해, 우리는 protein structure generative models을 protein의 3D space 내 layout으로 controlling하는 수단으로 ProtComposer를 개발합니다. 구체적으로, 우리는 protein의 대략적인 “sketch”를 제공하기 위해 annotations로 augmented된 3D ellipsoids를 통해 modular protein layouts를 설명합니다 (Figure 1). image generation에서의 blob 또는 bounding box representations와 유사하게, 이러한 ellipsoids는 data-level (즉, voxels) constraints와 global conditioning 사이의 intermediate 수준인 level of abstraction을 제공합니다. 이들은 diverse proteins의 generation을 control할 만큼 충분히 informative하지만, human-interpretable하고, easy-to-construct하며, protein structures의 low-level details를 constrain하지 않습니다. 따라서 3D ellipsoids는 complex protein designs가 수작업이나 heuristic algorithms를 통해 spatial sketches로 표현되고, deep learning models이 이러한 sketches를 high-quality하고 designable한 backbones로 "fill in"하는 two-stage paradigm을 facilitate합니다.
이 연구에서, 우리는 state-of-the-art designability를 가진 joint sequence-structure flow-matching model인 Multiflow를 secondary structure로 annotated된 3D ellipsoid layouts로 controlling하기 위해 우리의 philosophy를 적용합니다. Multiflow는 protein structures를 SE(3) 내 residue frames의 cloud로 represents하고, Invariant Point Attention으로 flow network를 parameterizes합니다. 이 network에 ellipsoid conditioning을 inject하기 위해, 우리는 3D space 내 ellipsoids와 residue frames 간의 message passing을 위한 equivariant mechanism을 개발하며, 이를 Invariant Cross Attention이라고 부릅니다. 그런 다음 우리는 이 cross attention을 사용하여 Multiflow를 ellipsoid layouts에 conditioned된 protein structure 및 sequence의 conditional model로 fine-tune합니다. 우리는 Multiflow를 위한 classifier-free guidance mechanism을 개발하여, ellipsoid-conditioned generation과 unconditional generation 사이의 interpolation을 enabling합니다. Empirically하게, 우리는 이 conditional generative distributions family가 세 가지 axes를 따라 protein generation의 state of the art를 advances시킨다는 것을 발견했습니다:
- Control—기존 models과 달리, 우리는 secondary structure ellipsoids의 (기존 또는 novel한) arrangements에 대해 우리의 method를 prompt할 수 있습니다. 우리는 ellipsoid conditioning에 대한 adherence를 measure하기 위한 metrics family를 개발하고, conditioning ellipsoid layouts와 generated backbones 사이에 strong consistency가 있음을 발견했습니다. 이 consistency는 training distribution을 훨씬 넘어서까지 persists하며, 일부 특히 impressive한 generations가 Figure 7에 나와 있습니다.
- Diversity and Novelty—간단한 statistical models family에서 drawn된 synthetic ellipsoid layouts에 conditioning함으로써 (Section 3.4), 우리는 generations의 diversity와 novelty를 significantly increase합니다. generated proteins의 designability에는 cost가 따르지만, 이 tradeoff에 따른 우리의 Pareto frontier는 Multiflow generations의 diversity를 controlling하기 위한 유일한 기존 option인 inference-time parameters 조정의 그것을 훨씬 surpasses합니다.
- Compositionality—비록 highly designable하지만, protein generations는 종종 낮은 수준의 architectural complexity를 exhibit합니다 (예: 단일 alpha-helix bundle로만 구성됨). 우리는 그러한 proteins가 high model likelihood를 가지지만 low information content를 가진 병적인 language model outputs (예: 많은 large language models 하에서 "and and and ..." 문장은 높은 likelihood를 가짐)와 analogous하다고 argue합니다. 우리는 이 phenomenon을 quantify하기 위해 compositionality metric을 introduce하고, ellipsoid conditioning이 generated structures의 complexity와 compositionality를 improve할 수 있음을 show합니다.
.
- 왼쪽 (Left): 모델 작동 방식
- 입력: 어지럽게 흩어진 단백질 구조 정보(noisy protein structure)와 함께, 단백질의 특정 부분(예: Sheet 또는 Helix 구조)을 나타내는 ellipsoids(타원체 설계도) 정보를 받습니다.
- 처리: 이 ellipsoids 정보는 **Invariant Cross Attention (ICA)**라는 특별한 메커니즘을 통해 Transformer 기반의 AI 모델에 주입됩니다. 모델은 이 '설계도' 정보를 참고하여 흩어진 구조 정보를 '정리'하거나 새로운 구조를 생성(denoises a noisy protein structure)합니다.
- 결과: ellipsoids로 지정된 spatial protein layout(공간 배치)에 맞는 단백질 구조가 generate(생성)됩니다.
- 가운데 (Middle): 통계 모델 활용 예시
- 입력: 사람이 직접 만드는 대신, simple statistical model(간단한 통계 모델)이 자동으로 생성한 novel(새롭고) diverse(다양한) ellipsoid layouts를 사용합니다. 그림 위쪽에 여러 형태의 ellipsoids 예시가 보입니다.
- 결과: 이 자동으로 생성된 '설계도'를 바탕으로 단백질 구조를 생성하면, 결과물 역시 novelty(참신함)와 diversity(다양성)가 높은 단백질들이 만들어집니다. 동시에 designability(실제로 안정적인 구조일 가능성)도 유지됩니다. 즉, 새롭고 다양한 단백질을 자동으로 탐색하고 생성하는 데 유용합니다.
- 오른쪽 (Right): 사용자 직접 설계 예시
- 입력: 사용자가 직접 배치한(hand-specified) ellipsoid layouts를 사용합니다. 그림에서는 여러 ellipsoids를 고리 모양이나 특정 형태로 배열한 예시를 보여줍니다.
- 결과: 사용자가 원하는 대로 '설계도'를 제공하면, 생성되는 단백질 구조 역시 그 layout을 따르게 됩니다. 이는 사용자가 원하는 구조를 controllable generation(제어 가능한 생성) 방식으로 만들 수 있음을 보여줍니다.
요약: 이 그림은 ProtComposer가 ellipsoids라는 '설계도'를 사용하여 단백질 구조 생성을 가이드하는 방법을 보여줍니다. 이 설계도는 (왼쪽) 모델 내부의 ICA 같은 기술을 통해 처리되며, (가운데) statistical model로부터 얻어 novel하고 diverse한 단백질을 만들거나, (오른쪽) 사용자가 직접 지정(hand-specified)하여 원하는 구조를 control하며 만들 수 있게 해줍니다.
ProtComposer Introduction: 정리 노트 (AI 연구자 대상)
핵심 문제 (Problem):
- 기존 ML 기반 protein structure generation 모델들은 unconditional generation 또는 inpainting 위주로, 생성될 단백질의 higher-level spatial layout (전체적 배치, substructure 위치/형태)에 대한 control 능력이 부재함.
- 이는 image generation 분야에서 흔히 제공되는 수준의 제어 능력과 비교했을 때, 생성되는 단백질의 diversity와 사용자 control을 제한함.
제안 솔루션 (Proposed Solution): ProtComposer
- Protein structure generative models을 protein's layout in 3D space를 이용하여 controlling하는 새로운 방법론 제시.
핵심 아이디어 (Key Idea): 3D Ellipsoids for Spatial Layout Representation
- Modular protein layouts를 annotations (e.g., secondary structure)로 augmented된 3D ellipsoids를 사용하여 "spatial sketch"로 표현.
- Ellipsoids는 data-level (voxels) constraints와 global conditioning 사이의 intermediate level of abstraction을 제공:
- Generation control에 충분히 informative함.
- Human-interpretable하고 easy-to-construct.
- Low-level details를 과도하게 constrain하지 않음.
- Two-stage paradigm 가능: 사용자가 spatial sketch (ellipsoids)를 정의 -> Deep learning model이 designable backbones로 "fill in".
기술적 핵심 (Technical Core):
- Target Model: Multiflow (state-of-the-art joint sequence-structure flow-matching model)를 control.
- Conditioning Mechanism: Invariant Cross Attention (ICA) 개발.
- Ellipsoids (in 3D space)와 residue frames 간의 message passing을 위한 equivariant mechanism.
- Invariant Point Attention 기반 flow network에 ellipsoid conditioning을 inject.
- Fine-tuning: Multiflow를 ICA를 사용하여 ellipsoid layouts에 conditioned된 conditional model로 fine-tune.
- Guidance: Classifier-free guidance mechanism을 적용하여 ellipsoid-conditioned / unconditional generation 간 interpolation 가능.
주요 기여 및 주장 (Contributions & Claims - 3 Axes):
- Control:
- 기존/신규 secondary structure ellipsoid 배열(arrangements)을 prompt로 사용하여 제어 가능.
- Ellipsoid conditioning에 대한 adherence(일치도) 측정 metrics 개발 및 높은 consistency 확인.
- Diversity & Novelty:
- Simple statistical models에서 생성된 synthetic ellipsoid layouts에 conditioning하여 diversity/novelty 크게 향상.
- Designability-diversity tradeoff에서 기존 inference-time parameter 조정 방식보다 우월한 Pareto frontier 달성.
- Compositionality:
- 기존 모델들이 highly designable하지만 architectural complexity가 낮은 (e.g., single alpha-helix bundle) 구조를 생성하는 문제 지적.
- Ellipsoid conditioning이 compositionality metric으로 측정한 생성 구조의 complexity와 compositionality를 향상시킴을 보임.
핵심 차별점: Ellipsoid 기반 spatial control 도입, ICA 메커니즘, Control/Diversity/Compositionality 측면에서의 명확한 개선점 제시.
쉬운 설명 :
현재 단백질을 AI로 만드는 연구들이 있지만, 대부분 전체적인 모양이나 각 부분들이 어떻게 배치될지 세밀하게 조종하기는 어려웠어요. 그냥 AI에게 "알아서 만들어봐" 하는 식이었죠. 이미지 생성 AI처럼 "이런 느낌으로 만들어줘"라고 지시하기가 힘들었습니다.
이 논문은 ProtComposer라는 새로운 방법을 제안합니다. 핵심은 AI가 단백질 구조를 만들기 전에, 우리가 원하는 단백질의 대략적인 3D '설계도' 또는 '스케치' 를 먼저 그려주는 거예요.
이 '설계도'는 3D ellipsoids라는 길쭉한 공 모양의 간단한 도형들로 만들어집니다. 이 도형들은 단백질의 특정 부분이 어디에 위치하고, 얼마나 크고, 어떤 방향이어야 하는지, 그리고 대략 어떤 구조(예: 나선형(helix))를 가져야 하는지를 표시해줍니다.
연구팀은 Multiflow라는 기존의 강력한 단백질 생성 AI 모델을 수정해서, 이 ellipsoid '설계도'를 이해할 수 있도록 만들었어요. 이때 Invariant Cross Attention (ICA) 라는 새로운 기술을 사용했습니다.
ProtComposer로 할 수 있는 것들:
- 원하는 대로 디자인 (Control): 사용자가 직접 ellipsoid 설계도를 그려서 원하는 특정 모양의 단백질을 만들도록 제어할 수 있어요.
- 새롭고 다양한 단백질 발견 (Diversity & Novelty): 프로그램(statistical model)이 자동으로 만들어주는 다양한 ellipsoid 설계도를 이용해서, 새롭고(novel) 다양한(diverse) 단백질 구조들을 탐색하고 생성할 수 있어요.
- 더 복잡한 구조 만들기 (Compositionality): 기존 AI들이 종종 너무 단순한 구조(예: 나선 구조만 뭉쳐놓은 것)만 만드는 경향이 있었는데, ProtComposer는 더 복잡하고 짜임새 있는(compositional) 구조를 만들도록 도와줘요.
결론적으로, ProtComposer는 과학자들이 최종 3D 단백질 구조를 훨씬 더 세밀하게 제어하면서 설계하고 생성할 수 있게 해주는 발전된 AI 도구입니다.
2 BACKGROUND AND RELATED WORK
Protein Structure Generation. Protein structure generative models의 주된 목표는 novel proteins의 computational design을 돕는 것입니다. 따라서, 우리는 종종 이미 existing folds나 secondary structure compositions를 넘어서 generate하고, design specifications를 만족시키도록 generations을 control하기를 원합니다. controllability를 다루기 위해, existing structural motifs를 scaffolding하는 것 외에 주로 탐색된 두 가지 경로가 있습니다. 첫째, block contact maps와 sequential secondary structure specifications에 conditioning하는 것입니다. 둘째, Chroma에서처럼, manifold로의 projections를 통하거나 arbitrary differentiable energy function으로부터의 forces를 통한 inference time conditioning입니다. 이러한 inference time control은 높은 generality를 가지는 반면, ProtComposer는 단일 유형의 shape 및 semantic conditioning에 대해 trained되므로, control에 대한 adherence가 향상됩니다 (Table 1).
Spatial Conditioning. image generative models의 경우, controllable generation은 인상적인 이미지를 generating하는 것을 넘어 그 utility의 상당 부분을 차지하는 새로운 applications를 가능하게 했습니다. protein generation을 위해 유사한 새로운 capabilities를 열기 위해, 우리는 crucial하다고 결정된 transferrable concepts를 따릅니다. technical 및 architectural level에서, 이는 initialized conditional model에서 original model의 output을 minimally perturbing하는 원칙에 따라 strong existing generative models을 fine-tuning하는 것을 포함합니다. conceptual level에서, 이는 올바른 input specification을 찾는 것을 의미합니다: 다른 tasks에 대해 다른 levels of granularity가 appropriate합니다. image domain에서는, semantic segmentation maps나 sketches와 같은 pixel-level specifications에서부터 proteins에 대한 우리의 aims과 유사한 bounding-boxes나 "blobs"와 같은 더 coarse-grained specifications에 이르기까지 다양합니다.
Flow models. Flow matching은 의 시작 시간부터 까지 integrated될 때 noise distribution 의 samples을 data distribution 으로 transports하는 time-dependent vector field 를 learn하는 것을 목표로 합니다. 를 train하기 위해, 우리는 및 을 satisfying하는 conditional probability path 로부터 partially noised data를 sample합니다. 일반적인 선택은 와 사이의 straight line을 따라가는 Dirac이거나 manifolds에서의 flow matching을 위한 geodesic입니다. sampled된 noisy datapoints 에서, 우리는 vector field 를 evaluate하고, continuity equation 를 통해 conditional probability path에 corresponds하는 conditional vector field 에 대해 regress합니다. convergence 시, 는 (gradients가 에 대해 regressing하는 것과 equivalent하기 때문에) prior 를 marginal probability path 를 통해 data distribution 으로 evolves시키는 marginal vector field 를 approximates합니다.
ProtComposer 배경 및 관련 연구: 정리 노트 (AI 연구자 대상)
Protein Structure Generation 맥락:
- 목표: Computational design을 통한 novel proteins 생성 지원. 이를 위해 기존 구조(existing folds/compositions)를 넘어서는 생성 및 design specifications 충족을 위한 control 필요.
- 기존 Control 연구 방향:
- Scaffolding: 기존 structural motifs 확장. (제한적)
- Specific Conditioning: Block contact maps, sequential secondary structure 등에 conditioning. (공간적 유연성 부족 가능성)
- Inference Time Conditioning (e.g., Chroma): Manifold projections, differentiable energy functions 활용.
- 장점: 높은 generality.
- 단점: Control signal에 대한 adherence(충실도)가 낮을 수 있음.
- ProtComposer의 위치:
- Ellipsoids를 이용한 shape & semantic conditioning에 특화되어 trained된 모델.
- Inference time control 대비 향상된 adherence 추구 (Table 1).
- 기존 specific conditioning 방식 대비 유연한 spatial control 제공 목표.
Spatial Conditioning (Image Generation에서의 영감):
- 동기: Image generation에서 controllable generation이 새로운 applications과 utility를 창출함. Protein generation에서도 유사한 capabilities 확보 목표.
- 차용한 개념:
- Technical/Architectural: 강력한 기존 generative model(여기서는 Multiflow)을 fine-tuning하되, original model's output을 최소한으로 변경(minimally perturbing)하는 원칙 적용 (ControlNet 등과 유사).
- Conceptual: 적절한 input specification granularity 탐색. Ellipsoids는 image domain의 coarse-grained specifications(bounding-boxes, blobs)에 상응하는, 제어와 유연성 사이의 균형을 맞춘 중간 수준의 명세로 제안됨.
Flow Models (기술 기반):
- 기반 Generative Framework: Flow matching 활용. (ProtComposer는 Multiflow 기반)
- 핵심 원리: Noise distribution()을 data distribution()으로 변환하는 time-dependent vector field() 학습.
- 학습 방식: Noise-data 샘플 쌍을 잇는 conditional probability path()에서 정의된 conditional vector field()를 타겟으로 를 regress.
- 연관성: ProtComposer는 flow matching model인 Multiflow 위에 구축되며, 제안하는 ellipsoid conditioning 메커니즘(ICA)을 이 flow 프레임워크 내에 통합함.
핵심: 기존 단백질 생성 제어 방식들의 한계(특정 조건 제약, 낮은 adherence 등)를 인지하고, 이미지 생성 분야의 성공적인 제어 개념(fine-tuning, 적절한 입력 명세)과 강력한 생성 모델링 프레임워크(flow matching)를 결합하여, ellipsoid라는 새로운 형태의 spatial conditioning을 제안하고 그 효과(특히 adherence)를 강조함.
쉬운 설명 :
이전 연구들은 어땠나? (단백질 생성 분야)
- 목표: 컴퓨터로 새로운 단백질을 설계하는 것을 돕는 것인데, 그러려면 AI가 만드는 단백질을 우리가 원하는 대로 제어할 수 있어야 했어요.
- 기존 제어 방법들:
- 이미 있는 단백질 조각에 새로운 조각을 이어 붙이는 방식 (Scaffolding).
- 단백질의 어떤 부분끼리 닿아야 하는지 같은, 아주 구체적이지만 공간적이지는 않은 정보를 주는 방식 (Block contact maps 등).
- AI가 단백질을 생성하는 중간 과정에 계속 개입해서 방향을 틀어주는 방식 (Inference time conditioning). 이건 유연하지만, 때로는 지시를 잘 못 따를 때가 있었어요.
- ProtComposer의 차이점: 이 논문의 방법(ProtComposer)은 처음부터 ellipsoid라는 '설계도'를 따르도록 특별히 **훈련(trained)**되었기 때문에, 지시를 더 정확하게 따를(adherence가 높을) 것이라고 기대해요.
이미지 AI에서 배운 점 (공간 제어)
- 요즘 그림 그려주는 AI들을 보면, 사용자가 "여기에 사람을, 저기에 나무를 그려줘" 하고 스케치나 네모 상자(bounding-boxes)로 위치를 지정해주면 그대로 그려주잖아요? 이런 제어 기능 덕분에 이미지 AI가 훨씬 유용해졌어요.
- ProtComposer도 여기서 아이디어를 얻었어요:
- 이미 성능 좋은 단백질 생성 AI(여기서는 Multiflow)를 가져와서, 원래 성능은 최대한 유지하면서(minimally perturbing) ellipsoid '설계도'를 이해하도록 살짝 **튜닝(fine-tuning)**해요.
- '설계도'를 얼마나 자세하게 줄지 정하는 게 중요한데, ellipsoids는 이미지 AI에서 '대략적인 덩어리(blobs)'나 '네모 상자'를 쓰는 것과 비슷하게, 너무 자세하지도 너무 추상적이지도 않은 적절한 수준(coarse-grained specifications)의 정보를 주는 방식이에요.
AI의 엔진 (Flow Models)
- ProtComposer가 기반으로 하는 Multiflow 모델은 Flow matching이라는 기술을 사용해요.
- 이 기술은 마치 아무 의미 없는 노이즈(noise) 상태에서 시작해서, AI가 배운 방향(vector field)을 따라 시간이 흐르면서 점차 실제 단백질 구조()로 부드럽게 변형시키는 방식이라고 생각할 수 있어요. ProtComposer는 이 변형 과정에 ellipsoid라는 길잡이를 추가해주는 거죠.
이 이미지(Figure 2)는 주어진 protein structure 데이터로부터 어떻게 ellipsoid representation(타원체 표현)을 추출하는지 3단계로 보여줍니다. 이 과정은 주로 AI 모델 training을 위해 필요합니다.
- 1단계 (Input protein):
- 분석하려는 input protein(입력 단백질) 구조에서 시작합니다. (위: 3D 리본 구조, 아래: 단순화된 연결 구조)
- 2단계 (Annotate residues and draw edges):
- 단백질을 구성하는 각 residue(아미노산 잔기)에 feature(특성)를 부여(annotate)합니다. 논문에서는 secondary structure type (-helix는 파란색, -sheet는 주황색 등)을 사용했습니다.
- 그 다음, 두 residue가 (1) 같은 feature를 가지고 (2) 서로 가까이 (within 5Å) 있으면, 그 사이에 edge(선)를 그어 연결합니다. 이렇게 하면 같은 종류의 구조이면서 서로 가까이 있는 residue들이 그룹으로 묶이기 시작합니다. (아래 그림에서 같은 색깔의 가까운 점들이 연결된 것을 볼 수 있습니다.)
- 3단계 (Fit Gaussians to connected components):
- 2단계에서 edge로 연결된 residue들의 묶음(connected components) 각각에 대해, 통계적인 Gaussian 분포를 계산하여 맞춥니다(fit). 이 Gaussian 분포는 각 묶음의 중심 위치(mean)와 퍼진 모양/방향(covariance) 정보를 담고 있으며, 이것이 바로 3D 공간에서의 ellipsoid로 시각화됩니다.
- 즉, 각 semantically coherent region(의미적으로 유사한 영역, 여기서는 같은 secondary structure 묶음)을 하나의 ellipsoid로 나타내는 것입니다.
요약: 이 그림은 복잡한 단백질 구조를 받아서 (1단계), 구조의 종류와 근접성을 기준으로 잔기들을 그룹화하고 (2단계), 각 그룹을 대표하는 간단한 3D ellipsoid로 변환하는 (3단계) 과정을 보여줍니다. 이렇게 추출된 ellipsoid 정보가 모델 학습이나 분석에 사용됩니다.
알고리즘 1과 2를 설명해 드리겠습니다. 이 알고리즘들은 Section 3.2에서 설명된 것처럼, ellipsoid conditioning(타원체 조건) 정보를 기존 Multiflow 모델에 어떻게 주입하는지를 보여줍니다.
Algorithm 1: Invariant Cross Attention
- 목적: 이 알고리즘은 논문에서 새로 제안한 핵심 메커니즘입니다. 각 residue(단백질 구성 요소)의 정보(residue tokens )가 ellipsoid(타원체 '설계도') 정보()를 참고하여 업데이트될 수 있도록 합니다. 중요한 점은 이 과정이 단백질 전체의 위치나 회전에 영향을 받지 않도록(SE(3)-invariant) 설계되었다는 것입니다.
- 입력:
- Residue tokens (): 각 residue의 현재 특징 벡터.
- Residue frames (): 각 residue의 3차원 위치 및 방향 정보.
- Ellipsoid parameters (): 각 ellipsoid 의 중심()과 모양/방향() 정보.
- 주요 단계:
- r_ik = ...: 각 ellipsoid 중심()이 해당 residue()의 로컬 좌표계 기준으로 어디에 있는지 상대 위치를 계산합니다. (전체 단백질이 회전/이동해도 이 값은 불변)
- C_ik = ...: Ellipsoid의 모양/방향 정보()를 residue의 로컬 좌표계 기준으로 변환합니다. (역시 불변성 유지)
- a_ik = s_i + ...: 원래 residue 정보()와 ellipsoid의 상대 위치 정보(PosEmbed(r_ik))를 결합합니다.
- a_ik += ...: 위 결과에 ellipsoid의 상대 모양/방향 정보(Flatten(C_ik))를 추가합니다. 이제 는 residue 와 ellipsoid 간의 관계 정보를 종합적으로 담습니다.
- q_i = ...: Residue 로부터 '질문(query)' 벡터 를 만듭니다.
- k_ik, v_ik = ...: 종합 정보 로부터 '키(key)' 와 '값(value)' 벡터를 만듭니다. 에는 ellipsoid 로부터 추출한, residue 에게 유용한 정보가 담겨있다고 볼 수 있습니다.
- s_i += Attention_k(...): Attention 메커니즘을 수행합니다. Residue의 질문()이 모든 ellipsoid의 키()들을 살펴보고, 관련성 높은 ellipsoid들의 값()들을 가중합하여 원래 residue 정보()에 더해줍니다. 이를 통해 residue token이 모든 ellipsoid로부터 관련 정보를 얻어 업데이트됩니다.
- 출력: Ellipsoid 조건 정보가 반영되어 업데이트된 residue tokens .
Algorithm 2: Update Block
- 목적: Algorithm 1(Invariant Cross Attention)이 Multiflow 모델의 기본 처리 단위인 update block 안에 어떻게 통합되는지를 보여줍니다. 이 블록은 단백질의 다양한 정보(residue tokens, pair representations, frames)를 업데이트합니다.
- 입력:
- Residue tokens ()
- Pair representations (): Residue 쌍 사이의 상호작용 정보.
- Residue frames ()
- Ellipsoid tokens (): Ellipsoid 자체의 특징 벡터 (Algorithm 1에서는 파라미터 를 사용했지만, 여기서는 토큰을 직접 사용).
- Ellipsoid parameters (): Algorithm 1 내부에서 사용.
- 주요 단계 (초록색으로 표시된 부분이 추가/수정된 부분):
- s += InvariantPointAttention(...): 기존 Multiflow처럼, **Invariant Point Attention(IPA)**을 통해 residue들이 서로 상호작용하여 를 업데이트합니다.
- s += InvariantCrossAttention(...): (새로 추가된 부분) Algorithm 1을 호출하여, ellipsoid 정보()를 이용해 를 추가로 업데이트합니다. 여기서 ellipsoid 조건이 주입됩니다.
- s = Concat(s, e): 업데이트된 residue tokens()와 ellipsoid tokens()를 하나로 합칩니다.
- s = Transformer(s): 합쳐진 토큰들을 Transformer 모델에 통과시킵니다. Transformer 내부의 self-attention을 통해 residue와 ellipsoid 정보가 서로 영향을 주고받으며 업데이트될 수 있습니다.
- s, e = Split(s): Transformer를 통과한 결과를 다시 residue tokens()와 ellipsoid tokens()로 분리합니다.
- T = RigidUpdate(...): 업데이트된 를 기반으로 residue의 3D 위치/방향()을 업데이트합니다.
- z = EdgeUpdate(...): 업데이트된 를 기반으로 residue 쌍 간의 상호작용 정보()를 업데이트합니다.
- 출력: 다음 레이어(block)로 전달될 업데이트된 .
요약: Algorithm 1은 residue와 ellipsoid 간의 SE(3)-invariant한 정보 교환을 위한 구체적인 계산 방법(새로운 cross attention)이고, Algorithm 2는 이 새로운 계산(Algorithm 1)을 기존 Multiflow 모델의 update block 구조 안에 어떻게 통합하여 (IPA, Transformer 등과 함께) 전체 정보를 업데이트하는지를 보여줍니다.
이 이미지(Figure 3)는 Section 3.3에서 설명된 classifier-free guidance의 효과를 보여줍니다. 특히 guidance parameter인 lambda () 값을 조절했을 때 단백질 생성 결과가 어떻게 변하는지를 보여줍니다.
- 두 가지 예시 (위/아래 행): 그림에는 두 가지 다른 목표 protein layout(배경에 희미하게 보이는 ellipsoids '설계도')에 대한 생성 결과가 각각 위쪽 행과 아래쪽 행에 나와 있습니다.
- Guidance 강도 조절 (): 각 행에서 왼쪽에서 오른쪽으로 갈수록 값이 0.0에서 1.0으로 증가합니다.
- (No guidance): AI가 ellipsoid '설계도'를 전혀 참고하지 않고, 학습된 대로 자유롭게 단백질 구조를 생성합니다 (unconditional generation). 생성된 구조(파란색/주황색 리본)가 배경의 ellipsoid와 잘 맞지 않는 것을 볼 수 있습니다.
- (Full guidance): AI가 ellipsoid '설계도'를 최대한 강하게 따르도록 지시받습니다 (ellipsoid-conditioned generation). 생성된 구조가 ellipsoid 모양과 배치에 매우 잘 맞춰진 것을 볼 수 있습니다.
- : 값이 커질수록 AI는 ellipsoid '설계도'를 더 중요하게 고려하여 구조를 생성합니다. 왼쪽에서 오른쪽으로 갈수록 생성된 구조가 점점 더 ellipsoid에 맞춰지는 것을 시각적으로 확인할 수 있습니다.
- 평가 지표 (Metrics): 각 생성된 구조 위에는 ellipsoid alignment metrics(타원체 정렬 지표) 값들이 표시되어 있습니다.
- Cov (coverage): 생성된 구조가 목표 ellipsoids를 얼마나 잘 덮는지.
- Acc (accuracy): 생성된 구조가 목표 ellipsoids 안에 얼마나 정확하게 들어맞는지.
- LL (likelihood): 모델이 해당 구조를 생성할 확률 또는 designability(설계 가능성) 점수일 수 있습니다 (값이 클수록 좋을 수 있음 - 여기서는 음수이므로 0에 가까울수록).
- Soft Acc (soft accuracy): 좀 더 유연하게 측정한 정확도.
- 경향: 일반적으로 가 증가함에 따라 Cov와 Acc (및 Soft Acc)가 증가하는 경향을 보여, ellipsoid 조건에 대한 adherence(충실도)가 높아짐을 나타냅니다. LL은 약간 감소할 수 있는데, 이는 조건을 강하게 따르면서 모델의 자체적인 생성 확률이나 designability와 약간의 trade-off가 발생할 수 있음을 시사합니다.
요약: 이 그림은 guidance parameter 를 조절함으로써, 사용자가 ellipsoid '설계도'를 얼마나 엄격하게 따를지 (adherence)와 AI 모델의 기본적인 생성 능력(designability/likelihood) 사이에서 trade-off를 조절할 수 있음을 보여줍니다. 를 높이면 ellipsoid 모양대로 더 정확하게 제어(control)된 구조를 얻을 수 있습니다.
3 METHOD
3.1 ELLIPSOID REPRESENTATION OF PROTEINS
Proteins는 compositional objects입니다—서로 다른 regions는 서로 다른 properties를 가지며, 우리는 generative model의 sampling을 control하기 위해 이 information을 succinctly describe할 language를 찾고자 합니다. 이를 위해, 우리는 protein의 spatial layout을 K개의 ellipsoids 세트를 사용하여 represent할 것을 제안하며, 각 ellipsoid는 protein의 semantically coherent region에 해당합니다. 각 ellipsoid는 associated region의 residues 수, categorical semantic feature, 그 position, 그리고 region 내 coordinates의 covariance matrix로 표현되는 shape을 기록합니다. 우리는 protein spatial layouts의 이러한 representation이 text prompt나 protein family와 같은 single global annotation과 meshes나 voxel grids와 같은 더 complex shape descriptors 사이에서 favorable tradeoff를 찾는다고 주장합니다. global annotation은 spatial layout에 대한 desired control을 provide하기에 insufficient할 수 있으며, 더 complex annotation은 additional model의 training 없이는 generate하거나 specify하기 어려울 수 있습니다. 반면, 3D ellipsoids는 expressive하고 precise하면서도 generate하고 manipulate하기 쉽습니다.
Mathematically하게, 우리는 K개의 ellipsoids로 구성된 protein spatial layout을 unordered set 로 define하며, 여기서 각 ellipsoid는 mean , covariance , count , 그리고 feature annotation (여기서 는 application-dependent feature space)를 가진 Gaussian으로 represented됩니다. Gaussian probability distributions로 Viewed될 때, 우리의 ellipsoids는 well-defined boundaries를 가지지 않습니다; 그러나, visualization 및 evaluation purposes를 위해, 우리는 ellipsoid boundary를 Mahalanobis distance
의 surface로 define합니다, 즉, [ \partial E_k = \left{ x \in \mathbb{R}^3 : \sqrt{(x - \mu_k)^T \Sigma_k^{-1} (x - \mu_k)} = \sqrt{5} \right} ] 이것은 conventional ellipsoid의 functional form입니다. distance
는 density의 83%가 surface 내부에 falls하도록 chosen되었으며, 이는 best visual results를 yields합니다 (Appendix Fig 12).
Ellipsoid Segmention. protein structure가 Provided되었을 때, (예: training purposes를 위해) 그것의 ellipsoid representation을 obtaining하는 것은 두 단계로 consists됩니다: protein을 semantically coherent regions으로 segmentation하는 것, 그리고 각 region에 대한 ellipsoid descriptions 를 extraction하는 것입니다. protein을 segment하기 위해, 우리는 두 residues가 spatially proximal하고 semantically similar할 경우에만 같은 region에 places하는 simple, non-learned segmentation algorithm을 consider합니다. 우리는 이러한 각 residues pair에 대해 edge를 drawing하여 segmentation graph를 construct하고 이 segmentation graph의 connected components list를 return합니다. 각 segmented region에 대해, 우리는 residue features를 aggregate하여 를 obtain하고 positions의 mean과 covariance를 compute합니다. 이러한 steps는 Figure 2에 illustrated되어 있으며 Appendix Algorithm 3에 더 detailed되어 있습니다. 우리는 이것이 예를 들어 K-means나 spectral clustering을 사용하는 더 sophisticated variants보다 more reliable하다는 것을 found했습니다.
이 연구에서는, -helices와 -sheets의 regions, 즉 secondary structure layouts를 specifying하는 3D ellipsoids에 focus합니다. 우리의 feature space는 따라서 secondary structures types의 two-class space이며, 입니다. 우리는 DSSP를 사용하여 residues를 featurize하고, 같은 secondary structure label을 가지며 내에 있는 amino acids 사이에 segmentation graph의 edges를 draw합니다. ellipsoid annotation 는 그 constituent residues의 label을 그대로 inherits합니다. 우리는 모든 loop residues와 다섯 개 이하의 residues를 가진 ellipsoids는 exclude합니다.
3.2 ELLIPSOID CONDITIONING
Unconditional model. 우리는 우리의 ellipsoid conditioning을 existing protein structure generative model에 inject하며, 이를 위해 sequence와 structure를 jointly generates하는 Multiflow를 choose합니다. 그들의 framework를 Following하여, 우리는 각 residue의 frame 이 associated translation 과 AlphaFold2를 following하여 backbone coordinates로부터 constructed된 rotation matrix 을 가지는 frames 의 array로 represented된 proteins을 generate합니다. Additionally, 각 residue는 amino acid type 를 가집니다. translations, rotations, amino acids를 jointly generate하기 위해, Multiflow는 세 가지 modalities 모두를 iteratively update하는 세 types의 flow matching procedures를 employs합니다. Translations는 Gaussian prior로부터의 linear flow matching으로 handled되고, rotations는 상에서의 Riemannian flow matching으로, residue types는 discrete flow matching으로 처리되어, 일 때 probability path 를 tracing out하면서 prior 에서 data distribution 로 transports하는 joint flow를 resulting합니다. flow는 translations, rotations 및 residue type inputs를 받는 single backbone architecture에 의해 parameterized되며, 이로부터 time dependent translation vector field , rotation vector field , 그리고 residue type updates를 dictating하는 rate matrix 를 predicts합니다. architecture는 여러 identical update blocks로 composed되어 있으며, 각 블록은 에 대한 d-dimensional residue representations , 에 대한 residue pair representations , 그리고 에 대한 residue frames 를 updates합니다. updates는 -equivariant하며 shallow transformers와 Invariant Point Attention (IPA)의 mixture로 accomplished됩니다; complete architectural details는 Campbell et al. (2024)를 refer합니다. 모든 update blocks 후, final residue tokens 와 frames 는 flow fields , , 를 parameterize하는 데 used됩니다.
Injecting Ellipsoid Conditioning. 우리는 이제 를 sample하도록 trained된 pre-trained unconditional Multiflow model을 fine-tune하여 ellipsoid conditioned density 를 sampling하고 ProtComposer를 obtain하는 것을 aim합니다. At inference time, ellipsoids는 manually specified되거나 novel하고 diverse ellipsoids의 second distribution (see Section 3.4)에서 sampled되어 density 를 target할 수 있습니다. For fine-tuning, 우리는 conditioning information을 additional input으로만 provide하며 - training loss는 Multiflow의 unchanged loss로 remains합니다. ellipsoid information을 inject하기 위해, 우리는 image diffusion models을 위한 semantic map conditioning의 best principles를 follow하고, initialization 시점에 unconditional model을 minimally perturb하는 architecture modifications를 design합니다. 즉, input으로 empty set의 ellipsoids가 주어지면, untrained conditional model은 unconditional model과 identical outputs을 produce해야 합니다. 이것은 initial residue representations 를 preserving하고, 아래에서 described되는 그들의 updates를 inform하기 위해 3D ellipsoids로부터 additional information만 supplying함으로써 accomplished됩니다.
우리는 residue tokens 와 same dimensionality 를 가지는 각 ellipsoid 에 대해 additional tokens 를 introduce합니다. 이 tokens는 ellipsoids의 모든 -invariant quantities—그들의 size , squared radius of gyration , 그리고 secondary structure type —의 embeddings로 initialized됩니다. Then, 각 model layer에서, 이 tokens는 두 가지 mechanisms를 via하여 residue representations 의 updates를 inform하고 (그리고 themselves updated됩니다):
- ellipsoids의 location과 shapes에 대한 information으로 residue tokens 를 update하기 위해, 우리는 novel invariant cross attention mechanism (Algorithm 1 and Figure 1)을 introduce하며, 이를 통해 values는 ellipsoid tokens로부터 -invariant manner로 aggregated됩니다. Similar to IPA, 이 mechanism은 invariance를 enforce하기 위해 residue local frames를 uses하지만, ellipsoid tokens는 themselves updated되지 않으며, 이는 Appendix B에서 further discuss합니다.
- residue와 ellipsoid tokens가 mutually update each other할 mechanism을 provide하기 위해, 우리는 Transformer stack 직전에 sequence dimension을 따라 tokens를 concatenate하고, 그 후에 sequence를 re-split합니다.
frame update 및 edge update layers와 같은 Multiflow update blocks의 All other aspects는 architecturally unmodified 상태로 remain합니다. Algorithms 1과 2에서는, modifications가 녹색으로 colored된 modified update block과 함께 new invariant cross attention을 outline합니다.
3.3 GUIDANCE FOR THE SELF-CONDITIONED AND JOINT FLOW
를 samples하는 base protein structure generative model을 fine-tuning하여 ProtComposer의 distribution 를 obtain한 후, 우리는 guidance parameter 에 의해 controlled되는 classifier-free guidance를 via하여 두 distributions 사이를 interpolate합니다. This enables 으로 recovered되는 의 designability와 에 corresponding하는 의 diversity, novelty, ellipsoid adherence 사이에서 trade off하기 위한 optimal 를 finding하는 것을 가능하게 합니다. individual samples에 대한 Interpolations는 Figure 3에 visualized되어 있습니다.
이러한 guidance를 Implementing하는 것은 우리가 (diffusion models에서처럼) score 대신 flow field를 model한다는 사실, ProtComposer의 conditional probability paths가 Gaussian이 아니라는 점 (Ho & Salimans (2022)의 guided flows가 directly applicable하지 않음), 그리고 우리가 self-conditioning을 employ한다는 점 때문에 complicated됩니다. self-conditioning difficulty에 대해 elaborating하기 전에, 우리는 각 inference step에서 그들의 flow fields를 separately interpolating함으로써 translations, rotations, discrete residue types에 대한 joint flow를 어떻게 guide하는지 lay out합니다:
- Translations: 우리는 unconditional vector field 와 conditioned version 를 로 interpolate합니다. Since translations에 대한 conditional probability paths는 Gaussian paths이므로, 이는 Zheng et al. (2023b)의 guided flows에 corresponds하며, 이는 guided diffusion models과 같이 conditional distribution에 의해 tilted된 unconditional distribution의 same approximation을 sample합니다, i.e., 만약 우리가 translations만 sample하는 models을 interpolate한다면 의 approximation입니다.
- Rotations: 상의 non-Gaussian paths에서는 Zheng et al. (2023b)의 results가 더 이상 hold하지 않지만, 우리는 translations와의 analogy를 사용하여, Yim et al. (2024b)를 follows하는 를 사용하는 것에서 empirical success를 find합니다.
- Discrete Flow: 우리는 discrete flow를 위한 rate matrix를 unconditional model's predictions와 ellipsoid conditioned model's predictions의 combination으로 obtain하는 denoised residues의 predicted probabilities에 대한 conditional rate matrix (see Campbell et al. (2024))의 expectation으로 construct합니다. Specifically, 우리는 ellipsoid conditioned probabilities (여기서 superscript는 denoising time을 denotes)에 의해 tilted된 unconditionally predicted probabilities를 use합니다.
Multiflow와 ProtComposer 모두 self-conditioning을 use하며, 이는 inference 중에 flow-model이 additional self-conditioning input으로 previous integration step의 output을 receives하는 것입니다. During inference, unconditional model 는 self-conditioning variable 를 produces하고, ellipsoid conditioned model 로부터 우리는 를 obtain합니다. unconditional model에 를, conditioned model에 를 supplying하는 Instead of, 우리는 모든 에 대해 better designability와 ellipsoid adherence를 achieves하는 를 for both 사용합니다. self-conditioning variants의 exploration 및 ablation은 Appendix C.2에 있습니다.
3.4 GENERATING NOVEL ELLIPSOIDS
ellipsoid conditioning mechanism을 designing함에 있어, 우리는 지금까지 inference 동안 provided된 ellipsoids의 sources에 대해 no assumptions를 made했습니다. manually specified ellipsoids를 사용하는 Next to, protein structures의 unconditional distribution을 로 factorized하여 sample하기 위해 additional generative model 로부터 synthetic ellipsoids를 sampling하는 opportunity도 있습니다. 에 대해 deep learning solution을 use하는 것이 tempting할 수 있지만, 우리는 purposefully 이를 avoid하고 factorization이 ellipsoid layouts에 대한 simple statistical model로 best leveraged된다고 argue합니다. training data와 similar한 layouts를 produce할 수 있는 deep learned Instead of, 에 대한 simple statistical model은 diverse하고 novel layouts를 sampling하는 것을 guarantees하며, 이는 로부터 더 diverse하고 novel protein structures로 lead합니다 - 이는 aim이 commonly novel designs를 produce하는 것인 protein design에 crucial한 properties입니다.
novel ellipsoid layouts를 generate하기 위해, 우리는 first K ellipsoids에 대한 means와 covariances를 sample한 다음 secondary structure와 residue count annotations를 assign합니다. means와 covariances에 대한 model은 [ p \left( { (\mu_k, \Sigma_k) }{k=1}^{K} \right) \propto \prod{k=1}^{K} \mathcal{N} \left( \mu_k; 0, \sigma^2 I_3 \right) \mathcal{W}3 \left( \Sigma_k; \psi^2 I_3, \nu \right) \exp \left( - U \left( { (\mu_k, \Sigma_k) }{k=1}^{K} \right) \right), \quad U \left( { (\mu_k, \Sigma_k) }{k=1}^{K} \right) = \sum{k \neq j} \frac{1}{\left[ (\mu_k - \mu_j)^T \Sigma_k^{-1} (\mu_k - \mu_j) \right]^2} . ] 입니다. 즉, ellipsoid means와 covariances는 각각 isotropic Gaussian 및 Wishart distributions로부터 i.i.d.하게 drawn되고, ellipsoid overlaps를 penalizes하는 energy function의 Boltzmann factor와 multiplied됩니다. Intuitively, 는 ellipsoid의 spread를 controls하고, 는 그들의 volume을 controls하며, 는 그들의 anisotropy 또는 "roundness"를 controls하고, 는 overlaps를 prevents합니다. energy 는 pairwise Mahalanobis distances에 based한 simple inverse square repulsion입니다. 우리는 rejection sampling을 via하여 이를 sample합니다, i.e., 를 sampling하고, 그들의 energy 를 evaluating하며, probability 로 rejecting합니다.
ellipsoid annotations를 choose하기 위해, 우리는 first 각 ellipsoid를 probability 로 로, probability 로 로 independently annotate합니다. We then observe that 주어진 choice에 대해, ellipsoid volume
가 simple linear fit (Appendix 10)에 의해 residue count를 strongly determines한다는 것을 observe합니다. Hence, 우리는 이를 independently modeling하는 instead of residues의 number를 assign하기 위해 이 linear fit을 use합니다.
ProtComposer 방법론: 정리 노트 (AI 연구자 대상)
3.1 Ellipsoid Representation of Proteins (타원체 표현법):
- 핵심: Protein spatial layout을 개의 annotated 3D ellipsoids 세트로 represent. 각 ellipsoid는 semantically coherent region(e.g., secondary structure)을 나타냄 (mean , covariance , feature , residue count ).
- 장점: Global annotation보다는 높은 control, mesh/voxel보다는 낮은 복잡도. Human-interpretable, easy-to-construct, low-level detail 과소 제약 방지.
- 추출 방법 (Training용):
- Simple graph-based segmentation: 공간적 근접성(<5Å) && semantic 유사성(secondary structure type from DSSP) 기준으로 residue 연결.
- Connected components 각각에 Gaussian() fitting.
- Loop 및 작은 segment(<=5 residues) 제외.
3.2 Ellipsoid Conditioning (타원체 조건 주입):
- 기반 모델: Multiflow (joint sequence-structure flow matching, SE(3) frames , amino acids , Transformer/IPA 기반).
- 조건 주입 전략:
- Pre-trained Multiflow ()를 fine-tuning하여 conditional density 학습 (ProtComposer).
- Minimal perturbation 원칙: Unconditional 모델 성능을 최소한으로 해치면서 조건 정보 주입 (Image diffusion 분야 참고).
- 구현:
- Ellipsoid tokens () 도입: Ellipsoid의 invariant 속성()으로 초기화.
- Invariant Cross Attention (ICA - Alg 1): 제안된 핵심 메커니즘. Residue tokens()가 ellipsoid parameters() 정보를 SE(3)-invariant 방식으로 attend하여 업데이트. 상대적 기하 정보()와 attention 활용.
- Transformer Integration (Alg 2): Multiflow의 update block 내에 ICA 통합. 와 token을 concatenate하여 Transformer 통과 (상호 업데이트), 이후 split. 기존 IPA, RigidUpdate 등은 구조 유지.
3.3 Guidance for the Self-Conditioned and Joint Flow (가이던스):
- 목표: Classifier-free guidance( 파라미터)를 이용, unconditional()과 conditional() 분포 사이 interpolation. Designability() vs Adherence/Diversity/Novelty() trade-off 조절.
- 구현: Flow field 모델링, Non-Gaussian path(rotation), self-conditioning 고려 필요.
- 각 modality별 flow field 선형 보간:
- Translations: . (Gaussian path 가정 시 표준 guided flow와 유사)
- Rotations: 유사하게 . (경험적 성공)
- Discrete Types: 예측 확률 보간 후 rate matrix 계산.
- Self-Conditioning 처리: Guidance 시 unconditional/conditional 예측 모두에 보간된 self-conditioning 변수 사용.
- 각 modality별 flow field 선형 보간:
3.4 Generating Novel Ellipsoids (신규 타원체 생성):
- 목표: Novel/diverse한 ellipsoid layouts () 생성을 통해 완전한 generative sampling 가능.
- 접근법: Deep learning 대신 simple statistical model 사용. (Diversity/Novelty 보장 목적).
- 모델:
- Means() / Covariances() 샘플링 + Overlap 방지 repulsion energy() 기반 rejection sampling().
- Feature() 랜덤 할당 (확률 ).
- Residue count()는 ellipsoid volume()과 기반 linear fit으로 결정.
핵심: Ellipsoid 표현법 정의 및 추출 -> ICA 메커니즘을 통한 Multiflow fine-tuning -> Guidance 기반 trade-off 조절 -> Statistical model 기반 novel layout 생성.
쉬운 설명 :
3.1 단백질을 '타원체(Ellipsoid)'로 표현하기:
- '설계도' 만들기: 단백질의 전체적인 모양과 배치를 간단하게 표현하기 위해, 주요 부분들을 3D ellipsoids(3차원 타원체)로 나타냅니다.
- 타원체 정보: 각 타원체는 해당 부분이 어디에 있는지(위치 ), 어떤 모양과 방향인지(형태 ), 어떤 종류의 구조인지(특징 , 예: 나선형/병풍형), 그리고 얼마나 많은 구성 요소(residues )로 이루어졌는지 정보를 담고 있어요.
- 왜 타원체? 너무 자세하지도, 너무 단순하지도 않아서 제어하기 좋고 사람이 이해하거나 만들기 쉬워요.
- 실제 단백질에서 타원체 얻기 (학습 데이터 준비): 간단한 프로그램을 이용해, 실제 단백질 구조에서 비슷한 종류(특징 )이면서 가까이 있는 부분들을 그룹으로 묶고, 각 그룹에 맞는 타원체를 계산해서 얻습니다.
3.2 AI에게 '타원체 설계도' 보는 법 가르치기:
- 기본 AI: Multiflow라는 강력한 단백질 생성 AI를 사용합니다.
- 튜닝: 이 Multiflow AI를 약간 수정(fine-tuning)해서, 우리가 준 타원체 '설계도'()를 이해하고 그에 맞춰 단백질을 만들도록 가르칩니다.
- 핵심 기술 (Invariant Cross Attention - ICA): AI의 단백질 구성 요소(residue tokens )가 타원체 설계도() 정보를 '참조'할 수 있게 하는 새로운 계산법(Algorithm 1)입니다. 단백질 전체가 돌아가거나 움직여도(SE(3)-invariant) 헷갈리지 않고 정보를 잘 받아들이도록 설계되었어요.
- AI 구조에 통합 (Algorithm 2): 이 새로운 계산법(ICA)을 Multiflow AI의 기존 처리 과정(Update Block)에 끼워 넣습니다. 단백질 내부 정보와 타원체 설계도 정보가 서로 영향을 주고받을 수 있도록 Transformer 등을 함께 사용합니다.
3.3 '설계도' 얼마나 따를지 조절하기:
- 조절 스위치 (Lambda - ): AI가 타원체 설계도를 얼마나 엄격하게 따를지 조절하는 스위치()를 도입합니다.
- : 설계도 무시. AI가 원래 배운 대로 자유롭게 생성 (더 자연스럽거나 안정적(designable)일 수 있음).
- : 설계도 최대한 따름. 설계대로 제어(control)가 잘 되고, 설계도가 새롭다면(novel/diverse) 결과물도 새로울 수 있음.
- : 둘 사이의 trade-off. 사용자가 원하는 균형점을 찾을 수 있습니다.
- 구현: 이 스위치를 AI가 다루는 여러 정보(위치, 방향, 아미노산 종류)에 각각 적용하는 방법과, self-conditioning이라는 AI 기법을 함께 사용할 때의 처리 방법을 설명합니다.
3.4 새로운 '설계도' 자동으로 만들기:
- 설계도 출처: 사용자가 직접 타원체 설계도를 그릴 수도 있지만, AI가 스스로 새롭고 다양한 설계도를 만들게 할 수도 있습니다.
- 방법: 복잡한 AI 대신, **간단한 통계 프로그램(statistical model)**을 사용해서 무작위적이면서도 그럴듯한 타원체 배치()를 만듭니다. (간단한 프로그램이 오히려 더 새롭고(novel) 다양한(diverse) 결과를 보장하기 때문입니다.)
- 프로그램 작동 방식: 무작위로 타원체들의 위치와 모양을 만들되, 너무 겹치지 않게 서로 밀어내도록 합니다. 각 타원체에 나선형/병풍형 특징()을 무작위로 부여하고, 크기에 맞춰 적절한 구성 요소 개수()를 추정합니다.
주인장 이해
- ProtComposer 모델 (Section 3.2): Ellipsoid 조건 를 입력받아, 현재 상태(초기 노이즈 또는 중간 생성물)에서 최종 구조로 어떻게 변해가야 하는지에 대한 방향(flow field)을 예측하도록 학습된 모델입니다.
- Guidance 메커니즘 (Section 3.3): 생성 과정의 각 단계에서 ellipsoid 조건 를 얼마나 강하게 따를지( 값) 조절하면서, ProtComposer 모델(조건부)과 Multiflow 모델(조건 없는)의 예측 flow field를 적절히 혼합합니다.
- Flow Matching Inference: 이 혼합된 flow field를 따라, 초기 노이즈 상태에서 시작하여 점진적으로(여러 시간 단계에 걸쳐) 구조와 서열을 업데이트해 나갑니다. 이 과정(ODE Solver 등을 이용한 적분 과정)을 끝까지 수행하면 최종적인 단백질 구조와 서열이 **생성(generate)**됩니다.
요약:
- 구조 '예측'보다는 '생성'에 가깝습니다.
- 실제 생성은 Inference 단계에서 일어납니다.
- ProtComposer 모델(3.2)이 생성 방향(flow field)을 예측하고,
- Guidance 메커니즘(3.3)이 이 방향을 조절하며,
- Flow matching 과정을 통해 노이즈로부터 최종 구조/서열을 점진적으로 만들어냅니다.
1. 단백질 구조, residue의 좌표 그리고 각 residue가 어떤 스트럭쳐에 속하는지 이는 DSSP라는 프로그램을 사용하여 정보를 얻음. 암튼 이걸 준비.
2. 각 residue를 node라고 보고 같은 secondary structure와 가까울때 edge로 연결한다.
3. 서로 연결된 residue들을 찾고 이를 connected components라고 부름 이는 의미적으로 일관된 영역임.
4. 암튼 그 식별된 region에 대해 4가지 정보를 계산한다. residue들의 공통된 secondary structure 라벨, residue 평균 위치, 공분산행렬, residue의 수
5. Pre-trained Unconditional Multiflow Model을 준비. 이는 단백질 구조를 곧 잘 생성하는 모델이며, sequence와 structure도 함께 생성할 수 있음.
6. residue tokens 과ellipsoid tokens 를 임시로 하나로 합쳐서 트랜스포머를 태워서 후에 다시 스플릿.
7. 그 정보를 가지고 올바른 단백질 구조 방향으로 flow field를 업데이트
Translation vector field Residue 위치를 어디로 이동시킬지 방향과 크기
Rotation vector field Residue 방향을 어떻게 회전시킬지
Rate matrix 아미노산 종류(sequence)를 어떻게 바꿀지
8. 모델 두개써서 사이의 값을 사용하게 끔 함 그래서 단백질 구조 예측에 대한 가이던스
9. 3.4는 자동 생성