기술 요약: 역강화학습을 통한 확산 모델의 샘플링 학습

문제 정의

확산 모델(Diffusion models)은 고차원 데이터 생성의 표준이 되었으나, 샘플링 절차는 여전히 중요한 병목 구간으로 남아 있다. 일반적으로 디노이징 신경망은 학습 후 고정되지만, 샘플링 알고리즘 자체는 하이퍼파라미터(예: 노이즈 스케줄, 가이던스 스케일, 확률성 프로파일)의 세심한 튜닝을 필요로 한다. 현재 이러한 튜닝은 비용이 많이 드는 경험적 그리드 탐색(grid search)을 통해 수행된다. ImageNet-64와 같은 복잡한 모델의 경우, 철저한 하이퍼파라미터 탐색에 드는 계산 비용은 디노이저를 학습시키는 비용에 근접하거나 이를 초과할 수 있다. 또한, 기존의 강화학습(RL) 또는 제어 접근 방식은 명시적인 보상 함수나 디노이저의 미세 조정(fine-tuning)을 요구하는 경우가 많은데, 이는 사전 학습된 생성기를 수정하지 않고 샘플러를 최적화하려는 목표와 상충된다.

방법론

저자들은 확산 샘플링을 **역강화학습(Inverse Reinforcement Learning, IRL)**을 통해 해결 가능한 순차적 의사결정 문제로 프레이밍할 것을 제안한다. 핵심 혁신은 전문가의 궤적(expert trajectories), 명시적 보상, 또는 디노이저의 재학습 없이도 타겟 상태 분포를 일치시키는 샘플링 정책을 학습하는 것이다.

1. MDP 정식화

확산 샘플링 과정은 유한 호라이즌 마르코프 결정 과정(finite-horizon MDP)으로 정식화된다:

• 상태 ($s_t$): $(x_t, \sigma_t)$로 정의되며, 여기서 $x_t$는 현재 샘플이고 $\sigma_t$는 노이즈 수준이다.
• 행동 ($a_t$): 가이던스 스케일 선택, 노이즈 주입 수준 결정, 또는 이전 노이즈 수준으로 리노이징(reenoise/restart) 여부 결정과 같이 샘플링 역학에 적용되는 이산적 제어 결정이다.
• **역학(Dynamics): 이산화된 확산 역학과 선택된 제어 행동에 의해 유도되는 전이 커널(transition kernel)이다. 프로세스는 노이즈 수준이 $\Sigma_0$(흡수 상태)에 도달할 때 종료된다.

2. 상태 전용 역강화학습 (State-Only IRL)

전통적인 IRL이 전문가의 행동이나 궤적을 가정하는 것과 달리, 이 방법은 상태 전용(state-only) 체제에서 작동한다. "전문가"는 선택된 노이즈 수준에서의 데이터 샘플의 주변 분포(marginal distributions)에 의해서만 정의된다.

• 목표: 정책 $\pi_\theta$의 상태 점유 측도(state occupancy measure) $\mu_\theta$(방문된 상태의 주변 분포)가 전문가의 점유 측도 $\mu_E$와 일치하도록 하는 것이다.
• 최적화: 보상 함수를 직접 학습하는 대신, 저자들은 전문가와 정책 간의 점유 측도 사이의 **$f$-발산(f-divergence)**을 직접 최소화한다:
  $$\min_\theta D_f(\mu_E \parallel \mu_\theta)$$
  논문에서는 특히 쿨백-라이블러(KL) 발산과 역-KL(Reverse-KL) 발산을 탐구하며, $f$-발산의 선택이 모드 커버링(mode-covering, 다양성)과 모드 시킹(mode-seeking, 충실도) 사이의 트레이드오프를 결정한다는 점에 주목한다.

3. 정책 경사 및 최적화

