Diffusion 모델 샘플링을 역강화학습으로 자동 최적화하기

초록

본 논문은 사전 학습된 디퓨전 모델의 denoiser를 재학습하지 않고, 샘플링 과정에서 사용되는 하이퍼파라미터와 동적 제어를 역강화학습(IRL) 프레임워크로 학습한다. 디퓨전 샘플링을 유한 시간 마코프 결정 과정(MDP)으로 정의하고, 전문가 상태 분포(노이즈 레벨별 데이터 샘플)와 정책이 생성하는 상태 점유 분포를 f‑divergence로 직접 매칭한다. 보상 함수를 설계하지 않고 정책 그라디언트를 이용해 최적화함으로써, stochastic injection, classifier‑free guidance, restart sampling 등 기존 휴리스틱을 자동으로 튜닝하고 샘플 품질을 향상시킨다.

상세 분석

이 연구는 디퓨전 모델의 샘플링 단계가 실제로는 고정된 denoiser 네트워크 위에 여러 가지 하이퍼파라미터(예: 가이드 스케일 ω, 노이즈 증폭 γ, 재시작 횟수 K 등)를 적용해 동적 시스템을 구성한다는 점에 착안한다. 기존 연구들은 이러한 파라미터를 수동으로 조정하거나, 별도의 보상 함수를 정의해 강화학습으로 최적화했지만, 보상 설계 자체가 복잡하고 목표와의 정합성이 떨어지는 문제가 있었다. 저자들은 이를 해결하기 위해 “전문가 상태 분포”라는 새로운 감독 신호를 도입한다. 여기서 전문가란 실제 데이터 분포와 그에 대응하는 여러 노이즈 레벨에서의 샘플을 의미하며, 이는 전통적인 전문가 정책(액션 시퀀스) 없이도 얻을 수 있다.

논문은 먼저 디퓨전 샘플링을 유한‑ horizon MDP로 모델링한다. 상태 sₜ는 (xₜ, σₜ) 쌍이며, σₜ는 현재 노이즈 레벨, xₜ는 현재 이미지 벡터이다. 행동 aₜ는 이산형으로, 각 행동은 특정 가이드 스케일 선택, 노이즈 증폭 정도, 혹은 재시작 여부와 같은 샘플링 동작을 지정한다. 전이 확률 P는 선택된 행동에 따라 기존 디퓨전 업데이트 연산자 H에 변형을 가한 결과로 정의된다. 특히 σₜ = Σ₀ 일 때는 흡수 상태로 설정해 이후 변동이 없도록 함으로써, 최종 이미지 품질을 직접 평가할 수 있게 만든다.

핵심은 정책 πθ가 생성하는 상태 점유 측도 μθ와 전문가 점유 측도 μᴇ 사이의 f‑divergence를 최소화하는 것이다. KL, reverse‑KL, TV 등 다양한 f‑divergence를 선택할 수 있으며, 저자는 KL과 reverse‑KL을 사용해 두 가지 형태의 손실을 명시적으로 분해한다.

• KL(μᴇ‖μθ) = KL(wᴇ‖wθ) + Σσ wᴇ(σ)·KL(pᴇ(·|σ)‖pθ(·|σ))
• reverse‑KL(μᴇ‖μθ) = KL(wᴇ‖wθ) + Σσ wθ(σ)·KL(pᴇ(·|σ)‖pθ(·|σ))

여기서 w(σ)는 각 노이즈 레벨에 할당된 방문 비중, p(x|σ)는 해당 레벨에서의 이미지 분포이다. 이 분해는 정책이 “노이즈 레벨 분포”와 “레벨별 이미지 품질”을 동시에 최적화하도록 만든다. 특히 wᴇ는 전문가가 빠르게 Σ₀에 도달하기를 바라는 가중치를 부여함으로써, 정책이 효율적인 샘플링 경로를 학습하도록 유도한다.

학습은 보상 함수를 명시하지 않고, 정책 그라디언트(예: REINFORCE)와 샘플 기반 추정으로 f‑divergence의 미분 가능한 근사치를 최소화한다. 이때 전문가 점유 측도 μᴇ는 실제 데이터와 노이즈 레벨별 샘플을 이용해 경험적으로 추정한다. 정책 파라미터는 가이드 스케일, 노이즈 증폭, 재시작 전략을 매 timestep마다 선택하도록 설계되며, 연속적인 정책 네트워크 대신 간단한 테이블이나 작은 MLP로 구현해도 충분히 학습된다.

실험에서는 사전 학습된 Stable Diffusion, EDM, DPM‑Solver 등 다양한 베이스 모델에 적용했으며, FID, IS, Precision‑Recall 등 정량 지표에서 기존 수동 튜닝 대비 5~15% 정도 개선을 보였다. 특히 복잡한 데이터셋(ImageNet)에서 non‑zero γ_EDM을 자동으로 학습해 샘플 다양성을 유지하면서도 품질을 높였고, classifier‑free guidance의 ω 값을 데이터에 맞게 조정함으로써 조건부 생성에서도 성능 향상을 확인했다. 또한 재시작(renoise) 전략을 정책이 동적으로 선택하도록 함으로써, 특정 단계에서만 재시작을 수행해 연산 비용을 크게 늘리지 않으면서도 로컬 최소함정에서 탈출할 수 있었다.

이 접근법의 장점은 (1) denoiser 재학습이 필요 없으므로 기존 대규모 모델을 그대로 활용 가능, (2) 전문가 행동이 필요 없고 데이터만 있으면 학습이 가능해 적용 범위가 넓다, (3) f‑divergence 기반 목표가 이론적으로 명확한 최적화 기준을 제공한다는 점이다. 한계로는 전문가 점유 측도 μᴇ를 정확히 추정하기 위해 충분한 노이즈 레벨별 샘플이 필요하고, 정책이 이산 행동 공간에 제한돼 연속적인 하이퍼파라미터 조정이 어려울 수 있다는 점을 들 수 있다. 향후 연구에서는 연속 행동 공간으로 확장하거나, 다중 목표(예: 속도‑품질 트레이드오프)를 동시에 고려하는 다중 목표 IRL 프레임워크를 탐색할 여지가 있다.