CREPE: 복제 교환으로 확산 모델 제어하기

초록

CREPE는 사전 학습된 확산 모델을 재학습 없이 다양한 제약을 만족하도록 제어하는 새로운 방법이다. 기존 SMC 기반 접근법의 메모리·다양성·재정제 한계를 극복하기 위해 병렬 템퍼링(Replica Exchange) 알고리즘을 확산 과정에 적용한다. 입자들을 여러 시점에 동시에 배치하고, 순차적으로 샘플을 생성·교환함으로써 높은 다양성을 유지하고, 연속적인 온라인 정제와 조기 종료가 가능하다. 온도 조절, 보상 기울기, 모델 합성, 클래스‑프리 가이드 디바이어싱 등 다양한 작업에서 SMC 대비 경쟁력 있는 성능을 보인다.

상세 분석

본 논문은 확산 모델의 추론 단계에서 제약을 부여하는 문제를 “샘플링” 관점에서 재정의한다. 기존에는 휴리스틱 가이던스나 Sequential Monte Carlo(SMC) 기반 방법이 주를 이루었지만, SMC는 대규모 파티클 유지가 필요하고, 파티클 다양성이 급격히 감소하며, 한 번 샘플링이 끝나면 재정제가 불가능하다는 구조적 한계가 있다. 저자들은 이러한 문제를 해결하기 위해 원래 물리학에서 다중 모드 분포를 탐색하기 위해 고안된 Parallel Tempering(PT), 즉 Replica Exchange MCMC를 확산 모델에 직접 적용한다. 핵심 아이디어는 “시간”과 “병렬성”을 뒤바꾸는 것이다. SMC가 모든 파티클을 동일한 시간 단계에서 진행하면서 가중치를 재조정한다면, CREPE는 서로 다른 디노이징 단계(t₀, t₁,…,t_M)를 갖는 체인을 병렬로 실행하고, 각 체인에서 순차적으로 새로운 파티클을 생성한다. 이렇게 하면 파티클 수에 비례한 메모리 사용이 크게 감소하고, 각 체인 간 교환(swap) 단계에서 Radon‑Nikodym Derivative(RNE)를 이용해 정확한 메트로폴리스 수용률을 계산한다. 교환은 인접한 시간 단계 사이에만 수행되며, 교환 확률은 앞서 정의된 RNE와 보상·온도·합성 등 사용자가 정의한 잠재 함수들의 비율로 구성된다.

CREPE는 세 가지 구성 요소를 명시한다. 1) Annealing Path: 목표 분포 π₀와 쉬운 샘플링이 가능한 기준 분포 π₁ 사이를 연결하는 연속적인 경로를 정의한다. 여기서는 기존 확산 모델의 마진 밀도 p_t와 온도 β, 보상 r(x) 등을 결합해 π_t(x)∝p_t(x)·β·exp(r_t(x)) 형태로 설계한다. 2) Communication Move: 인접 체인 간 교환을 수행하며, 교환 확률 α_{t,t’}은 제안 경로의 전·후방 확률밀도 비율을 이용해 메트로폴리스를 만족한다. 3) Local Exploration: 각 체인 내부에서는 기존 확산 모델의 역방향 SDE(또는 CTMC)를 이용해 MCMC 스텝을 수행한다. 이 단계는 기존 확산 모델의 스코어 네트워크를 그대로 활용하므로 추가 학습이 필요 없다.

이론적으로 CREPE는 기존 PT와 동일한 수렴 보장을 갖지만, 목표 밀도가 명시적으로 주어지지 않은 상황에서도 Radon‑Nikodym Estimator(RNE)를 통해 교환 확률을 정확히 계산한다는 점이 혁신적이다. 실험에서는 이미지(Imagenet‑512), 비디오, 텍스트 등 다중 모달리티에 대해 온도 annealing, 보상‑tilting, 모델 composition, 클래스‑프리 가이드 디바이어싱을 수행하였다. SMC 대비 파티클 다양성 지표(Frechet Inception Distance, Diversity Score 등)에서 현저히 우수했으며, 메모리 사용량은 2~3배 절감되었다. 특히 “온라인 정제” 실험에서는 초기 샘플이 만족스럽지 않을 경우 추가 PT iteration만으로 품질을 향상시킬 수 있음을 보였다.

한계점으로는 (1) 교환 빈도가 인접 시간 단계 간 겹침 정도에 크게 의존해, 매우 급격한 스케줄에서는 교환 효율이 떨어질 수 있다. (2) 현재 구현은 동일한 하드웨어에서 병렬 체인을 실행하도록 설계돼, 분산 환경에서의 스케일링은 추가 연구가 필요하다. (3) 복잡한 보상 함수가 포함된 경우 RNE 계산이 수치적으로 불안정해질 가능성이 있다. 향후 연구에서는 적응형 annealing schedule, 교환 제안의 비대칭 설계, 그리고 대규모 클러스터에서의 효율적인 파이프라인 구축이 기대된다.