Energy-Based Flow Matching for Generating 3D Molecular Structure 정리

논문 핵심 정리 및 블로그용 요약 문서

논문 정보

• 제목: Energy-Based Flow Matching for Generating 3D Molecular Structure
• 저자: Wenyin Zhou, Christopher Iliffe Sprague, Vsevolod Viliuga, Matteo Tadiello, Arne Elofsson, Hossein Azizpour
• 학회/연도: ICML 2025 (PMLR 267)
• 주제: 3D 분자 구조 생성, 분자 도킹, 단백질 백본 생성, Flow Matching, Energy-Based Model

---

TL;DR

이 논문의 핵심은 flow matching을 “에너지 최소화” 문제로 다시 해석했다는 점이다.
저자들은 중간 상태에서 한 번 예측하고 끝내는 대신, 예측한 구조를 다시 같은 네트워크로 한 번 더 정제(refine) 하도록 학습한다. 이때 모델은 f(f(x)) ≈ f(x) 형태의 idempotent flow map을 배우며, 결과적으로 샘플러이면서 동시에 리파이너가 된다.
실험에서는 이 아이디어가 단일/다중 리간드 도킹과 단백질 백본 생성 모두에서 유효했고, 같은 수준의 계산 예산에서 기존 flow matching 및 diffusion 계열 대비 더 좋은 성능을 보였다.

---

이 논문을 한 문장으로 요약하면

“Flow matching의 출력을 최종 답이 아니라, 에너지를 더 낮출 수 있는 중간 추정치로 보고 반복 정제를 학습시킨 논문”이다.

---

왜 이 논문이 중요한가

3D 분자 구조 생성에서는 “그럴듯한 하나의 구조”보다 여러 가능한 구조 분포를 잘 샘플링하는 능력이 중요하다.
그래서 diffusion과 flow matching이 많이 쓰이는데, 기존 flow matching은 보통 중간 샘플 x_t로부터 목표 구조 x_1을 한 번에 예측하는 방식이다. 이 접근은 빠르고 단순하지만, 다음과 같은 질문이 남는다.

1. 예측된 샘플 자체를 더 좋은 방향으로 다듬을 수는 없을까?
2. 생성된 샘플이 정말 낮은 에너지 영역(= 그럴듯한 구조 영역)으로 가고 있는가?
3. 샘플링과 refinement를 하나의 네트워크로 합칠 수는 없을까?

이 논문은 바로 이 지점을 건드린다.

---

배경: 기존 Flow Matching을 어떻게 보고 있나

논문은 $x_1$ parameterization을 쓰는 flow matching을 바탕으로 설명한다.
여기서 네트워크 $f_{\theta,t}(x)$는 중간 상태 $x_t$에서 최종 데이터 $x_1$을 직접 예측하는 함수처럼 동작한다.
가장 기본적인 학습 목표는 다음처럼 볼 수 있다.

$$\mathcal{L}_{\text{CFM}} = \mathbb{E} \left\| f_{\theta,t}(x) - x_1 \right\|_2^2$$

 

즉, “지금 보이는 중간 상태에서 정답 구조를 맞혀라”에 가깝다.
문제는 이 손실이 정답 x_1 쪽으로는 끌어당기지만, 모델이 스스로 만든 예측 샘플 \hat{x}_1을 더 좋은 구조로 재정렬하는 메커니즘은 약하다는 점이다.

---

핵심 아이디어 1: Flow Matching을 에너지 기반으로 다시 해석하기

저자들은 생성 구조의 품질을 “에너지”로 보자는 관점을 제안한다.
직관적으로는,

• 좋은 구조일수록 에너지가 낮고
• 나쁜 구조일수록 에너지가 높다

라는 전형적인 energy-based model 관점이다.

이 논문에서 제안하는 아주 단순한 에너지 형태는 재구성 오차(reconstruction error) 다.

$$E_\theta(\hat{x}_1) = \left\| G(\hat{x}_1) - \hat{x}_1 \right\|_2^2$$

여기서 G는 예측 구조를 더 좋은 구조 쪽으로 밀어주는 refiner 역할을 한다.
논문의 핵심은 이 G를 별도 네트워크로 두기보다, flow map 자체가 refiner 역할도 하도록 만들자는 데 있다.

---

핵심 아이디어 2: Idempotent Flow Map

이 논문의 가장 중요한 개념은 다음 한 줄로 요약된다.

$$f_{\theta,t}(f_{\theta,t}(x)) \approx f_{\theta,t}(x)$$

즉, 한 번 더 적용해도 거의 바뀌지 않는 함수를 배우게 하겠다는 것이다.
이 성질을 논문은 idempotent flow map이라고 부른다.

이게 왜 중요한가?

