DiffusionART가 확산 모델 강화학습의 불안정성을 구조적으로 제거한다

무엇인가

확산(diffusion)·플로우 매칭 모델은 시각 품질, 모션 품질, 인간 선호 같은 다운스트림 기준이 아니라 노이즈/속도 회귀 손실로 학습되기 때문에, 보상 기반 파인튜닝이 최근 이미지·비디오 생성의 핵심 작업이 됐다. 이 논문이 다루는 것은 온라인 확산 강화학습이다. 미분 가능한 보상 모델이 필요 없고 블랙박스 보상을 쓸 수 있지만, 보상은 완성된 출력 하나에만 주어지므로 값비싼 출력 수준 피드백을 수많은 디노이징 상태에 배분하는 신용 할당(credit assignment) 문제가 생긴다. 기존 역방향(reverse-process) 방식인 Flow-GRPO 계열은 ODE→SDE 변환이 필요하고, 각 전이를 그 전이 자체의 샘플링 노이즈 방향으로 업데이트하면서 궤적 수준 어드밴티지로 가중해 신용 할당이 나빠 효율이 제한된다. 전방(forward-process) 방식인 DiffusionNFT는 전이 우도가 필요 없어 매우 효율적이지만, 공동 양·음 학습과 태스크마다 튜닝해야 하는 EMA 롤아웃 정책이라는 두 안정화 장치 없이는 불안정하다.

어떻게 동작하나

논문의 첫 번째 진단은 NFT 목적함수에 숨어 있는 전방-역방향 일관성 힘(consistency force)이다. Lemma 1은 이상적인 NFT 가이던스가 상태 조건부 보상-속도 공분산 g와 같고 양·음 브랜치가 같은 방향을 가리킨다는 것을 보인다. 그런데 실제 역방향 샘플링 장은 자기 엔드포인트가 유도하는 이상적 전방 속도와 다르며, 그 차이 e = v_fwd_old − v_rvs_old가 업데이트에 섞여든다. Proposition 1은 순수 온폴리시 지점에서 NFT의 정규화된 출력 공간 힘이 F_NFT = 2g + (2q−1)e임을 유도한다. 각 브랜치 힘은 F+ = g + qe, F− = g − (1−q)e로, 두 브랜치를 함께 학습하면 g 방향으로는 서로 보강하지만 e 방향으로는 서로 상쇄되어 (2q−1)e만 남는다. 문제는 이 상쇄가 모든 상태에서 q = 1/2일 때만 정확하다는 점이다. 프롬프트 수준 보상 중심화는 Σ(2r_i−1) ≈ 0을 만들 뿐 상태별 균형이나 Σ(2r_i−1)e_i ≈ 0을 보장하지 않는다. t=1에서는 상태가 엔드포인트 정보를 담지 않아 q = 1/2로 잔차가 사라지지만, 노이즈가 줄수록 q가 극단화되어 |2q−1| 가중치가 커진다. 즉 공동 학습은 일관성 힘을 상태별이 아니라 샘플 간 집계 수준에서 부분적으로만 지운다.

무엇과 다른가

두 번째 진단은 EMA가 왜 안정화하는가다. Proposition 2는 NFT 손실을 L_NFT = E[−2β(2r−1)⟨d, h⟩ + β²‖d‖²]로 분해한다. d = v_θ − v_old는 현재 정책이 롤아웃 정책에서 벗어난 드리프트, h = u_t − v_old는 샘플 학습 잔차다. 선형항이 보상 가이던스와 일관성 힘을 함께 나르고, 이차항이 v_θ를 v_old 쪽으로 되돌리는 복원(restoration)을 담당한다. d = 0이면 복원 그래디언트가 사라지므로 순수 온폴리시 NFT에는 복원이 없다. 저자들은 두 가지 절제 실험으로 이를 확인한다. EMA 롤아웃은 유지하되 손실의 v_old 앵커를 v_θ로 바꿔 그래디언트에서 제거하면 여전히 붕괴한다. 즉 오프폴리시 데이터 자체는 안정화 요인이 아니다. 반대로 EMA 없이 롤아웃 배치를 두 개의 서로소 미니배치로 나눠 각각 한 번씩 업데이트하면 관측 구간에서 붕괴를 피하지만 학습이 15.2배 느려진다. 흥미롭게도 롤아웃 0~17 평균에서 복원의 출력 힘 RMS는 가이던스보다 9.4배 작은데, 집계된 파라미터 그래디언트 노름은 4.9배 크다. 상태별로 이질적인 가이던스는 집계 과정에서 상당 부분 상쇄되는 반면 복원은 살아남아 업데이트를 지배하고 보상 최적화를 늦춘다는 해석이다. 저자들은 이 현상의 상세 분석은 향후 과제로 남긴다.

어떻게 쓰나

