OpenAI最近推出了在大語言模型LLM上的強化微調(diào)（Reinforcement Finetuning，ReFT），能夠讓模型利用CoT進(jìn)行多步推理之后，通過強化學(xué)習(xí)讓最終輸出符合人類偏好。

無獨有偶，齊國君教授領(lǐng)導(dǎo)的MAPLE實驗室在OpenAI發(fā)布會一周前公布的工作中也發(fā)現(xiàn)了圖像生成領(lǐng)域的主打方法擴(kuò)散模型和流模型中也存在類似的過程：模型從高斯噪聲開始的多步去噪過程也類似一個思維鏈，逐步「思考」怎樣生成一張高質(zhì)量圖像，是一種圖像生成領(lǐng)域的「圖像鏈CoT」。

與OpenAI不謀而和的是，機器學(xué)習(xí)與感知（MAPLE）實驗室認(rèn)為強化學(xué)習(xí)微調(diào)方法同樣可以用于優(yōu)化多步去噪的圖像生成過程，論文指出利用與人類獎勵對齊的強化學(xué)習(xí)監(jiān)督訓(xùn)練，能夠讓擴(kuò)散模型和流匹配模型自適應(yīng)地調(diào)整推理過程中噪聲強度，用更少的步數(shù)生成高質(zhì)量圖像內(nèi)容。

圖片

論文地址：https://arxiv.org/abs/2412.01243

研究背景

擴(kuò)散和流匹配模型是當(dāng)前主流的圖像生成模型，從標(biāo)準(zhǔn)高斯分布中采樣的噪聲逐步變換為一張高質(zhì)量圖像。在訓(xùn)練時，這些模型會單獨監(jiān)督每一個去噪步驟，使其具備能恢復(fù)原始圖像的能力；而在實際推理時，模型則會事先指定若干個不同的擴(kuò)散時間，然后在這些時間上依次執(zhí)行多步去噪過程。

這一過程存在兩個問題：

1. 經(jīng)典的擴(kuò)散模型訓(xùn)練方法只能保證每一步去噪能盡可能恢復(fù)出原始圖像，不能保證整個去噪過程得到的圖像符合人類的偏好；

2. 經(jīng)典的擴(kuò)散模型所有的圖片都采用了同樣的去噪策略和步數(shù)；而顯然不同復(fù)雜度的圖像對于人類來說生成難度是不一樣的。

如下圖所示，當(dāng)輸入不同長度的prompt的時候，對應(yīng)的生成任務(wù)難度自然有所區(qū)別。那些僅包含簡單的單個主體前景的圖像較為簡單，只需要少量幾步就能生成不錯的效果，而帶有精細(xì)細(xì)節(jié)的圖像則需要更多步數(shù)，即經(jīng)過強化微調(diào)訓(xùn)練后的圖像生成模型就能自適應(yīng)地推理模型去噪過程，用盡可能少的步數(shù)生成更高質(zhì)量的圖像。

值得注意的是，類似于LLM對思維鏈進(jìn)行的動態(tài)優(yōu)化，對擴(kuò)散模型時間進(jìn)行優(yōu)化的時候也需要動態(tài)地進(jìn)行，而非僅僅依據(jù)輸入的prompt；換言之，優(yōu)化過程需要根據(jù)推理過程生成的「圖像鏈」來動態(tài)一步步預(yù)測圖像鏈下一步的最優(yōu)去噪時間，從而保證圖像的生成質(zhì)量滿足reward指標(biāo)。

方法

MAPLE實驗室認(rèn)為，要想讓模型在推理時用更少的步數(shù)生成更高質(zhì)量的圖像結(jié)果，需要用強化微調(diào)技術(shù)對多步去噪過程進(jìn)行整體監(jiān)督訓(xùn)練。既然圖像生成過程同樣也類似于LLM中的CoT：模型通過中間的去噪步驟「思考」生成圖像的內(nèi)容，并在最后一個去噪步驟給出高質(zhì)量的結(jié)果，也可以通過利用獎勵模型評價整個過程生成的圖像質(zhì)量，通過強化微調(diào)使模型的輸出更符合人類偏好。

圖片

OpenAI的O1通過在輸出最終結(jié)果之前生成額外的token讓LLM能進(jìn)行額外的思考和推理，模型所需要做的最基本的決策是生成下一個token；而擴(kuò)散和流匹配模型的「思考」過程則是在生成最終圖像前，在不同噪聲強度對應(yīng)的擴(kuò)散時間（diffusion time）執(zhí)行多個額外的去噪步驟。為此，模型需要知道額外的「思考」步驟應(yīng)該在反向擴(kuò)散過程推進(jìn)到哪一個diffusion time的時候進(jìn)行。

為了實現(xiàn)這一目的，在網(wǎng)絡(luò)中引入了一個即插即用的時間預(yù)測模塊（Time Prediction Module, TPM）。這一模塊會預(yù)測在當(dāng)前這一個去噪步驟執(zhí)行完畢之后，模型應(yīng)當(dāng)在哪一個diffusion time下進(jìn)行下一步去噪。

具體而言，該模塊會同時取出去噪網(wǎng)絡(luò)第一層和最后一層的圖像特征，預(yù)測下一個去噪步驟時的噪聲強度會下降多少。模型的輸出策略是一個參數(shù)化的beta分布。

由于單峰的Beta分布要求α>1且β>1，研究人員對輸出進(jìn)行了重參數(shù)化，使其預(yù)測兩個實數(shù)a和b，并通過如下公式確定對應(yīng)的Beta分布，并采樣下一步的擴(kuò)散時間。