• 첫 번째 적용: 샘플을 대략 맞는 구조 쪽으로 보낸다.
• 두 번째 적용: 그 예측값을 한 번 더 정제한다.
• 충분히 잘 학습되면: 더 이상 크게 바뀌지 않는 안정적인 저에너지 영역에 머무른다.

이 관점은 분자 구조 생성에서 매우 자연스럽다.
좋은 구조란 결국 작은 perturbation을 더해도 크게 흔들리지 않는 안정 상태여야 하기 때문이다.

---

Figure 1. 논문의 아이디어를 가장 잘 보여주는 그림

 

원논문 Figure 1 (p.2) cropped.
왼쪽은 일반적인 flow matching이고, 오른쪽은 IDFlow다.
핵심 차이는 예측값을 다시 네트워크에 넣어 refinement를 수행한다는 점이다.
즉, IDFlow는 “예측기(preditor)”와 “정제기(refiner)”를 분리하지 않고, 동일한 맵을 반복 적용하는 방식으로 학습한다.

---

핵심 아이디어 3: Predictor-Refiner Sampler

논문은 학습된 모델을 predictor-refiner sampler로 해석한다.
이는 diffusion에서의 predictor-corrector sampler와 닮았지만, 여기서는 같은 네트워크가 샘플링과 정제를 동시에 맡는다는 점이 다르다.

실제 추론 시 각 step은 대략 이렇게 진행된다.

1. 현재 샘플 $x_t$에서 $\hat{x}1 = f{\theta,t}(x_t)$를 예측한다.
2. 예측값을 다시 넣어 $f_{\theta,t}(\hat{x}_1)$로 한 번 더 정제한다.
3. 그 방향을 이용해 ODE step을 진행한다.

논문은 계산량을 맞추기 위해 step 수는 줄이고, step당 refinement는 1회로 제한한다.
즉, “무한히 refine할 수 있다”는 이론을 실제 추론에서는 한 번 정제라는 현실적인 형태로 쓴다.

---

학습은 어떻게 바뀌는가

학습 손실은 크게 두 축이다.

1) 기본 flow matching 손실

$$\mathcal{L}_{\text{CFM}} = \mathbb{E} \left\| f_{\theta,t}(x) - x_1 \right\|_2^2$$

2) refinement 손실

먼저 모델이 예측한 $\hat{x}1 = f{\theta,t}(x)$를 만들고 ,
그다음 정제 결과가 정답 $x_1$에 더 가까워지도록 학습한다.

$$\mathcal{L}_R = \mathbb{E} \left\| f_\theta(\hat{x}_1.\text{detach()}) - x_1 \right\|_2^2$$

직관적으로는:

• 기본 손실은 “정답 방향으로 맞혀라”
• refinement 손실은 “네가 만든 예측값도 다시 정답 쪽으로 끌어와라”

라고 볼 수 있다.

이 두 손실이 합쳐지면 모델은 단순 예측기에서 벗어나 자기 출력을 더 낮은 에너지 방향으로 정렬하는 맵을 배우게 된다.

---

이 논문의 포인트를 더 쉽게 말하면

포인트 1. “출력값 자체”를 학습 대상으로 다시 삼는다

기존 flow matching은 x_t → x_1 대응에 집중한다.
이 논문은 한 걸음 더 나아가 x_t → \hat{x}_1 → 더 나은 \hat{x}_1 로 본다.

포인트 2. 샘플링과 refinement를 하나의 네트워크에 묶는다

별도의 relaxation 단계나 confidence model 없이도,
모델 내부에서 예측 → 정제 구조를 만들 수 있다.

포인트 3. 안정성과 고정점 관점이 들어간다

좋은 생성 모델은 단지 그럴듯한 샘플을 뱉는 것을 넘어서,
반복 적용해도 무너지지 않는 안정적 맵이어야 한다는 시각을 제시한다.

포인트 4. Euclidean / Riemannian 둘 다 통한다

분자 도킹은 보통 유클리드 좌표계에서,
단백질 백본 생성은 SE(3)^N 같은 리만 기하 위에서 다루는데,
이 논문은 두 설정 모두에서 효과가 있음을 보였다.

---

실험 1: 단일 리간드 도킹

논문은 PDBBind에서 single-ligand docking을 평가했다.
지표는 주로 % < 2Å, median RMSD 등이다.

핵심 결과만 뽑으면

설정	HarmonicFlow	IDFlow	차이
Sequence similarity split, Distance-pocket, %<2Å	30.1	35.6	+5.5
Sequence similarity split, Radius-pocket, %<2Å	20.5	21.0	+0.5
Time split, Distance-pocket, %<2Å	42.8	44.3	+1.5
Time split, Radius-pocket, %<2Å	28.3	34.7	+6.4

