提速4.6倍！NVIDIA×港大×MIT聯(lián)手重新定義擴散模型訓(xùn)練速度上限

當(dāng)強化學(xué)習(xí)后訓(xùn)練的大規(guī)模 rollout 已經(jīng)被證明能夠提升圖像生成模型的偏好對齊能力，推理負(fù)擔(dān)就成了制約訓(xùn)練速度的核心瓶頸。來自 NVIDIA、港大和 MIT 的團隊提出的 Sol-RL，通過「FP4 先探索、BF16 再訓(xùn)練」的后訓(xùn)練框架，將達(dá)到等效 reward 水平的收斂速度最高提升到 4.64x，在訓(xùn)練速度與對齊效果之間給出了一條更具工程可行性的解法。

Sol-RL（Speed-of-light RL）是一種將 NVFP4 推理融合進(jìn) Diffusion 強化學(xué)習(xí)微調(diào)的高效訓(xùn)練框架。該方法并不是用量化推理結(jié)果直接訓(xùn)練，而是讓 NVFP4 rollout 承擔(dān)高吞吐的大規(guī)模探索任務(wù)，先在海量候選中篩出最有對比性的樣本的初始噪聲集合，再讓 BF16 對這些關(guān)鍵初始噪聲進(jìn)行高保真再生成并完成策略優(yōu)化。實驗結(jié)果表明，該設(shè)計在 SANA、FLUX.1 和 SD3.5-L 模型上都帶來了明顯收益：在相同 GPU-hour 預(yù)算下，達(dá)到等效 reward 水平的收斂速度最高提升 4.64x，同時基本保持了 BF16 高精度 pipeline 的訓(xùn)練保真度。本文將深入探討 Sol-RL 的核心思路、方法設(shè)計、實驗結(jié)果與實際意義。

• 論文名稱：FP4 Explore, BF16 Train: Diffusion Reinforcement Learning via Efficient Rollout Scaling
• 論文地址：https://arxiv.org/abs/2604.06916
• 項目主頁：https://nvlabs.github.io/Sana/Sol-RL/
• 代碼鏈接：https://github.com/NVlabs/Sana/

效果展示

Sol-RL 在速度和性能上表現(xiàn)出色，下圖展示了經(jīng)過 Sol-RL 優(yōu)化后 SANA 模型的整體效果，以及與基礎(chǔ)模型相比在復(fù)雜細(xì)節(jié)、語義一致性和整體審美上的改善。

研究背景

在文生圖模型的后訓(xùn)練階段，強化學(xué)習(xí)正在成為提升人類偏好對齊能力的重要路徑。現(xiàn)有研究不斷表明，擴大 rollout，也就是為同一個 prompt 生成更多候選圖像，再從中篩選高對比樣本（例如其中表現(xiàn)最好的 n 個和最差的 n 個構(gòu)成的 2n 個圖像）進(jìn)行優(yōu)化，能夠顯著提升模型的對齊效果。對于基于 GRPO 的 Diffusion RL 來說，更多候選意味著更強的對比信號，也意味著更可靠的梯度更新。

但問題在于：rollout 一旦增大，訓(xùn)練瓶頸很快就會從參數(shù)更新轉(zhuǎn)移到海量候選樣本生成。尤其在 FLUX.1、SD3.5-L 這類參數(shù)量較大、需要多次迭代推理的文生圖 Diffusion 模型上，想真正把 rollout scaling 做起來，rollout 階段生成候選樣本的計算成本會迅速上升。為了提高海量 rollout 效率，團隊引入 nvfp4 量化推理，但論文進(jìn)一步指出，如果直接把低比特 rollout 樣本當(dāng)作訓(xùn)練目標(biāo)，訓(xùn)練穩(wěn)定性和最終效果都會受到明顯影響。因此，關(guān)鍵并不只是「能不能用 FP4」，而是「FP4 應(yīng)該在訓(xùn)練流程里扮演什么角色」。

核心創(chuàng)新

Sol-RL 的核心創(chuàng)新在于其雙階段解耦框架。與直接把低精度計算貫穿整個訓(xùn)練流程不同，Sol-RL 把 rollout 樣本的「探索」和「生成」拆開處理，讓不同精度承擔(dān)不同職責(zé)。

