AI바라기의 인공지능
Diffusion : 논문리뷰 : Coarse-to-Fine: a Hierarchical Diffusion Model for Molecule Generation in 3 본문
Diffusion : 논문리뷰 : Coarse-to-Fine: a Hierarchical Diffusion Model for Molecule Generation in 3
AI바라기 2025. 4. 3. 11:55쉬운 설명:
이 논문은 3D 분자를 레고 블록처럼 조립하는 새로운 방법을 제안합니다. 기존 방식은 아주 작은 원자(atom) 단위로 하나씩 붙여나가서 모양(특히 반지 모양 같은 복잡한 부분)이 쉽게 망가졌습니다. HierDiff는 먼저 큼직한 블록(fragment) 들의 대략적인 위치와 모양을 diffusion이라는 방법으로 정합니다(coarse-grained). 그런 다음, 이 큼직한 블록들을 더 정교하게 다듬고(fine-grained) 서로 어색하지 않게 잘 연결하는 과정을 거칩니다. 이때 Iterative Refinement라는 특별한 단계를 두어 블록들이 서로 부딪히거나 이상하게 연결되지 않도록 계속 수정합니다. 이렇게 하면 처음부터 작은 원자 하나하나 신경 쓰는 것보다 훨씬 안정적이고 제대로 된 모양의 3D 분자를 만들 수 있습니다.
HierDiff 논문 학습 노트
Purpose of the Paper:
- 기존 atom-level 3D 분자 생성 모델들의 한계점 극복:
- 특히 큰 분자 생성 시, ring과 같은 국소적 구조(local structure)의 화학적 유효성(validity) 저하 문제 해결 (Fig 1).
- Euclidean geometric constraint 부재로 인한 비현실적 구조 생성 문제 해결.
- Fragment-based 생성 방식의 장점 활용 + 3D non-autoregressive 생성의 어려움 해결:
- Fragment-based 방식은 유효한 국소 구조 보존에 유리하지만, 3D non-autoregressive 방식에 적용 시 fragment 간 연결 충돌(conflict) 문제와 이로 인한 combinatorial optimization 문제 발생 (Fig 2, Sec 4.2.1).
- 핵심 목표: Coarse-to-fine 전략을 통해, non-autoregressive 방식으로 fragment 기반 3D 분자 생성을 수행하면서도 국소 구조 유효성을 보장하고 fragment 연결 충돌 문제를 효과적으로 해결하는 새로운 hierarchical diffusion model (HierDiff) 제안.
Key Contributions & Novelty:
- Contribution 1: 3D 분자 생성을 위한 최초의 Hierarchical Diffusion model (HierDiff) 제안.
- Novelty: 분자 구조를 coarse-grained (fragment) 레벨과 fine-grained (atom) 레벨로 나누어 계층적으로 생성하는 diffusion model 접근 방식을 이 분야에 처음 도입.
- Contribution 2: 2단계 생성 프로세스 설계:
- Coarse-grained Fragment Diffusion (Sec 5.1):
- SE(3)-equivariant diffusion process를 사용하여 fragment의 화학적/위치적 representation (Hf, Hp)과 3D 좌표 생성.
- 화학적으로 의미 있는 두 가지 coarse feature (Property-based, Element-based) 제안 (Fig 4).
- Novelty: Fragment 레벨에서 SE(3)-equivariant diffusion을 통해 coarse geometry를 직접 생성.
- Fine-grained Fragment Generation (Sec 5.2):
- Coarse node 정보를 바탕으로 equivariant message-passing network (EGNN 변형)를 이용해 fine-grained fragment type과 연결 관계 예측.
- Fragment 연결 시 발생할 수 있는 편향(bias)과 충돌(conflict)을 수정하기 위한 새로운 Iterative Refinement 모듈 도입 (Fig 6).
- Novelty: Non-autoregressive 방식에서 fragment 조립 시 발생하는 제약 조건을 만족시키기 위한 Iterative Refinement 메커니즘. Coarse/Fine 양 단계 모두에서 equivariant network 활용.
- Coarse-grained Fragment Diffusion (Sec 5.1):
- Contribution 3: Fragment 생성을 제약 조건 만족 문제(constraint generation problem) 로 정의하고, 이를 hierarchical latent variable (H) 도입을 통해 해결 (Sec 4.2.2, Eq 5).
- Novelty: Fragment 충돌 문제를 확률 모델 내에서 명시적으로 정의하고 계층적 분해(decomposition)를 통해 해결 방안 제시.
Experimental Highlights:
- Datasets: GEOM-DRUG, CrossDocked2020 (주요 평가), QM9.
- Baselines: EDM, G-SphereNet.
- Metrics: Drug-likeness (QED, RA, MCF, SAS, ΔLogP, ΔMW), Conformation Quality (Cov, Mat - atom/fragment level, MD simulation 기반), Validity, Stability, Diversity, Energy (MFF MMD).
- 주요 결과:
- Drug-likeness (Table 1): HierDiff가 EDM, G-SphereNet 대비 모든 drug-likeness 지표에서 significantly 우수한 성능 달성. 특히 G-SphereNet이 약한 큰 분자 생성에서 강점 (ΔMW).
- Conformation Quality (Table 2, Fig 7): HierDiff가 EDM 대비 atom 및 fragment 레벨 모두에서 더 높은 Cov, Mat 달성. 생성된 분자 구조가 시각적으로 더 안정적임.
- Validity/Stability (Table 5, 6): HierDiff가 GEOM-DRUG에서 EDM보다 높은 validity 및 stability 달성. QM9에서도 경쟁력 있는 성능.
- Iterative Refinement 효과 검증 (Table 9): Ablation study를 통해 Iterative Refinement 단계가 더 복잡하고 현실적인 분자 생성에 필수적임을 입증.
- Energy Distribution (Fig 9): HierDiff가 생성한 분자들의 MFF energy 분포가 EDM, JT-VAE보다 ground truth에 더 근접함.
- Error Accumulation (Table 12): 분자 크기 증가 시 autoregressive 모델(G-SphereNet) 대비 오류 누적 현상이 적음을 보임.
Limitations and Future Work:
- Limitations:
- 사전 정의된 fragment vocabulary 및 분해 방식(JT-VAE 방식 기반)에 의존적이며, vocabulary 품질이 성능에 영향 미칠 수 있음 (Appendix A.1, C.7).
- Diffusion model 자체의 샘플링 속도 + fine-grained decoding 과정의 계산 비용 (Appendix C.8).
- 일반적인 ML 모델의 한계인 일반화 성능 및 해석 가능성 부족 (Appendix B).
- Future Work:
- 다양한 fragment 분해 전략 탐색 (Appendix C.7).
- 샘플링 효율성 개선 (diffusion step 수 감소 등) (Appendix C.8).
- 제안된 hierarchical framework를 다른 제약 조건 하 구조 생성 문제 (e.g., protein alignment)에 적용 (Sec 7).
- 조건부 생성(conditional generation) 능력 심화 및 실제 신약 개발 적용 가능성 탐색 (Appendix B, C.6).
Overall Summary:
이 논문은 3D 분자 생성 시 기존 atom-level 모델의 국소 구조 유효성 문제를 해결하고, fragment-based 방식의 장점을 non-autoregressive diffusion 모델에 효과적으로 접목하기 위해 Hierarchical Diffusion model (HierDiff) 를 제안한다. Coarse-grained fragment geometry를 먼저 생성하고, 이를 equivariant message passing과 Iterative Refinement를 통해 fine-grained fragment로 변환 및 조립하는 coarse-to-fine 전략을 사용한다. 실험 결과, HierDiff는 기존 SOTA 모델 대비 drug-likeness, conformation quality, validity 측면에서 우수한 성능을 보였으며, 특히 크고 복잡한 분자 생성에서 강점을 나타냈다. 이 연구는 non-autoregressive 방식의 장점을 살리면서도 화학적으로 유효한 3D 분자 구조 생성을 위한 새로운 방향을 제시한다.
Abstract
원하는 3D 분자 구조를 Generating하는 것은 신약 개발에 있어 근본적인 문제입니다. 상당한 progress를 이루었음에도 불구하고, 기존 methods는 일반적으로 분자를 원자 해상도로 generate하며 고리와 같은 본질적인 local structures를 무시하여, 생성된 구조의 quality가 저하되는 경향이 있습니다. 특히 큰 분자를 Generating할 때 더욱 그렇습니다. Fragment-based molecule generation은 유망한 전략이지만, 조합 최적화 문제로 인해 3D non-autoregressive generations에 적용하기는 non-trivial합니다.
본 paper에서는 이 problem을 해결하기 위해 coarse-to-fine strategy를 활용하며, 이를 위해 Hierarchical Diffusion-based model (즉, HierDiff)을 제안하여 autoregressive modeling에 의존하지 않고 local segments의 유효성을 보존합니다.
Specifically, HierDiff는 먼저 equivariant diffusion process를 통해 coarse-grained 분자 기하 구조를 generates합니다. 여기서 각 coarse-grained node는 분자 내의 fragment를 반영합니다. 그런 다음 coarse-grained nodes는 message-passing process와 새로 설계된 iterative refined sampling module에 의해 fine-grained fragments로 decoded됩니다. 마지막으로, fine-grained fragments를 조립하여 완전한 원자 분자 구조를 도출합니다.
광범위한 experiments는 HierDiff가 기존 methods에 비해 molecule generation의 quality를 일관되게 향상시킨다는 것을 보여줍니다.

