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

텍스트-이미지 확산 모델을 인간의 선호도에 맞추는 강화 학습 과정에서 대규모 샘플 생성(Rollout)은 막대한 비용이 든다. 이 논문은 저정밀도 FP4 연산으로 후보를 빠르게 탐색하고 고정밀도 BF16으로 핵심 샘플만 다시 생성해 학습 효율과 품질을 동시에 잡는 새로운 패러다임을 제시한다.

확산 모델의 강화 학습은 모델이 생성한 여러 이미지 중 좋은 것을 골라 학습하는 과정을 반복한다. 이때 더 좋은 이미지를 찾으려면 한 번에 수백 장의 이미지를 뽑아봐야(Rollout Scaling) 하는데, FLUX.1 같은 거대 모델에서는 이 과정이 너무 느려 학습의 병목 현상이 된다.

이 문제를 해결하기 위해 연구진은 '대충 빠르게 훑어보고 중요한 것만 제대로 다시 보기' 전략을 취한다. 4비트 부동소수점(FP4)이라는 매우 낮은 정밀도를 사용하면 연산 속도는 4배 빨라지지만 이미지 품질은 떨어진다. 하지만 연구진은 FP4로 만든 '못생긴 이미지'라도 그들 사이의 상대적인 등수(Ranking)는 고정밀도 이미지와 거의 일치한다는 사실을 발견했다.

결국 FP4로 96장의 후보를 순식간에 뽑아 등수를 매긴 뒤, 그중 가장 점수가 높은 것과 낮은 것 24장만 골라낸다. 이 24장만 원래의 고정밀도(BF16)로 다시 깨끗하게 생성하여 학습에 사용함으로써, 전체 연산량은 줄이면서도 대규모 샘플링이 주는 학습 효과는 그대로 누릴 수 있게 된다.

Sol-RL은 크게 두 단계의 파이프라인으로 구성된다. 첫 번째 단계인 'FP4 Exploration'에서는 NVIDIA Transformer Engine을 통해 모델 가중치를 NVFP4로 양자화하고, 단 6단계의 적은 추론 스텝으로 96개의 후보 샘플을 생성한다. 이때 생성된 샘플들의 보상 점수를 계산하여 상위 K/2개와 하위 K/2개의 노이즈 시드(Seed)를 추출한다.

두 번째 단계인 'BF16 Re-generation'에서는 선택된 K개의 시드만을 사용하여 BF16 정밀도와 표준 추론 스텝(10단계 이상)으로 고품질 이미지를 재생성한다. 이후 DiffusionNFT 또는 GRPO 목적 함수를 사용하여 정책 네트워크를 업데이트한다. 업데이트된 가중치는 다시 NVFP4로 양자화되어 다음 반복(Iteration)의 탐색 단계에 투입된다.

수학적으로는 FP4 양자화 오차를 유계 섭동(Bounded Perturbation)으로 정의하고, 극값 이론(Extreme Value Theory)을 통해 샘플링 규모 N이 커질수록 FP4 탐색으로 얻는 이득이 양자화 노이즈로 인한 손실을 압도함을 증명했다. 구체적으로 기대 보상 범위 E[W]가 2σ√(2 log N) - 4Δ 이상임을 보여, N을 늘릴수록 강력한 학습 신호를 확보할 수 있음을 이론적으로 뒷받침했다.

FLUX.1 모델 실험에서 Sol-RL은 BF16 기반의 나이브 스케일링 대비 롤아웃 시간을 2.33배 단축했으며, 전체 학습 시간은 1.62배 가속화했다. SD3.5-Large에서는 롤아웃 2.41배, 전체 시간 1.61배의 가속을 달성했다. 특히 동일한 GPU 시간 예산 내에서 DiffusionNFT 대비 ImageReward 점수 수렴 속도가 최대 4.64배 빨랐다.

정성적 평가에서도 Sol-RL로 학습된 모델은 프롬프트에 대한 의미론적 정렬(Semantic Alignment)이 더 뛰어났으며, 세부 묘사와 예술적 일관성 측면에서 베이스 모델 및 기존 강화 학습 방법론들을 압도했다. Ablation Study 결과, FP4 탐색 스텝을 6단계로 설정했을 때 보상 순위의 신뢰도가 포화 상태에 도달하여 효율성이 극대화됨을 확인했다.

Sol-RL은 알고리즘과 하드웨어의 시너지를 극대화한 설계가 특징이다. NVFP4(NVIDIA 4-bit Floating Point) 포맷을 사용하여 BF16 대비 4배 높은 TFLOPs 처리량을 확보했다. 특히 탐색 단계에서 ODE Solver의 스텝 수를 줄여도 초기 노이즈가 전체적인 구조와 보상 수준을 결정한다는 '결정론적 샘플링의 특성'을 활용해 탐색 효율을 높였다.

학습 안정성을 위해 정책 업데이트는 오직 BF16 재생성 샘플에서만 수행함으로써, 저정밀도 샘플이 직접 그래디언트에 노출될 때 발생하는 분포 편향(Distribution Shift) 문제를 원천 차단했다. 또한 LoRA(rank=32)를 적용해 파라미터 효율성을 높였으며, 업데이트된 가중치를 실시간으로 재양자화하여 컴파일 오버헤드 없이 다음 롤아웃에 즉시 반영하는 최적화된 루프를 구현했다.

FP4 탐색의 신뢰도는 ODE Solver의 결정론적 특성에 의존하므로, 확률적 샘플링(SDE) 비중이 높은 환경에서는 순위 보존 능력이 저하될 가능성이 있다. 또한 NVFP4 하드웨어 가속이 지원되지 않는 구형 GPU 아키텍처에서는 본 논문에서 제시한 수준의 가속 효과를 기대하기 어렵다.