NeurIPS 2025 Oral
출발 분포가 남긴 흔적을 따로 배워 지우면 확산 샘플러는 아무 데서나 출발할 수 있다
확산 샘플러, 슈뢰딩거 다리, 볼츠만 분포, 확률 최적 제어
한 줄 요약
에너지 함수만 주어진 볼츠만 분포에서 표본을 뽑는 확산 샘플러(diffusion sampler)는 지금까지 출발 분포를 한 점에 묶거나 잡음을 크게 키워야 했다. 이 논문의 ASBS는 문제를 임의의 출발 분포와 목표를 잇는 슈뢰딩거 다리(Schrödinger bridge)로 다시 쓰고, 회귀 손실 두 개를 번갈아 풀어 그 제약을 없앴다.
무엇이 달라졌나
분자의 열평형 분포 \(\nu(x)\propto e^{-E(x)}\)에는 정답 표본이 없다. 데이터로 학습하는 볼츠만 생성기와 달리 확산 샘플러는 에너지만 보고 출발 분포 \(\mu\)에서 시작하는 확률 미분 방정식의 표류(drift) \(u\)를 배워 시간 1에 \(\nu\)에 닿게 한다(1장, 식 2). 수반 샘플링(Adjoint Sampling, AS)은 모델이 직접 뽑은 표본만으로 회귀하는 수반 매칭(adjoint matching)을 써서 중요도 가중치 없이 이 일을 해냈다. 대신 기준 과정의 시작점과 끝점이 서로 독립이라는 무기억 조건(memoryless condition)이 필요했고, 그래서 출발 분포가 원점 한 점이었다(2장, 표 1).
확률 최적 제어(stochastic optimal control, SOC)의 최적 분포는 기준 과정을 끝점 비용 \(g(X_1)\)과 초기 가치 함수 \(V_0(X_0)\)로 기울인 꼴이어서, 무기억이 아니면 \(V_0\)가 끝점 분포에 흔적을 남긴다(식 5). 저자들은 슈뢰딩거 다리의 최적 표류가 끝점 비용 \(g=E+\log\hat\varphi_1\)인 SOC 문제의 해와 같음을 보였다(정리 3.1, 식 9). 여기서 \(\nabla\log\hat\varphi_1\)이 출발 분포의 흔적을 지우는 보정자(corrector)다. 출발점이 원점 한 점이면 이 식은 AS의 손실로 돌아간다(3.2절).
어떻게 보였나
보정자를 알면 표류를 배울 수 있고, 표류를 알면 보정자를 배울 수 있다. 그래서 둘을 번갈아 푼다(알고리즘 1). 수반 매칭 단계는 직전 보정자 \(h\)를 넣은 목표 \(-\sigma_t(\nabla E+h)(X_1)\)로 표류를 회귀하고(식 14), 보정자 매칭 단계는 현재 표류로 뽑은 \((X_0,X_1)\) 쌍에서 \(\nabla_{x_1}\log p^{\text{base}}(X_1\mid X_0)\)를 회귀한다(식 15). 어느 쪽도 정답 표본이 필요 없다. 앞 단계는 출발 쪽, 뒤 단계는 도착 쪽 반쪽 다리를 풀기 때문에 전체가 반복 비례 맞춤(iterative proportional fitting, IPF)이 되고, 단계마다 임계점까지 풀리면 슈뢰딩거 다리의 해로 수렴한다(정리 3.2, 4.1, 4.2).
(가) 단계를 거듭하며 닫히는 치우침
(나) 잡음과 수송 거리
그림은 이 교대를 1차원 격자에서 정확히 계산한 결과다. 보정 없이 시작한 1단계는 끝점 분포가 목표에서 벗어나고, 단계를 거듭하면 치우침이 닫힌다. 잡음을 키우면 보정 없이도 목표에 가까워지지만, 같은 목표에 닿는 다리의 이동 거리가 늘어난다. 블로그의 해석으로는 무기억 조건이 잡음으로 치우침을 덮는다면, 이 방법은 치우침을 보정자로 따로 계산해 잡음과 출발 분포를 고를 자유를 준다.
실험은 세 갈래다. 합성 에너지 네 개에서는 모든 지표가 가장 좋았다. 4입자 이중 우물(DW-4)의 \(\mathcal W_2\)는 AS의 0.62에서 0.43으로, 5차원 다중 우물(MW-5)의 싱크혼 거리는 0.32에서 0.15로 줄었다. 다만 55입자 레너드존스(LJ-55)에서는 4.04 대 4.00으로 차이가 작다(표 2). 알라닌 다이펩타이드에서는 두 주사슬 이면각 결합 분포의 \(\mathcal W_2\)가 0.65에서 0.25로 줄었다(표 3). 처음 보는 분자의 형태 이성질체(conformer)를 만드는 과제에서는 조화 사전 분포(harmonic prior)에서 출발한 모델이 SPICE 재현율(1.0Å 기준)을 AS의 56.75%에서 74.05%로 올렸고, RDKit 표본으로 미리 학습한 AS의 72.21%보다도 높았다(표 4). 학습 갱신 한 번에 드는 에너지 평가는 PIS, DDS, iDEM보다 훨씬 적고 AS와 비슷하다(그림 4).
눈여겨볼 점
에너지 평가를 아끼는 힘은 중요도 가중치 없는 회귀에서 오고, 이는 AS도 가졌다. 일반화가 더한 것은 출발 분포와 잡음의 선택지다. 표 4에서 가우스 출발 73.04%와 조화 사전 분포 74.05%의 차이는 작으니, 이득의 대부분은 영역 지식보다 한 점 출발을 벗어난 데서 왔다고 보는 편이 자연스럽다(블로그의 해석). 그런데 표 4의 ASBS는 20층 EGNN, AS는 12층이다. 같은 12층으로 맞춘 부록 표 6에서도 70.70% 대 56.75%로 차이는 남지만, 표에 적힌 표준편차가 30%포인트를 넘어 평균만으로 판단하기는 이르다. DW-4와 LJ의 기준선 값은 AS 논문에서 옮겨 왔다(부록 D.1.2).
수렴 정리는 단계마다 임계점에 닿는다고 가정하지만, 실제 학습은 교대를 3번에서 20번만 돌리고 재생 버퍼로 지난 표본을 섞는다(부록 C, 표 5). 내부 좌표 설정이 쉽다는 리뷰어 지적에 저자들이 OpenReview 토론에서 낸 데카르트 좌표 예비 결과는 \((\phi,\psi)\) \(\mathcal W_2\)가 1.30으로 내부 좌표의 0.25보다 크게 나빴다. 밀도가 낮은 모드를 놓치는 성향도 저자들이 한계로 적었다(6장, 그림 5).
열린 질문
- 출발 분포를 목표에 가깝게 고를수록 수송 거리는 줄지만 보정자가 지울 흔적의 모양도 바뀐다. 출발 분포의 선택, 교대 횟수, 보정자의 근사 오차가 끝점 분포의 오차에 어떻게 함께 쌓이는지 가늠할 수 있을까.
이해 확인
ASBS에서 보정자 \(h\approx\nabla\log\hat\varphi_1\)이 맡는 일은?
- 에너지 기울기의 크기를 잘라 학습을 안정시킨다
- 출발 분포가 끝점 분포에 남기는 치우침을 지워, 무기억이 아닌 과정으로도 목표에 닿게 한다
- 표본마다 중요도 가중치를 매겨 목표 분포의 표본을 추정한다
- 출발 분포를 원점 한 점으로 바꿔 무기억 조건을 맞춘다
정리 3.1에서 끝점 비용에 더해지는 \(\log\hat\varphi_1\)이 초기 가치 함수의 흔적을 상쇄해 끝점 분포가 \(\nu\)와 같아진다(3.1절). 1번의 기울기 자르기는 구현의 보조 장치이고(부록 C), 3번은 이 방법이 피하려는 방식이다. 4번은 AS가 무기억 조건을 맞추던 방식이다.
논문 정보
- 제목: Adjoint Schrödinger Bridge Sampler
- 저자: Guan-Horng Liu (FAIR at Meta), Jaemoo Choi (Georgia Institute of Technology), Yongxin Chen (Georgia Institute of Technology), Benjamin Kurt Miller (FAIR at Meta), Ricky T. Q. Chen (FAIR at Meta)
- 학회: NeurIPS 2025 구두 발표(Oral)
- 링크: arXiv, OpenReview, 코드