1. introduction
Deep generative models는 복잡한 graph-like molecular structures를 modeling하는 데 특히 큰 가능성을 보여주었으며, 이는 molecular atom-bond graph 생성에서부터 graph로부터 molecular conformations 생성에 이르기까지 다양합니다. 상당한 progress가 이루어졌음에도 불구하고, 이 분야에서 여전히 중요하게 남아있는 연구 방향은 3D space에서의 de novo design of drug molecules입니다. 3D information을 molecule design process에 통합하는 것은 structure-based drug design, molecular dynamic simulation, 3D similarity searching과 같은 많은 중요한 applications에서 topological information만 사용하는 것에 비해 여러 장점을 가집니다.
3D molecule generation에 대한 일부 초기 studies는 일반적으로 autoregressive approach를 채택했는데, 이는 atoms에 인위적인 order를 도입하고 language generation 방식처럼 atoms를 하나씩 generates합니다. 그러나 분자는 3D에서 자연적인 geometric structure를 가지고 있습니다. 게다가 이러한 models는 scale 및 error accumulation problem으로 어려움을 겪습니다. 이러한 problems를 tackle하기 위해, non-autoregressive models가 이 area에 도입되어 인상적인 results를 얻고 있습니다. 예를 들어, text and image generation에서 성공적인 diffusion model에 inspired 받아, Hoogeboom 등은 molecule generation을 위한 첫 번째 diffusion model을 proposed하여 generated molecules의 validity를 크게 향상시켰습니다.
그럼에도 불구하고, 이러한 works에서의 atom-level generation 방식은 각 atom을 배치하는 데 더 높은 flexibility를 제공하지만, 특히 큰 분자를 generating할 때 신뢰할 수 있는 molecule structures를 얻기 위한 필요한 constraints가 부족합니다. Figure 1에서 볼 수 있듯이, modeling process에 Euclidean geometric constraints를 부과하지 않으면, generated 3D 방향족 고리(aromatic rings)는 기본적인 화학 규칙을 심각하게 위반할 수 있습니다.
본 paper에서는 위의 problems를 tackle하기 위해 coarse-to-fine approach를 propose합니다. 기본적인 idea는 먼저 분자의 coarse-grained structure를 generate하고 (여기서 각 node는 fragments의 cluster를 represents합니다), 그런 다음 coarse-grained structure를 fine-grained fragments로 decoded하여 atomic molecule structure를 assemble하는 것입니다. 이러한 방식으로, 계산 단위를 atoms에서 fragments로 대체함으로써 valid local structures가 preserved될 것입니다. 그러나 이러한 process는 generated된 이웃 fragments가 atom-bonds conflicts를 겪어 서로 connected되는 것을 방해할 수 있기 때문에 non-trivial합니다.
이 problem을 tackle하기 위해, 우리는 3D molecule generation을 constraint generation problem으로 취급하고 새로운 Hierarchical Diffusion-based model (HierDiff)을 propose합니다. coarse-grained phase에서, 우리의 method는 deterministic fragment 대신 fragment representation을 generates합니다. Specifically, 우리는 fragments를 representing하기 위한 화학적으로 해석 가능한 features를 얻는 두 가지 다른 방법을 introduce합니다. 그런 다음 우리는 이러한 fragment representations와 그들의 Cartesian coordinates를 효율적인 non-autoregressive manner로 generate하기 위해 geometric diffusion model을 propose합니다. fine-grained phase에서는, connectivity를 보장하기 위해 equivariant message-passing network를 사용하고 bias를 correct하기 위해 iterative refinement module을 utilize합니다. 마지막으로, decoded 3D fragment graph를 기반으로 atom-level 3D structure를 construct합니다.
제안된 coarse-to-fine approach는 미리 정의된 functional group database에서 fragments를 결합함으로써 화학 전문가의 약물 설계 process를 잘 mimics합니다. 이러한 방식으로, 이 area의 중요한 inductive biases가 우리 model에 encoded됩니다. Furthermore, machine learning의 관점에서 볼 때, 분자의 fragment-based representation은 atom-based methods에서의 불필요한 degrees of freedom을 크게 줄여주므로, global optimum convergence와 더 나은 generalization ability로 이어질 것입니다.
어려운 과제인 약물 크기 분자 generating에 대해 우리 model을 test하기 위해 광범위한 experiments가 수행되었습니다. baseline model과 비교했을 때, HierDiff는 더 나은 약물 유사 특성을 가진 현실적인 분자와 ground truth conformations에 훨씬 더 가까운 conformations를 모두 generate할 수 있습니다. Visualized results는 또한 HierDiff가 더 안정적인 substructures를 가진 high-quality molecules를 generating할 수 있음을 demonstrate합니다.
정리 노트 (AI 연구자용)
핵심 문제:
- 기존 3D 분자 생성 models (Autoregressive 및 Non-autoregressive Diffusion 포함)은 주로 atom-level에서 작동.
- 이 방식은 특히 큰 분자에서 local structures(e.g., 방향족 고리)의 화학적/기하학적 타당성(validity)을 보장하기 어려움 (Euclidean geometric constraints 부족). 결과적으로 비현실적인 구조 생성 가능.
- Autoregressive 방식은 추가로 error accumulation 및 scale 문제 보유.
제안된 해결책: HierDiff (Hierarchical Diffusion-based model)
- 접근법: Coarse-to-fine strategy 채택.
- 핵심 아이디어: 계산 단위를 atom에서 fragment (화학적으로 의미있는 local structure 단위)로 변경하여 local structure의 타당성을 내재적으로 보존.
- Coarse Phase:
- Geometric diffusion model (non-autoregressive)을 사용하여 결정론적 fragment가 아닌, 화학적으로 해석 가능한 features로 표현된 fragment representation과 그 Cartesian coordinates를 생성.
- Fine Phase:
- 생성된 fragment representations를 fine-grained fragments로 decode.
- Equivariant message-passing network를 사용하여 fragment 간의 연결성(connectivity) 보장 (연결 시 atom-bonds conflicts 문제 해결).
- Iterative refinement module을 사용하여 bias 수정.
- 최종적으로 decoded 3D fragment graph로부터 atom-level 3D structure 구성.
- 부차적 이점:
- 화학 전문가의 fragment-기반 설계 process를 모방하여 화학적 inductive biases를 model에 encode.
- Fragment-based representation은 atom-based methods 대비 불필요한 degrees of freedom을 줄여 global optimum convergence 및 generalization ability 향상 가능성 제시.
주요 특징 요약: Hierarchical, Coarse-to-fine, Fragment-based representation, Geometric diffusion (Coarse), Equivariant message passing & Iterative refinement (Fine), Non-autoregressive (Coarse 단계).
(Note: 위 내용은 Introduction 섹션에서 주장하는 바를 요약한 것입니다.)
쉬운 설명:
AI로 3D 신약 분자를 새로 디자인하는 건 중요하지만 어려운 일입니다. 기존 AI들은 보통 분자를 원자(atom) 하나하나씩 순서대로 또는 한꺼번에 배치해서 만들려고 했습니다 (마치 레고 아주 작은 기본 블록만 써서 복잡한 구조를 만들려는 것처럼요). 하지만 이 방식은 특히 분자가 커지거나, 벤젠 고리처럼 특정 모양을 꼭 유지해야 하는 '부분 구조'(local structures)를 만들 때, 화학적으로 말이 안 되는 이상한 모양이나 불안정한 구조를 만들기 쉬웠습니다. 최근의 diffusion model들도 원자 단위 생성에서는 이런 문제가 있었습니다.
이 논문(HierDiff)은 다른 접근법을 제안합니다. 원자 하나하나에 집중하기 전에, 먼저 분자를 구성하는 좀 더 큰 의미있는 '부품' 또는 '조각'(fragments - 예를 들어 고리 구조 전체)들을 대략적인 위치에 배치합니다 (Coarse 단계). 이때 똑똑한 AI 기법(diffusion model)을 사용해서 이 조각들의 정보와 3D 위치를 만듭니다. 그런 다음, 이 조각들을 정교하게 다듬고 서로 화학적으로 올바르게 연결해서 최종적인 원자 수준의 3D 분자 구조를 완성합니다 (Fine 단계).
이렇게 하면, 중요한 '부분 구조'들(고리 등)이 처음부터 깨지지 않고 그 형태가 잘 유지되기 때문에, 더 현실적이고 안정적인 3D 분자 구조를, 특히 큰 분자에 대해서도 더 잘 만들 수 있다고 주장합니다. 마치 레고를 만들 때, 미리 만들어진 특수 블록(창문 블록, 바퀴 블록 등)을 먼저 활용해서 전체 구조를 잡는 것과 비슷하다고 생각할 수 있습니다.
2. Related work
Molecule generation은 신약 개발의 근본적인 문제 중 하나입니다. 초기 works에서는, molecule generation tasks가 분자의 sequential representations, 즉 SMILES를 generating함으로써 해결되었습니다. graph neural networks의 발전과 함께, 연구자들은 molecular topological structures를 generate하기 위해 graph-based generative model을 활용하기 시작했고 큰 progress를 이루었습니다. 그러나 sequence 기반 model이나 graph-based models 모두 3D geometric information을 capture하지 못하며, 이는 molecule property prediction 및 protein-ligand docking과 같은 다양한 분자 applications에 매우 중요합니다.
최근에는 3D molecule generation이 이 area에서 떠오르는 핫 토픽이 되었으며, 이 problem을 tackle하기 위해 다양한 deep generative models가 proposed되었습니다. 예를 들어, Gschnet은 atoms와 bonds를 반복적으로 sample하기 위해 Schnet을 갖춘 autoregressive process를 사용합니다. G-spherenet은 invariant geometric features를 autoregressively하게 generate하기 위해 discrete flows를 적용했습니다. EnFlow는 valid molecules를 sample하기 위해 continuous time normalizing flows를 utilizes합니다. EDM은 강력한 diffusion model을 이 area에 처음으로 적용하여 추가적인 improvements를 얻었습니다. 그러나 EDM은 large molecules에 대해 training할 때 종종 unrealistic ring systems와 broken molecules를 generates합니다.
molecule generation의 관련 분야 중 하나는 hierarchical graph generation입니다. 대부분의 이전 works는 어떤 intrinsic rules에 기반하여 hierarchical structure를 도출합니다. 예를 들어, 일부 work는 다른 hierarchies를 construct하기 위해 atom-motif 또는 node-edge와 같은 다른 granularity levels를 사용합니다. Zhou 등과 Chauhan 등은 다른 nodes를 다른 levels로 구별하기 위해 predefined rules를 사용합니다. Mi 등은 hierarchy를 정의하기 위해 natural graphical topology를 employs합니다. Geng 등은 hierarchical structure를 형성하기 위해 connection information을 수집합니다. Kuznetsov & Polykovskiy는 model의 다른 layers에 latent variables를 추가하여 hierarchy를 obtain합니다. 반면 우리의 method에서는 semantic-guided hierarchy를 근사하기 위해 learnable decoding module을 사용합니다.
정리 노트 (AI 연구자용)
배경:
- 초기 분자 생성: Sequence (SMILES) 또는 Graph (GNNs) 기반. 3D 정보 부재.
- 최근 동향: 3D Molecule Generation이 중요해짐 (Property prediction, Docking 등).
기존 3D 생성 연구:
- 접근법: Autoregressive (Gschnet, G-spherenet), Normalizing Flows (EnFlow), Diffusion (EDM) 등 다양한 deep generative models 제안됨.
- 주요 한계점 (본 논문이 주목하는): 특히 diffusion model 기반의 state-of-the-art인 EDM이 large molecules에 대해 training될 경우, 화학적으로 비현실적인 ring systems나 broken molecules를 생성하는 경향이 있음. 이는 atom-level generation의 한계와 연결됨 (Introduction에서 언급된 내용).
관련 연구 분야: Hierarchical Graph Generation:
- 개념: 분자(Graph)를 계층적으로 생성.
- 기존 방식: 대부분 고정된 규칙(intrinsic/predefined rules), 구조적 특징(granularity levels, topology), 연결 정보(connection info), 또는 model 구조(latent variables in layers)에 기반하여 hierarchy를 정의함.
- 본 논문(HierDiff)과의 차별점: 고정된 규칙 대신, learnable decoding module을 사용하여 semantic-guided hierarchy를 근사(approximate)하는 방식을 채택. 즉, 학습을 통해 의미론적으로 적절한 계층 구조를 파악하려는 접근.
(Note: 위 내용은 Related Work 섹션에서 파악된 정보를 요약한 것입니다.)
쉬운 설명:
AI로 분자를 만드는 연구는 꽤 오래됐어요. 처음에는 분자를 텍스트(SMILES)나 평면적인 연결도(graph)로 만들려고 했죠. 그런데 이런 방식은 분자의 중요한 3D 입체 구조 정보를 담지 못했어요. 약효 예측이나 단백질 결합 같은 걸 하려면 3D 구조가 꼭 필요하거든요.
그래서 요즘엔 AI로 바로 3D 분자 구조를 만들려는 연구가 활발해요. 여러 가지 방법들이 나왔는데, 원자를 하나씩 순서대로 붙이는 방식(autoregressive), 좀 다른 수학적 기법을 쓰는 방식(normalizing flows), 그리고 요즘 강력하다고 알려진 diffusion model을 쓰는 방식(EDM) 등이 있어요.
그런데 문제가 있어요. 현재 좋은 성능을 보이는 diffusion model(EDM)도 큰 분자를 만들려고 하면 종종 실패해요. 화학적으로 말이 안 되는 이상한 모양의 고리 구조를 만들거나 분자가 중간에 끊어진 것처럼 생성될 때가 있다는 거죠.
한편, 관련 연구로 '계층적 생성'이라는 아이디어가 있어요. 분자를 한 번에 다 만드는 게 아니라, 큰 덩어리부터 만들고 점점 세부적으로 만들어가는 방식이죠. 기존 연구들은 대부분 미리 정해진 규칙(예: '이런 모양은 상위 레벨', '저런 연결은 하위 레벨' 등)에 따라 이 '덩어리'나 '단계'(hierarchy)를 나눴어요.
이 논문(HierDiff)은 여기서 차별점을 갖는데요. 이 논문도 계층적으로 분자를 만들긴 하는데(앞서 말한 coarse-to-fine 방식), 미리 정해진 규칙 대신 AI가 스스로 학습해서('learnable') 의미상 적절한('semantic-guided') 계층 구조를 파악하게 만들어요. 즉, 좀 더 똑똑하게 분자의 구조적 단계를 이해해서 생성하겠다는 접근 방식입니다.