해석은 분명하다.
같은 계열의 baseline(HarmonicFlow)보다 더 잘 맞는 도킹 포즈를 더 자주 생성한다.

---

실험 2: 다중 리간드 도킹

Binding MOAD의 multi-ligand docking에서는 다음이 눈에 띈다.

모델	%<2Å	%<5Å	Median RMSD
EigenFold Diffusion	39.7	73.5	2.4
HarmonicFlow	44.4	75.0	2.2
IDFlow	43.8	83.1	2.1

여기서는 %<2Å는 baseline과 거의 비슷하지만,
%<5Å가 75.0 → 83.1로 크게 오른다.
즉, IDFlow는 정답 근처 모드를 더 넓게 커버하는 성향을 보인다고 읽을 수 있다.

---

실험 3: 단백질 백본 생성

단백질 백본 생성에서는 FrameFlow를 기반으로 IDFlow를 만들었다.
핵심 지표는 designability / diversity / novelty다.

• Designability: 생성된 backbone에 대해 서열을 설계하고 다시 접었을 때, 원래 backbone으로 잘 돌아오는가
• Diversity: 생성 구조끼리 얼마나 다양한가
• Novelty: 기존 자연 단백질과 얼마나 다른가

PDB 결과 핵심 비교

모델	NFEs	Designability	Diversity	Novelty
FrameFlow	200	0.824 ± 0.037	0.35 ± 0.01	0.70 ± 0.00
IDFlow	100	0.871 ± 0.013	0.35 ± 0.00	0.70 ± 0.02
IDFlow	200	0.927 ± 0.020	0.35 ± 0.01	0.72 ± 0.01

이 표에서 특히 중요한 해석은 두 가지다.

1. 같은 200 NFE에서 designability가 크게 오른다.
2. 100 NFE만으로도 baseline 200 NFE보다 더 좋다.

즉, IDFlow는 품질이 올라갈 뿐 아니라, 저 step regime에서도 상대적으로 강하다는 메시지를 준다.

---

Figure 2. 생성된 단백질 백본 예시

 

원논문 Figure 2 (p.8) cropped.
길이 100~300 범위에서 생성한 designable backbone 예시다.
단순히 헬릭스만 반복적으로 만드는 것이 아니라, 길이에 따라 다양한 topology를 샘플링하고 있다는 점이 보인다.

---

Figure 3. 테스트 시 refinement 횟수 ablation

 

원논문 Figure 3 (p.8) cropped.
테스트 시 refinement를 너무 많이 하면 오히려 성능이 떨어진다.
논문은 이를 ODE discretization error 때문이라고 해석한다.
즉, “정제를 많이 할수록 무조건 좋다”가 아니라, 적절한 1회 refinement가 가장 실용적이라는 결론이다.

블로그 포인트:
이 논문은 “반복 정제”를 말하지만, 실제 최적점은 무한 반복이 아니라 1회 refinement에 가깝다.

---

Figure 4. 학습 수렴 속도

 

원논문 Figure 4 (p.20) cropped.
IDFlow가 HarmonicFlow보다 validation RMSD<2 지표를 더 빨리 끌어올린다.
즉, 이 방법은 최종 성능뿐 아니라 수렴 과정 자체도 더 안정적임을 시사한다.

---

Figure 5. 샘플링 중 에너지 감소

 

원논문 Figure 5 (p.21) cropped.
이 그림은 이 논문의 주장을 가장 잘 뒷받침한다.
IDFlow는 샘플링 step이 진행될수록 재구성 기반 에너지(L2 error)가 더 빠르고 더 낮게 감소한다.
저자들의 핵심 주장이 “flow matching을 에너지 기반으로 보자”는 것이므로, 이 결과는 단순 성능표 이상으로 의미가 있다.

---

Figure 6. 적은 NFE에서도 강한가?

 

원논문 Figure 6 (p.22) cropped.
SCOPe 데이터에서 NFEs를 줄였을 때, IDFlow는 FrameFlow보다 designability와 diversity 측면에서 더 강한 robustness를 보인다.
즉, 모델이 더 안정적인 sampling dynamics를 배웠다는 해석이 가능하다.

---

Figure 7. 도킹 예시 시각화

 

원논문 Figure 7 (p.23) cropped.
파란색이 ground truth, 초록색이 IDFlow 생성 결과다.
예시마다 차이는 있지만, 전반적으로 binding pocket 내부에서 리간드 배치가 꽤 근접하게 복원되는 것을 볼 수 있다.
정량 지표가 실제 구조 예시에서도 어느 정도 설득력을 가진다는 점을 보여준다.

---

Figure 8. 생성된 백본 갤러리

 

