기술 요약: 에너지 기반 모델을 통한 물리적으로 일관된 분자 생성

문제 정의

열평형 상태의 분자 상태는 볼츠만 분포 $p(x) \propto \exp(-E(x)/(k_BT))$에 의해 지배되며, 여기서 구성 확률은 그 위치 에너지에 의해 결정됩니다. 분자 생성용 표준 데이터셋 (예: QM9, GEOM-Drugs) 은 국소 에너지 최소값으로 완화된 기하학적으로 최적화된 구조를 포함하고 있지만, 볼츠만 분포에서 추출된 것은 아닙니다. 이들의 상대적 빈도는 열역학적 가중치보다는 열거 전략을 반영합니다. 따라서, 최소값이 안정적인 구성과 일치하는 스칼라 에너지 지형이 자연스러운 모델링 대상입니다.

그러나 이러한 지형을 학습하고 샘플링하는 것은 여전히 열린 과제입니다. 전통적인 에너지 기반 모델 (EBM) 은 훈련을 위해 볼츠만 분포 데이터를 요구하거나, 노이즈와 유효 구조 사이의 거칠고 비물리적인 영역을 통과하는 데 어려움을 겪는 비싼 마르코프 체인 몬테카를로 (MCMC) 방법에 의존합니다. 반면, 최첨단 확산 및 흐름 매칭 모델은 시간 조건부 운송 장을 학습함으로써 이러한 문제를 우회합니다. 효과적이지만, 이러한 모델들은 에너지 귀납 편향을 상실하여 추론 시 구성을 평가, 순위 매기거나 조합할 수 있는 명시적 에너지 함수를 제공하지 않습니다.

방법론: EBMol

저자들은 훈련 중 명시적 시뮬레이션 없이 에너지 귀납 편향을 복원하는 3D 분자 생성용 에너지 기반 모델인 EBMol을 소개합니다. 이 프레임워크는 두 가지 주요 구성 요소로 이루어져 있습니다.

1. 훈련을 위한 복원장 매칭 (Restoring Field Matching, RFM)

볼츠만 분포 데이터를 사용하지 않고 다루기 쉬운 에너지 지형 $E_\theta(x)$를 학습하기 위해, 저자들은 **복원장 매칭 (RFM)**을 제안합니다.

• 개념: 표준 흐름 매칭에서와 같이 데이터 포인트를 궤적 끝점으로 취급하는 대신, RFM 은 보간 구간을 $t \in [-1, 1]$로 확장하여 $t=0$에서 데이터를 통과하도록 합니다.
• 목표 장 구성: 모델은 에너지의 음의 기울기 $-\nabla E_\theta(x_t)$를 목표 "복원장" $u$와 일치하도록 훈련됩니다. 이 장은 매니폴드의 양쪽 (보간 $t \in [0, 1]$ 및 외삽 $t \in [-1, 0]$) 에서 데이터 $x_0$를 향합니다.
• 정상성: 스무딩 함수 $b(|t|)$가 목표 장을 스케일링하여 데이터 포인트 ($t=0$) 에서 사라지도록 하여, 데이터 포인트가 학습된 에너지 지형의 국소 최소값이 되도록 합니다.
• 정규화: 원자별 정규화 항이 깨끗한 데이터 포인트의 에너지 값을 0 을 향해 고정시켜 샘플러에서의 인위적 깊이 편향을 방지합니다.
• 상태 공간 처리: 이 방법은 혼합 연속 - 이산 상태 ($x = (c, p)$, 여기서 $c$는 좌표이고 $p$는 원자 유형) 를 처리하기 위해 별도의 보간자를 정의합니다. 좌표에는 최적 수송을 통해 정렬된 선형 궤적을, 범주형 특성에는 심플렉스 내에 머무르기 위한 반사 기반 궤적을 사용합니다.

2. 샘플링을 위한 미러 - 랑주뱅 알고리즘 (Mirror-Langevin Algorithm, MLA)

학습된 분포 $p_\theta(x) \propto \exp(-E_\theta(x)/\tau)$에서 샘플링하기 위해, 저자들은 미러 - 랑주뱅 알고리즘을 적용합니다.

• 통합 업데이트: 표준 랑주뱅 역학은 원자 유형을 확률 심플렉스 밖으로 밀어냅니다. MLA 는 제약이 없는 쌍대 공간에서 업데이트를 수행하고 미러 맵을 통해 투영함으로써 이를 해결합니다.
  • 연속 좌표의 경우, 유클리드 포텐셜은 MLA 를 수정되지 않은 랑주뱅 알고리즘 (Unadjusted Langevin Algorithm) 으로 축소합니다.
  • 범주형 원자 유형의 경우, 음의 엔트로피 포텐셜이 심플렉스를 로짓 공간으로 매핑하며, 소프트맥스 함수가 다시 투영합니다.