Figure 2. fragment conflicts의 예시: One-hop conflict는 연결된 두 fragments가 valid edge를 형성하기 위해 공유하는 elements가 없음을 의미합니다. Two-hop conflict는 연결된 fragments가 동일한 atom/bond를 공유하여 edges를 형성할 수는 있지만, valence가 위반될 때 conflicts가 발생하는 것을 나타냅니다.
3. Backgrounds
3.1. Denoising Diffusion Probabilistic Model
Denoising diffusion probabilistic model (DDPM)은 diffusion process를 역으로 진행함으로써 강력한 generative modeling 도구를 제공합니다. 더 구체적으로 말하면, diffusion process는 noise를 ground truth data로 project하고, generative process는 그 과정을 역으로 되돌리는 것을 learns합니다. 이 두 process는 latent variable model을 내포하며, 여기서 x1, · · · , xt−1는 latent variables입니다. forward process는 고정된 근사 사후 분포(fixed approximate posterior distribution)로 볼 수 있습니다: [ q(x_{1:T}|x_{0}) = \prod_{t=1}^{T} q(x_{t}|x_{t-1}) q(x_{t}|x_{t-1}) = \mathcal{N}(x_{t}; \sqrt{1 - \beta_{t}} x_{t-1}, \beta_{t} I) ] 여기서 β1, · · · , βT는 고정된 variance schedule에 해당합니다. 단순화를 위해 αt = 1 − βt 그리고 ¯αt = Qt i=1 αi로 두면, 임의의 time step에 대한 forward pass는 해석적인 형태(analytic form)를 가집니다. 즉, q (xt | x0) = N (xt; √¯αtx0, (1 − ¯αt) I) 입니다. generative process는 Markov chains의 transition kernel Pθ (xt−1|xt)를 parameterized하며, 해당 likelihood function은 다음과 같이 유도될 수 있습니다: [ P_\theta (x_{t-1} | x_t) = \mathcal{N} (x_{t-1}; \mu_\theta (x_t, t), \sigma^2_t I) \ P_\theta (x_0) = \int p(x_T) P_\theta (x_{0:T-1} | x_T) dx_{1:T} ] 여기서 μθ는 parameterized means function을 의미하고 σ2 t는 predefined variance입니다. 초기 분포(initial distribution) Pθ (xT )에 대해서는, equivariant coordinates를 위해 invariant base distribution을 선택합니다.
3.2. Equivariance and SE(3)-invariant Density Estimation
Equivariance는 physical world, 특히 atomic systems에서 널리 존재합니다. 예를 들어, atomic forces의 vector fields는 molecule의 3D positions과 함께 상응하여 rotate하거나 translate해야 합니다. 따라서 이러한 inductive bias를 function modeling에 integrating하는 것은 매력적인 properties를 가지며 널리 explored되었습니다. 더 구체적으로, space X와 Y에 작용하는 두 transformation Tg와 Sg가 주어졌을 때, 만약 다음 식이 만족되면 function f는 group G에 대해 equivariant하다고 간주됩니다: [ f \circ T_g(x) = S_g \circ f(x) ] 이 task에서는 주로 SE(3) group, 즉 3D space에서의 rotation 및 translation group에 초점을 맞춥니다. 3D molecule graph의 generative modeling을 위해, model distribution Pθ (.)의 density function은 SE(3)-invariant여야 합니다. 즉, Pθ (x) = Pθ (Tg (x)) 입니다. 이를 위해, 이전 methods는 bond angles과 같은 invariant components를 직접 model하거나, invariant base distribution을 사용하고 equivariant neural network에 의해 transformation을 model했습니다. HierDiff는 후자를 확장하여 fragment coordinates를 model하고 bond lengths를 predefined rules에 fit시키는 equivariant hierarchical framework를 사용하며, 이는 Sec. 5에서 논의될 것입니다.
정리 노트 (AI 연구자용)
3.1. Denoising Diffusion Probabilistic Model (DDPM):
- 개요: Ground truth data에 점진적으로 noise를 추가하는 forward process와, noise로부터 원본 data를 복원하도록 학습하는 generative process (reverse)를 기반으로 하는 generative model. Latent variable model로 해석 가능.
- 수식: Forward process q(xt|xt-1) 및 q(xt|x0) (Gaussian), Reverse process Pθ(xt-1|xt) (parameterized Gaussian μθ, σt) 정의.
- HierDiff 관련성: HierDiff의 Coarse Phase에서 fragment representation 및 coordinates 생성에 geometric diffusion model 사용의 이론적 기반. Equivariant coordinates 생성을 위해 invariant base distribution (e.g., noise) 사용 언급.
3.2. Equivariance and SE(3)-invariant Density Estimation:
- 개념: Equivariance는 입력의 transformation (e.g., rotation, translation)에 따라 출력도 예측 가능하게 변환되는 성질 (f ∘ Tg = Sg ∘ f). SE(3)는 3D rotation + translation group.
- 3D 생성 요구사항: 3D 분자 생성 model의 density function Pθ(.)는 SE(3)-invariant여야 함 (Pθ(x) = Pθ(Tg(x))). 즉, 생성된 분자를 회전/이동해도 그럴듯함(확률)은 변하지 않아야 함.
- 기존 Invariance 구현 방식:
- Invariant components (e.g., bond angles, 거리) 직접 model.
- Invariant base distribution (e.g., noise) 사용 + equivariant neural network로 transformation model.
- HierDiff의 접근 방식 (핵심 연결고리): HierDiff는 두 번째 방식(invariant base + equivariant network)을 equivariant hierarchical framework로 확장함. 이를 통해 다음을 수행:
- Diffusion으로 생성된 fragment coordinates를 SE(3) equivariant하게 model.
- Bond lengths를 predefined rules에 맞도록 fit (Sec. 5에서 상세 설명 예정).
(Note: 위 내용은 Backgrounds 섹션의 핵심 내용을 요약한 것입니다.)
쉬운 설명:
이 섹션에서는 HierDiff 모델을 이해하는 데 필요한 두 가지 배경 지식을 설명합니다.
1. Diffusion Models (DDPM): 노이즈 제거해서 그림 그리기
- 깨끗한 이미지에 점점 노이즈(지직거림)를 추가해서 완전한 노이즈로 만드는 과정(diffusion process)을 생각해 보세요.
- Diffusion 모델은 이 과정을 거꾸로 학습합니다. 완전한 노이즈에서 시작해서, 단계적으로 노이즈를 조금씩 제거하면서 원래의 깨끗한 이미지(generative process)를 만들어내는 방법을 배웁니다.
- HierDiff는 이 아이디어를 사용해서, 분자의 대략적인 fragment(조각)들의 3D 위치를 생성하는 첫 단계를 수행합니다.
2. Equivariance (SE(3) 불변성): 돌려도 똑같아야 해!
- 분자는 3D 공간에 존재합니다. 실제 분자를 잡고 돌리거나 위치를 옮겨도 그 분자 자체는 변하지 않죠. 그리고 분자의 에너지 같은 물리적 속성도 예측 가능한 방식으로 함께 변합니다.
- AI가 3D 분자를 만들 때도 이런 물리 법칙을 존중해야 합니다. AI가 만든 분자를 프로그램 상에서 돌려보거나 위치를 옮겼을 때, AI가 "이건 완전 다른 분자야!" 라고 인식하거나 구조가 망가지면 안 되겠죠.
- Equivariance: 분자를 돌리면 예측값(예: 힘 벡터)도 똑같이 돌아가는 성질.
- Invariance: 분자를 돌리거나 옮겨도 그 분자가 생성될 확률값 자체는 변하지 않는 성질 (SE(3)는 3D 회전/이동에 대한 불변성).
- AI 모델이 이런 성질을 갖도록 특별한 architecture(구조)를 사용합니다 (equivariant neural network).
HierDiff가 이 두 가지를 어떻게 사용할까?
- HierDiff는 diffusion model(노이즈 제거 아이디어)을 사용해서 분자 fragment들의 3D 좌표를 생성합니다.
- 이때 그냥 생성하는 게 아니라, 3D 공간의 회전/이동을 이해하는 똑똑한 equivariant 방식을 사용합니다 (equivariant hierarchical framework).
- 이렇게 하면 생성된 fragment들의 3D 위치가 물리적으로 더 타당해지고, 나중에 이 fragment들을 조립할 때도 bond length(결합 길이) 등을 미리 정해진 규칙에 맞게 조절하는 데 이 framework를 활용합니다 (자세한 내용은 나중에 설명).

