FP4 탐색, BF16 학습: 효율적인 롤아웃 확장을 통한 확산 강화 학습

초록

강화 학습 기반 사후 학습은 최근 텍스트-이미지 확산 모델을 인간 선호도에 정렬시키는 유망한 패러다임으로 부상하고 있습니다. 최근 연구에 따르면 롤아웃 그룹 크기를 증가시키면 성능이 현저히 개선되어, 추가 정렬 효과를 얻을 수 있는 상당한 여지가 있음이 나타났습니다. 그러나 대규모 기초 확산 모델(예: FLUX.1-12B)에서 롤아웃 규모를 확장하는 것은 막대한 계산 부담을 초래합니다. 이러한 병목 현상을 완화하기 위해 우리는 Diffusion RL 롤아웃에 FP4 양자화 통합을 탐구합니다. 그러나 순진한 양자화 파이프라인은 본질적으로 성능 저하 위험을 수반함을 확인했습니다. 효율성과 학습 무결성 사이의 딜레마를 극복하기 위해 우리는 새로운 FP4 기반 2단계 강화 학습 프레임워크인 Sol-RL(Speed-of-light RL)을 제안합니다. 첫째, 높은 처리량의 NVFP4 롤아웃을 활용하여 대규모 후보 풀을 생성하고 높은 대비를 가진 하위 집합을 추출합니다. 둘째, 선택된 샘플들을 BF16 정밀도로 재생성하고 이를 배타적으로 사용하여 정책을 최적화합니다. 후보 탐색과 정책 최적화를 분리함으로써 Sol-RL은 롤아웃 스케일링의 알고리즘 메커니즘과 NVFP4의 시스템 수준 처리량 이점을 통합합니다. 이 상호 보완적인 알고리즘-하드웨어 설계는 롤아웃 단계를 효과적으로 가속화하면서 최적화를 위한 고품질 샘플을 확보합니다. 우리는 본 프레임워크가 BF16 정밀도 파이프라인의 학습 무결성을 유지하면서 FP4 연산이 가능하게 하는 처리량 이점을 완전히 활용함을 실증적으로 입증합니다. SANA, FLUX.1, SD3.5-L에 대한 광범위한 실험을 통해 우리의 접근 방식이 여러 메트릭에 걸쳐 우수한 정렬 성능을 제공하는 동시에 훈련 수렴 속도를 최대 4.64배까지 가속화하여, 훨씬 낮은 비용으로 대규모 롤아웃 스케일링의 힘을 활용할 수 있음을 확인했습니다.

English

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.64times, unlocking the power of massive rollout scaling at a fraction of the cost.

FP4 탐색, BF16 학습: 효율적인 롤아웃 확장을 통한 확산 강화 학습

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

초록

Support