[논문리뷰] FewStepRL: Fine-tuning Few-step Diffusion Models with Reinforcement Learning

링크: 논문 PDF로 바로 열기

저자: Haitao Lin, Minsi Ren, Yang Li, Xingtong Ge, Zheyuan Hu, Peiyan Hu, Tailin Wu (Tencent Hunyuan, Westlake University, UCAS)

1. Key Terms & Definitions (핵심 용어 및 정의)

• Few-step Model (Average Velocity Model): 짧은 구간의 instantaneous velocity가 아니라 구간 [r, t] 전체의 average velocity u_θ(x_t, t, r)를 학습합니다. 그래서 1~4회의 network evaluation만으로 sample을 생성합니다. 대표 예는 MeanFlow와 DMD, AnyFlow입니다.
• x-space Transformation: linear noise schedule에서는 velocity가 noise level 사이의 affine map을 정의합니다. 이 성질을 이용해 density, velocity, reward 기반 RL objective를 공통 policy prediction(기본값은 clean end x̃_θ = x̃_t − t·v_θ)으로 표현하는 변환입니다.
• Decoupled Loss Decomposition: RL loss의 gradient를 두 항의 합으로 나누는 분해입니다. 하나는 reference 쪽으로 당기는 self-regression 항이고, 다른 하나는 reward를 담은 reward augmentation 항입니다. 이렇게 나누면 각 method의 차이는 gradient level에서 anchors와 weights의 선택으로 정리됩니다.
• Anchors (x̃_reg, x̃_base, x̃_rwd): self-regression의 목표가 되는 regression anchor, 그리고 reward force의 방향을 정하는 baseline anchor와 rewarded anchor를 말합니다.
• Group-advantage vs. Reward-gradient methods: 앞의 것은 prompt group 안에서 scalar reward로 advantage를 계산하는 방식으로, FewStepGRPO, AWM 계열, NFT 계열이 여기에 속합니다. 뒤의 것은 ∇R을 직접 사용하는 방식으로, FewStepOPSD와 ReFL이 해당합니다.

2. Motivation & Problem Statement (연구 배경 및 문제 정의)

이 논문은 instantaneous velocity와 many-step sampler를 전제로 만들어진 기존 diffusion RL objective를, average velocity를 모델링하는 few-step model에 체계적으로 적용하는 방법이 없다는 문제를 다룹니다. 기존 RL objective는 매 iteration마다 current policy로 rollout을 해야 합니다. many-step sampler(보통 32~64 step)에서는 이 rollout이 학습 time과 memory 대부분을 차지하며, 긴 video sequence에서는 그 비중이 더 커집니다. few-step model은 이 비용을 줄여 주지만, 기존 objective를 그대로 적용하면 몇 가지가 정해지지 않은 채 남습니다. 어떤 network evaluation을 미분할지, reference policy를 어떻게 평가할지, 구간의 어느 부분을 학습할지가 그것입니다. 선행 연구는 이런 선택을 objective마다 따로 해결했습니다. 실제로 distilled model에 average velocity를 그대로 대입해 FlowGRPO, AWM, DiffusionNFT를 적용하면, 세 경우 모두 이 논문의 adapted objective보다 best GenEval2가 낮았습니다 [Figure 1].

3. Method & Key Results (제안 방법론 및 핵심 결과)

저자들은 Transformation → Decomposition → Adaptation의 3단계 framework를 제안합니다 [Figure 2].

• Transformation: 모든 objective를 x-space의 공통 policy prediction으로 표현해 서로 비교할 수 있게 만듭니다.
• Decomposition: 각 objective의 gradient를 self-regression과 reward augmentation의 합으로 분해합니다. 그러면 method 간 차이는 anchors와 weights의 차이로 축소됩니다.
• Adaptation: 각 anchor를 few-step model에서 계산합니다. 계산 위치는 transport의 끝 지점(u-prediction) 또는 trajectory의 tangent(v-prediction)입니다. instantaneous velocity는 network 자체에서 복원하거나(MeanFlow 방식), auxiliary network v_ϕ로 얻습니다(DMD 방식).