이 그림은 **HierDiff 모델의 전체적인 작동 방식인 Coarse-to-Fine Approach를 세 단계로 나누어 보여주는 개요(overview)**입니다.
- Coarse-grained Diffusion (왼쪽 단계):
- 시작: 3D 공간에 무작위로 흩어진 점들(Random Noise)에서 시작합니다.
- 과정: Coarse-grained Diffusion 과정을 거칩니다. (앞서 Background에서 설명한 Diffusion 모델 활용)
- 결과: 각 fragment의 대략적인 3D 위치(Coordinates)와 그 fragment가 어떤 종류인지에 대한 정보(High-level Feature)를 생성합니다. 이미지에서는 위치가 정해진 약간 더 큰 구체들로 표현되어 있습니다. 아직 원자 수준의 상세 정보는 없습니다.
- Fine-grained Generation (중간 단계):
- 시작: 이전 단계에서 얻은 High-level Feature + Coordinates를 입력으로 받습니다.
- 과정: Fine-grained Generation 과정을 거칩니다. (앞서 텍스트에서 언급된 equivariant message passing networks, iterative refinement process 등이 여기서 사용될 것으로 추정됩니다).
- 결과: Coarse 정보로부터 실제 fragment들을 decode하고, 이 fragment들이 어떻게 서로 연결되는지를 결정하여 3D 공간 상의 Fragment 3D Graph를 만듭니다. 이미지에서는 구체(fragment)들이 선(edge)으로 연결된 그래프 형태로 표현되어 있습니다.
- Conformation Assembling (오른쪽 단계):
- 시작: 이전 단계에서 얻은 Fragment 3D Graph를 입력으로 받습니다.
- 과정: Conformation Assembling 과정을 거칩니다.
- 결과: 연결된 fragment 정보를 바탕으로 완전한 원자(atom) 수준의 3D 분자 구조 (Drug-like 3D Molecules)를 최종적으로 조립(assemble)합니다. 이미지에서는 실제 분자처럼 보이는 ball-and-stick 모델로 표현되어 있습니다.
요약하면, 이 그림은 HierDiff가 처음에는 noise로부터 대략적인 fragment 위치와 정보(Coarse)를 생성하고, 점차적으로 세부적인 연결 정보(Fine)를 추가하여 최종적인 3D 분자 구조를 만들어가는 계층적인(hierarchical) 생성 과정을 시각적으로 보여줍니다.
4. Coarse-to-Fine Approach
이 섹션에서는 problem formulation과 coarse-to-fine 정의를 포함하여 우리의 coarse-to-fine approach를 소개합니다.
4.1. Problem Formulation and Notations
G를 3D graphs의 공간이라고 합시다. 여기서 각 3D graph G는 fragment set V와 edge set E로 구성됩니다. 더 구체적으로, 모든 fragment V ∈ V는 여러 atoms와 bonds의 조합을 represents합니다. 예를 들어, benzene ring은 여섯 개의 carbon atoms와 aromatic bonds를 포함하는 fragment가 될 수 있습니다. 이전 works에서처럼 edges를 사용하여 chemical bonds를 나타내는 대신, 우리는 edge Eij ∈ E를 사용하여 두 fragments Vi와 Vj가 공유하는 bond/atom이 있음을 나타냅니다. 이러한 종류의 정의는 합리적인 크기의 fragment vocabulary를 사용하여 Fig. 5에서처럼 fragment sphere 상의 tangent condition을 이용해 molecule geometry를 model할 수 있게 합니다. 따라서, 3D generation model의 목표는 3D molecule graphs의 empirical distribution을 model하는 probabilistic model Pθ (V, E)를 learn하는 것이며, 이는 새로운 molecules를 sample하는 데에도 사용될 수 있습니다.
4.2. Coarse-to-Fine Framework
4.2.1. CHALLENGES OF NON-AUTOREGRESSIVE METHODS
autoregressive approach와 비교할 때, non-autoregressive generative models는 global modeling ability라는 자연스러운 장점 때문에 3D molecule generation에 더 유망합니다. 경험적으로(Empirically), previous work는 일관된 observations를 보여주었습니다. 몇 가지 매력적인 properties가 있음에도 불구하고, fragment level에서의 non-autoregressive model은 실제로 hard constraints 하에서 다음과 같은 structure generation procedure를 내포합니다. [ G \sim P_{\theta}(V, E), \text{ s.t. } G_s(V_i) \in W, \forall i = 1, \ldots, n, ] 여기서 Gs(Vi)는 Vi와 그 neighbors로 구성된 substructure를 의미하고, W는 모든 valid substructures의 집합을 의미합니다. 직관적으로, valid substructures는 다음 conditions를 만족합니다: neighbors는 ensembled되기 위해 matched components(atoms/bonds)를 가져야 합니다; node가 multiple neighbors를 갖는 경우에는 이 node에도 모든 neighbors를 match시키기에 충분한 matched components가 있어야 합니다. Eq. (4)의 constraints를 더 잘 설명하기 위해 Fig. 2를 제공합니다.
제한된 chemical valency는 fragment generation에서 conflicts를 정말 흔하게 만듭니다. fragment conflict를 avoiding하는 problem은 multi-hop conflicts 때문에 high complexity를 가지며 소위 ”combinatorial exploding” issues를 야기한다는 점에 유의해야 합니다. non-autoregressive modeling fashion의 경우, complexity는 structure size에 따라 지수적으로 increases합니다. 이 problem은 다른 tasks에서도 discovered됩니다. 예를 들어, Dispatching Route Generation, Optimal Experiment Design, 그리고 Protein Alignment Generation이 있습니다. fragment-based molecule generation에서는 realistic drug-size molecules를 ordering-agnostic way로 generate하는 것이 challenging합니다.
4.2.2. SOLUTIONS FOR AVOIDING CONFLICTS
hard constraints가 있는 distribution에서 sample하는 것은 difficult합니다. previous work에서 adopted된 한 가지 direct solution은 rejection sampling을 수행하는 것입니다. 즉, connectable molecules만 accept하는 것입니다. 그럼에도 불구하고, rejection sampling은 극도로 낮은 acceptance rate 때문에 fragment-based methods에서는 practice에서 applicable하지 않습니다; 대안으로, hard constraint를 approximate하기 위해 learnable model, Pϕ(Vi|N (Vi))를 사용할 수 있습니다. 여기서 N (Vi)는 Vi의 neighbors를 의미합니다. 상응하여, generative distribution Pθ,ϕ(V, E)는 다른 target distribution Pθ (V, E) Q 1≤i≤n Pϕ(Vi|N (Vi))로 쓰일 수 있습니다. 불행히도, 이러한 distribution에서 sampling을 수행하려면 Gibbs Sampling과 같은 Markov chain Monte Carlo (MCMC) sampling이 필요하며, 이는 여전히 efficiency issues로 어려움을 겪습니다.
generated samples가 filter나 refinement를 통해 constraint를 satisfy하도록 만드는 대신, 우리는 constraint를 hierarchical fashion 내에서 model phase에 직접 decompose하고 embed하려고 시도합니다.
더 구체적으로, 우리는 variable (H)를 latent variable로 설계하고, probabilistic model은 Pθ,ϕ(V, E) = Pθ (H)Pϕ(V, E|H)로 expressed될 수 있습니다. 결과적으로, 우리는 logarithm의 concavity를 이용하여 maximum likelihood objective의 lower bound를 다음과 같이 obtain할 수 있습니다: [ E ( V,E ) \sim P_{\text{data}} \log \sum_{H \in \mathcal{H}} P_0(H) P_{\phi}(V, E | H) \geq E ( V,E ) \sim P_{\text{data}} E_{H \sim Q(H | V,E)} \left[ \log P_0(H) \text{ Coarse-grained Diffusion} + \log P_{\phi}(V, E | H) \text{ Fine-Grained Generation} - \log Q(H | V, E) \text{ Constant Term} \right] ] 여기서 H는 H의 possible support를 의미합니다. 공식적으로, Eq. (5)의 Q(H|V, E)는 chemical features를 extracting하고 atom coordinates를 averaging함으로써 implemented됩니다. Q(H|V, E)가 constant term으로 설정되므로, objective는 앞의 두 terms, 즉 log Pθ (H)와 log Pϕ(V, E|H)만 include할 것입니다. 첫 번째 objective log Pθ (H)는 coarse-grained fragment diffusion model에 의해 approximated될 수 있고, 두 번째 objective log Pϕ(V, E|H)는 equivariant message passing networks와 iterative refinement process에 의해 modeled됩니다.
직관적으로, coarse-to-fine process는 먼저 각 component의 position과 function을 determining한 다음, 작은 subsets에서 connectable fragments를 finding하고 그것들을 assembling하는 것과 같습니다. 따라서, HierDiff는 non-autoregressive methods의 global modeling property를 maintain할 수 있으며, 또한 connectable fragments를 finding하는 complexity를 significantly reduce할 수 있습니다.
정리 노트 (AI 연구자용)
4.1 Problem Formulation:
- Representation: 3D 분자를 Graph G=(V, E)로 정의.
- V: Fragment들의 집합 (atoms+bonds의 조합, e.g., 벤젠 고리). 원자 단위가 아님.
- E: Fragments 간 공유되는 bond/atom을 나타내는 edge. 직접적인 화학 결합이 아님.
- 이 정의는 fragment sphere에서의 tangent condition을 이용한 geometry modeling을 가능하게 함 (Fig. 5 참조).
- 목표: Probabilistic model Pθ(V, E)를 학습하여 3D molecule graphs의 empirical distribution을 모델링하고 sampling.
4.2 Coarse-to-Fine Framework:
- 4.2.1 Challenges (Non-autoregressive Fragment Generation):
- Fragment level에서의 non-autoregressive 생성은 hard constraints 문제에 직면.
- Validity Constraint: 생성된 substructure Gs(Vi)가 유효(W)하려면, neighbors 간 matched components (atoms/bonds)가 필요 (Fig. 2 conflicts 참조).
- Complexity: 제한된 chemical valency로 conflicts (특히 multi-hop)가 빈번하며, non-autoregressive 방식에서는 combinatorial exploding 문제 발생 (복잡도 지수적 증가). Ordering-agnostic 방식으로 큰 분자 생성 challenging.
- 4.2.2 Solution (HierDiff's Approach):
- 기존 해결책 한계: Rejection sampling (낮은 acceptance rate), MCMC 기반 learnable constraint model (efficiency issues).
- HierDiff 전략: Constraint를 filter/refinement 단계에서 처리하는 대신, hierarchical fashion으로 model phase 내에 decompose 및 embed.
- Latent Variable (H): 모델을 P(V, E) = Pθ(H)Pϕ(V, E|H)로 인수분해. H는 coarse 상태 (High-level features + Coordinates).
- Objective (ELBO Optimization, Eq. 5):
- log Pθ(H) (Coarse Phase): Coarse-grained fragment diffusion model로 근사. Latent H 생성.
- log Pϕ(V, E|H) (Fine Phase): Equivariant message passing networks + iterative refinement process로 모델링. H를 조건으로 Fragment Graph(V, E) 생성 (Constraint 처리 내재화).
- 핵심 이점: Hard constraint 문제를 계층적으로 분해. Non-autoregressive 장점 유지하며 connectable fragments 탐색의 complexity를 크게 감소시킴.
(Note: 위 내용은 Section 4의 핵심 내용을 요약한 것입니다.)
쉬운 설명:
이 섹션에서는 HierDiff 모델의 핵심 아이디어인 Coarse-to-Fine (대략적인 것에서 세밀한 것으로) 접근 방식을 설명합니다.
1. HierDiff는 분자를 어떻게 볼까?
- 원자 하나하나보다는, **미리 정의된 좀 더 큰 화학적 '부품' 또는 '조각' (fragment)**들로 분자를 본다고 생각합니다. 예를 들어 벤젠 고리 전체가 하나의 fragment가 될 수 있습니다. 그리고 이 fragment들이 서로 어떻게 연결되는지에 집중합니다.
2. 뭐가 그렇게 어려운 문제일까?
- 그냥 AI보고 이런 fragment들을 만들어서 3D 공간에 배치하고 연결하라고 하면 문제가 많이 생깁니다. 특히 non-autoregressive 방식 (한 번에 또는 순서 없이 생성하는 방식)은 더 어렵습니다.
- 연결 불가: 만들어진 fragment 조각들이 화학적으로 서로 연결될 수 없는 경우가 많습니다 (앞서 본 Figure 2의 conflicts). 원자들이 맞지 않거나, 결합 규칙(원자가 가질 수 있는 팔 개수 - valence)을 위반하게 됩니다.
- 복잡도 폭발: 특히 분자가 커지면, 어떤 조각들을 어떻게 연결해야 규칙에 맞는지 모든 경우의 수를 따지는 것이 너무 복잡해집니다 (combinatorial exploding).
- 기존 해결책의 비효율: 오류가 생기면 버리고 다시 만드는 (rejection sampling) 방식은 성공률이 너무 낮고, 다른 복잡한 샘플링 방법(MCMC)은 너무 느립니다.
3. HierDiff의 똑똑한 해결책: 2단계 접근!
- HierDiff는 이 어려운 문제를 두 단계로 나누어 풉니다.
- 1단계: 밑그림 그리기 (Coarse Phase): 먼저, diffusion model(노이즈 제거 아이디어)을 이용해서 **각 fragment 조각이 대략 어디쯤 위치해야 하고 어떤 종류인지 (High-level Feature + Coordinates)에 대한 '밑그림'(Latent Variable H)**을 그립니다. 아직은 대략적인 정보입니다.
- 2단계: 세밀하게 조립하기 (Fine Phase): 앞에서 그린 '밑그림'(H) 정보를 바탕으로, 실제 fragment 조각들을 선택하고 화학적으로 올바르게 연결하여 Fragment 3D Graph를 만듭니다. 이 단계에서는 equivariant message passing network 같은 똑똑한 도구를 사용해서 앞서 말한 연결 오류(conflicts)가 발생하지 않도록 조심스럽게 조립합니다.
- 최종적으로 이 Fragment 3D Graph로부터 완전한 원자 구조를 만듭니다.
왜 이게 더 좋을까?
- 어려운 문제를 한 번에 풀려고 하지 않고, (1) 쉬운 밑그림 그리기 -> (2) 밑그림 기반으로 조심스럽게 조립하기로 나누었기 때문에 훨씬 효율적입니다.
- 특히 복잡한 연결 규칙(hard constraints)을 조립 단계(Fine Phase)에서 효과적으로 처리할 수 있게 되어, combinatorial exploding 문제를 피하면서도 크고 현실적인 분자를 생성할 가능성이 높아집니다.


Figure 4: Feature 추출 예시
- 목적: Coarse-grained 단계에서 사용되는 두 종류의 chemical feature (Hf)가 3D Benzaldehyde 분자로부터 어떻게 추출되는지를 보여줍니다.
- 내용:
- Property-Based Coarse Feature: 분자의 특정 표면적이나 화학적 속성(수소 결합 수, 고리 수 등)을 계산하여 8차원 벡터로 만듭니다. 이는 분자의 전반적인 물리화학적 특성을 나타냅니다.
- Element-Based Coarse Feature: 분자를 구성하는 원소(예: C, H, O)의 개수를 세어 3차원 벡터(원소 빈도 히스토그램)로 만듭니다. 이는 분자의 기본적인 원소 구성을 나타냅니다.
- 요점: 이 두 가지 feature를 합쳐서 Coarse-grained Diffusion 모델의 입력 또는 중간 representation으로 사용합니다 (Hf 부분).
Figure 5: Positional Feature (Center Coordinate)의 정당성
- 목적: Coarse-grained 단계에서 fragment의 3D 위치 정보(Hp)로 center coordinates(중심 좌표)만 사용하는 것이 왜 충분한지를 설명합니다.
- 내용:
- 각 coarse-grained node(fragment)는 화학 규칙에 따라 결정될 수 있는 반경(radius)을 가진 구(sphere)로 생각할 수 있습니다.
- 두 fragment가 연결된다는 것은 이 두 구가 접하는(tangent) 조건으로 표현될 수 있습니다.
- 구들의 중심 좌표와 반경 정보만 있어도, 구들이 접하는 지점(tangent space)에서 원자들의 가능한 위치를 샘플링(Sample from Tangent Space)하여 전체 원자 구조(full atom conformation)를 재구성할 수 있습니다.
- 요점: 따라서 fragment의 복잡한 3D 형태 대신 center coordinate만 positional feature로 사용해도 나중에 원자 수준 구조를 복원하기에 충분하며, 이것이 좋은 representation임을 보여줍니다.
Figure 6: Fine-grained 생성 과정 상세
- 목적: Coarse-grained 단계 이후 진행되는 Fine-grained fragment generation 과정의 4단계를 시각적으로 보여줍니다. (Section 5.2 설명)
- 내용: (왼쪽에서 오른쪽으로 진행)
- Choose a Focal Node (ϕfocal): 이미 fine-grained 상태인 노드들 중에서 다음 연결을 시작할 중심 노드(Focal Node)를 선택합니다 (빨간색 상자).
- Predict a New Edge (ϕedge): Focal Node와 연결될 새로운 edge(결합)를 아직 coarse-grained 상태인 노드들 중에서 예측하여 선택합니다 (점선 중 하나 선택).
- Choose the Fragment (ϕnode): 새로 연결된 coarse-grained node의 정확한 fine-grained fragment type을 결정합니다 (예: 노드 레이블이 'CO'로 바뀜).
- Iterative Refine (ϕrefine): 새로 추가된 정보를 바탕으로 기존 fine-grained node들의 타입을 반복적으로 재평가하고 수정하여 전체적인 구조의 정확도를 높입니다 (bias 보정).
- 범례: 메시지 전달 방향, Focal Node, Fine-grained Node, Coarse-grained Node를 구분하여 보여줍니다.
- 요점: 이 그림은 coarse 정보로부터 어떻게 점진적으로 fine-grained fragment들의 타입과 연결 관계를 결정하고, iterative refinement를 통해 오류를 수정하며 최종 Fragment 3D Graph를 만들어나가는지를 단계별로 보여줍니다. 각 단계에서 어떤 neural network module(ϕfocal, ϕedge, ϕnode, ϕrefine)과 message passing 방식(EGNN)이 사용되는지도 나타냅니다.
5. HierDiff: hierarchical diffusion-based model
이 섹션에서는 Fig. 3에 설명된 대로, coarse-grained fragment generation, fine-grained fragment generation, 그리고 atom conformation assembling을 포함하여 제안된 HierDiff model을 자세히 소개합니다. 이는 또한 Eq. (5)의 parameterized terms에 해당합니다.
5.1. Coarse-Grained Fragment Diffusion
우리는 coarse nodes의 representation을 H = [Hf , Hp]로 정의합니다. 여기서 Hf는 invariant chemical features를, Hp는 equivariant positional features를 의미합니다. 공식적으로, Eq. (5)의 Q(H|V, E)는 chemical features를 extracting하여 Hf를 obtain하고 모든 atom coordinates를 averaging하여 Hp를 obtain함으로써 implemented됩니다. 구체적으로(Specifically), property-based features는 fragment V와 attachment E, 즉 neighbor fragments와의 connection 모두에 의존할 수 있습니다.
coarse-grained phase에서는 diffusion model이 log Pθ (H)를 approximate하기 위해 proposed됩니다. diffusion model에서 sampling할 때, 우리는 먼저 training set에서 calculated한 histogram으로부터 coarse nodes의 number를 sample합니다.
5.1.1. CHEMICAL FEATURE
우리는 서로 다른 chemical 및 geometrical properties를 가진 fragments와 molecules에 대해 충분히 discriminative하도록 features를 신중하게 설계했으며, 이를 통해 우리의 domain knowledge를 inductive bias로 model에 쉽게 integrate할 수 있습니다. 그리고 우리는 구체적으로(specifically) 두 종류의 features를 employ합니다:
- Property-based Coarse Feature: 우리는 drug discovery에서 널리 사용되는 몇 가지 중요한 properties를 8차원 vector로 summarize합니다. 여기에는 hydrogen bonds 및 rings의 number, 서로 다른 surfaces의 area 등이 포함됩니다.
- Element-based Coarse Feature: 우리는 또한 동일한 수의 valence electrons를 가진 elements가 보통 동일한 properties를 share한다는 fact에 inspired 받아, element frequency의 histogram, 즉 3차원 vector를 feature representation에 include합니다.
자세한 implementations는 Fig. 4, Table 3, Table 4를 참조하십시오.
5.1.2. POSITIONAL FEATURE
fragment level에서 3D conformation systems을 represent하는 데에는 몇 가지 가능한 방법이 있습니다. 예를 들어, neighbor fragments 사이의 dihedral angle, 그리고 distance matrix가 있습니다. 이 paper에서는 coarse node의 positional feature로 단순히 center coordinates를 simply use합니다. 왜냐하면 이 information만으로도 RDkit ETKDG module에서 가져온 bond length와 bond angles의 predefined vocabulary를 사용하여 atom resolution에서 conformation을 determine하기에 충분하다는 것을 발견했기 때문입니다. coarse node의 center position은 rotation에 대한 degree of freedom으로부터 generated된 모든 possible conformations를 includes하는 conformation sphere의 center로 seen될 수 있다는 점에 유의하십시오. connected fragments는 실제로 degree of freedom을 eliminates하는 tangent condition에 correspond하며, 이는 Fig. 5에 illustrated되어 있습니다.
5.1.3. DIFFUSION PROCESS
다음으로, Hf와 Hp의 modeling을 개별적으로 introduce합니다. Hf는 step t > 0에 대해 Gaussian noise를 사용하는 전형적인(typical) diffusion model에 의해 modeled될 수 있습니다. 그러나 우리는 continuous feature Hint f와 Hcont f에 대한 0번째 term, 즉 L0는 유사하게 관찰된 바와 같이 신중하게 designed되어야 함을 발견했습니다. 이 paper에서는 더 나은 empirical performance를 보여준 다음 form을 사용합니다. [ L_0(H_{f}^{int}, H_{f}^{cont}) = -\log \left[ \int_{H_{f}^{int} - \frac{1}{2}}^{H_{f}^{int} + \frac{1}{2}} \mathcal{N}(u | x_{0}^{(H_{f}^{int})}, \sigma_{0}) , du \right] - \log \mathcal{N} \left( H_{f}^{cont} \bigg| \frac{x_{0}^{(H_{f}^{cont})}}{\alpha_{0}} - \frac{\sigma_{0}}{\alpha_{0}} \hat{\epsilon}{0}, \frac{\sigma{0}^2}{\alpha_{0}^2} I \right). ] 다음으로, Hp의 generation을 설명합니다. Eq. (2)의 likelihood function을 SE(3)-invariant하게 만들기 위해, 우리는 질량 중심(center of mass, CoM)이 0인 systems 하에서 initial distribution을 set합니다. 즉, CoM-free Gaussian을 applying합니다: [ N(H_p | 0, \sigma^2 I) = \left( \sqrt{2 \pi \sigma} \right)^{-(M - 1) \cdot n} \exp\left( -\frac{1}{2 \sigma^2} | H_p |^2 \right). ] 여기서 Hp는 space RM ×n에 속하며, M은 fragment nodes의 number이고 n은 coordinate dimension과 같습니다. 게다가, 널리 applied되는 noise parameterization 하에서 equivariant Markov transition kernel이 constructed됩니다. [ \mu_0 (H_t^p, t) = \frac{1}{\sqrt{\alpha_t}} \left( H_t^p - \frac{\beta_t}{\sqrt{1 - \bar{\alpha}_t}} \epsilon_0 (H_t^p, t) \right) \cdot ] 만약 ϵθ가 SE(3)-equivariant networks에 의해 parameterized된다면, transitional kernel Pθ (Ht−1 p |Ht−1 p ) 또한 SE(3)-equivariant입니다. 즉, Pθ Ht−1 p | Ht p = Pθ Tg Ht−1 p | Tg Ht p 입니다. 자세한 proof는 Appendix A.2에 남겨둡니다.
5.2. Fine-Grained Fragment Generation
[Hf , Hp]가 diffusion model에 의해 generated된 후, Fig. 3에 illustrated된 바와 같이 3D space의 coarse-grained nodes 집합이 obtained됩니다. 이제 coarse-grained nodes에 conditioned하여 fine-grained fragment types와 edges를 generating하는 detailed process를 introduce합니다. 이는 term Pϕ(V, E|H)에 correspond합니다.
먼저 여기서 decoding logic을 간략하게 introduce합니다. 각 decoding step에서, fine-grained generation process는 네 가지 stages를 contains합니다. 첫째, 우리는 parameterized neural network module ϕfocal을 사용하여 모든 existing fine-grained nodes 중에서 focal node를 select합니다. 다음으로, link prediction network, 즉 ϕedge를 utilize하여 모든 나머지 coarse-grained nodes 중에서 focal node에 linked될 수 있는 new node를 identify합니다. 그런 다음 우리는 다른 network ϕnode의 도움으로 위에서 새로 linked된 coarse-grained node의 fine-grained fragment type을 obtain합니다. 마지막으로, 새로 determined된 fragment type에 based on, fine-grained nodes의 bias를 correct하기 위해 iterative refinement process가 conducted됩니다. beginning에는 모든 nodes가 coarse-grained이므로, node를 randomly select하고 직접 ϕnode를 use하여 fragment type을 predict한다는 점에 유의해야 합니다. 위의 procedure는 Fig. 6에 illustrated되어 있습니다. 우리는 다음에서 우리의 assembling module의 몇 가지 key elements를 emphasize합니다:
5.2.1. ITERATIVE REFINEMENT
이제 iterative refinement process의 details를 자세히 살펴봅니다. 이 iterative refinement process를 introducing하는 motivation은 existing fine-grained nodes의 bias를 correct하여, global view에서 더 realistic molecules를 generate하는 ability를 enhance하는 것입니다. 구체적으로(Specifically), 우리는 모든 coarse-grained nodes와 다른 fine-grained nodes에 conditioned된 각 decoded fine-grained fragment의 probability를 approximate하기 위해 mask prediction model ϕrefine을 design합니다. target은 다음과 같이 fine-grained nodes의 joint probability를 maximize하는 것입니다: [ P_{\text{target}} = \prod_{V_i \in T_f} P_{\phi \text{ refine}} (f(V_i) | T_c, T_f \setminus V_i), ] 여기서 Tf는 모든 existing fine-grained nodes의 set이고, Tc는 coarse-grained nodes의 set이며, f는 특정 node에 대한 fragment type을 return하는 function입니다. 위의 target distribution에서 sample하기 위해, 우리는 node type replacement를 state transition으로 defined하는 Markov Chain을 defined했고, sampling process를 conduct하기 위해 early-stopping Monte Carlo sampling strategy를 adopted했습니다. detailed algorithm은 Appendix A.4에 illustrated되어 있습니다. 우리는 또한 제안된 iterative refinement process의 effectiveness를 prove하기 위해 ablation study를 conduct했으며, 이는 Appendix C.7에 shown되어 있습니다.
5.2.2. MESSAGE PASSING NEURAL NETWORKS
여기서는 fine-grained process에서 앞서 언급된 models, 즉 ϕfocal, ϕedge, ϕnode 그리고 ϕrefine을 introduce합니다. Fig.2와 같은 disconnectivity problem을 avoid하기 위해, 우리는 이 models를 정교하게 design해야 합니다.
구체적으로(Specifically), input은 chemical features와 positional features 모두로 represented되는 3D fragments의 set입니다. 또한, 모든 fine-grained nodes에 대한 fragment types를 indicate하기 위해 one-hot vectors를 add합니다. initial stage에서는 input이 fully connected 3D graph로 treated되고, vanilla EGNN이 모든 links와 nodes에 대한 initial embeddings를 extracts합니다. fine-grained generation process에 따라, ϕfocal은 단순히 fine-grained nodes 사이에서 information을 passes합니다. ϕedge는 그 다음 focal node가 tree structure의 root로 treated되는 tree bottom-up pattern으로 모든 fine-grained nodes의 information을 focal node로 aggregates합니다. new edge가 predicted된 후, network는 new edge의 addition을 tree top-down pattern으로 모든 fine-grained nodes에 broadcasts합니다. 마지막으로, ϕnode는 fine-grained fragment type을 decoding하기 위해 bottom-up pattern으로 모든 fine-grained nodes에서 new node로 information을 aggregates합니다. mask prediction module ϕrefine에 대해서는, ϕnode와 similar한 bottom-up EGNN을 utilize합니다. information은 Eq. (6)에서와 같이 target distribution을 compute하기 위해 모든 fine-grained nodes에서 masked nodes로 aggregated됩니다. message-passing process의 illustration은 Fig. 6에서 found될 수 있습니다. 우리는 EGNN의 model-level modification을 Appendix A.3에서 discussed했습니다.
5.2.3. TRAINING
training 중에는, 먼저 각 step에서 connected subgraph를 randomly sampling하는 것으로 start합니다. 그런 다음 랜덤한 leaf node가 picked되고, 우리는 이 node의 fine-grained generation을 simulate합니다. 선택된 node를 제외한 subgraph의 모든 fine-grained nodes와 edges는 kept됩니다. 다른 nodes에 대해서는 coarse features와 그들의 position만 maintain합니다. 그런 다음 ϕfocal은 위의 feature에 based on fine-grained subgraph 내의 nodes 중에서 선택된 node의 parent의 probability를 maximize하도록 trained됩니다. ϕedge는 다른 모든 coarse-grained nodes 중에서 focal node와 선택된 node 사이의 edge link의 probability를 maximize하도록 trained됩니다. ϕnode는 선택된 node의 fine-grained fragment type을 output하도록 trained됩니다. iterative refinement part에 대해서는, 우리는 그냥 subgraph에서 node의 fine-grained feature를 randomly mask하고, ϕrefine은 masked fragment type을 reconstruct하도록 trained됩니다. Detailed implementations와 objectives는 Appendix A.4에서 found될 수 있습니다.
5.3. Assembling to Atom Conformation
fine-grained generation process에서 determined된 모든 fine-grained nodes와 link relations가 Given되면, 우리는 atom-level conformation을 construct하기 위해 두 개의 linked fragments 내에서 어떤 atoms가 merged될 수 있는지 decide해야 합니다. 이 process를 conduct하기 위해, 우리는 먼저 fragment를 randomly choose하고, 그것의 neighbor fragments에 대한 모든 possible attachments를 enumerate하며, coarse-grained fragment generation에서 우리의 generated positional features와 가장 가까운 fragment center geometric을 가진 것을 select합니다. 구체적으로(Specifically), 우리는 각 candidate attachment에 대한 local conformation을 generate하기 위해 RDkit을 use하며, fragment center coordinates 간의 difference를 measure하기 위해 root-mean-square deviation (RMSD)를 apply합니다. 위의 process는 모든 local connections가 determined될 때까지 neighboring structure를 following하며 continued될 것입니다. 그런 다음 우리는 각 atom의 coordinates를 generate합니다. local conformations를 각 molecule coordinate system에 plug하기 위해, 우리는 rotation matrix (R)와 translation vector (t)를 determine해야 합니다. 여기서 우리는 fragment level에서 Kabsch Algorithm을 사용하여 generated local coordinates와 RDkit predicted local coordinates 사이의 R과 t를 compute합니다. 그런 다음 우리는 obtained된 R, t를 RDkit generated atom level coordinates에 applied하여 RDkit generated local geometry를 우리의 sampled center positions에 align합니다. 이 process는 randomly selected fragment와 그 neighbors에 의해 constructed된 subgraph에서 starts하며 full atom conformation이 derived될 때까지 successively conducted됩니다. RDkit은 conformation generation에서 local geometry를 generating하기 위해서만 utilized된다는 점에 Noted합니다.
full detailed algorithm은 Appendix A.5에 introduced되어 있습니다.
정리 노트 (AI 연구자용)
5.1 Coarse-Grained Fragment Diffusion (log Pθ(H) 모델링):
- 출력: Coarse latent state H = [Hf, Hp] 생성 (Hf: invariant chemical features, Hp: equivariant positional features).
- 시작: 생성할 coarse nodes 개수를 training set 분포에서 sampling.
- Chemical Features (Hf):
- Property-based: 8-dim (H-bond 수, ring 수, surface area 등).
- Element-based: 3-dim (원소 빈도 히스토그램). (Fig 4 참조)
- 표준 Diffusion 사용, continuous feature에 대한 L0 (t=0 항) 특별 처리.
- Positional Features (Hp):
- Fragment center coordinates 사용. (Fig 5 정당성: tangent condition으로 원자좌표 재구성 가능).
- SE(3)-Equivariance 확보: CoM-free Gaussian 초기 노이즈 분포 + SE(3)-equivariant network로 parameterize된 diffusion kernel ϵθ 사용.
- 구현: Q(H|V, E)는 feature 추출 및 좌표 평균으로 구현.
5.2 Fine-Grained Fragment Generation (log Pϕ(V, E|H) 모델링):
- 입력: H = [Hf, Hp]. 출력: Fragment Graph (V, E).
- Iterative Decoding Process (Fig 6): 4단계 반복:
- ϕfocal: Fine-grained nodes 중 focal node 선택.
- ϕedge: Focal node와 연결될 coarse-grained node 및 edge 예측.
- ϕnode: 새로 연결된 node의 fine-grained fragment type 결정.
- ϕrefine: Iterative Refinement - 기존 fine-grained nodes의 bias 수정.
- Iterative Refinement (5.2.1):
- Mask prediction model (ϕrefine) 사용, MCMC 샘플링 (node type replacement)으로 joint probability (Eq. 다음 식) 최적화하여 전역적(global) 사실성 증진.
- Message Passing Networks (5.2.2):
- ϕfocal, ϕedge, ϕnode, ϕrefine 모두 EGNN 변형 사용 (3D 구조 및 연결성 제약 처리).
- 특정 message passing 패턴 (bottom-up, top-down) 활용 (Fig 6 참조).
- Training (5.2.3):
- Connected subgraph 샘플링 후 leaf node 생성 시뮬레이션. 각 ϕ 모듈을 해당 예측/재구성 작업에 맞게 학습.
5.3 Assembling to Atom Conformation:
- 입력: Fine-grained Graph (V, E). 출력: 원자 수준 3D conformation.
- 프로세스:
- 연결된 fragments 간 최적의 원자 결합 지점(attachment) 결정: RDkit으로 생성한 후보 local conformations 중, Coarse 단계에서 생성된 center coordinates(Hp)와 가장 유사한 fragment center geometric을 갖는 것을 RMSD 기준으로 선택.
- 원자 좌표 배치: 선택된 RDkit local geometry와 목표 위치(sampled center positions) 간의 Rotation(R) 및 Translation(t)을 fragment level에서 Kabsch Algorithm으로 계산.
- 계산된 R, t를 RDkit 원자 좌표에 적용하여 최종 원자 conformation 생성 및 정렬.
- 참고: RDkit은 최종 조립 단계에서 local geometry 후보 생성에만 제한적으로 사용됨.
(Note: 위 내용은 Section 5의 핵심 구현 내용을 요약한 것입니다.)
쉬운 설명:
이번 섹션에서는 HierDiff 모델이 실제로 어떻게 작동하는지, 각 단계를 자세히 설명합니다. (Figure 3의 각 단계를 떠올려 보세요.)
1단계: 밑그림 그리기 (Coarse-Grained Fragment Diffusion)
- 무엇을 만드나? 분자를 구성할 fragment(부품)들의 목록을 만듭니다. 이 목록에는 각 부품이 어떤 종류인지(Chemical Features: 원소 종류 개수, 수소 결합 수, 고리 수 등)와 3D 공간에서 대략 어디쯤 있는지(Positional Features: 부품의 중심 좌표) 정보가 담깁니다.
- 어떻게 만드나?
- Diffusion (노이즈 제거) 방식을 사용합니다.
- 특히 3D 위치를 다룰 때는 분자를 마음대로 돌리거나 움직여도 괜찮도록(SE(3)-Equivariance), **3D 공간을 이해하는 특별한 AI 신경망(SE(3)-equivariant network)**을 사용합니다. (마치 로봇 팔이 물체의 방향과 상관없이 물체를 잡는 것과 비슷합니다.)
2단계: 세밀하게 조립하기 (Fine-Grained Fragment Generation)
- 무엇을 하나? 1단계에서 만든 대략적인 밑그림(부품 목록과 위치 정보 H)을 바탕으로, **실제 부품(fragment)들을 결정하고 어떻게 연결할지 정해서 'Fragment 3D Graph'(부품 연결도)**를 만듭니다.
- 어떻게 하나? (조심스럽게 하나씩 만들고 확인하는 과정, Figure 6 참조)
- 중심 부품 선택 (ϕfocal): 이미 놓인 부품 중 하나를 중심으로 잡습니다.
- 새 연결 예측 (ϕedge): 중심 부품에 연결될 새로운 부품(아직 종류가 미정인 상태)을 선택하고 연결선을 긋습니다.
- 부품 종류 결정 (ϕnode): 새로 연결된 부품이 정확히 어떤 종류의 fragment인지 결정합니다.
- 반복 검토 및 수정 (ϕrefine - Iterative Refinement): 이 부분이 중요합니다! 새 부품을 추가한 후, 기존에 놓았던 부품들이 여전히 최적인지 다시 확인하고 필요하면 수정합니다. (마치 퍼즐 조각을 맞추다가 전체 그림을 보고 잘못된 부분을 고치는 것과 같습니다.) 이 과정을 통해 전체 분자 구조가 더 자연스러워지도록 합니다.
- 이 모든 과정에서 3D 구조와 연결 정보를 잘 처리하는 똑똑한 신경망(EGNN 기반 message passing networks)을 사용합니다.
3단계: 최종 완성 (Assembling to Atom Conformation)
- 무엇을 하나? 2단계에서 만든 'Fragment 3D Graph'(부품 연결도)를 보고, 실제 원자들을 모두 배치하여 최종적인 3D 분자 모델을 만듭니다.
- 어떻게 하나?
- 연결된 두 fragment 부품 사이를 어떻게 '이어붙일지' 결정합니다. 여러 가능한 연결 방법 중, 1단계에서 예측된 중심 위치 정보와 가장 잘 맞는 방식을 선택합니다 (이때 후보 모양 생성에 RDkit이라는 화학 도구를 잠시 활용).
- 각 부품의 세부 원자들을 3D 공간에 정확히 배치합니다. (수학적인 방법 - Kabsch Algorithm - 을 사용해서 위치와 방향을 맞춥니다.)
이렇게 3단계를 거쳐 HierDiff는 크고 복잡한 분자도 화학적으로 더 정확하고 안정적으로 생성하려고 시도합니다.
주인장 이해
1. 랜덤노이즈 에서 시작. Positional feature의 경우 CoM free Gaussian 노이즈를 사용 그 이유는 이동 불변성(translation invariance)을 만족시키기 위함.
2. 생성 전에 fragment의 개수를 미리 sample 해야함 통계 정보를 바탕으로 뽑는다고 했는데 어떻게 뽑는지?
3. Hf(화학적 특성)와 Hp(3D 위치)는 각각 다른 방식으로 모델링 하나의 디퓨전 내에서 네트워크가 다르게 움직임
4. 아웃풋으로 H가 나오게 됨 .
H를 바탕으로 fragment를 결정하고 Fragment 3D graph를 만드는 것이 목표
ϕnode 모듈을 사용해 fine grained node를 1차 예측
거기서 ϕfocal을 사용해서 중심점 역할을 할 focal node 선택 (fragment)
ϕedge를 사용해서 focal node와 연결될 후보가 되는 coarse-grained nodes 중에서, 가장 적절한 연결(edge)을 예측하고 coarse-grained node를 선택
ϕnode를 사용해서 fragment 를 결정하고 fine-grained node가 최종 구성
ϕrefine을 사용해서 fine-grained node를 마지막으로 재검토
rdkit를 사용해서 일단 스케치를 함.
fragment의 목표 중심 위치(Hp) 정보와 스케치와 비교해서 계산하여 가장 가까운 연결 방식을 최종적으로 선택
Kabsch 알고리즘으로 회전, 이동을 계산하고
암튼 계산이 끝나면 R와 t값을 모든 원자들의 좌표에 똑같이 적용
