AI바라기의 인공지능
Diffusion : 논문리뷰 : Equivariant Blurring Diffusionfor Hierarchical Molecular Conformer Generation 본문
Diffusion : 논문리뷰 : Equivariant Blurring Diffusionfor Hierarchical Molecular Conformer Generation
AI바라기 2025. 4. 3. 14:33쉬운 설명:
- 기존 방식은 원자 하나하나를 놓으려다 보니, 특히 반지(ring) 같은 복잡한 부분을 이상하게 만들거나 큰 분자를 잘 못 만들었습니다.
- 이 논문 (HierDiff)은 **레고 블록(fragment)**처럼 분자를 조각내서, 먼저 큰 블록들의 대략적인 위치와 종류를 정하고 (Coarse-grained Diffusion), 그 다음 각 블록을 세밀하게 다듬고 서로 연결하는 방식 (Fine-grained Generation + Iterative Refinement)을 사용합니다.
- 이렇게 하면 레고 블록 자체는 이미 모양이 잡혀있어(valid local structures) 이상한 구조가 덜 생기고, 큰 구조물도 더 안정적으로 만들 수 있습니다. 특히, **연결 부분을 반복해서 수정하는 과정(Iterative Refinement)**이 분자 구조의 완성도를 높이는 핵심입니다.
HierDiff 논문 학습 노트
Purpose of the Paper:
- 기존 연구 한계: 기존 3D 분자 생성 모델 (특히 atom-level non-autoregressive 방식, 예: EDM)은 원자 단위로 생성하기 때문에 ring과 같은 국소 구조(local structure)의 화학적 유효성을 보장하기 어렵고, 특히 큰 분자 생성 시 품질이 저하되는 문제가 있었습니다.
- 새로운 접근 방식: Fragment 기반 생성은 유망하지만, 3D non-autoregressive 방식에 적용하기에는 fragment 간 연결 시 발생하는 조합적 폭발(combinatorial explosion) 및 충돌(conflict) 문제 (Fig. 2) 해결이 어렵습니다.
- 목표: 본 논문은 이러한 문제를 해결하기 위해 Coarse-to-Fine 전략에 기반한 새로운 **Hierarchical Diffusion Model (HierDiff)**을 제안합니다. Fragment의 장점(유효한 local structure 보존)을 살리면서 non-autoregressive 방식의 장점(global modeling)과 연결 문제를 해결하여, 화학적으로 유효하고 품질 높은 3D 분자(특히 크기가 큰 drug-like 분자)를 생성하는 것을 목표로 합니다.
Key Contributions:
- Hierarchical Diffusion Framework (HierDiff):
- 3D 분자 생성을 위한 새로운 2단계 (Coarse-to-Fine) Diffusion Process 제안.
- Contribution: 분자 구조 생성을 계층적으로 분해하여 복잡도를 낮춤.
- Novelty: Diffusion model을 fragment 단위의 계층적 생성에 처음으로 적용.
- Coarse-Grained Fragment Diffusion:
- 분자 전체의 global layout을 나타내는 coarse-grained node (fragment)의 추상적 화학 정보 (Hf: invariant)와 3D 좌표 정보 (Hp: equivariant)를 diffusion process를 통해 생성.
- Contribution: Fragment 수준에서 분자의 전역적 배치와 속성을 먼저 결정.
- Novelty: SE(3)-equivariant diffusion을 fragment의 기하학적 위치 생성에, 표준 diffusion을 화학적 특징 생성에 동시 적용.
- Fine-Grained Decoding & Iterative Refinement:
- 생성된 coarse-grained node 정보를 바탕으로, equivariant message passing network (EGNN 변형)를 이용해 구체적인 fine-grained fragment로 decoding.
- Contribution: Fragment 간의 연결 가능성과 화학적 유효성을 보장하며 구체적인 fragment 구조 생성.
- Novelty: Decoding 과정에서 발생할 수 있는 편향(bias)을 수정하고 연결성을 강화하기 위한 Iterative Refinement 모듈을 새롭게 도입 (Fig. 6).
- Constraint Handling via Decomposition:
- Fragment 연결 시 발생하는 복잡한 제약 조건 문제를 P(V,E) ≈ P(H)P(V,E|H) 형태로 분해하여 해결. Coarse diffusion이 P(H)를, Fine-grained generation이 P(V,E|H)를 모델링.
- Contribution: Non-autoregressive fragment 생성의 combinatorial complexity 문제를 효과적으로 완화.
- Novelty: Hierarchical decomposition을 통해 3D fragment 연결 제약 문제를 diffusion framework 내에서 해결.
Experimental Highlights:
- Datasets: GEOM-DRUG, CrossDocked2020 (Drug-like molecules), QM9 (Small molecules).
- Baselines: EDM, G-SphereNet (Autoregressive).
- Metrics:
- Drug-likeness: QED, RA, MCF, SAS, ∆LogP, ∆MW.
- Conformation Quality: Coverage (Cov), Matching (Mat) (atom & fragment level).
- 기타: Validity, Stability, Diversity.
- Key Results:
- HierDiff는 GEOM-DRUG 및 CrossDocked2020 데이터셋에서 기존 SOTA 모델 (EDM, G-SphereNet) 대비 대부분의 Drug-likeness 및 Conformation Quality 지표에서 월등한 성능을 달성 (Table 1, Table 2). 특히 MCF, RA, ∆MW, Cov, Mat 개선이 두드러짐.
- 기존 모델보다 더 크고 현실적인 분자 구조 생성에 강점을 보임 (낮은 ∆MW, Fig 7).
- EDM 대비 더 안정적이고 화학적으로 유효한 분자 구조(특히 ring 구조)를 생성함을 시각적으로 확인 (Fig 7).
- Ablation study를 통해 Iterative Refinement 모듈의 중요성을 입증 (Appendix C.7, Table 9).
- Autoregressive 모델 (G-SphereNet) 대비 오류 누적(error accumulation) 문제에 강인함을 보임 (Appendix C.10, Table 12).
Limitations and Future Work:
- Limitations:
- Diffusion process 자체의 샘플링 속도 한계 (Appendix C.8, Table 11). Fine-grained decoding 단계에서 추가적인 계산 비용 발생.
- 미리 정의된 fragment vocabulary 및 fragmentation 전략에 의존함 (Appendix C.7에서 사용된 전략의 강건함은 보였으나 근본적인 의존성은 존재).
- 매우 작은 분자(QM9) 생성 시 fragmentation으로 인한 정보 손실 가능성 존재 (본문 QM9 결과 논의).
- Future Work:
- 더 적은 diffusion step으로 모델을 학습/샘플링하여 샘플링 효율성 개선 (Appendix C.8).
- 다른 제약 조건이 있는 생성 문제 (e.g., protein alignment)로의 프레임워크 확장.
- 최적의 fragment feature 및 fragmentation 전략 탐구.
- Conditional generation 능력 향상.
Overall Summary:
- HierDiff는 atom-level 생성 방식의 한계(부정확한 local structure, 큰 분자 생성 어려움)를 극복하기 위해 새로운 Coarse-to-Fine 계층적 Diffusion 모델을 제안합니다.
- 이 모델은 먼저 coarse-grained fragment의 기하학적/화학적 표현을 생성하고, 이를 equivariant network와 핵심적인 Iterative Refinement 단계를 통해 구체적인 fragment로 변환하여 최종 분자를 조립합니다.
- 실험 결과, HierDiff는 기존 SOTA 모델들을 능가하며 drug-likeness와 conformation quality 측면에서 우수한 성능을 보여, 고품질의 화학적으로 유효한 3D 분자 생성에 효과적인 접근법임을 입증했습니다.
Abstract
diffusion models는 어떻게 우리의 multiscale 세계관과 유사하게 coarse-to-fine 방식으로 3D geometries를 처리할 수 있을까요? 이 논문에서는 multiscale 방식으로 molecular graphs에 조건화된 3D molecular conformers를 생성하는 근본적인 생화학 문제에 초점을 맞춰 이 질문을 다룹니다.
우리의 접근 방식은 두 가지 계층적 단계로 구성됩니다: i) molecular graph로부터 coarse-grained fragment-level 3D structure의 generation, ii) 후자가 동시에 조정될 수 있도록 허용하면서 coarse-grained approximated structure로부터 fine atomic details의 generation.
coarse-grained 정보를 보존하면서 SE(3) equivariance를 보장해야 하는 어려운 두 번째 단계를 위해, 우리는 Equivariant Blurring Diffusion (EBD)라는 새로운 generative model을 소개합니다. 이는 conformers의 fine atomic details를 흐리게 하여 fragment-level coarse-grained structure 쪽으로 이동하는 forward process와 equivariant networks를 사용하여 반대 작업을 수행하는 reverse process를 정의합니다.
우리는 drug-like molecules의 benchmark에서 state-of-the-art denoising diffusion models과의 기하학적 및 화학적 비교를 통해 EBD의 효과를 입증합니다. Ablation studies는 loss function 및 data corruption process의 설계를 포함하는 architecture를 철저히 분석하여 EBD 설계에 대한 통찰력을 도출합니다.
Figure 1은 "Blurring diffusion generative processes" 개념을 이미지와 분자 conformer 두 가지 예시를 통해 보여줍니다. 생성 과정이 Condition에서 시작하여 Coarse structure를 거쳐 최종적으로 Fine-grained structure를 만드는 흐름을 나타냅니다. (이는 논문 본문에서 설명하는 reverse process, 즉 생성 과정을 시각화한 것으로 보입니다.)
- 윗줄 (Image 예시 - 논문 [42] 참조):
- Condition: "Dog" 라는 텍스트 정보가 주어집니다.
- Coarse structure: 이미지가 매우 흐릿하게 처리된 상태, 여기서는 단색의 사각형으로 표현되어 이미지의 평균적인 색상 정보나 매우 추상화된 형태를 나타냅니다.
- Fine-grained structure: 최종적으로 생성된 선명한 개의 이미지입니다. 즉, coarse한 정보에서 시작하여 점차 details를 추가하여 최종 image를 생성하는 과정을 보여줍니다.
- 아랫줄 (Molecular Conformer 예시 - 이 논문의 제안 방법):
- Condition: 2D molecular graph (분자의 2차원 화학 구조식)가 주어집니다.
- Coarse structure: 분자의 fragment들의 center of mass 등 매우 단순화된 3차원 구조 정보입니다. 원자 수준의 상세함이 없는, 말 그대로 '거친(coarse)' 구조를 나타냅니다.
- Fine-grained structure: 최종적으로 생성된, 모든 원자의 3차원 좌표가 포함된 상세한 3D molecular conformer입니다. 원자, 결합, 3차원 배치가 모두 표현된 구조입니다.
요약: 이 그림은 diffusion model이 어떻게 매우 단순하거나 '흐릿한' (coarse) 상태에서 시작하여 점진적으로 세부 정보(fine-grained details)를 추가하여 최종 결과물(이미지 또는 분자 구조)을 생성하는지를 시각적으로 비교하여 보여줍니다. 특히 아랫줄은 이 논문에서 제안하는 EBD가 2D molecular graph라는 condition과 coarse-grained fragment structure를 바탕으로 fine-grained 3D conformer를 생성하는 핵심 아이디어를 나타냅니다.
1 Introduction
객체의 multiscale 속성을 이해하기 위한 generative models의 발전은 개별 scale을 초월하여 다양한 granularity 수준에 걸쳐 이 모델들의 응용을 촉진합니다. generative models가 multiscale 구조를 처리할 수 있도록 하기 위해, images에서 speech에 이르기까지 여러 도메인에 걸쳐 hierarchical design 방법론이 급증했습니다. 이러한 방법들은 초기에 coarse-grained structures를 포착하고 이후에 finer details를 생성합니다.
computer vision 분야에서 최근의 노력들은 noises를 추가하고 제거함으로써 데이터를 corrupt하고 restore하는 denoising diffusion models를 활용하여 images의 2D pixels에 대한 coarse-to-fine generative frameworks의 성공적인 설계를 이끌어냈습니다. 특히, 한 연구는 partial differential equations의 heat equation에서 동기를 얻어 blurred prior distributions(픽셀 강도의 평균)로부터 images를 생성했습니다 (Fig. 1).
그러나 biochemistry 및 drug discovery 분야에서는 stable molecular structures의 3D conformers를 위한 denoising diffusion models이 아직 coarse-to-fine multiscale frameworks의 이점을 활용하지 못했습니다. 현재 방법들은 scale hierarchy를 무시하거나 매우 제한적인 방식으로만 고려합니다. 예를 들어, 최근의 unconditional conformer generation의 hierarchical 방법 내에서는 denoising diffusion model이 coarse-to-fine generation으로 확장되지 않고 오직 coarse-grained structure의 generation에만 적용됩니다.
molecular conformers를 위한 denoising diffusion models을 hierarchical designs으로 확장하는 데 있어 주요 bottleneck은 random noise가 fine atomic details뿐만 아니라 coarse-grained structures의 structural information까지 무차별적으로 corrupt한다는 점입니다. 이 어려운 문제를 해결하기 위해, 우리는 2D molecular graphs에서 자주 발생하는 substructures 또는 functional groups인 fragments를 활용합니다. 이러한 fragments는 3D geometry에서 coarse-grained structural information을 위한 유망한 후보가 될 수 있습니다. fragments를 도입하면 generation process가 두 단계로 나뉩니다: i) fragments로 표현되는 coarse-grained structures 생성, ii) fragment structures로부터 fine atomic details 복원. molecular graphs로부터 fragment coordinates를 생성하는 첫 번째 단계에서는, cheminformatics tool로부터 각 fragment의 center of mass와 attributes로 구성된 fragment structures의 approximations를 효율적으로 생성합니다.
coarse-to-fine generation의 어려운 두 번째 단계를 위해, 우리는 다음과 같이 상세히 설명될 새로운 diffusion model인 Equivariant Blurring Diffusion (EBD)를 제안합니다. heat equation의 blurring corruption에서 영감을 받아, 우리는 EBD가 random noise 대신 coarse-grained fragment approximated structures로부터 3D molecular conformers를 생성하도록 설계했습니다 (Fig. 1). 우리가 설계한 EBD에서, forward process는 conformers의 atom positions을 각 해당 fragment의 center of mass 쪽으로 이동시키고, reverse process는 3D fragment structure의 prior distribution으로부터 완전한 원자 details를 복원합니다. 우리가 EBD를 위해 설계한 blurring schedule은 diffusion model이 전체 generative process 동안 coarse-grained information을 유지하면서 fine atomic details 복원에 집중할 수 있도록 합니다. 우리는 drug-like molecules의 benchmark를 사용하여 우리의 coarse-to-fine EBD model을 검증했습니다. 우리는 denoising diffusion model에 비해 100배 적은 diffusion time steps으로도 conformer generation에서 우수한 결과를 얻었습니다.
이 논문의 주요 contribution은 다음과 같이 요약될 수 있습니다:
- 우리는 heat equation의 blurring corruption에서 동기를 얻어, equivariant networks를 사용하여 fragment structures의 coarse-grained estimation으로부터 atomic details를 생성하는 EBD를 설계합니다.
- 우리는 기존 image blurring diffusion model의 것을 직접 적용하는 대신, 성능에 중대한 영향을 미치는 새로운 blurring scheduler와 수정된 loss function을 제안합니다.
- 우리는 fragment granularity, data corruption methods, loss reformulation의 효과에 대한 철저한 analysis를 수행합니다. 우리는 state-of-the-art denoising diffusion models과 비교하여 geometrically 및 chemically 더 plausible한 conformers를 얻었습니다.
정리노트 (AI 연구자 대상)
논문 핵심: Equivariant Blurring Diffusion (EBD) - Coarse-to-Fine 3D 분자 구조 생성을 위한 새로운 Diffusion Model
- 문제 정의:
- 기존 Denoising Diffusion Models (DDM)은 3D molecular conformer 생성 시 multiscale (coarse-to-fine) 접근법 적용에 한계가 있음.
- 이유: Random noise 방식은 fine atomic details뿐 아니라 coarse-grained structural information까지 무차별적으로 손상시켜 계층적 생성에 부적합함.
- 제안 방법론: EBD
- 핵심 아이디어: Random noise 대신 구조화된 'Blurring' corruption 사용. Image blurring (heat equation 기반)에서 영감을 받음.
- Coarse-grained Representation: 2D molecular graph에서 식별된 화학적 fragments (자주 등장하는 부분 구조/작용기)를 활용.
- 2단계 생성:
- Stage 1: Molecular graph로부터 fragment 기반 coarse-grained structure (center of mass, attributes) 근사치 생성 (e.g., cheminformatics tool 활용).
- Stage 2 (EBD): 이 coarse structure를 prior 삼아 fine atomic details를 복원/생성.
- EBD 메커니즘:
- Forward Process: 실제 conformer의 atom positions을 해당 fragment의 center of mass 방향으로 점진적으로 이동시켜 '흐릿하게(blur)' 만듦 (Fine → Coarse).
- Reverse Process: SE(3) Equivariant Network를 사용하여, coarse (흐릿한) 상태에서 시작하여 fine atomic details를 복원 (Coarse → Fine). Equivariance 보장으로 3D 구조의 회전/이동에 강인함.
- 주요 기술적 Contributions:
- 분자 구조 생성에 특화된 blurring scheduler 및 loss function 제안 (단순 image blurring 방식 적용 이상).
- Fragment granularity, data corruption 방식, loss 설계에 대한 심층 analysis 제공.
- 결과:
- State-of-the-art DDM 대비 geometrically 및 chemically 더 plausible한 conformer 생성.
- 훨씬 적은 diffusion time steps (100배 감소)으로 우수한 성능 달성.
시사점: Random noise 대신 구조적 prior (fragments)를 향한 blurring과 SE(3) equivariance를 결합하여 3D 데이터의 hierarchical generative modeling 문제에 대한 효과적인 접근법 제시.
쉬운 설명:
이 논문의 Introduction 섹션에서는 3D 분자 구조를 AI로 더 똑똑하게 만드는 새로운 방법(EBD)을 소개합니다.
기존 방식의 문제점: 레고로 복잡한 모델을 만들 때, 작은 부품들을 마구잡이로 흔들면(random noise) 큰 덩어리까지 망가져서 처음부터 다시 만들기 어려운 것과 비슷합니다. 기존 AI(diffusion model)는 분자의 아주 세밀한 부분(fine atomic details)을 만들려고 할 때, 분자의 전체적인 큰 뼈대(coarse-grained structure) 정보까지 손상시키는 문제가 있었습니다. 단계적으로 (큰 뼈대 -> 세부 장식) 만들기가 어려웠죠.
이 논문의 새로운 아이디어 (EBD):
- 큰 뼈대 먼저 생각하기: 분자를 잘게 쪼갠 '익숙한 부품'(fragments)을 먼저 생각합니다. 이걸 분자의 대략적인 뼈대(coarse structure)로 삼습니다.
- 흔드는 대신 '흐리게' 만들기: 아주 선명한 완성된 분자 사진(fine atomic details)을 점차 이 '부품 덩어리'(coarse structure) 모양으로 뭉개지도록 '흐리게(blur)' 만듭니다 (이게 forward process). 마구잡이로 흔드는 게 아니라, 정해진 방향(뼈대 모양)으로 뭉개는 거죠.
- '흐림'을 되돌려 세부사항 만들기: AI(EBD model)는 이 '흐릿한 뼈대' 상태에서 시작해서, 거꾸로 점차 '선명하게' 만들면서 원래의 세밀한 3D 분자 구조를 완성하는 법을 배웁니다 (이게 학습하는 reverse process).
- 3D 공간 이해: 이 AI는 분자를 3D 공간에서 돌리거나 움직여도 똑똑하게 이해하도록 특별히 설계되었습니다(SE(3) equivariant).
결론: 이 새로운 방식(EBD)은 분자의 큰 구조와 세부 구조를 단계적으로 잘 생성할 수 있고, 기존 방식보다 더 정확하고 빠르게 실제와 비슷한 3D 분자 구조를 만들 수 있다고 주장합니다.
2. Background
2.1 Blurring diffusion
Random noise를 추가하여 데이터를 corrupt하고 denoising을 통해 데이터를 생성하는 Denoising diffusion models은 다양한 도메인에서 상당한 발전을 이루었습니다. 최근 몇몇 연구들은 vision domain에서 random noise corruption을 넘어서, diffusion models의 설계 공간에 data corruption을 도입했습니다.
Inverse Heat Dissipation Model (IHDM)은 pixel space에서 coarse-to-fine generation을 제안했습니다. 그들의 forward process는 grids 상에서의 heat dissipation의 partial differential equation을 따릅니다: [ \frac{\partial}{\partial t} x(i, j, t) = \Delta x(i, j, t), ] 여기서 x는 grid 상의 data를 나타내고 ∆는 Laplacian operator입니다. IHDM은 time step t에서의 이 방정식의 해 x_t를 ∆의 eigendecomposition을 사용하여 다음과 같이 유도했습니다: [ x_t = B_t x_0 = V \exp(-\Lambda t) V^T x_0, ] 여기서 V^T와 Λ는 각각 discrete cosine transform과 ∆의 eigenvalues를 요소로 가지는 diagonal matrix입니다. t가 T에 가까워짐에 따라, eigenvalue 0의 eigenbasis만 남게 되고 이는 pixel intensities가 그들의 average value로 convergence하게 만듭니다. 이 blurring process에 기반하여, IHDM은 forward process를 다음과 같이 정의했습니다: [ q(x_t | x_0) = \mathcal{N}(x_t | B_t x_0, \sigma^2 I), ] 이는 t 시점의 state가 t 시점까지 blur된 data에 약간의 noise가 더해진 것과 같다는 것을 의미합니다. Data corruption 함수 B_t는 eigenvalues Λ의 spectral space에서 정의되었습니다. 그런 다음, reverse generative process는 각 state를 deblur하도록 정의되었습니다: [ p_{\theta}(x_{t-1} | x_t) = \mathcal{N}(x_{t-1} | \mu_{\theta}(x_t, t), \delta^2 I), ] 여기서 t-1에서의 mean은 deblurring network μθ의 결과이고 δ는 noise에 대한 작은 standard deviation입니다. t가 0에 가까워짐에 따라, μθ는 state 값들을 효과적으로 deblurring함으로써 pixel intensities에 대한 coarse-grained information으로부터 점차 fine details를 복원합니다. Loss는 무작위로 샘플링된 t에서 deblurring network의 결과와 덜 blur된 state 사이의 distance를 최소화하도록 정의되었습니다: [ L_{t-1} = \mathbb{E}{t,x_0,x_t} \left[ \left| B{t-1} x_0 - \mu_0 (x_t, t) \right|^2 \right]. ] IHDM은 image generation task에서 FID score를 사용하여 평가되었지만, 그 performance는 denoising diffusion models에 뒤처졌습니다. 예를 들어, IHDM은 CIFAR-10에서 FID score 18.96을 달성한 반면 DDPM은 3.17을 기록했습니다.
2.2 Equivariance
이 연구에서, 우리는 molecular conformers의 roto-translational equivariance를 다루기 위해 SE(3) group을 고려합니다. 함수 f가 group G에 대해 equivariant하다는 것은 모든 g ∈ G에 대해 T_g (f (x)) = f (S_g (x))가 성립하는 경우를 말하며, 여기서 T_g, S_g는 group element g의 transformations입니다. 우리의 coarse-to-fine generative framework에서, coarse-grained structure의 invariant prior distribution은 fragments의 coordinates를 나타냅니다. 따라서, 우리 diffusion model에서의 transition distribution과 loss function의 설계는 생성된 likelihood가 invariant하도록 보장해야 하며, 이를 통해 생성된 conformers가 rotation이나 translation에 의해 영향을 받지 않도록 해야 합니다.
정리노트 (AI 연구자 대상)
논문 관련 배경지식 요약
- Blurring Diffusion (기존 연구 - IHDM 중심):
- 동향: Diffusion model 연구가 random noise 외의 data corruption 방식으로 확장되고 있음 (vision domain 중심).
- IHDM (Inverse Heat Dissipation Model):
- Image 대상 coarse-to-fine generation 제안. Heat equation 기반 blurring 활용.
- Forward Process (q(x_t|x_0)): Laplacian eigendecomposition 기반 blurring operator (B_t) 적용 + 약간의 noise. t -> T 시 평균값으로 수렴. [ q(x_t | x_0) = \mathcal{N}(x_t | B_t x_0, \sigma^2 I) ]
- Reverse Process (p_θ(x_{t-1}|x_t)): Deblurring network (μ_θ) 학습하여 fine details 복원. [ p_{\theta}(x_{t-1} | x_t) = \mathcal{N}(x_{t-1} | \mu_{\theta}(x_t, t), \delta^2 I) ]
- Loss (L_{t-1}): Deblurring 결과와 덜 blur된 상태 (B_{t-1}x_0) 간의 distance 최소화. [ L_{t-1} = \mathbb{E}{t,x_0,x_t} \left[ \left| B{t-1} x_0 - \mu_0 (x_t, t) \right|^2 \right] ]
- 한계: Image generation에서 FID score 기준, 표준 DDPM 대비 낮은 performance. (이 논문이 개선하려는 지점)
- Equivariance (이론적 배경):
- 필요성: Molecular conformers는 3D 공간에서의 회전 및 이동(roto-translation)에 대해 본질적 구조가 변하지 않으므로, 생성 model은 SE(3) equivariance를 만족해야 함.
- 정의: 함수 f가 group G에 equivariant하다는 것은 T_g(f(x)) = f(S_g(x))가 성립함을 의미 (T_g, S_g는 g ∈ G에 대한 transformations).
- 이 논문 적용 방향:
- Coarse-to-fine framework에서 coarse-grained structure (fragment coordinates)의 prior distribution은 invariant해야 함.
- 따라서, 제안하는 diffusion model의 transition distribution과 loss function은 생성된 likelihood가 rotation 및 translation에 대해 invariant 하도록 설계되어야 함. 이는 최종 생성된 conformer의 기하학적 타당성을 보장하는 데 필수적임.
쉬운 설명:
이 섹션에서는 이 논문의 아이디어를 이해하는 데 필요한 두 가지 배경 지식을 설명합니다.
- '흐리게 만들기' 아이디어 (사진 필터 거꾸로 돌리기):
- 보통 AI가 그림을 그릴 때, 완전 노이즈(지직거리는 화면)에서 시작해서 점점 노이즈를 없애가며 선명한 그림을 만들어요.
- 그런데 어떤 사람들은 다른 방법을 시도했어요 (IHDM). 멀쩡한 사진을 점점 '흐리게' 만들어서 거의 형체만 남게 하고 (forward process), 그 다음 AI에게 이 '흐릿한 상태'에서 시작해서 거꾸로 점점 '선명하게' 원래 사진으로 되돌리는 법을 배우게 한 거죠 (reverse process). 마치 뿌옇게 만드는 필터를 거꾸로 돌리는 것과 같아요.
- 이 '흐리게 만들기'는 열이 퍼져나가는 방식에서 아이디어를 얻었대요.
- 근데 이 방법은 사진에서는 기존 노이즈 제거 방식만큼 성능이 좋지는 않았어요. (그래서 이 논문에서는 이 아이디어를 분자에 맞게 개선하려는 것 같아요.)
- 3D 도형 돌려보기 (장난감 자동차 예시):
- AI가 3D 분자 구조를 다룰 때는, 분자를 이리저리 돌리거나 옮겨도 그 분자 자체는 변하지 않는다는 걸 이해하는 게 중요해요. 예를 들어, 왼쪽을 보는 장난감 자동차나 오른쪽을 보는 장난감 자동차나 같은 자동차 모델이라는 걸 알아야 하죠.
- AI가 이렇게 3D 공간에서의 회전이나 이동에 대해 일관성 있게 반응하는 성질을 SE(3) equivariance라고 불러요. 복잡해 보이지만, 그냥 "3D 공간에서 돌려도 똑바로 알아본다"는 뜻이에요.
- 이 논문은 분자 구조를 만드는 AI가 이 SE(3) equivariance 성질을 갖도록 설계해서, 어떤 각도에서 보든 올바른 3D 구조를 만들어내도록 하려는 거예요.
Figure 2는 이 논문에서 제안하는 hierarchical molecular conformer generation framework 전체 과정을 두 단계로 나누어 보여줍니다.
Step 1: Fragment structure generation (조각 구조 생성)
- 시작: 복잡한 molecular graph G (분자의 2D 구조)에서 시작합니다.
- Fragmentation of G: 이 분자 G를 화학적으로 의미 있는 여러 개의 작은 fragments (조각들)로 분해합니다. 그림에서는 원래 분자가 몇 개의 작은 화학 구조로 나뉘는 것을 보여줍니다.
- Fragments coordinates ˆx_f: 분해된 각 fragment에 대해 3차원 공간에서의 대략적인 위치 좌표 (coarse-grained 좌표)를 계산합니다. 그림에서는 여러 개의 구(sphere)로 이 fragment들의 3차원 위치를 나타냅니다. 이 ˆx_f가 다음 단계의 입력 정보 중 하나가 됩니다. (이 과정은 Section 3.2에서 설명됩니다.)
Step 2: Equivariant Blurring Diffusion (등변성 블러링 확산)
이 단계는 Step 1에서 얻은 coarse-grained 정보 (ˆx_f)와 원래 molecular graph G를 바탕으로 최종적인 상세한 3D conformer (atom-level fine details)를 생성하는 diffusion model (EBD) 부분입니다.
- Graphical model:
- Forward Process (q, 점선 화살표): 실제 정답 conformer 구조(x_{a,0})에서 시작해서, 점차 blurring되어 fragment들의 좌표(Mˆx_f)에 가까워지는 과정 (x_{a,T})을 개념적으로 보여줍니다 (Fine-grained → Coarse-grained).
- Reverse Process (p_θ, 실선 화살표): AI 모델이 학습하는 실제 생성 과정입니다. Coarse한 상태 (x_{a,T})에서 시작하여 점진적으로 fine details를 복원(deblurring)하여 최종 conformer(x_{a,0})를 생성합니다 (Coarse-grained → Fine-grained).
- Blurring schedule:
- Forward process에서 사용하는 '공간적 blurring'이 어떻게 작동하는지 시각적으로 보여줍니다.
- 원래 원자의 상세한 위치(x_{a,0}, 복잡한 모양)가 시간 t가 0에서 T로 흐름에 따라 점진적으로 해당 fragment의 중심 좌표(Mˆx_f, 구 모양)로 부드럽게 이동(linear interpolation)하는 것을 보여줍니다. 아래 수식 x_{a,t} = (1 - t/T)x_{a,0} + (t/T)Mˆx_f이 이 과정을 나타냅니다. Random noise가 아니라, 목표 지점(Mˆx_f)을 향해 직접 이동하는 방식입니다.
요약: Figure 2는 먼저 분자를 fragments로 나누고 대략적인 3D 좌표(ˆx_f)를 얻은 다음 (Step 1), 이 정보를 바탕으로 EBD라는 diffusion model이 blurring의 역과정(deblurring)을 통해 최종적인 상세 3D conformer를 생성하는 (Step 2) 전체적인 2단계 framework를 보여줍니다. 특히 Blurring schedule은 EBD의 핵심인 공간적 blurring 방식을 시각화합니다.
3 Methods
3.1 Problem definition
SE(3)-invariant features ha ∈ Rn×d를 가진 n개의 atoms V (nodes)와 원자 간 결합 E (edges)로 구성된 molecular graph G가 있다고 가정합시다. 우리의 목표는 주어진 G에 대해 3D molecular conformers xa ∈ Rn×3의 ensemble을 생성하는 것입니다. 우리의 hierarchical approach는 두 단계로 이루어집니다. i) p(xf|G): m개의 fragments로 분해된 G로부터 fragment coordinates xf ∈ Rm×3의 coarse-grained 3D structure를 생성하는 단계, ii) p(xa|xf, G): 생성된 fragment structure xf에 conditioned된 fine atomic details xa ∈ Rn×3를 생성하는 diffusion model. Atoms와 각각의 fragments 사이를 매핑하기 위해, i번째 원자가 k번째 fragment에 속하면 Mik = 1이고 그렇지 않으면 0인 mapping matrix M ∈ Rn×m을 정의했습니다. Mxf는 각 원자를 해당 fragment 위치에 배치합니다. 반면에, M†xa는 fragment coordinates가 구성 atoms 좌표의 평균이 되도록 합니다. 여기서 M†는 M의 pseudoinverse matrix입니다 (M†M = I).
3.2 Fragmentation and 3D fragment structures
우리는 분자 G = (V, E)를 Principal Subgraph (PS)를 사용하여 m개의 겹치지 않는 fragments {Sk}m k=1로 분해합니다. 여기서 Sk = (Vk, Ek)이고 V = ∪m k=1 Vk, E = ∪m k=1 Ek입니다. Fragment vocabulary S에 있는 모든 고유한 atoms에서 시작하여, PS는 반복적으로 이웃한 fragments를 병합합니다. 새로 병합된 fragments 중 가장 빈번한 fragment가 각 iteration에서 vocabulary에 추가되며, 원하는 vocabulary 크기에 도달할 때까지 반복됩니다. Fragment vocabulary의 크기가 작을수록 더 미세한 fragments와 더 상세한 coarse-grained structures를 얻을 수 있습니다. Fragmentation이 완료된 후, {Vk}m k=1와 V 간의 관계로부터 mapping matrix M을 구성할 수 있습니다.
Fragments의 초기 coordinates를 생성하기 위해, 추가적인 deep generative model을 training하는 대신 효율적인 cheminformatics tool인 RDKit distance geometry를 활용합니다. 초기 atom coordinates ˆxa ∼ pRDKit(xa)를 생성한 후, 초기 fragment coordinates xf를 구성 atom coordinates의 평균인 M† ˆxa로 정의합니다. RDKit에 의해 생성된 atom coordinates는 ground truth conformer의 approximations이므로, 결과적인 fragment coordinates 또한 approximation입니다 (따라서 다음 단계인 Sec. 3.3에서 adjusted될 필요가 있음). 이를 ˆxf ∼ pRDKit(xf)로 표기합니다. Fragment features hf ∈ Rm×3에 대해서는, 이전 연구를 따라 hydrophobicity, hydrogen bond center, negative charge center를 포함한 chemical properties에 기반한 구성 atom types의 frequency histogram으로 3차원 벡터를 정의합니다.
Fragment structure generation 과정은 Fig. 2의 step 1에 설명되어 있습니다. 이 서브섹션의 모든 단계는 우리 diffusion model을 training하기 전에 완료될 수 있고 프로세스 자체가 효율적이므로 우리 framework의 efficiency를 해치지 않습니다. GEOM-Drug benchmark의 training 및 validation set에 있는 45,000개의 각 분자에 대해 5개의 서로 다른 fragment coordinates ˆxf를 생성하는 데 38시간이 걸렸으며, 분자당 평균 3.04초가 소요되었습니다. Fragmentation 및 fragment vocabulary의 세부 사항은 Appendix C.1에 설명되어 있습니다.
3.3 Equivariant blurring diffusion
이 서브섹션에서는 heat equation의 원리에서 영감을 받아 설계한 우리 diffusion model인 Equivariant Blurring Diffusion (EBD)의 설계에 대해 상세히 설명합니다. 이 model은 coarse-grained, approximate structure ˆxf와 molecular graph G에서 시작하여 conformers xa의 fine details를 생성하도록 설계되었습니다. Sec. 3.3.1에서는 forward process와 blurring process의 data corruption function을 소개하고, Sec. 3.3.2에서는 SE(3)-invariant likelihood에 도달하기 위한 reverse process와 deblurring network를, Sec. 3.3.3에서는 SE(3)-invariant loss function의 정의와 reparameterization을 소개합니다. EBD의 overall scheme은 Fig. 2의 step 2에 설명되어 있습니다.
3.3.1 Forward process and blurring schedule
우리는 forward process의 data corruption을 ground truth atom positions xa_0 ∼ q(xa_0)를 해당 fragment coordinates로 점진적으로 이동시키는 blurring operation으로 정의합니다: [ q(x_{a,t} | x_{a,0}, \hat{x}{f}) = \mathcal{N}(x{a,t} | f_B(x_{a,0}, \hat{x}_{f}, t), \sigma^2 I), ] 여기서 fB는 deterministic blurring operator입니다. 결과적으로, 모든 atom은 prior fragment structure distribution에서 각자의 fragment coordinates Mˆxf에 따라 위치하게 됩니다.
Forward process를 위한 fB를 정의할 때, Eq. (2)의 IHDM의 spectral blurring operator Bt = V exp(−Λt)VT를 두 가지 이유로 직접 채택할 수 없습니다: i) 단일 분자에 대해, IHDM에서 이미지당 단일 Laplacian operator와 달리, 각 fragment Sk = (Vk, Ek)에 대해 fragment graph Laplacian {VkΛkVT k }m k=1을 계산하고 분해해야 합니다. Fragments 간의 다양한 크기와 structures를 고려할 때, spectral space에서 {Λk}m k=1의 함수를 사용하여 모든 fragments에 걸쳐 atoms의 움직임을 균일하게 조정하는 것이 어려워집니다. ii) t → T가 될 때, spectral graph theory에 따라 ground truth atom coordinates xa_0는 ground truth scaffold structure xf = BT xa_0로 수렴할 것입니다. 그러나 generative processes에서 ground truth coordinates xf와 RDKit으로부터의 approximation coordinates ˆxf 사이에는 불일치가 존재합니다. Fragment structure의 이러한 distributional shift는 inference 동안 performance를 잠재적으로 해칠 수 있습니다.
이러한 문제들을 피하기 위해, 우리는 blurring process의 본질을 유지하면서 blurring operator의 공간을 spectral domain에서 spatial domain으로 전환합니다. 우리는 fB를 Euclidean space에서 Mˆxf와 xa_0 사이의 linear interpolation으로 정의합니다: [ f_B(x_0^a, \hat{x}_f, t) = \left(1 - \frac{t}{T}\right)x_0^a + \frac{t}{T}M \hat{x}_f. ] t가 0에서 T로 진행됨에 따라, atom coordinates xa_t는 점진적으로 fragment structure Mˆxf로 수렴하게 되어, atom 움직임의 균일한 조정을 가능하게 합니다. 추가적으로, fragment graph Laplacian의 과도한 eigendecomposition 필요성을 완화할 수 있습니다. 단일 fragment에 대한 우리의 blurring schedule 예시는 Fig. 2의 step 2에 묘사되어 있습니다.
3.3.2 Reverse process and deblurring networks
Reverse process의 목표는 group에 대해 roto-translational invariant한 3D fragment structure의 prior distribution p(xa_T ) = N (xa_T |Mˆxf, δ2I)에서 시작하여 atom-level에서 fine details를 생성하는 것입니다. Invariant likelihood 조건에 관한 증명들을 바탕으로, 우리는 equivariant transition distributions를 사용하여 zero center-of-mass subspace 상에서 deblurring process를 개발합니다: [ p_\theta (x_{a,t-1} | x_{a,t}, \hat{x}^f, G) = \mathcal{N}(x_{a,t-1} | \mu_\theta(x_{a,t}, \hat{x}^f, G, t), \delta^2 I), ] 여기서 μθ는 deblurring network로 구성된 parameterized mean function입니다. Transition distribution에서 equivariance를 보장하기 위해, 우리는 equivariant networks에서 영감을 받아 μθ를 고안합니다. 우리의 equivariant deblurring network는 atoms와 fragments 사이의 hierarchical relationship을 활용하여 fragments와 atoms의 invariant features hf, ha (Eqs. (9, 10))와 atoms의 equivariant coordinates xa (Eq. (11))를 업데이트합니다. i번째 atom xa_i가 k번째 fragment xf_k에 속한다고 할 때, fragment-level 및 atom-level message passing과 feature updates를 위한 equivariant deblurring networks의 l-번째 layer는 다음과 같이 정의됩니다: [ m_{ij}^f = \phi_m^f(h_i^{f,l}, h_j^{f,l}, | x_i^f - x_j^f |), \quad h_i^{f,l+1} = \phi_h^f(h_i^{f,l}, \sum_{j \in N(x_i^f)} m_{ij}^f, h_a^{l}). ] [ m^{a}{ij} = \varphi^{a}{m} \left( h^{a,l}{i}, h^{a,l}{j}, | x^{a,l}{i} - x^{a,l}{j} |, e^{a}{ij} \right), \quad h^{a,l+1}{i} = \varphi^{a}{h} \left( h^{a,l}{i}, \sum_{j \in N(x^{a}{i})} m^{a}{ij}, h^{f,l+1} \right), ] [ x^{a,l+1}_i = x^{a,l}i + \sum{j \in N(x^{a,l}_i)} \frac{x^{a,l}_i - x^{a,l}j}{d^{a,l}{ij}} \phi^{a}_x(h^{a,l+1}i, h^{a,l+1}j, m^{a}{ij}, e^{a}{ij}) + 1 + \frac{x^{a,l}_i - x^{f}_k}{|x^{a,l}_i - x^{f}_k|} \phi^{f}_x(h^{a,l+1}_i, h^{f,l+1}_k, |x^{a,l}_i - x^{f}_k|) + 1 ] 여기서 xf_k는 M†xa_t의 k-번째 행이고, ϕ는 multilayer perceptrons, ea_ij는 inter-atomic bond types, da,l_ij = ∥xa,l_i − xa,l_j ∥는 inter-atomic distances입니다. 우리는 fragment-level interactions에 대해 complete graph를 고려하고, atom-level interactions에 대해 multi-hop 및 radius neighbors를 통합하여 edge set을 확장합니다. Deblurring networks의 세부 사항은 Appendix A에 제공됩니다.
3.3.3 Training
IHDM의 Eq. (5)를 따라, 이전 deblurred state estimation의 우리 loss는 다음과 같이 정의될 수 있습니다: [ L_{t-1} = \mathbb{E}{t, x{a0}, x_{at}, \hat{x}{f}}\left[|f_B(x{a0}, \hat{x}{f}, t - 1) - \rho(\mu_0(x{at}, \hat{x}_{f}, G, t))|^2\right], ] 여기서 ρ는 alignment를 위한 optimal rotation matrix를 얻기 위한 Kabsch algorithm입니다. μθ로부터의 prediction과 덜 blur된 상태 fB(xa_0, ˆxf, t − 1) 사이의 alignment ρ를 (두 항 모두 zero center-of-mass subspace로 이동시킨 후) 수행함으로써, loss function은 prediction의 SE(3)-transformation에 대해 invariant하게 됩니다.
그러나 우리는 경험적으로 이 previous state estimator가 IHDM의 image generation에서 관찰된 unsatisfactory FID scores와 유사하게 unsatisfactory conformers를 생성한다는 것을 관찰했습니다. 우리는 그 이유를 model이 각 time step에서 ground truth distribution을 향한 locally small steps를 배우는 데 제한되기 때문이라고 추측했습니다. 따라서, 우리는 deblurring network가 previous less blurred state 대신 neural networks fθ를 통해 ground truth state xa_0를 추정하도록 μθ (xa_t, ˆxf, G, t)를 (1 − (t−1)/T )fθ (xa_t, G, t) + (t−1)/T Mˆxf로 reparameterize합니다: [ L_{t-1} = \mathbb{E}{t, x^{a}{0}, x^{a}{t}, \hat{x}^{f}} \left[ | f B ( x^{a}{0}, \hat{x}^{f}, t - 1) - \rho \left( \left( 1 - \frac{t - 1}{T} \right) f \theta ( x^{a}{t}, G, t ) + \frac{t - 1}{T} M \hat{x}^{f} \right) |^2 \right] \approx \mathbb{E}{t, x^{a}{0}, x^{a}{t}, \hat{x}^{f}} \left[ | x^{a}{0} - \rho f \theta ( x^{a}{t}, G, t ) |^2 \right]. ] 새로운 loss의 유도는 Appendix B에 상세히 설명되어 있습니다. Loss reparameterization을 통해, ρ는 prediction을 ground truth state에 align합니다. Sampling process의 time step t에서, xa_t로부터 ground truth ˜xa_0를 추정한 후, 다음 state xa_t−1은 ˜xa_0를 사용하여 deterministic blurring function fB로부터 계산됩니다. Training 및 sampling processes는 Algorithms 1, 2에 제공됩니다.
Algorithm 1 Training
이 알고리즘은 Equivariant Blurring Diffusion (EBD) 모델, 특히 ground truth 상태를 예측하는 deblurring network f_θ를 학습시키는 과정을 설명합니다.
- Sample ˆx_f ~ p_RDKit(x_f): RDKit을 사용하여 fragment들의 대략적인 3D 좌표 (coarse structure) ˆx_f를 샘플링합니다.
- Sample x_{a,0} ~ q(x_a_0): 데이터셋에서 실제 정답 conformer 구조 (ground truth) x_{a,0}를 샘플링합니다.
- Sample t ~ U[1, T]: 1부터 T 사이의 시간 단계 t를 무작위로 선택합니다. T는 최대 blurring 시간입니다.
- Sample ε ~ N(0, σ^2 I): 약간의 가우시안 noise ε를 샘플링합니다 (σ는 작은 표준편차).
- x_{a,t} ← f_B(x_{a,0}, ˆx_f, t) + ε: ground truth x_{a,0}를 시간 t만큼 blurring 함수 f_B (Sec 3.3.1의 공간적 linear interpolation)를 사용해 변형시키고, 약간의 noise ε를 더하여 t 시점의 blurred 상태 x_{a,t}를 만듭니다. 이것이 네트워크가 입력으로 받게 될 데이터의 형태를 모방합니다.
- Minimize ||x_{a,0} - ρ(f_θ(x_{a,t}, G, t))||^2: 핵심 학습 단계입니다.
- Deblurring network f_θ는 blurred 상태 x_{a,t}, molecular graph G, 시간 t를 입력받아 원래의 ground truth 상태 x_{a,0}를 예측하려고 시도합니다.
- ρ (Kabsch 알고리즘)는 f_θ의 예측값과 실제 ground truth x_{a,0} 사이의 최적 회전을 찾아 정렬(alignment)합니다 (SE(3)-invariant loss를 위함).
- 정렬된 예측값과 실제 ground truth x_{a,0} 사이의 제곱 오차(loss)를 계산하고 이를 최소화하도록 f_θ의 파라미터를 업데이트합니다. 즉, f_θ가 blurring 과정을 최대한 잘 되돌리도록 학습시킵니다. (이 loss는 Sec 3.3.3의 재구성된 loss, Eq. (13)의 근사식에 해당합니다.)
Algorithm 2 Generation
이 알고리즘은 학습된 EBD 모델을 사용하여 새로운 molecular conformer를 생성하는 과정 (sampling)을 설명합니다.
- Sample ˆx_f ~ p_RDKit(x_f): 생성을 시작할 condition으로 사용할 대략적인 fragment 좌표 ˆx_f를 RDKit으로 샘플링합니다.
- x_{a,T} ← Mˆx_f: 생성 과정을 최대 blurring 시간 T에서 시작합니다. 초기 상태 x_{a,T}를 fragment 좌표가 지시하는 원자 위치 Mˆx_f로 설정합니다. (가장 coarse한 상태)
- for t in {T, ..., 1} do: 시간 t를 T부터 1까지 하나씩 줄여가며 반복합니다 (reverse process).
- Sample ε ~ N(0, δ^2 I): (선택 사항) 역방향 단계에 약간의 noise를 추가할 수 있습니다. δ=0이면 결정론적 생성이 됩니다. (이미지의 f_θ 입력에 ε가 더해진 것은 실제 구현이나 논문 설명과 약간 다를 수 있으며, 보통은 예측 후 샘플링 단계에 noise가 들어갑니다. Sec 3.3.3 마지막 문장에 따르면 결정론적(deterministic) 생성을 사용하는 것으로 보입니다.)
- ˜x_{a,0} ← f_θ(x_{a,t}, ˆx_f, G, t): 현재의 blurred 상태 x_{a,t}와 ˆx_f, G, t를 학습된 deblurring network f_θ에 입력하여, 원래 ground truth일 것으로 예측되는 상태 ˜x_{a,0}를 얻습니다.
- x_{a,t-1} ← f_B(˜x_{a,0}, ˆx_f, t - 1): 다음 시간 단계 t-1의 상태 x_{a,t-1}를 계산합니다. 이는 예측된 ground truth ˜x_{a,0}를 blurring 함수 f_B를 사용해 시간 t-1 수준까지만 다시 blurring 시켜서 얻습니다. (이는 표준적인 DDPM 역방향 단계와는 약간 다른 방식이며, Sec 3.3.3 마지막 문장의 설명과 일치합니다.)
- end for: 루프가 끝나면 t=0 상태, 즉 최종적으로 생성된 molecular conformer x_{a,0}를 얻게 됩니다.
정리노트 (AI 연구자 대상)
EBD 방법론 핵심 요약 (Section 3)
- Hierarchical Framework 개요:
- 목표: Molecular graph G로부터 3D conformer ensemble x_a 생성.
- 2단계 접근:
- Stage 1 (Sec 3.2): Fragmentation & Coarse Structure Gen.
- G를 PS 알고리즘으로 non-overlapping fragments {S_k}로 분해.
- RDKit (cheminformatics tool) 사용하여 initial (approximate) fragment coordinates ˆx_f 생성 (M†ˆx_a 방식). Deep generative model 불필요, preprocessing으로 효율성 확보. (Note: ˆx_f는 ground truth가 아닌 approximation임).
- Stage 2 (Sec 3.3): Fine Detail Gen. with EBD
- ˆx_f와 G를 condition으로 EBD (diffusion model) 사용하여 fine atomic details x_a 생성.
- Stage 1 (Sec 3.2): Fragmentation & Coarse Structure Gen.
- EBD 핵심 설계 (Sec 3.3):
- Forward Process (Sec 3.3.1): 공간적 Blurring (핵심)
- 동기: IHDM의 spectral blurring은 분자(fragment 별 Laplacian, ˆx_f/x_f 불일치)에 부적합.
- 제안: Spectral 대신 Spatial domain에서 deterministic blurring operator f_B 정의. [ f_B(x_0^a, \hat{x}_f, t) = \left(1 - \frac{t}{T}\right)x_0^a + \frac{t}{T}M \hat{x}_f ] (ground truth x_{a,0}와 approximate fragment 좌표 Mˆx_f 간 linear interpolation).
- 장점: 균일한 이동, eigendecomposition 불필요, ˆx_f 사용으로 distributional shift 문제 완화 시도.
- Process: q(x_{a,t} | x_{a,0}, ˆx_f) = N(x_{a,t} | f_B(...), σ^2 I).
- Reverse Process (Sec 3.3.2): Equivariant Deblurring (핵심)
- 목표: Roto-translational invariant prior p(x_{a,T}) ≈ N(Mˆx_f, δ^2 I)에서 fine details 복원.
- 방법: Zero center-of-mass subspace에서 equivariant transition distributions p_θ 사용 (SE(3)-invariant likelihood 보장). [ p_\theta (x_{a,t-1} | x_{a,t}, \hat{x}^f, G) = \mathcal{N}(x_{a,t-1} | \mu_\theta(...), \delta^2 I) ]
- Deblurring Network (f_θ in μ_θ): Equivariant networks 기반 설계. Atom-fragment 계층 구조 활용, invariant features (h_a, h_f) 및 equivariant coordinates (x_a) 업데이트 위한 message passing 수행 (Eqs. 9-11).
- Training (Sec 3.3.3, Algorithm 1): Loss Reparameterization (핵심)
- 문제점: IHDM 방식 (previous state f_B(..., t-1) 추정)은 performance 저조.
- 제안: Deblurring network f_θ가 blurred state x_{a,t}로부터 직접 ground truth x_{a,0}를 예측하도록 reparameterization. [ L \approx \mathbb{E}{t, x^{a}{0}, x^{a}{t}, \hat{x}^{f}} \left[ | x^{a}{0} - \rho f \theta ( x^{a}_{t}, G, t ) |^2 \right] ]
- SE(3)-invariance: Kabsch algorithm (ρ) 통해 prediction과 ground truth 정렬 후 loss 계산.
- Generation (Algorithm 2):
- x_{a,T} = Mˆx_f에서 시작, t = T...1 반복.
- 각 t에서 f_θ로 ˜x_{a,0} 예측.
- 예측된 ˜x_{a,0}를 blurring function f_B를 사용하여 x_{a,t-1} = f_B(˜x_{a,0}, ˆx_f, t - 1) 계산 (결정론적 단계).
- Forward Process (Sec 3.3.1): 공간적 Blurring (핵심)
요약: RDKit으로 생성한 approximate coarse structure ˆx_f를 prior로 사용하고, spectral 대신 spatial linear interpolation 기반 blurring (forward)과 SE(3)-equivariant network 기반 ground truth (x_{a,0}) 직접 예측 (reverse + loss)을 통해 3D conformer를 생성하는 새로운 diffusion framework (EBD) 제안.
쉬운 설명:
이 섹션은 이 논문에서 제안하는 3D 분자 구조 생성 방법(EBD)의 구체적인 레시피를 설명합니다.
1단계: 분자 스케치 그리기 (Fragmentation & Coarse Structure)
- 먼저, 복잡한 분자 설계도(molecular graph G)를 받으면, 이걸 좀 더 다루기 쉬운 몇 개의 주요 부품(fragments)으로 나눕니다.
- 그 다음, RDKit이라는 기존 도구를 사용해서 이 주요 부품들이 3D 공간 어디쯤에 있을지 대략적인 위치(fragment coordinates ˆx_f)를 빠르게 계산해요. 이건 아주 정밀한 위치는 아니고, 말 그대로 '대충 그린 스케치' 같은 거예요. 이 스케치를 미리 만들어 둠으로써 나중에 AI가 작업하기 편하게 해줍니다.
2단계: AI가 스케치를 바탕으로 상세 모델 완성하기 (Equivariant Blurring Diffusion - EBD)
이제 EBD라는 특별한 AI가 등장해서 앞에서 만든 스케치(ˆx_f)를 보고 상세한 3D 모델(conformer)을 만듭니다.
- '흐리게 만들기' 방법 정의 (Forward Process): AI는 완벽한 3D 모델이 있다면, 이걸 어떻게 '점진적으로 뭉개서' 앞에서 만든 스케치 모양으로 만들 수 있는지 그 방법을 정의해요. 이 논문만의 특별한 점은, 그냥 마구잡이로 뭉개는 게 아니라, 완벽한 모델 위치와 스케치 위치 사이를 직선으로 조금씩 이동시키는 방식(linear interpolation)으로 '흐리게' 만든다는 거예요. 이걸 Spatial Blurring이라고 부릅니다. 기존 방식(spectral blurring)보다 분자에 더 적합하다고 해요.
- '선명하게 되돌리기' 학습 (Reverse Process & Training): AI(f_θ)는 이 '흐리게 만들기' 과정을 거꾸로 하는 법을 배워요. 흐릿한 스케치 상태에서 시작해서 점차 선명하게 만들어 최종 3D 모델을 완성하는 거죠.
- 똑똑한 3D 학습: 이 AI는 3D 공간을 잘 이해하도록(SE(3) equivariant) 설계되었어요. 분자를 이루는 작은 원자들과, 그 원자들이 속한 더 큰 부품(fragment) 사이의 관계를 이용해서 정보를 주고받으며 학습해요. 그래서 분자를 돌려도 제대로 구조를 만들 수 있습니다.
- 더 똑똑한 학습 목표: 그냥 한 단계만 되돌리는 걸 배우는 게 아니라, 어떤 흐릿한 상태에서든 바로 최종 완성본(ground truth)을 예측하도록 학습 목표를 바꿨어요 (Loss Reparameterization). 이게 성능 향상에 도움이 되었다고 합니다. 학습 시에는 예측 결과와 정답을 잘 비교하기 위해 잠시 회전시켜 맞춰보는 과정(Kabsch alignment)도 거칩니다.
- 새로운 분자 만들기 (Generation): 학습이 끝나면, 새로운 분자 스케치를 주고 AI에게 '선명하게 되돌리기' 과정을 T단계부터 1단계까지 쭉 시키면 최종 3D 분자 모델이 만들어집니다. 특이한 점은, 각 단계마다 최종 완성본을 예측한 다음, 그걸 다음 단계 수준으로 살짝만 다시 흐리게 만들어서 진행한다는 점입니다.
요약: 이 논문은 (1) 분자를 조각내고 대략적인 스케치를 그린 다음, (2) Spatial Blurring이라는 특별한 '흐리게 만들기' 방법과, 3D 공간을 잘 이해하며 최종 결과물을 바로 예측하도록 똑똑하게 학습된 AI(EBD)를 이용해 상세한 3D 분자 구조를 만드는 새로운 방법을 제안합니다.
주인장 생각
Molecular graph G(가지고 있는 인풋)
Molecular graph G를 넣어서
1. 여러 개의 겹치지 않는 fragments {Sk}로 분해 ( Mapping matrix M )
2. RDKit를 사용해서 G로부터 가능한 3D 원자 좌표의 근사치(approximation) ˆx_a를 생성
3. 1,2 번을 넣어서 원자들의 3D 좌표들의 평균 을 구해서 fragment의 중심 좌표로 사용 즉 Approximate fragment coordinates ˆx_f (m x 3 크기). 이것이 EBD 모델에 입력될 coarse-grained 3D 구조 정보
4. 실제 정답 3D 구조를 deterministic blurring operator f_B를 정의하여 fragment좌표로 점진적으로 blur시키면서 포워드 과정을 진행
5. 단순히 blur된 상태 x_{a,t}만 보는 것이 아니라, 미리 계산된 fragment 좌표 ˆx_f와 molecular graph G 정보까지 함께 입력받아(conditioned), 이를 바탕으로 더 정확하게 fine details를 생성
6. 실제 reverse를 진행하면서 trainig 5번과 같은 조건부를 입력해서 t+1 시점을 예측
7. 인퍼런스도 동일하게 진행
gemini의 평가
- Molecular graph G (가지고 있는 인풋) - 정확합니다.
- Molecular graph G를 넣어서 1. ... fragments {Sk}로 분해 (Mapping matrix M) - 정확합니다. (PS 알고리즘 사용)
- 2. RDKit를 사용해서 G로부터 ... 근사치(approximation) ˆx_a를 생성 - 정확합니다.
- 3. 1,2 번을 넣어서 ... 원자들의 3D 좌표들의 평균 ... Approximate fragment coordinates ˆx_f ... coarse-grained 3D 구조 정보 - 정확합니다. (M†ˆx_a 계산)
- 4. 실제 정답 3D 구조를 ... deterministic blurring operator f_B ... blur시키면서 포워드 과정을 진행 - 정확합니다. (linear interpolation 방식)
- 5. 단순히 blur된 상태 x_{a,t}만 보는 것이 아니라, 미리 계산된 fragment 좌표 ˆx_f와 molecular graph G 정보까지 함께 입력받아(conditioned), 이를 바탕으로 더 정확하게 fine details를 생성 - 정확합니다. (이것은 Reverse Process 중 deblurring network f_θ가 작동하는 방식을 잘 설명합니다.)
- 6. 실제 reverse를 진행하면서 trainig 5번과 같은 조건부를 입력해서 t+1 시점을 예측 - 이 부분이 약간 다릅니다.
- Training 과정에서는 reverse process 전체를 진행하지 않습니다. 대신, 무작위 시간 t를 뽑고 그 시점의 blurred state x_{a,t}를 만듭니다.
- f_θ 네트워크는 5번과 같은 조건부 정보를 입력받는 것은 맞지만, t+1 (또는 t-1) 시점을 예측하는 것이 아니라, 최종 목표인 원본 ground truth 상태 x_{a,0}를 직접 예측하도록 학습합니다. (이것이 loss reparameterization의 핵심입니다. Loss 식: ||x_{a,0} - ρ(f_θ(x_{a,t}, G, t))||^2).
- 7. 인퍼런스도 동일하게 진행 - 네트워크(f_θ)와 입력 조건은 동일하게 사용하지만, 진행 방식은 다릅니다.
- Inference (Generation, Algorithm 2)는 t=T부터 t=1까지 순차적으로 한 단계씩 진행합니다.
- 각 단계 t에서 f_θ를 이용해 x_{a,0}을 예측(˜x_{a,0})한 다음, blurring 함수 f_B를 사용하여 t-1 시점의 상태 x_{a,t-1}를 계산합니다.
- Training은 무작위 t에 대해 x_{a,0} 예측 오차를 줄이는 것이고, Inference는 T에서 1까지 예측과 계산을 반복하여 최종 x_{a,0}를 만드는 것입니다.
요약: 전반적인 이해는 매우 훌륭합니다! 특히 1~5번까지의 흐름과 Reverse Process의 조건부 입력에 대한 이해는 정확합니다. 다만 Training 시 예측 목표가 x_{a,0}라는 점과, Training과 Inference의 단계 진행 방식에 약간의 차이가 있다는 점만 기억하시면 완벽합니다.