• 병렬 템퍼링: 거친 에너지 지형을 통과하는 어려움을 해결하기 위해, 저자들은 서로 다른 온도의 여러 체인을 사용하는 병렬 템퍼링을 적용합니다. 주기적인 메트로폴리스 - 헤이스팅스 스왑은 "뜨거운" 체인이 광범위하게 탐색하고 "차가운" 체인이 정제할 수 있게 하여 추론 시 계산 확장성을 가능하게 합니다.

주요 기여

1. 복원장 매칭 (RFM): 데이터 포인트를 국소 최소값으로 갖는 스칼라 원자 가산 에너지 지형을 형성하는 시뮬레이션 없는 훈련 목적 함수. 저자들은 해의 존재성과 유일성에 대한 이론적 보장을 제공하고 다양한 정규화 강도 하에서의 거동을 특징화합니다.
2. 미러 - 랑주뱅 샘플링 프레임워크: 단일 단계 내에서 좌표 업데이트와 심플렉스 제약 원자 유형 업데이트를 결합하고 병렬 템퍼링을 보강한, 최초의 통합 혼합 기하학 랑주뱅 프레임워크로, 결합된 연속 - 이산 딥 생성 모델링을 가능하게 합니다.
3. 최첨단 성능: EBMol 은 QM9 및 GEOM-Drugs 벤치마크에서 최첨단 결과를 달성하여, 에너지 기반 모델링이 확산/흐름 모델의 생성 품질과 일치하면서도 추가적인 기능을 제공할 수 있음을 보여줍니다.

실험 결과

이 모델은 QM9(134k 개의 작은 분자) 와 GEOM-Drugs(430k 개의 약물 유사 분자) 에서 평가되었습니다.

• 생성 품질: QM9 에서 EBMol(1810 개의 네트워크 함수 평가) 은 원자 및 분자 안정성에서 모든 기준선을 능가하여 각각 99.76% 및 97.09% 를 달성했습니다. GEOM-Drugs 에서 EBMol 은 강력한 기준선인 SLDM 을 크게 능가하여 높은 계산 예산 하에서 98.68% 의 원자 안정성과 61.42% 의 분자 안정성을 달성했습니다.
• 물리적 일관성: 학습된 에너지 지형은 원칙적인 품질 지표 역할을 합니다. EBMol 은 GEOM-Drugs 에서 중앙값 완화 에너지를 1.78 kcal/mol로 달성했는데, 이는 이전 결합 암시적 모델보다 한 자릿수 낮으며 명시적 결합 감독이 없는 모델임에도 불구하고 결합 명시적 모델 (Megalodon) 과 경쟁 가능한 수준입니다 (18 배 적은 매개변수).
• 물리적 에너지와의 상관관계: 학습된 에너지는 GFN2-xTB 완화 시 에너지 감소와 유의미하게 상관관계 ($r=0.65$) 를 보여, 모델이 물리적으로 긴장된 기하학에 대한 경량 필터로 작용할 수 있게 합니다.
• 제어 가능한 생성:
  • 형태 유도: 학습된 포텐셜을 미분 가능한 형태 포텐셜 (선형, 평면, 구형) 과 결합함으로써, 모델을 재훈련하지 않고도 과소 표현된 기하학 방향으로 생성을 유도할 수 있습니다.
  • 제로샷 링커 설계: 에너지의 가산성으로 인해 분자 조각을 고정함으로써 조건부 생성이 가능합니다. EBMol 은 작업별 훈련 없이 다양한 복잡성의 링커 설계 작업에서 73.9% 의 성공률을 달성했습니다.

중요성 및 주장

이 논문은 EBMol 이 3D 분자 생성을 위한 다목적 프레임워크로서 에너지 기반 모델의 잠재력을 보여준다고 주장합니다. 주요 중요성은 시간 조건부 모델이 상실한 에너지 귀납 편향을 복원함으로써 다음을 가능하게 한다는 점에 있습니다:

• 물리적으로 의미 있는 점수 매기기: 명시적 에너지 함수를 기반으로 생성된 샘플을 순위 매기고 필터링할 수 있는 능력.
• 조합 가능한 생성: 포텐셜 조합 (예: 형태 유도) 을 통한 생성 유도 및 재훈련 없이 제로샷 조건부 생성 (예: 링커 설계) 수행 능력.
• 고품질 생성: 명시적 결합 감독이 있는 모델과 경쟁하는 물리 국소 최소값에 놀라울 정도로 가까운 구조를 생성하면서 최첨단 안정성과 유효성을 달성.

저자들은 샘플링 알고리즘이 하이퍼파라미터에 민감하다는 점과 최적의 결과를 얻기 위해 현재 상당한 계산 예산 (NFEs) 이 필요하다는 한계를 지적합니다. 향후 작업은 물리량 (DFT/GFN2-xTB) 에 대한 에너지 보정 및 더 정교한 샘플러의 통합에 초점을 맞출 수 있다고 제안합니다.