기술 요약: 이산 확산 샘플러 및 브리지

문제 정의

본 논문은 정규화 상수 $Z$가 계산 불가능한 에너지 함수 $E$는 알려져 있으나, 이산 상태 공간 $\mathcal{S} = \{1, \dots, C\}^d$ 위의 분포 $p(x) \propto e^{-E(x)}$로부터 샘플링하는 문제를 다룹니다. 확산 샘플러는 확률적 최적 제어 및 경로 공간 매칭을 사용하여 연속 공간에 성공적으로 적용되어 왔으나, 이산 공간에서의 적용은 아직 미흡한 상태입니다. 기존 이산 확산 샘플러는 연속 환경에서 표준적인 오프-정책 강화 학습 (RL) 기법을 활용하지 못해, 특히 다중 모드 분포에서 모드 커버리지와 샘플링 품질이 최적화되지 못하는 경우가 많습니다. 또한, 두 개의 임의 분포를 연결하는 문제 (슈뢰딩거 브리지) 는 두 분포 모두로부터 샘플이 이용 가능한 설정 (데이터 - 대 - 데이터) 에 국한되어 왔으며, 한 분포가 정규화되지 않은 밀도로만 정의되는 "데이터 - 대 - 에너지" 설정은 거의 탐구되지 않았습니다.

방법론

저자들은 오프-정책 훈련 전략을 이산 확산 샘플러에 통합하고, 이러한 방법을 이산 슈뢰딩거 브리지로 일반화하는 프레임워크를 제안합니다.

1. 이산 확산 샘플러를 위한 오프-정책 훈련

핵심 방법론은 목표 분포와 일치하도록 순방향 커널 $\overrightarrow{p}_\theta$를 훈련하는 것으로, 순방향 과정 궤적 분포와 역방향 (노이즈 추가) 과정 궤적 분포 간의 발산을 최소화합니다. 연구된 목적 함수는 2 차 모멘트 발산이며, 구체적으로 다음과 같습니다:

• 궤적 균형 (Trajectory Balance, TB): 정규화 상수를 학습 가능한 매개변수로 취급합니다.
• 로그 분산 (Log-Variance, LV): 순방향 및 역방향 확률 비율의 로그 분산을 최소화합니다.

탐색과 모드 커버리지를 개선하기 위해 저자들은 RL 및 생성 흐름 네트워크 (GFlowNets) 에서 차용한 세 가지 오프-정책 전략을 도입합니다:

• 리플레이 버퍼: 샘플링된 상태 $X_N$을 저장하고, 훈련을 위해 역방향 궤적 $\overleftarrow{p}^{\otimes N}$을 샘플링하여 재사용합니다.
• 중요도 가중 버퍼: 궤적을 생성한 중요도 가중치에 기반하여 버퍼 요소에 샘플링 확률을 할당하여, 현재 샘플러가 놓칠 수 있는 고확률 영역을 우선시합니다.
• MCMC 탐색: 훈련에 사용되기 전에 리플레이 버퍼의 샘플을 마르코프 연쇄 몬테 카를로 (MCMC) 단계 (예: Swendsen-Wang 또는 Metropolis-Hastings) 를 사용하여 정제합니다. 이는 모델 롤아웃 없이 에너지 함수를 직접 활용하여 샘플러를 고밀도 영역으로 유도합니다.

2. 이산 데이터 - 대 - 에너지 슈뢰딩거 브리지

본 논문은 두 개의 임의 분포 $p_0$와 $p_N$ 사이의 슈뢰딩거 브리지 (SB) 를 찾는 샘플링 문제를 일반화합니다. 두 끝점에 대한 샘플을 가정하는 기존 이산 SB 연구와 달리, 본 작업은 $p_N$이 에너지 함수로만 정의되는 데이터 - 대 - 에너지 설정을 다룹니다.

• 이 방법은 비례 적합 (Iterative Proportional Fitting, IPF) 절차를 사용합니다.
• $p_0$의 샘플에 대한 최대 가능도를 통해 역방향 커널을 피팅하고, 에너지로 정의된 $p_N$의 샘플에 대한 로그 분산 손실을 통해 순방향 커널을 피팅하는 과정을 번갈아 수행합니다.
• 에너지로 정의된 분포의 탐색을 용이하게 하기 위해 오프-정책 기법 (버퍼 및 MCMC) 이 IPF 루프에 통합됩니다.

3. 잠재 공간 적용

