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

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

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

• 論文名稱：FP4 Explore, BF16 Train: Diffusion Reinforcement Learning via Efficient Rollout Scaling
• 論文地址：https://arxiv.org/abs/2604.06916
• 項(xiàng)目主頁：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)練階段，強(qiáng)化學(xué)習(xí)正在成為提升人類偏好對(duì)齊能力的重要路徑。現(xiàn)有研究不斷表明，擴(kuò)大 rollout，也就是為同一個(gè) prompt 生成更多候選圖像，再從中篩選高對(duì)比樣本（例如其中表現(xiàn)最好的 n 個(gè)和最差的 n 個(gè)構(gòu)成的 2n 個(gè)圖像）進(jìn)行優(yōu)化，能夠顯著提升模型的對(duì)齊效果。對(duì)于基于 GRPO 的 Diffusion RL 來說，更多候選意味著更強(qiáng)的對(duì)比信號(hào)，也意味著更可靠的梯度更新。

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

核心創(chuàng)新

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

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

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

方法概述

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

實(shí)驗(yàn)結(jié)果

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

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

結(jié)論與展望

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