그래프 에너지 매칭: 그래프 생성 모델을 위한 운송 정렬 기반 에너지 모델링

Graph Energy Matching: Transport-Aligned Energy-Based Modeling for Graph Generation

그래프와 같은 이산 영역에 대한 에너지 기반 모델은 상대적인 가능성을 명시적으로 표현하여, 조건부 생성 또는 테스트 시 제약 조건 적용과 같은 구성 가능한 확률 추론 작업을 가능하게 합니다. 그러나 이산 에너지 기반 모델은 일반적으로 효율적이고 고품질의 샘플링에 어려움을 겪는데, 이는 오프-서포트 영역에 종종 존재하며, 샘플러를 가두고 학습 불안정을 유발하는 가짜 지역 최소값을 포함하기 때문입니다. 이러한 문제는 이산 확산 모델과 비교하여 충실도 격차를 초래했습니다. 본 논문에서는 이러한 격차를 해소하는 그래프 생성 프레임워크인 Graph Energy Matching (GEM)을 소개합니다. Jordan-Kinderlehrer-Otto (JKO) 스킴의 운송 맵 최적화 관점에서 영감을 받아, GEM은 순열 불변 포텐셜 에너지를 학습하여, 노이즈로부터 데이터 방향으로 운송 정렬된 지침을 제공하고, 동시에 데이터 가능성이 높은 영역 내에서 샘플을 개선합니다. 또한, 우리는 에너지 기반 스위치를 활용하여 다음과 같은 작업을 원활하게 수행하는 샘플링 프로토콜을 도입합니다: (i) 고확률 영역으로의 빠른, 기울기 기반 운송, (ii) 학습된 그래프 분포의 탐색을 위한 혼합 모드. 분자 그래프 벤치마크에서 GEM은 강력한 이산 확산 모델의 성능에 필적하거나 능가합니다. 샘플 품질 외에도, 상대적인 가능성의 명시적인 모델링은 추론 시의 표적 탐색을 가능하게 하여, 구성 가능한 생성, 속성 제약 샘플링 및 그래프 간의 지오데식 보간을 용이하게 합니다.

Energy-based models for discrete domains, such as graphs, explicitly capture relative likelihoods, naturally enabling composable probabilistic inference tasks like conditional generation or enforcing constraints at test-time. However, discrete energy-based models typically struggle with efficient and high-quality sampling, as off-support regions often contain spurious local minima, trapping samplers and causing training instabilities. This has historically resulted in a fidelity gap relative to discrete diffusion models. We introduce Graph Energy Matching (GEM), a generative framework for graphs that closes this fidelity gap. Motivated by the transport map optimization perspective of the Jordan-Kinderlehrer-Otto (JKO) scheme, GEM learns a permutation-invariant potential energy that simultaneously provides transport-aligned guidance from noise toward data and refines samples within regions of high data likelihood. Further, we introduce a sampling protocol that leverages an energy-based switch to seamlessly bridge: (i) rapid, gradient-guided transport toward high-probability regions to (ii) a mixing regime for exploration of the learned graph distribution. On molecular graph benchmarks, GEM matches or exceeds strong discrete diffusion baselines. Beyond sample quality, explicit modeling of relative likelihood enables targeted exploration at inference time, facilitating compositional generation, property-constrained sampling, and geodesic interpolation between graphs.