FP4로 탐색하고, BF16으로 학습: 효율적인 Rollout 확장 기반 확산 강화 학습

FP4 Explore, BF16 Train: Diffusion Reinforcement Learning via Efficient Rollout Scaling

최근 강화 학습 기반의 후속 학습은 텍스트-이미지 확산 모델을 인간의 선호도에 맞추는 유망한 방법으로 부상했습니다. 최근 연구에서 Rollout 그룹 크기를 늘리면 성능이 크게 향상되어 추가적인 정렬 개선의 여지가 충분히 있음을 보여줍니다. 그러나 대규모 기본 확산 모델(예: FLUX.1-12B)에서 Rollout을 확장하면 상당한 계산 부담이 발생합니다. 이러한 병목 현상을 완화하기 위해, 우리는 확산 강화 학습 Rollout에 FP4 양자화를 통합하는 방법을 탐구합니다. 그러나, 단순하게 양자화된 파이프라인은 성능 저하의 위험을 내포하고 있음을 확인했습니다. 효율성과 학습의 무결성 사이의 이러한 딜레마를 해결하기 위해, 우리는 FP4 기술을 활용한 새로운 2단계 강화 학습 프레임워크인 Sol-RL (Speed-of-light RL)을 제안합니다. 먼저, 고처리량 NVFP4 Rollout을 사용하여 대규모 후보 풀을 생성하고, 매우 차별적인 부분집합을 추출합니다. 두 번째로, 선택된 샘플을 BF16 정밀도로 다시 생성하고, 정책을 오직 이 샘플에 대해서만 최적화합니다. Sol-RL은 후보 탐색과 정책 최적화를 분리함으로써, Rollout 확장의 알고리즘적 메커니즘과 NVFP4의 시스템 수준 처리량 향상을 결합합니다. 이러한 시너지 효과를 내는 알고리즘-하드웨어 설계는 Rollout 단계를 효과적으로 가속화하면서도, 고품질 샘플을 최적화를 위해 보존합니다. 우리의 프레임워크가 BF16 정밀도 파이프라인의 학습 무결성을 유지하면서 FP4 연산이 제공하는 처리량 향상을 최대한 활용한다는 것을 경험적으로 입증했습니다. SANA, FLUX.1 및 SD3.5-L에 대한 광범위한 실험 결과는, 우리의 접근 방식이 여러 지표에서 우수한 정렬 성능을 제공하며, 최대 4.64배까지 학습 수렴 속도를 가속화하여, 막대한 Rollout 확장의 잠재력을 훨씬 낮은 비용으로 실현할 수 있음을 보여줍니다.

Reinforcement-Learning-based post-training has recently emerged as a promising paradigm for aligning text-to-image diffusion models with human preferences. In recent studies, increasing the rollout group size yields pronounced performance improvements, indicating substantial room for further alignment gains. However, scaling rollouts on large-scale foundational diffusion models (e.g., FLUX.1-12B) imposes a heavy computational burden. To alleviate this bottleneck, we explore the integration of FP4 quantization into Diffusion RL rollouts. Yet, we identify that naive quantized pipelines inherently introduce risks of performance degradation. To overcome this dilemma between efficiency and training integrity, we propose Sol-RL (Speed-of-light RL), a novel FP4-empowered Two-stage Reinforcement Learning framework. First, we utilize high-throughput NVFP4 rollouts to generate a massive candidate pool and extract a highly contrastive subset. Second, we regenerate these selected samples in BF16 precision and optimize the policy exclusively on them. By decoupling candidate exploration from policy optimization, Sol-RL integrates the algorithmic mechanisms of rollout scaling with the system-level throughput gains of NVFP4. This synergistic algorithm-hardware design effectively accelerates the rollout phase while reserving high-fidelity samples for optimization. We empirically demonstrate that our framework maintains the training integrity of BF16 precision pipeline while fully exploiting the throughput gains enabled by FP4 arithmetic. Extensive experiments across SANA, FLUX.1, and SD3.5-L substantiate that our approach delivers superior alignment performance across multiple metrics while accelerating training convergence by up to $4.64\times$, unlocking the power of massive rollout scaling at a fraction of the cost.