원논문 Figure 8 (p.24) cropped.
PDB에서 학습한 뒤 생성한 다양한 길이의 backbone 모음이다.
이 그림의 포인트는 “잘 접히는가”뿐 아니라 생성물의 형태적 다양성이다.
IDFlow가 한 가지 구조 패턴으로 붕괴하지 않고 여러 topology를 만들어내고 있음을 시각적으로 전달한다.

---

이 논문이 기존 방법과 다른 지점

1. 외부 에너지 최소화가 아니라, 모델 내부에서 정제를 학습한다

분자 구조 생성에서는 종종 생성 후 relaxation을 추가한다.
이 논문은 그런 후처리를 별도로 두기보다, flow map 안에 refinement 성질을 내장한다.

2. confidence model을 따로 두는 대신, 맵 자체를 안정화한다

샘플을 뽑은 뒤 “무엇이 좋은 샘플인지”를 따로 판별하는 대신,
애초에 좋은 샘플 쪽으로 고정점을 갖는 함수를 만들려 한다.

3. AlphaFold의 recycling과 닮았지만, generative sampling에 붙였다

저자들도 구조 refinement와의 연결성을 강조한다.
이 점은 단순히 flow matching을 조금 고친 것이 아니라, 구조생성 모델을 안정적 고정점 학습으로 보는 관점 전환에 가깝다.

---

장점

• 아이디어가 단순하다.
  복잡한 새로운 sampler를 만드는 대신, 기존 flow map을 한 번 더 쓰게 만든다.
• 개념적으로 설득력 있다.
  생성된 구조는 낮은 에너지의 안정 상태여야 한다는 분자 과학의 직관과 잘 맞는다.
• 도킹과 단백질 생성 양쪽에서 통한다.
  유클리드 공간과 리만 공간 모두에서 효과를 보였다.
• 낮은 step 수에서도 강하다.
  실제 활용 측면에서 중요한 포인트다.

---

한계

논문이 직접 인정한 한계도 분명하다.

1. 학습 비용이 증가한다.
   refinement를 위해 추가 forward pass가 필요하다.
2. 정제를 너무 많이 하면 오히려 discretization error가 커질 수 있다.
   즉, refinement는 “많을수록 좋다”가 아니다.
3. 에너지 정의가 아직 물리 기반 에너지 자체는 아니다.
   현재는 주로 reconstruction error를 기반으로 한다.
   향후에는 biophysics-informed energy와 결합하는 방향이 남아 있다.

---

블로그에 바로 넣기 좋은 핵심 문장

핵심 문장 1

Flow matching의 출력을 최종 예측치가 아니라 반복적으로 정제 가능한 구조 추정치로 본 것이 이 논문의 핵심이다.

핵심 문장 2

IDFlow는 f(f(x)) ≈ f(x) 라는 idempotent 성질을 학습해, 샘플링과 refinement를 하나의 네트워크 안에 통합한다.

핵심 문장 3

실험적으로는 도킹과 단백질 백본 생성 모두에서, 더 좋은 구조 품질과 더 안정적인 샘플링 동역학을 보여준다.

핵심 문장 4

특히 적은 NFE에서도 성능이 덜 무너진다는 점은, 이 방법이 단순한 성능 향상보다 더 좋은 생성 경로를 학습했다는 증거에 가깝다.

---

개인적인 해석 대신, 이 논문에서 읽어야 할 메시지

이 논문이 주는 가장 큰 메시지는 다음이다.

좋은 생성 모델은 “정답을 한 번에 맞히는 모델”이 아니라,
자기 출력을 더 안정적인 데이터 매니폴드 방향으로 반복적으로 정제할 수 있는 모델일 수 있다.

즉, 생성 문제를 단순한 one-shot regression처럼 다루지 않고,
점진적으로 에너지를 낮추는 동역학으로 이해하려는 시도라는 점에서 의미가 크다.

---

결론

이 논문은 flow matching을 조금 튜닝한 수준의 작업이 아니다.
보다 정확히는, flow matching을 energy minimization과 fixed-point refinement 관점으로 재정의한 작업에 가깝다.

실험 결과만 보면 “baseline보다 몇 포인트 좋아졌다”로 읽을 수 있지만,
정작 더 중요한 것은 다음 두 가지다.

1. 생성된 샘플을 다시 입력해도 무너지지 않는 안정적 맵을 학습한다는 발상
2. 그 발상이 실제 도킹과 단백질 생성에서 모두 통한다는 점

그래서 이 논문은 구조 생성 쪽에서 앞으로도 자주 인용될 만한 포인트를 갖고 있다.
특히 샘플링-정제 통합, idempotency, 에너지 기반 해석이라는 세 키워드로 기억해두면 좋다.

---

참고

• 원논문: Energy-Based Flow Matching for Generating 3D Molecular Structure (ICML 2025)
• 본 문서는 업로드된 논문 PDF를 바탕으로 한국어로 재구성한 요약/블로그 초안이다.