Reward-Tilted Distribution Matching Distillation으로 Few-step 생성기 강화

왜 중요한가

소수 단계 diffusion/flow 기반 생성기에서 인간 선호를 반영하기 어렵다. 기존 방법은 교사 분포와 인간 보상의 간극을 효과적으로 맞추지 못했고, 중간 시간단계의 노이즈와 이동 타깃 문제로 학습 신호가 불안정했다. RTDMD는 distribution matching과 reward 최적화를 하나의 프레임워크로 결합해 학습 안정성과 샘플 품질을 동시에 개선한다.

핵심 기여

Reward-Tilted Distribution Matching Distillation (RTDMD) 제안

교사 분포 pψ에 보상 tilt를 적용한 p̃ψ를 정의하고, DKL(pθ ∥ p̃ψ)을 최소화하는 두 단계 프레임워크를 제시한다. 이로써 distribution matching과 reward maximization이 자연스럽게 분해되어 학습 효율과 품질 향상의 균형을 달성한다.

Ambient-Consistent Distribution Matching Distillation (AC-DMD)

AC-DMD는 4단계 생성기의 각 하위 구간에서 분포 매칭을 수행하고, 가짜 score 모델의 예측이 시간에 따라 변하는 분포를 추적하도록 consistency regularizer를 추가한다. 이로써 한정된 업데이트 하에서도 generator distribution의 이동을 안정적으로 따라간다.

보상 최적화를 위한 하이브리드 정책 그래디언트

확률적 중간 전이는 GRPO 스타일의 estimator로 다루고, 최종 결정적 단계는 역전파를 통해 직접 보상에 연결하는 하이브리드 그래디언트를 도출한다. 또한 SubGRPO를 도입해 샘플 간 분산을 감소시키고, 선택된 단계의 신호만으로 학습이 집중되게 한다.

Step-subset GRPO (SubGRPO) 도입

트레이닝 예산 하에서 전체 경로의 모든 스텝을 독립적으로 샘플하지 않고, S ⊂ {1,...,K−1}의 부분 집합만 독립 노이즈를 사용하도록 하고 나머지 스텝은 그룹 노이즈를 공유한다. 이를 통해 크로스 스텝 크레딧 배분의 분산을 줄인다.

대규모 벤치마크에서 SOTA 달성

SD3-M, SD3.5-M, FLUX.2 4B에서 4-NFE로 학습한 RTDMD가 기존 최첨단 방법들을 상회하며, 4단계 추론으로도 강력한 품질을 달성한다. 4B FLUX.2 모델은 9B(50-step) 대비 다수 지표에서 우수한 성능을 보인다. 코드 및 모델은 GitHub에서 공개된다.

핵심 아이디어 이해하기

기본 아이디어는 두 가지 목표를 하나의 최적화 문제로 연결하는 것이다. 먼저 pθ를 pψ와 최대한 가깝게 만드는 분포 매칭이 필요하다. 그러나 pψ 자체가 사람의 선호를 반영하진 않으므로, 보상 r(x)를 이용해 보상 강도를 조정한 reward-tilted 분포 p̃ψ를 도입한다. KL(pθ ∥ p̃ψ)를 최소화하면, (i) 분포 매칭 항과 (ii) 보상 최대화 항의 두 가지 항으로 분해된다. 이 두 항을 각각 다른 방식으로 최적화하여 학습 신호의 안정성과 샘플 품질을 동시에 확보한다.

방법론

단계 1: Ambient-Consistent Distribution Matching Distillation (AC-DMD) 아키텍처를 도입한다. K=4인 생성기에서 각 단계 tk(= 1, 0.75, 0.5, 0.25)마다 xˆtk를 생성하고, CPS(Coefficient-Preserving Sampling) 방식으로 노이즈 수준 tk에서 샘플을 재구성한다. tk>0에서의 중간 샘플 xˆtk에 대해, p(k)θ,t를 구하고 교사 분포 pψ,t와의 역 KL을 최소화한다. 가짜 score 모델 sϕ를 사용해 p(k)θ,t의 근사치를 얻고, 상태별 DSM(denoising score matching) 손실 L(k)fake(ϕ)로 학습한다. 또한 L(k)cons(ϕ)를 도입해인접 timesteps 간 self-consistency를 강제하여 변동성을 감소시킨다. 이 두 항은 L(k)fake-total(ϕ)=L(k)fake(ϕ)+γ L(k)cons(ϕ)로 합쳐진다.