量化 rollout 不適合直接訓(xùn)練：論文首先指出，直接將 FP4 量化推理樣本用作優(yōu)化目標(biāo)，會帶來訓(xùn)練退化和不穩(wěn)定性的隱患，如下圖實驗中紅色曲線所示。

FP4 適合作為探索代理：盡管 FP4 樣本在像素層面會引入偏差，但研究發(fā)現(xiàn)，給定相同的初始噪聲，在同組候選樣本的相對 reward 排序上，NVFP4 推理樣本依然保持了和 BF16 推理樣本足夠高的一致性。因此，它非常適合用來承擔(dān)「大規(guī)模探索」和「候選篩選」的任務(wù)，也就是在大量初始噪聲中快速識別出哪些會產(chǎn)生「最好」或者「最壞」的最終圖像。

方法概述

Sol-RL 的整體流程可以概括為兩個階段。第一階段，框架使用 NVFP4 rollout 和更少的采樣步數(shù)快速生成一個大規(guī)模候選池，并根據(jù) reward 對候選樣本做排序，篩選出對應(yīng)圖像得分最高 / 最低的初始噪聲種子集合。第二階段，框架并不會直接拿這些低精度樣本做訓(xùn)練，而是只保留第一階段篩出的關(guān)鍵種子，再用 BF16 精度重新生成高保真樣本，并僅基于這些高保真樣本完成策略優(yōu)化。這樣一來，F(xiàn)P4 負(fù)責(zé)快速找方向，BF16 負(fù)責(zé)對其中一小部分高對比度樣本做高質(zhì)量生成用于訓(xùn)練，效率和穩(wěn)定性被重新組織到同一個框架中。這樣的流程讓高成本 BF16 計算只集中在真正會影響梯度更新的部分，而不是平均浪費在大量最終不會參與訓(xùn)練的候選圖像上。

實驗結(jié)果

實驗結(jié)果表明，Sol-RL 在多個基礎(chǔ)模型和多個 reward 指標(biāo)上都展現(xiàn)出明顯優(yōu)勢。在相同 GPU-hour 預(yù)算下，Sol-RL 在 SANA、FLUX.1 和 SD3.5-L 上持續(xù)優(yōu)于基線方法，并將達(dá)到等效 reward 水平的收斂速度最高提升至 4.64x。如下圖所示，Sol-RL 在相同 wall-clock 預(yù)算下能夠更快達(dá)到基線性能，在有限時間內(nèi)達(dá)到更高對齊質(zhì)量。

進(jìn)一步做時間拆解，相對于直接使用高精度進(jìn)行 rollout scaling，Sol-RL 在 rollout 階段的加速最高達(dá)到 2.41x，訓(xùn)練迭代時間最高提升 1.62x。Sol-RL 通過探索 - 重生成的兩階段設(shè)計，顯著緩解了完全使用 BF16 rollout scaling 帶來的計算瓶頸，并且避免了直接使用低精度樣本進(jìn)行訓(xùn)練帶來的不穩(wěn)定性，相對于全程 NVFP4 低精度推理版本僅帶來約為 2% 的額外開銷。

結(jié)論與展望

當(dāng)大規(guī)模 rollout 已經(jīng)被證明能夠持續(xù)提升生成模型偏好對齊能力，接下來的關(guān)鍵問題就是如何以更低成本釋放這種擴展帶來的收益。Sol-RL 給出的答案是：讓低精度負(fù)責(zé)探索，讓高精度負(fù)責(zé)優(yōu)化。這也意味著，F(xiàn)P4 在生成式模型后訓(xùn)練中的角色被重新定義了。它不再只是一個推理加速工具，而是開始成為強化學(xué)習(xí)探索階段的有效代理。對于文生圖后訓(xùn)練、偏好對齊、低比特量化以及系統(tǒng)級優(yōu)化方向的研究者和工程團隊來說，這個方案為生成模型后訓(xùn)練提供了一條更現(xiàn)實的落地路徑，具備持續(xù)關(guān)注的價值。