AIxiv專欄是機器之心釋出學術、技術內容的欄目。過去數年，機器之心AIxiv專欄接收報導了2000多篇內容，覆蓋全球各大高校與企業的頂級實驗室，有效促進了學術交流與傳播。如果您有優秀的工作想要分享，歡迎投稿或者聯絡報導。投稿郵箱：liyazhou@jiqizhixin.com；zhaoyunfeng@jiqizhixin.com
本論文作者主要包括澳大利亞科學院院士、歐洲科學院外籍院士、IEEE Fellow陶大程，現任南洋理工大學傑出教授；新加坡工程院院士、IEEE Fellow文勇剛，現為南洋理工大學電腦科學與工程學院校長講席教授，同時擔任IEEE Transactions on Multimedia主編；張森，曾在悉尼大學從事博士後研究工作，現任TikTok機器學習工程師；詹憶冰，京東探索研究院演算法科學家。本文的通訊作者是武漢大學計算機學院教授、博士生導師、國家特聘青年專家羅勇。第一作者為張子屹，目前在武漢大學計算機學院攻讀博士二年級，研究方向為強化學習、擴散模型和大模型對齊。
OpenAI 最近釋出的 o1 模型在數學、程式碼生成和長程規劃等複雜任務上取得了突破性進展，據業內人士分析披露，其關鍵技術在於基於強化學習的搜尋與學習機制。透過迭代式的自舉過程，o1 基於現有大語言模型的強大推理能力，生成合理的推理過程，並將這些推理融入到其強化學習訓練過程中。強化學習技術能夠顯著提升模型解決複雜問題的能力，尤其是在 o1 所採用的細粒度獎勵機制的加持下。這種獎勵機制為模型的每一步推理提供細粒度的反饋，而不僅是依賴最終答案的正確性來評估模型的表現。透過精細化的控制，使模型能夠不斷最佳化其推理路徑，有效應對複雜任務中的長程依賴問題，確保模型在推理過程中保持邏輯一致性。前不久，來自武漢大學、悉尼大學、京東探索研究院和南洋理工大學的研究人員便在 ICML 2024 上發表了一項 “擴散模型對齊” 方向的研究，所提出的名為 TDPO-R 的強化學習演算法與 o1 不謀而合地也採用了細粒度獎勵機制。- 論文連結：https://openreview.net/forum?id=v2o9rRJcEv
- 程式碼連結：https://github.com/ZiyiZhang27/tdpo
TDPO-R 在強化學習演算法中引入了時間差分獎勵機制，對文生圖擴散模型的每一步生成過程提供細粒度的反饋，從而有效緩解了在擴散模型對齊時常見的獎勵過最佳化問題。這項研究證實了細粒度獎勵機制在擴散模型對齊中的關鍵性，而 o1 的最新技術同樣揭示了這一機制在大模型領域中的廣泛應用前景，有望推動生成模型在多樣化、複雜任務中的持續發展與最佳化。擴散模型（Diffusion Model）作為一種先進的生成式模型，透過學習並模擬自然界中的擴散過程來合成新資料，尤其在生成高質量逼真影像方面相較於其他模型架構具有顯著優勢。因此，儘管近年來也有基於 GPT 等其它架構的優秀影像、影片生成模型不斷湧現，基於擴散模型的架構依然是當前大多數影像、影片生成任務的主流選擇。近期，為了確保生成的影像能夠符合下游任務目標以及使用者偏好，許多研究者開始探索如何對預訓練擴散模型進行獎勵或偏好驅動的微調訓練 —— 這就是所謂的擴散模型對齊（Diffusion Model Alignment）。目前，最流行的對齊方法之一便是源於大語言模型領域的一項常用技術 —— 基於人類反饋的強化學習（Reinforcement Learning from Human Feedback, RLHF），即以強化學習方式依據人類反饋來不斷地對已有的預訓練模型引數進行迭代最佳化。另外，也有其它一些不涉及強化學習的對齊方法，但它們大多也都需要以特定獎勵函式作為最佳化目標，即將所對齊的下游任務目標（如影像美學質量、文圖一致性等）構建為一個獎勵函式 / 模型，並以最大化該獎勵函式為目標對擴散模型進行迭代最佳化。還有一些方法，比如直接偏好最佳化（Direct Preference Optimization，DPO），雖然不顯式地使用獎勵函式來對齊模型，但其所採用的訓練資料本身往往也隱含了某種獎勵標準，間接引導模型朝著某種期望的目標最佳化。由於受獎勵目標驅動，這些擴散模型對齊方法常常面臨一個核心挑戰 —— 獎勵過最佳化（Reward Overoptimization），即經過微調後的模型可能會過度偏向於某一獎勵目標，導致生成的影像喪失個性化和多樣性、視覺保真度降低，最終偏離人類真實的審美偏好。如上圖所示，當我們使用美學分數（Aesthetic Score）獎勵模型，對 Stable Diffusion v1.4 模型進行對齊微調時，“獎勵過最佳化” 這一現象在生成影像上肉眼可見為：影像風格趨於統一、與文字提示詞的匹配度下降。如何更加通俗地理解獎勵過最佳化問題？想象你是一位學生，目標是透過考試拿到高分。你的老師設計了一套嚴格的評分標準（就像是獎勵模型），目的是測試你的知識深度和理解能力。這套標準非常全面，涵蓋了大量的知識點和難題，理論上應該能夠準確衡量你的學習水平。然而，你發現了一些 “捷徑”：考試中雖然有很多問題，但也有著某些 “答題套路”，比如背誦標準答案，或者用某些關鍵詞和答題模板。於是你並不真正花時間理解知識，而是把大量時間用在研究考試技巧上，專注於如何根據評分標準去 “迎合” 考試。結果，你每次考試都能拿到高分（就像模型獲得了高獎勵），但實際上你對知識的理解很膚淺。老師看到你的分數，誤以為你學得非常好，但你自己知道，雖然分數高了，但你並沒有真正掌握知識，甚至可能過於執著答題模板而錯誤地理解了一些問題，換另一位老師出題的話便可能原形畢露。如何對獎勵過最佳化問題進行定量分析？TDPO-R 演算法的研究者發現，當擴散模型過度最佳化某一特定獎勵函式時，其在域外獎勵函式上的泛化能力隨之下降。這種現象可以透過定量分析表現為：在以某一獎勵函式為最佳化目標時，採用其他域外獎勵函式（即不同於最佳化目標的獎勵標準）對生成樣本進行評分，會發現隨著目標獎勵逐漸上升，域外獎勵評分的增長速度明顯放緩，甚至出現域外評分逐漸下降的情況。這表明，模型在一個特定獎勵函式上取得高分時，可能會以犧牲其在其他任務上的表現為代價，最終導致生成質量的不平衡。因此，研究者採用了這種跨獎勵泛化度量，來定量地對比不同演算法的獎勵過最佳化程度。研究者進一步指出，相較於基於監督學習的方法，基於強化學習的擴散模型對獎勵過最佳化問題表現出更高的抗性。這是因為強化學習能夠根據生成過程中的反饋訊號不斷調整策略，避免模型在單一目標上過度擬合。然而，這種抗性並非完全免疫。強化學習的這種優勢往往是以犧牲樣本效率為代價，即需要更多的樣本和訓練時間來達到較為平衡的效果。因此，在設計新的擴散模型對齊方法時，必須在提高樣本效率與緩解獎勵過最佳化之間做出權衡。受到深度強化學習理論的啟發，研究者首先從擴散模型對齊中的歸納偏置角度出發，分析了加重獎勵過最佳化的潛在因素。歸納偏置（Inductive Bias）指的是機器學習演算法在從有限的訓練資料中學習時，基於特定假設提煉出的一般性模式。在深度強化學習的背景下，演算法所採用的歸納偏置與其解決任務的契合程度，直接影響模型的泛化能力。簡單來說，如果模型的偏置過於集中在某個特定目標上，就容易在泛化到新任務時出現問題。然而，現有的擴散模型對齊方法往往過度依賴稀疏獎勵 —— 也就是隻根據最終生成影像的質量進行最佳化。這種做法隱含著一種假設，即最終影像的質量是最重要的評估標準，忽略了生成過程中的中間步驟（即每一步的去噪影像）。模型只在最後一步表現良好時被獎勵，導致訓練過程中沒有充分利用多步去噪過程中的中間資訊，而這部分資訊恰恰包含了對生成過程的更細粒度反饋。因此，這種忽略中間步驟的做法，實際上引入了一種錯位的歸納偏置 —— 即模型在訓練時過於注重終極輸出的獎勵訊號，而忽視了去噪過程中可能幫助模型提高質量的階段性訊號。這種偏置錯位便是加重獎勵過最佳化的潛在因素之一。為了更好理解這種偏置錯位，可以結合之前老師與學生的比喻：想象你是一名學生，你的老師只根據期末考試的成績來評價你整個學期的表現（類似於稀疏獎勵）。這種情況下，學生的學習過程（就像擴散模型中的中間去噪過程）沒有被重視或衡量，學生很容易找到捷徑（如透過背題庫而非理解知識）來獲得高分。這就類似於模型在訓練中透過某種捷徑，在最後生成的影像中獲得高分，而忽視了生成過程中的細節和影像質量的演變。然而，如果老師在整個學期中對你的每一次小測驗、作業、甚至課堂表現都進行細粒度的評分和反饋（相當於細粒度獎勵），那麼你在每一步的學習過程中都會受到監控和指導。你無法依靠最後的考試來矇混過關，必須在每個階段都真正掌握知識。這與擴散模型對齊時類似，細粒度獎勵機制能夠讓模型在每個生成步驟中都得到反饋，避免模型透過過度最佳化某一單一目標（如最終影像質量）來走捷徑，確保生成過程中的每一步都保持質量和一致性。面對以上問題，TDPO-R 透過引入時間差分獎勵機制，為擴散模型的每一步去噪操作提供實時的獎勵反饋，從而修正這種偏置錯位，緩解獎勵過最佳化問題。這種獎勵機制將擴散過程中的每個時間步視為一個馬爾科夫決策過程中的狀態，模型在每一步的動作（即去噪操作）後，都會獲得相應的獎勵值。為了降低時間差分獎勵函式的學習成本，TDPO-R 透過一個時間差分評判器（Temporal Critic）來近似估計每個時間步的獎勵。這一評判器透過強化學習的方式與擴散模型一齊同步進行訓練，逐步學習如何為每個時間步估算合適的獎勵。具體來說，模型在每個時間步 t 做出去噪操作後，評判器會基於當前狀態（中間去噪影像）與目標狀態的差異，給出即時獎勵反饋。在策略更新時，TDPO-R 採用類似於策略梯度（Policy Gradient）的方法。特別是，得益於時間差分獎勵機制提供了每個時間步的即時反饋，TDPO-R 在擴散模型去噪的每一步都可以立即進行策略更新，而無需等待所有步驟結束。這種即時更新讓模型能夠及時修正去噪過程中出現的偏差，避免問題在後續步驟中積累。同時，由於每一步都進行即時更新，模型可以在更短的時間內獲得更好的最佳化效果，減少了傳統全域性更新方式下的延遲和不必要的計算開銷，提高了訓練的樣本效率。 接著，從首要偏置（Primacy Bias）的角度出發，TDPO-R 的研究者進一步分析了獎勵過最佳化的問題。首要偏置描述的是深度強化學習模型在訓練中傾向於過度擬合早期的訓練經驗，從而影響後續訓練階段的最佳化效果的一種現象。他們發現，在擴散模型的對齊任務中，首要偏置與獎勵過最佳化也有著密切的聯絡，特別是當模型在訓練過程中過度傾向於某個早期的策略時，後續的生成質量可能會出現下降。為了深入理解首要偏置與獎勵過最佳化之間的關係，研究者引入了神經網路中神經元啟用狀態的概念。通常，深度神經網路中的神經元可以處於兩種狀態：- 活躍神經元：在訓練過程中頻繁被啟用的神經元，通常與當前任務的主要模式相關。
- 休眠神經元：在訓練中很少被啟用的神經元，通常被認為對當前任務貢獻較少，甚至被認為會阻礙模型的學習能力。
在以往的研究中，活躍神經元往往被視為是模型的核心，因為它們反映了模型對當前資料的主要適應能力。而休眠神經元則通常被認為是冗餘的或不必要的，因為它們沒有參與主要的學習過程。而在 TDPO-R 的研究中，研究者驚訝地發現，休眠神經元實際上在應對獎勵過最佳化方面起到了重要作用。具體而言：- 活躍神經元更易受到首要偏置的影響：由於活躍神經元頻繁參與早期訓練任務，它們傾向於過擬合早期的學習模式。因此，這些神經元在模型的最佳化過程中可能會過度強化某一特定獎勵目標，導致獎勵過最佳化問題的加劇。例如，模型可能因為過度最佳化早期階段的獎勵訊號，導致後續的生成結果表現不佳，甚至在其他獎勵函式上的泛化能力下降。
- 休眠神經元作為自適應正則化手段：與活躍神經元相反，休眠神經元由於較少參與訓練任務，反而可以作為一種對抗獎勵過最佳化的自適應正則化手段。這是因為，休眠神經元的低啟用狀態意味著它們沒有過度依賴早期的獎勵訊號，從而在模型後期的訓練中可以起到平衡的作用。另外，它們還能夠提供一種類似於 “潛力儲備” 的功能，在獎勵函式發生變化或模型過擬合早期獎勵訊號時，休眠神經元可以重新被啟用，以補充模型的適應能力，增強模型的泛化效能。
為了應對首要偏置，TDPO-R 引入了一種神經元重置機制。這種機制透過定期重置模型中的活躍神經元，將其恢復到類似 “休眠” 狀態，從而打破首要偏置的影響，重新激發模型的學習能力。具體來說：- 週期性神經元重置：在訓練過程中，TDPO-R 會定期對評判器（critic）模型中過度活躍的神經元進行重置，降低它們的啟用頻率，從而打破它們在早期階段對特定獎勵訊號的過擬合。透過這一操作，模型可以避免過度強化某一獎勵目標，確保生成過程的多樣性和泛化能力。
- 重新啟用休眠神經元：隨著活躍神經元被重置，模型的其他神經元，包括那些此前處於休眠狀態的神經元，會被啟用，以參與新的學習任務。這種神經元的 “輪替” 確保了模型的學習能力不會因為早期訓練經驗的固定化而受到限制，從而緩解了獎勵過最佳化的問題。
研究者使用提出的 TDPO-R 對 Stable Diffusion v1.4 模型進行微調訓練，並採取不同文字提示詞集和獎勵函式進行評估實驗。為了更直觀地展現跨獎勵泛化度量的效果，研究者透過圖表展示了在最佳化某一特定目標獎勵函式（如 Aesthetic Score）時，模型在其他多種獎勵函式（如 ImageReward、HPSv2 和 PickScore）上的得分如何變化。如上圖所示，其中每幅圖表的橫座標代表訓練儲存的中間模型在目標獎勵函式上的得分，而縱座標代表對應的中間模型在域外獎勵函式上的得分。透過這些圖表，可以清楚地觀察到 TDPO-R 和其他擴散模型對齊方法（如 DDPO 和 AlignProp）在跨獎勵泛化能力上的差異。在泛化較差的目標獎勵函式（如 Aesthetic Score）上，TDPO-R 能夠保持域外獎勵函式的得分沒有顯著下滑，而其它方法則表現為明顯的下降趨勢；而在泛化較好的目標獎勵函式（如 HPSv2 和 PickScore）上，TDPO-R 能夠使域外獎勵函式的得分更快地上漲。這表明 TDPO-R 在最佳化特定目標時，展現出更強的跨獎勵泛化能力，有效緩解了獎勵過最佳化的問題。除了定量指標的對比，研究者還透過視覺化影像樣本的對比展示了 TDPO-R 相較於其他方法的優勢。如下圖所示對比了 TDPO-R 與其他方法，皆將目標獎勵函式（Aesthetic Score）的得分提升至同一臨界值時所生成的樣本影像。可以看出，在其它方法中，不同文字提示詞對應的影像樣本風格趨同，尤其是在光線、色調和構圖風格方面，缺乏多樣性，甚至還存在影像失真的情況。相比之下，TDPO-R 不僅能夠生成更具視覺自然感和細節保真的影像，還保留了風格的多樣性，具有更平衡的美學表現，直觀地體現了 TDPO-R 緩解獎勵過最佳化的有效性。如下圖所示，在生成包含複雜場景描述的影像時，其它方法存在影像細節不準確的問題，忽略了關鍵的語義元素（如物體的數量、型別以及場景地點等），表現出典型的文圖語義不一致問題。而 TDPO-R 生成的影像不僅視覺豐富，還能夠準確反映文字描述的內容。這體現了 TDPO-R 的跨任務泛化能力，即在提升影像美學表現的同時，還能夠有效保留文圖的語義一致性。除了上述實驗結果之外，研究者還展示了其他一些主要實驗內容，包括 “演算法樣本效率對比”、“對未見提示詞的泛化能力”、“神經元狀態的影響” 以及 “應對過最佳化的其它替代策略” 等方面。對於這些補充實驗以及 TDPO-R 方法的更多技術細節，感興趣的讀者可以參閱論文原文進一步瞭解。