저자들은 이러한 샘플러를 "아웃소싱"된 사후 샘플링에 적용합니다. 사전 훈련된 벡터 양자화 변분 오토인코더 (VQ-VAE) 의 이산 잠재 공간에서 확산 샘플러를 훈련합니다. 목표는 편향되지 않은 조건부 샘플을 관찰하지 않고도 사후 분포 $p(z|y) \propto p_{latent}(z)p(y|f(z))$로부터 샘플링하는 것으로, 잠재 공간에서 샘플링을 학습함으로써 생성 모델을 조건화하는 것입니다.

주요 기여

1. 이산 샘플러를 위한 오프-정책 알고리즘: 리플레이 버퍼, 중요도 가중치, MCMC 탐색과 같은 오프-정책 RL 기법을 이산 확산 샘플러 훈련에 도입했습니다. 논문은 이러한 기법이 온-정책 베이스라인보다 성능을 크게 향상시킨다는 것을 입증합니다.
2. 데이터 - 대 - 에너지 슈뢰딩거 브리지: 샘플 기반 분포와 에너지 정의 분포 간의 운송 지도 구축을 가능하게 하는 이산 확산 샘플러를 데이터 - 대 - 에너지 슈뢰딩거 브리지 문제로 처음 일반화했습니다.
3. 잠재 공간 사후 샘플링: 이미지 생성 모델의 이산 잠재 공간에서 사후 추론을 위한 이산 확산 샘플러를 시연하여, 생성 모델을 재훈련하지 않고도 클래스 레이블 (예: 홀수/짝수 숫자) 에 따라 조건부 생성이 가능함을 보여주었습니다.

실험 결과

저자들은 합성 벤치마크와 실제 응용 분야에서 방법을 평가했습니다:

• 이징 및 포트스 모델: $16 \times 16$ 격자에서 오프-정책 방법 (특히 버퍼 + MCMC 가 적용된 TB) 은 온-정책 방법 및 베이스라인 마스크드 확산 신경 샘플러 (MDNS) 일관적으로 우월한 성능을 보였습니다. 특히 모드 붕괴가 빈번한 저온 (높은 $\beta$) 에서 싱크혼 거리 및 자화/상관 오류를 더 낮게 달성했습니다.
• 이산화 합성 밀도: 그레이 코드를 통해 이산화된 복잡한 다중 모드 분포 (40GMM 및 ManyWell) 에서 오프-정책 방법은 모든 모드를 성공적으로 발견한 반면, 온-정책 방법과 MDNS 는 전체 분포를 커버하지 못하는 경우가 많았습니다.
• 슈뢰딩거 브리지: 단순 분포와 다중 모드 분포 간의 브리지 작업에서 오프-정책 방법은 온-정책 샘플러가 붕괴되는 복잡한 목표에 효과적으로 확장되었습니다.
• 아웃소싱 샘플링: 잠재 공간에서 생성된 MNIST 숫자의 시각화는 온-정책 및 오프-정책 샘플러 모두 원하는 클래스 (예: 홀수 대 짝수) 의 숫자를 성공적으로 생성할 수 있음을 보여주어, 조건부 생성을 위한 접근법의 유효성을 입증했습니다.

중요성 및 주장

본 논문은 오프-정책 훈련 기법을 통합하는 것이 이산 확산 샘플러의 중요한 발전 단계이며, 모드 붕괴 문제를 해결하고 고차원 다중 모드 이산 공간에서의 샘플링 효율성을 향상시킨다고 주장합니다. 이산 확산과 오프-정책 RL 간의 간극을 메움으로써, 저자들은 이러한 샘플러가 기존 방법보다 복잡한 통계 물리 문제 (이징/포트스) 와 조합 최적화 작업을 더 효과적으로 처리할 수 있음을 보여줍니다.

또한, 이 작업은 슈뢰딩거 브리지의 이론적 범위를 이산 영역의 데이터 - 대 - 에너지 설정으로 확장하여, 한 분포가 에너지 함수로 정의될 때 분포 간 운송을 위한 새로운 도구를 제공합니다. 마지막으로, 잠재 공간에 대한 적용은 전체 분포에 접근하거나 생성 모델 자체를 재훈련할 필요 없이 사전 훈련된 생성 모델을 조건화하는 실용적인 경로를 보여줍니다.

저자들은 노이즈 커널을 고정해야 한다는 점과 궤적 저장의 계산 비용과 같은 한계를 인정하지만, 연속 시간 이론 확장에 의해 지원되는 더 거친 시간 척도에서의 훈련이 일부 비용을 완화할 수 있다고 제안합니다. 또한, 데이터로부터 에너지 함수를 학습하고 이러한 방법을 더 큰 이미지 생성 모델로 확장하는 것과 같은 향후 방향을 제시합니다.