이 framework로 FewStepGRPO, MeanFlowAWM/DMDAWM, MeanFlowNFT/DMDNFT, FewStepOPSD, ReFL을 하나의 modular codebase에 구현했습니다. 학습 step은 rollout, reward/advantage, anchors/weights, loss의 stage hook으로 나뉘어 있습니다.

실험 설정은 다음과 같습니다.

• Image: AnyFlow-distilled SD3.5-Medium, 4-step, LoRA rank 32, 8× H20
• Video: AnyFlow-distilled Wan2.1-T2V-1.3B

주요 결과는 다음과 같습니다.

• Adaptation 효과: mixed reward(CLIPScore + HPSv2 + PickScore) 설정에서 FlowGRPO는 held-out HPSv3를 10.30에서 9.62로 떨어뜨렸습니다. 반면 FewStepGRPO는 11.58로 높였습니다. MeanFlowNFT는 DiffusionNFT baseline보다 training reward(1.509 vs 1.440)와 HPSv3(13.23 vs 11.79)가 모두 높았습니다.
• Reward-gradient methods 우위 (image): FewStepOPSD와 ReFL은 training reward를 1.450에서 각각 1.641, 1.623까지 올렸습니다. group-advantage 계열 5개 method는 1.461.52에 머물렀습니다. best GenEval2도 각각 0.463, 0.465로 가장 높았습니다. iteration당 비용도 268s, 219s로 가장 낮았는데, 다른 method는 304435s였습니다.
• OCR setting (non-differentiable reward): MeanFlowNFT는 OCR accuracy를 0.498에서 iteration 100 만에 0.85까지 올렸고, 최고치는 약 0.945였습니다. DMDAWM은 0.936에 도달했지만 HPSv3가 12.86에서 5.05로 떨어지는 등 held-out 지표의 degradation이 컸습니다.
• Video: group-advantage 계열에서는 DMDNFT가 training reward를 16.54에서 17.93으로 가장 크게 올렸습니다. 하지만 VBench overall은 83.38로 7개 method 중 최하위였습니다. FewStepGRPO는 15.89로 initialization보다 오히려 낮아졌습니다. video에서는 reward gradient 계산 비용 때문에 FewStepOPSD와 ReFL이 GPU-hours 기준으로 가장 비쌌고, DMD 계열도 MeanFlow 계열보다 느렸습니다(821s vs 734~737s).
• Test-time scaling: 학습 때보다 많은 step(8/16/32)으로 sampling해도 aggregate 지표가 체계적으로 좋아지지는 않았습니다. 일부 선택된 output에서만 2~4배 step으로 시각적 개선이 보였습니다.

4. Conclusion & Impact (결론 및 시사점)

이 논문은 x-space 변환과 decoupled loss decomposition을 통해 다양한 RL objective를 "anchors와 weights의 선택"이라는 하나의 관점으로 통합했습니다. 이를 바탕으로 few-step model에 체계적으로 적용하는 방법을 제시했습니다. 실무 지침은 다음과 같이 정리됩니다.

• reward가 differentiable하고 gradient 계산 비용이 감당 가능하면 reward-gradient method가 가장 효과적이며, ReFL 단독으로도 강력한 baseline입니다.
• 그렇지 않은 경우에는 DiffusionNFT-style method가 비교적 적은 iteration으로 높은 reward에 도달합니다.
• DMD 기반 variant의 속도 우위는 forward/backward 비용 비율에 따라 달라집니다(image에서는 빠르고 video에서는 느림). 또 held-out degradation이 큰 경우가 많았습니다.

이 연구는 rollout 비용이 큰 video와 고해상도 생성에서 few-step model 기반 RL post-training의 실용성을 높입니다. 통합 codebase를 공개했기 때문에 새 RL variant를 anchors와 weights 설계만으로 빠르게 정의하고 공정하게 비교할 수 있는 기반도 마련했습니다. 다만 reward hacking 방지와 test-time scaling은 여전히 열린 문제로 남아 있습니다.

⚠️ 알림: 이 리뷰는 AI로 작성되었습니다.