圖片

圖片

在強化微調(diào)的訓(xùn)練過程中，模型會在每一步按輸出的Beta分布隨機采樣下一個擴(kuò)散時間，并在對應(yīng)時間執(zhí)行下一個去噪步驟。直到擴(kuò)散時間非常接近0時，可以認(rèn)為此時模型已經(jīng)可以近乎得到了干凈圖像，便終止去噪過程并輸出最終圖像結(jié)果。

通過上述過程，即可采樣到用于強化微調(diào)訓(xùn)練的一個決策軌跡樣本。而在推理過程中，模型會在每一個去噪步驟輸出的Beta分布中直接采樣眾數(shù)作為下一步對應(yīng)的擴(kuò)散時間，以確保一個確定性的推理策略。

設(shè)計獎勵函數(shù)時，為了鼓勵模型用更少的步數(shù)生成高質(zhì)量圖像，在獎勵中綜合考慮了生成圖像質(zhì)量和去噪步數(shù)這兩個因素，研究人員選用了與人類偏好對齊的圖像評分模型ImageReward（IR）用以評價圖像質(zhì)量，并將這一獎勵隨步數(shù)衰減至之前的去噪結(jié)果，并取平均作為整個去噪過程的獎勵。這樣，生成所用的步數(shù)越多，最終獎勵就越低。模型會在保持圖像質(zhì)量的前提下，盡可能地減少生成步數(shù)。

圖片

將整個多步去噪過程當(dāng)作一個動作進(jìn)行整體優(yōu)化，并采用了無需值模型的強化學(xué)習(xí)優(yōu)化算法RLOO [1]更新TPM模塊參數(shù)，訓(xùn)練損失如下所示：

圖片

在這一公式中，s代表強化學(xué)習(xí)中的狀態(tài)，在擴(kuò)散模型的強化微調(diào)中是輸入的文本提詞和初始噪聲；y代表決策動作，也即模型采樣的擴(kuò)散時間；

代表決策器，即網(wǎng)絡(luò)中A是由獎勵歸一化之后的優(yōu)勢函數(shù)，采用LEAVE-One-Out策略，基于一個Batch內(nèi)的樣本間獎勵的差值計算優(yōu)勢函數(shù)。

通過強化微調(diào)訓(xùn)練，模型能根據(jù)輸入圖像自適應(yīng)地調(diào)節(jié)擴(kuò)散時間的衰減速度，在面對不同的生成任務(wù)時推理不同數(shù)量的去噪步數(shù)。對于簡單的生成任務(wù)（較短的文本提詞、生成圖像物體少），推理過程能夠很快生成高質(zhì)量的圖像，噪聲強度衰減較快，模型只需要思考較少的額外步數(shù)，就能得到滿意的結(jié)果；對于復(fù)雜的生成任務(wù)（長文本提詞，圖像結(jié)構(gòu)復(fù)雜）則需要在擴(kuò)散時間上密集地進(jìn)行多步思考，用一個較長的圖像鏈COT來生成符合用戶要求的圖片。

圖片

通過調(diào)節(jié)不同的γ值，模型能在圖像生成質(zhì)量和去噪推理的步數(shù)之間取得更好的平衡，僅需要更少的平均步數(shù)就能達(dá)到與原模型相同的性能。

圖片

同時，強化微調(diào)的訓(xùn)練效率也十分驚人。正如OpenAI最少僅僅用幾十個例子就能讓LLM學(xué)會在自定義領(lǐng)域中推理一樣，強化微調(diào)圖像生成模型對數(shù)據(jù)的需求也很少。不需要真實圖像，只需要文本提詞就可以訓(xùn)練，利用不到10,000條文本提詞就能取得不錯的明顯的模型提升。

經(jīng)強化微調(diào)后，模型的圖像生成質(zhì)量也比原模型提高了很多。可以看出，在僅僅用了原模型一半生成步數(shù)的情況下，無論是圖C中的筆記本鍵盤，圖D中的球棒還是圖F中的遙控器，該模型生成的結(jié)果都比原模型更加自然。

圖片

針對Stable Diffusion 3、Flux-dev等一系列最先進(jìn)的開源圖像生成模型進(jìn)行了強化微調(diào)訓(xùn)練，發(fā)現(xiàn)訓(xùn)練后的模型普遍能減少平均約50%的模型推理步數(shù)，而圖像質(zhì)量評價指標(biāo)總體保持不變，這說明對于圖像生成模型而言，強化微調(diào)訓(xùn)練是一種通用的后訓(xùn)練（Post Training）方法。

圖片

結(jié)論

這篇報告介紹了由MAPLE實驗室提出的，一種擴(kuò)散和流匹配模型的強化微調(diào)方法。該方法將多步去噪的圖像生成過程看作圖像生成領(lǐng)域的COT過程，通過將整個去噪過程的最終輸出與人類偏好對齊，實現(xiàn)了用更少的推理步數(shù)生成更高質(zhì)量圖像。

在多個開源圖像生成模型上的實驗結(jié)果表明，這種強化微調(diào)方法能在保持圖像質(zhì)量的同時顯著減少約50%推理步數(shù)，微調(diào)后模型生成的圖像在視覺效果上也更加自然。可以看出，強化微調(diào)技術(shù)在圖像生成模型中仍有進(jìn)一步應(yīng)用和提升的潛力，值得進(jìn)一步挖掘。

參考資料：

https://arxiv.org/abs/2412.01243