제안 방법 DiffusionART는 이 두 실패 모드를 각각 구조적으로 없앤다. 먼저 롤아웃 분포를 p_τ(x0|c) ∝ p_old(x0|c)exp(τR(x0,c))로 지수 틸팅하면, Lemma 2에 따라 τ=0에서 틸팅된 엔드포인트 평균의 변화율이 Cov(R, x0|c)와 같아진다. rectified flow에서 u_t = x0 − ε이고 ε이 x0와 독립이므로 이 공분산은 모든 노이즈 레벨에서 Cov(R, u_t|c)와 동일하다. 구현에서는 프롬프트당 K개 롤아웃의 엔드포인트와 보상으로 Cov̄(R, x0|c) = (1/K)Σ(R_i − R̄_c)(x0_i − x̄_0,c)를 추정하고 RMS로 정규화해 G(c)를 만든다. RMS 정규화는 보상 스케일·분산 차이를 없애 가이던스 강도 계수를 따로 둘 필요가 없게 한다. 이 G(c)를 해당 프롬프트의 모든 역방향 롤아웃 상태에 공유 적용하고, 타깃을 v_tgt^ART(s_t) = v_θ^(0)(s_t) + G(c)로 고정한다. 타깃에 전방 속도가 전혀 들어가지 않으므로 일관성 힘이 샘플 간 상쇄가 아니라 구조적으로 제거되고, 각 상태는 스칼라 어드밴티지 대신 명시적인 보상 개선 방향을 받는다. 안정화는 롤아웃 배치를 둘로 나눠 각각 MSE 업데이트하는 방식으로 얻는다. 두 업데이트 모두 같은 동결 타깃을 쓰는데, 두 번째 업데이트의 그래디언트는 G(c) − d(s_t)가 되어 드리프트에 비례하는 복원이 가이던스에 더해진다. 드리프트가 작으면 가이던스가 지배하고, 크면 복원이 커지는 적응적 구조이며 EMA 스케줄이 필요 없다. 여기에 ODE 롤아웃용 신뢰영역 마스킹을 추가한다. ODE 전이는 결정론적이라 유한 KL을 잴 수 없으므로, Flow-GRPO의 ODE→SDE 정식화를 따라 현재·이전 정책의 참조 가우시안 전이를 정의하고 KL 프록시 m(s_t) = ‖μ(s_t, v_θ) − μ(s_t, v_θ^(0))‖² / (2Dσ_ref²(t))를 쓴다. 드리프트가 임계값 δ = 10⁻⁵ 이상이면서 두 번째 업데이트의 하강 방향과 변위의 내적이 양수인(외향) 업데이트만 마스킹하고, 계수 t로 고노이즈 상태를 강조한다. 또한 DiffusionART-Pairwise는 프롬프트당 K개 엔드포인트를 K/2개의 서로소 쌍으로 묶어 승/패(1,0) 또는 무승부(1/2,1/2) 라벨을 만들고, 스칼라 보상 대신 이 라벨을 같은 공분산 파이프라인에 넣는다.

전제와 한계

텍스트-이미지 실험은 SD3.5-M에서 LoRA(rank 32, α=64), 512×512, CFG와 참조 정책 KL 정규화 없이, 학습률 3×10⁻⁴, 10스텝 롤아웃으로 수행했다. GenEval, OCR, PickScore 단일 보상과 GenEval·PickScore·CLIPScore·HPSv2.1 다중 보상 모두에서 DiffusionART는 태스크별 EMA 튜닝 없이 DiffusionNFT(온폴리시)의 심각한 붕괴를 피하면서 DiffusionNFT(EMA)에 필적하는 효율을 보였고, Flow-GRPO의 500스텝 보상 수준에 각각 10.0배, 5.2배, 9.6배 적은 스텝으로 도달했다. 12B FLUX.1-dev의 GenEval에서는 6스텝 롤아웃과 내장 가이던스 임베딩(스케일 3.5)을 쓰고 DiffusionNFT는 SD3.5-M용 β=1, EMA 감쇠 min(0.001n, 0.5)를 그대로 재사용했는데, DiffusionNFT는 보상이 크게 요동치는 반면 DiffusionART가 가장 빠르게 학습했다. 가이던스·복원 절제에서는 1회 업데이트 변형이 빠르지만 점차 나빠지고, 2회 업데이트가 장기 안정성을 제공하며, β=0.1로 튜닝한 NFT 2회 업데이트는 잠깐 거의 0으로 붕괴했다가 회복했다.

