KL을 넘어서는 확산 정렬: 분산 최소화 기반 정책 최적화

초록

본 논문은 사전 학습된 확산 모델을 보상에 따라 편향된 분포로 정렬하는 과정을 순차적 몬테카를로(SMC) 관점에서 재해석하고, KL 발산 대신 로그 중요도 가중치의 분산을 최소화하는 Variance Minimisation Policy Optimisation(VMPO) 프레임워크를 제안한다. 이론적으로 VMPO는 목표 보상-틸트 분포를 최적해로 갖고, 온‑policy 샘플링 하에서는 KL 기반 정렬과 동일한 그래디언트를 제공한다. 다양한 잠재 함수와 분산 최소화 전략을 선택함으로써 기존 DDPO, DPOK, Flow‑GRPO 등 여러 방법을 재현하고, 새로운 설계 방향도 제시한다. 실험에서는 Stable Diffusion 1.5·3.5을 인간 선호 점수(HPSv2)와 기타 보상 함수로 미세조정하여, VMPO‑Diff 변형이 기존 방법보다 높은 보상 점수를 달성함을 보인다.

상세 분석

본 연구는 확산 정렬을 “제안 분포 = 디노이징 모델, 보상 가중치 = 중요도 가중치”라는 SMC 구조로 모델링한다. 이때 목표는 보상에 따라 기울어진 목표 분포 p_tilt(x_{t‑1}|x_t)∝p_ref·exp(r/β)를 샘플링하는 것이며, 전통적인 접근은 KL 발산 KL(p_θ‖p_tilt)를 최소화하는 것이었다. 저자들은 KL 대신 로그 중요도 가중치 log w_t의 분산 Var_h(log w_t)를 최소화하는 목표 L_h^Var(t;θ)=½Var_h(log w_t) 를 정의하고, 이를 h라는 임의의 기준 분포(보통 현재 정책 p_θ) 위에서 기대값을 취한다. Proposition 1에 따르면 L_h^Var의 최적 θ*는 정확히 p_tilt와 일치하고, ∇θ L_h^Var|{h=p_θ}=∇_θ KL(p_θ‖p_tilt) 가 된다. 즉, 온‑policy 상황에서는 KL 기반 정책 그래디언트와 완전히 동일하므로 기존 PPO·GRPO와 호환된다.

핵심 기술적 기여는 두 가지이다. 첫째, 분산 최소화 목표를 Monte‑Carlo 샘플링으로 직접 추정하거나, 기대값 E_h