주요 결과

단계 2: Reward tilt를 적용한 두 번째 단계에서, θ의 업데이트는 분포 매칭과 보상 최대화의 합으로 구성된다. 생성기는 K−1개의 확률적 전이와 하나의 결정적 최종 매핑으로 구성된 하이브리드 정책 πθ( xˆtk | xˆtk−1 )를 따른다. ∇θ E[r(xˆ0)]는 두 부분으로 분해되어, (i) 확률적 중간 전이의 파라미터 의존성을 다루는 REINFORCE류의 기대값, (ii) 결정적 최종 매핑에 대한 경사(경로상 차단의 직접 미분)로 구성된다. SubGRPO를 통해 선택된 M개의 스텝만 독립 노이즈를 사용하고 나머지는 그룹 노이즈를 공유하여 분산을 줄인다. 전체 업데이트는 ∇θLtotal = ∇θLAC-DMD − β ∇θL SubGRPO stoc + ∇θLdet의 합으로 구성된다.

기술 상세

아키텍처: RTDMD는 2단계 구조로 작동한다. Stage I은 Ambient-Consistent Distribution Matching Distillation으로, 각 타임 서브인터벌 [tk, 1]에서 pψ,t와 p(k)θ,t를 비교하는 분포 매칭을 수행한다. 더불어 sϕ를 학습시키는 DSM 손실 L(k)fake(ϕ)와 consistency를 위한 L(k)cons(ϕ)를 합친 L(k)fake-total(ϕ)를 사용한다. CPS(코드 C)로 각 단계의 샘플링을 구현하고 η를 통해 노이즈를 주입한다. Stage II는 Reward-tilted KL의 gradient를 직접 최적화한다. xˆt0 → xˆt1 → ... → xˆtK−1 → xˆ0의 K-스텝 경로에서 K−1개의 확률적 전이와 하나의 결정적 최종 매핑으로 구성된 정책을 사용한다. ∇θE[r(xˆ0)]는 (i) 각 stochastic 단계의 로그 확도에 대한 기여, (ii) 최종 결정 매핑의 경로-적분 경향에 대한 기여로 구성된다. SubGRPO는 선택된 S의 스텝에서만 독립 노이즈를 사용하고 나머지는 그룹 노이즈를 공유하여 분산을 감소시킨다. 학습 세부사항으로는 CPS 파라미터 η, 가중치 γ, β 및 LoRA(α, r) 설정이 포함된다.

한계점

본 연구는 텍스트-투-이미지 생성에 초점을 맞춘다. 비디오 합성 및 이미지 편집과 같은 시퀀스/타임-일관성이 중요한 응용에서 보상 설계와 시간적 일관성 문제가 추가적으로 제시될 수 있다. 계산적 비용은 Stage II의 RL 기반 학습으로 인해 증가할 수 있으며, 보상 신호의 설계 품질에 따라 결과가 좌우될 여지가 있다.

실무 활용

RTDMD는 DMD 기반의 몇 단계 생성기를 reward-guided RL로 정렬하는 두 단계 프레임워크이다. AC-DMD로 cold-start를 안정화하고, 이후 보상 기반 최적화를 통해 인간 선호에 더 부합하는 샘플을 생성한다.

저지연 텍스트-투-이미지 생성에서 4-step 샘플링으로 품질과 일치도 향상
다양한 보상 신호를 결합한 다중 목표 최적화(예: CLIP, PickScore, OCR 등)에서의 효과 비교
책임 있는 이미지 합성 및 컨트롤 가능성 향상(랜더링 품질 관리, 스타일 일관성 개선)

코드 공개 여부: 공개

코드 저장소 보기

키워드

diffusion distillationreward-guided reinforcement learningdistribution matching distillationAmbient-Consistent Distribution Matching DistillationGRPOSubGRPOhybrid policy gradientRTDMD