텍스트-비디오는 CFG 증류된 HunyuanVideo-13B 전체 파라미터를 학습률 10⁻⁵, KL 정규화 없이 파인튜닝했다. 롤아웃당 32개 프롬프트, 프롬프트당 8개 비디오, 53프레임 480×480 8FPS 16스텝, VideoReward를 4FPS로 채점, TA:VQ:MQ = 7:2:1 가중치를 썼다. DiffusionART는 학습 VQ 6과 MQ 3을 각각 29스텝, 46스텝에 도달했고, Flow-GRPO-Fast·DanceGRPO·DiffusionNFT(EMA)는 기록된 궤적 안에서 이 수준에 도달하지 못했다. VQ 5.5 기준으로 DiffusionART는 24스텝·63.4 GPU시간이 필요한 반면 DiffusionNFT(EMA)는 273스텝·1,401.4 GPU시간이 걸려 스텝 기준 11.4배, GPU시간 기준 22.1배 적다. 400스텝까지 이득을 유지하며 평가 VQ 8.970, MQ 4.224를 기록해 최강 베이스라인 대비 VQ 31.4%, MQ 57.4% 향상됐다. 초록이 제시한 VQ:MQ = 1:1 조건에서는 DiffusionNFT와 Flow-GRPO-Fast의 후반부 평균 보상 수준에 각각 12.02배, 11.72배 적은 GPU시간으로 도달하면서 평균 학습 보상은 28.8%, 53.0% 더 높았다. 신뢰영역 마스킹을 제거하면 VQ:MQ = 1:1에서 200스텝 부근에 하락했다가 회복하는 반면, 마스킹을 유지하면 꾸준히 개선됐다. 페어와이즈 피드백에서는 그룹 크기 K에 둔감했고, 400스텝 예산에서 K=2의 최종 VQ가 12.235로 K=8의 12.092보다 낮지 않았는데, K=2는 롤아웃 비디오와 학습 샘플을 4분의 1만 써서 400스텝 동안 260.8 GPU시간으로 K=8의 883.3 GPU시간보다 3.4배 적게 소모했다. 이미지-투-비디오는 LightX2V가 공개한 CFG·스텝 증류 Wan2.1-I2V-14B를 32프롬프트, 프롬프트당 8비디오, 53프레임 400×640 8FPS 4스텝, VQ:MQ = 1:1로 파인튜닝해 MQ 6을 45.8 GPU시간에 달성했고, Flow-GRPO-Fast의 220.3 GPU시간 대비 4.8배 줄었다. DiffusionNFT는 일찍 정체됐는데, 저자들은 이를 그래디언트 집계 후 남은 자기 일관성 힘이 보상 개선 가이던스를 희석한 결과로 해석한다.

개발자 관점에서 이 논문의 실용적 의미는 세 가지다. 첫째, 미분 가능한 보상 모델이 필요 없고 블랙박스 스칼라 보상이나 페어와이즈 선호 피드백만 있으면 된다. 둘째, 태스크마다 EMA 감쇠 스케줄과 β를 다시 튜닝하는 작업이 사라진다. DiffusionNFT가 OCR에서 최대 유지율 0.999를 요구하는 등 태스크별 민감도가 컸던 것과 대비된다. 셋째, 프롬프트당 샘플 수 K에 둔감해 K=2(프롬프트당 한 쌍)만으로도 안정적 최적화가 가능해 롤아웃 비용을 크게 줄일 수 있다. 다만 도입 시 확인할 점은 명확하다. RMS 정규화가 보상 스케일 차이를 흡수하므로 보상 모델을 바꿔도 가이던스 강도 계수를 새로 잡을 필요는 없지만, 신뢰영역 임계값 δ는 전 실험에서 10⁻⁵로 고정돼 있어 다른 해상도·스텝 수에서도 그대로 통하는지는 검증이 필요하다. 또한 KL 프록시는 ODE를 SDE 참조 전이로 근사해 계산하므로 롤아웃 자체는 ODE 궤적을 유지한다는 점, 그리고 실험 범위가 SD3.5-M LoRA와 FLUX.1-dev, HunyuanVideo-13B, Wan2.1-I2V-14B 파인튜닝에 한정된다는 점을 감안해야 한다.

저자가 명시한 한계는 다음과 같다. 복원 항이 출력 힘 기준으로는 가이던스보다 9.4배 작으면서 집계 파라미터 그래디언트는 4.9배 크게 지배하는 현상에 대해 논문은 상세 분석을 향후 과제로 남겨두었다. EMA가 태스크별 튜닝을 요구한다는 점은 문제로 지적했지만, DiffusionART 역시 드리프트 임계값 δ 같은 상수를 선택해야 하며 이 값은 10⁻⁵로 고정해 실험했다. 신뢰영역 마스킹은 ODE 전이가 결정론적이라 유한 KL을 정의할 수 없다는 전제 위에서 SDE 참조 전이 기반 KL 프록시로 근사한 것이므로, 이 근사가 실제 정책 변화를 얼마나 잘 반영하는지는 별도 검증 대상이다. 또한 온폴리시 초기화와 롤아웃 정책 동결 등 NFT 계열의 기본 가정을 그대로 물려받는다.