저자들은 $f$-발산을 직접 최소화하기 위한 정책 경사 정리(policy gradient theorem)를 도출한다. 경사는 다음과 같이 표현된다:
$$\nabla_\theta L_f(\theta) = \mathbb{E}_{s_{1:T} \sim p_\theta} \left[ \frac{1}{T} \sum_{t=1}^T \nabla_\theta \log \pi_\theta(a_t|s_t) \sum_{t' \ge t} h_f\left(\frac{\mu_E(s_{t'})}{\mu_\theta(s_{t'})}\right) \right]$$
여기서 $h_f$는 $f$-발산의 생성자로부터 유도된 함수이다.

• 밀도 비율 추정(Density Ratio Estimation): $\mu_E$와 $\mu_\theta$를 알 수 없으므로, 전문가 상태(노이즈가 섞인 데이터)와 정책 생성 상태를 구별하도록 훈련된 판별기(discriminator)를 사용하여 밀도 비율을 추정한다.
• 안정화: 경사 추정치의 높은 분산을 처리하기 위해, 이 방법은 PPO 스타일의 클리핑된 목적 함수, 중요도 샘플링 보정, 그리고 글로벌 배치 평균 베이스라인을 사용한다.

4. 평가된 샘플링 전략

이 프레임워크는 세 가지 특정 추론 시점 개입에 적용된다:

1. 적응형 클래시파이어 프리 가이던스(Adaptive Classifier-Free Guidance): 상태 의존적 가이던스 스케일 $\omega(x, \sigma)$를 학습한다.
2. 확률성 주입(Stochasticity Injection): 상태 의존적 노이즈 주입 수준 $\gamma_{EDM}(x, \sigma)$를 학습한다.
3. 리노이징(Renoising/Restart Sampling): 잘못된 로컬 생성을 탈출하기 위해 프로세스를 이전 노이즈 수준으로 되돌릴 시점을 학습한다.

주요 기여

1. 샘플링을 위한 MDP 정식화: 추론 시점에 디노이저를 재학습하지 않고도 샘플러를 적응시키기 위해 확산 샘플링을 MDP로 구성하였다.
2. 직접적인 상태-주변 분포 매칭: 명시적인 보상 모델링이나 전문가 행동 데이터 없이도 전문가와 정책의 상태 점유 측도 간의 $f$-발산을 최소화하는 정책 경사 목적 함수를 도출하였다.
3. 해석 가능한 트레이드오프: $f$-발산(KL vs rKL)의 선택을 통해 샘플 다양성(Recall)과 충실도(Precision) 사이의 제어 가능한 트레이드오프를 보여준다.
4. 상태 의존적 제어: 학습된 정책은 고정된 글로벌 스케줄을 적용하는 대신, 특정 노이즈 수준이나 특정 샘플 유형에 대해 선택적으로 개입하는(예: 특정 노이즈 수준에서 높은 가이던스를 적용) 추론 시점 제어로 작용한다.

실험 결과

본 방법은 사전 학습된 EDM 모델을 사용하여 CIFAR-10, FFHQ, ImageNet-64에서 평가되었다.

• 성능: 학습된 정책은 FID, Precision, Recall 지표에서 가장 잘 튜닝된 베이스라인(그리드 탐색을 통해 발견됨)과 일치하거나 이를 능가한다.
  • 적응형 가이던스: 고정된 가이던스에 비해 품질-다양성 트레이드오프를 개선하였다. KL 발산은 Recall(다양성)을 향상시킨 반면, rKL은 Precision(충실도)을 향상시켰다.
  • 확률성 주입: 상태 의존적 프로파일은 최적의 EDM 스케줄의 질적 패턴(궤적 끝부분에서 낮은 확률성)을 복구하였으며, ImageNet과 같은 복잡한 데이터셋에서 고정된 베이스라인보다 우수한 성능을 보였다.
  • 리노이징: 표준 리스타트 전략과 비교하여 FID를 유지하거나 개선하면서도 기능 평가 횟수(NFE)를 크게 줄였다.
• 계산 효율성:
  • 학습: 단 한 번의 학습 실행이 방대한 그리드 탐색을 대체한다. ImageNet-64에서, 이 방법은 최적화 예산을 최대 9배 절감한다(3-하이퍼파라미터 그리드 탐색 시 약 1437 EFLOPs에서 IRL 학습 시 약 161 EFLOPs로 감소).
  • 추론: 정책은 디노이저 단독 대비 **16%에서 40%**의 완만한 FLOPs 오버헤드를 추가하며, 이는 수동 튜닝 제거를 고려할 때 무시할 만한 수준이다.
• 어블레이션(Ablation) 결과:
  • 조건화(Conditioning): 전체 상태 $(x, \sigma)$에 조건화된 정책는 노이즈 수준 $\sigma$에만 조건화된 정책보다 유의미하게 우수하며, 명시적 클래스 레이블 $(x, \sigma, c)$를 가진 경우와 거의 유사한 성능을 보인다. 이는 정책이 이미지 콘텐츠로부터 클래스 판별적 특징을 암시적으로 학습함을 시사한다.
  • 발산 선택: KL과 rKL이 정밀도-재현율 트레이드오프를 제어하는 데 가장 효과적임을 보여준 반면, 대칭적 발산(TV, Hellinger)은 덜 해석 가능한 결과를 냈다.

의의 및 주장

본 논문은 이 프레임워크가 확산 샘플링을 위한 체계적이고 자동화된 수동 하이퍼파라미터 튜닝의 대안을 제공한다고 주장한다. 샘플링을 IRL 문제로 다룸으로써, 수동으로 설계하기 어려운 복잡하고 상태 의dependent한 샘플링 전략을 최적화할 수 있게 한다.

저자들은 이 방법이 특히 per-trial 평가 비용이 막대한 대규모 모델에서 가치가 있다고 강조한다. 이 접근 방식은 생성 모델 자체를 수정하지 않고, 대신 추론 궤적을 위한 최적의 "제어 정책"을 학습함으로써, 재학습이나 철저한 탐색의 계산적 부담 없이 샘플 품질과 다양성을 개선할 수 있는 실용적인 경로을 제시한다. 본 논문은 이 혜택이 데이터셋에 따라 달라질 수 있으며, 현재 프레임워크가 이산 행동 공간에 의존한다는 점을 언급하며, 연속 행동 공간 및 더 강력한 초기화에 대한 향후 연구 가능성을 열어두며 겸허한 입장을 유지한다.