[Paper] Flow Sampling: 비정규화 밀도에서 샘플링을 학습하기 위한 디노이징 조건부 프로세스

Source: arXiv - 2605.03984v1

개요

이 논문은 Flow Sampling이라는 새로운 프레임워크를 소개한다. 이 프레임워크는 신경망이 정규화되지 않은 확률 밀도—에너지(또는 스코어) 함수만으로 정의되고 학습 데이터가 전혀 없는 분포—에서 샘플을 추출하도록 학습한다. 확산 모델과 흐름 매칭 아이디어를 결합함으로써, 저자들은 구나 쌍곡선 다양체와 같은 곡면 공간에서도 작동하는 빠르고 데이터가 필요 없는 샘플링을 구현한다.

핵심 기여

• Data‑free diffusion training: 목표 에너지에서 직접 유도된 디노이징 드리프트에 회귀하는 조건부 목표로, 데이터 샘플에 의존하는 일반적인 “노이징” 드리프트가 아니라 데이터‑프리 방식으로 확산을 학습합니다.
• Interpolant process: 학습 중에 필요한 비용이 큰 에너지‑함수 평가 횟수를 크게 줄여주는 영리한 기법입니다.
• Riemannian extension: 일정한 곡률을 갖는 다양체(예: 초구, 쌍곡공간)에 대한 닫힌 형태의 조건부 드리프트를 제공하여, 유클리드 공간을 넘어선 확산 기반 샘플링을 가능하게 합니다.
• Scalable implementation: 합성 벤치마크, 펩타이드 컨포머 생성, 대규모 암오티제이션 분자 컨포머 생성, 구형 분포 등에서 입증되었으며, 최첨단 결과를 달성했습니다.
• Amortized sampler: 한 번 학습된 모델은 추론 시에 소수의 신경망 평가만으로도 많은 독립 샘플을 생성할 수 있습니다.

방법론

1. Problem setup – 우리는 에너지 함수 (E(x))가 비정규화된 밀도 (\tilde{p}(x) \propto e^{-E(x)})를 정의한다는 전제 하에 있다. 목표는 정규화 상수를 전혀 계산하지 않고 정규화된 분포 (p(x))에서 샘플을 생성하는 것이다.

2. Diffusion backbone – 표준 diffusion 모델은 노이즈가 섞인 데이터 포인트를 데이터 매니폴드로 복원하는 역방향 확률 미분 방정식(SDE)을 학습한다. Flow Sampling은 조건을 뒤바꾼다: 모델은 노이즈 샘플 (z)를 입력받고, 그 노이즈 상태에서 (\tilde{p})의 고확률 영역으로 입자를 이동시키는 디노이징 드리프트를 예측하도록 학습한다.

3. Conditional drift regression – 학습 손실은
   [ \mathcal{L} = \mathbb{E}{t\sim[0,1],,z\sim\mathcal{N}(0,I)}\big|,\mathbf{v}\theta(t,z) - \underbrace{\big[-\nabla_x E(x_t) + \text{diffusion term}\big]}{\text{target drift}}\big|^2, ]
   여기서 (x_t)는 기준점(보통 원점)과 미지의 목표 샘플 사이의 보간된 상태이며, (\mathbf{v}\theta)는 신경망이 추정한 드리프트이다.

4. Interpolant process – 매 diffusion 단계마다 (E(\cdot))를 반복해서 평가하는 대신, 저자들은 알려진 앵커와 무작위 노이즈 샘플 사이의 선형(또는 측지) 보간을 구성한다. 이는 학습 예시당 단 하나의 에너지 평가만 필요하게 하여 비용을 한 차례 정도 절감한다.

5. Riemannian manifolds – 곡률이 일정한 매니폴드 (K)에서 보간은 측지 곡선을 따른다. 저자들은 매니폴드의 계량을 존중하는 조건부 드리프트에 대한 폐쇄형 식을 유도하여, 동일한 학습 파이프라인이 구( (K>0) )나 쌍곡선 공간( (K<0) )에서도 작동하도록 한다.

6. Inference – 테스트 시에는 학습된 드리프트 필드를 (예: Euler‑Maruyama) 통합하여 노이즈 샘플에서 시작해 목표 분포에서 최종 샘플을 생성한다. 드리프트가 사전 계산되어 있기 때문에 수천 개의 샘플을 생성하는 비용이 매우 저렴하다.

결과 및 발견

요약: Flow Sampling은 전통적인 MCMC 및 diffusion‑only 샘플러와 비교해 동일하거나 더 나은 성능을 보이며, 훈련 중 10‑100× 적은 에너지 평가와 추론 시 수십 배 이상의 계산량 감소를 달성합니다.

실용적 함의

• 분자 설계 파이프라인 – 연구자들은 Flow Sampling을 기존 생성 화학 스택에 연결하여 실시간으로 현실적인 컨포머를 생성할 수 있으며, 비용이 많이 드는 에너지 최소화 작업에 소요되는 시간을 크게 줄일 수 있습니다.
• 물리 기반 시뮬레이션 – 유체, 재료, 로봇공학을 모델링하는 엔지니어들은 사용자 정의 에너지 함수(예: 제약조건, 퍼텐셜)를 정의하고, 손수 MCMC 커널을 설계하지 않아도 빠른 샘플러를 얻을 수 있습니다.
• 기하학적 딥러닝 – 다양체 위에서 수행되는 작업(예: 방향 데이터, 쌍곡선 공간상의 그래프 임베딩)은 이제 곡률이 있는 공간을 떠나지 않고 확산 스타일 생성 모델을 사용할 수 있어 내재된 기하학을 보존합니다.
• 베이지안 모델을 위한 암시적 추론 – 사후분포가 정규화되지 않은 밀도 형태로만 알려진 경우, Flow Sampling은 반복적인 비용이 큰 그래디언트 계산을 피할 수 있는 즉시 배포 가능한 암시적 샘플러를 제공합니다.

개발자에게 이 방법은 표준 PyTorch 모듈로 구현되어 있으며(저자들이 코드를 공개함), 에너지 함수와 몇 가지 하이퍼파라미터(노이즈 스케줄, 통합 단계)만 있으면 됩니다. 기존 파이프라인에 최소한의 리팩터링으로 바로 적용할 수 있습니다.

제한 사항 및 향후 작업

• Energy evaluation cost는 매우 고차원 시스템(예: 대형 단백질)에서 여전히 병목 현상입니다; 보간 함수가 샘플당 비용을 줄이긴 하지만, 평가 횟수 자체는 여전히 많을 수 있습니다.
• Choice of noise schedule와 통합 단계 크기는 아직 다소 경험적이며, 최적이 아닌 설정은 샘플 품질이나 안정성을 저하시킬 수 있습니다.
• 현재 이론은 smooth 에너지 함수를 가정합니다; 비미분 가능한 제약(예: 강한 입체 충돌)은 추가적인 처리가 필요합니다.
• 저자들이 언급한 향후 연구 방향에는 적응형 단계 크기 스킴, Flow Sampling을 학습된 대리 에너지 모델과 결합하는 것, 그리고 가변 곡률을 갖는 다양체(예: 학습된 잠재 다양체)로 리만 형식을 확장하는 것이 포함됩니다.

저자

• Aaron Havens
• Brian Karrer
• Neta Shaul

논문 정보

• arXiv ID: 2605.03984v1
• 분류: cs.LG, cs.AI
• 출판일: 2026년 5월 5일
• PDF: PDF 다운로드