文章鏈接：https://arxiv.org/pdf/2412.01243

亮點直擊

• 提出了時間預(yù)測擴散模型（TPDM），該模型可以在推理過程中自適應(yīng)地調(diào)整噪聲調(diào)度，實現(xiàn)圖像質(zhì)量和模型效率之間的平衡。
• 為了訓練TPDM，通過強化學習最大化圖像質(zhì)量，并根據(jù)去噪步驟數(shù)折扣，直接優(yōu)化最終的性能和效率。
• 模型在多個評估基準上表現(xiàn)優(yōu)越，在減少推理步數(shù)的同時取得了更好的結(jié)果。

總覽全文

擴散模型和流模型在文本到圖像生成等多種應(yīng)用中取得了顯著成功。然而，這些模型通常在推理過程中依賴于相同的預(yù)定去噪調(diào)度策略，這可能限制了推理效率以及在處理不同提示時的靈活性。本文認為，最優(yōu)的噪聲調(diào)度應(yīng)該適應(yīng)每個推理實例，并提出了時間預(yù)測擴散模型（TPDM）來實現(xiàn)這一目標。

TPDM采用了一個即插即用的時間預(yù)測模塊（TPM），該模塊在每個去噪步驟中根據(jù)當前的隱空間特征預(yù)測下一個噪聲水平。使用強化學習來訓練TPM，目標是最大化一個獎勵，該獎勵通過去噪步驟的數(shù)量來折扣最終圖像質(zhì)量。通過這種自適應(yīng)調(diào)度器，TPDM不僅生成與人類偏好高度一致的高質(zhì)量圖像，還能動態(tài)調(diào)整去噪步驟的數(shù)量和時間，從而提升性能和效率。

在多個擴散模型基準上訓練了TPDM。在Stable Diffusion 3 Medium架構(gòu)下，TPDM實現(xiàn)了5.44的美學評分和29.59的人類偏好評分（HPS），同時使用大約50%更少的去噪步驟，取得了更好的性能。

動機

以下圖4中的幾張圖像為例，右側(cè)的圖像內(nèi)容豐富，需要更多的去噪步驟來捕捉更細致的細節(jié)。相比之下，左側(cè)的圖像相對簡單，可以使用較少的步驟生成，而不影響質(zhì)量。此外，Karras也證明了不同的噪聲調(diào)度對生成質(zhì)量有很大影響。

那么，不禁要問：是否可以在推理過程中自適應(yīng)地調(diào)整去噪步驟的數(shù)量和每個步驟的噪聲水平，而無需用戶的手動干預(yù)？

?

基于此，本文提出了時間預(yù)測擴散模型（TPDM），該模型能夠在推理過程中自適應(yīng)地調(diào)整去噪步驟的數(shù)量和去噪強度。具體而言，TPDM通過實現(xiàn)一個即插即用的時間預(yù)測模塊（TPM），該模塊可以根據(jù)當前步驟的隱空間特征預(yù)測下一個擴散時間，從而使噪聲調(diào)度能夠動態(tài)調(diào)整。

TPM通過強化學習進行訓練。將多步驟的去噪過程視為一個完整的軌跡，并將圖像質(zhì)量（經(jīng)過去噪步驟數(shù)的折扣）作為獎勵進行優(yōu)化。圖像質(zhì)量通過與人類偏好對齊的獎勵模型進行衡量。

TPM可以輕松集成到任何擴散模型中，幾乎不增加額外的計算負擔，并使模型能夠自動調(diào)整超參數(shù)，例如樣本步驟和每個步驟的噪聲水平，從而在不需要人工干預(yù)的情況下實現(xiàn)圖像質(zhì)量和效率之間的平衡。此外，在訓練過程中，模型的擴散過程與推理過程中保持一致，直接優(yōu)化推理性能并減少去噪步驟。

在多個最先進的模型上實現(xiàn)了TPDM，包括Stable Diffusion和Flux。通過自適應(yīng)噪聲調(diào)度，模型在生成圖像時平均使用了50%更少的步驟，且圖像質(zhì)量與Stable Diffusion 3持平或略有提高（0.322 CLIP-T，5.445美學評分，22.33選擇評分，29.59 HPSv2.1）。這些結(jié)果表明，TPDM具有在追求高質(zhì)量圖像生成和提高模型效率之間找到平衡的潛力。

方法

在本節(jié)中，我們首先簡要回顧擴散模型的基本原理，然后介紹時間預(yù)測模塊（TPM），最后詳細說明TPM的訓練算法。

基礎(chǔ)知識

擴散模型通過一個反向過程學習生成圖像，該過程逐漸去除樣本中的噪聲。實現(xiàn)這一反向過程的主要范式是流匹配（Flow Matching）。因此，在此介紹流匹配模型的工作原理以及當前最先進模型的詳細結(jié)構(gòu)。

目前，許多最先進的擴散模型都基于DiT架構(gòu)，它們只關(guān)注條件圖像生成，并采用調(diào)制機制來條件化網(wǎng)絡(luò)，既考慮擴散去噪步驟中的擴散時間，也考慮文本提示。

使用這樣的模型，可以根據(jù)隱特征和當前噪聲水平執(zhí)行單個去噪步驟。

時間預(yù)測擴散模型（TPDM）

如前所述，需要一系列去噪步驟才能使用訓練好的擴散模型生成圖像。通常，在這個過程中會對所有提示應(yīng)用固定的噪聲調(diào)度，為每個步驟分配一個預(yù)定的噪聲水平。

TPDM僅需要對原始擴散模型進行最小的修改：添加一個輕量級的時間預(yù)測模塊（TPM），如下圖3所示。該模塊將transformer塊前后的隱特征作為輸入，從而考慮到當前噪聲輸入和本步驟的預(yù)測結(jié)果。經(jīng)過幾層卷積后，TPM將隱空間特征池化成一個單一的特征向量，并通過兩個線性層預(yù)測a和b。我們還在TPM中使用了自適應(yīng)歸一化層，使模型能夠感知當前的時間嵌入。

在訓練過程中，凍結(jié)原始的擴散模型，只更新新引入的TPM。因此，模型在保持原有圖像生成能力的同時，學習預(yù)測下一個擴散時間。

訓練算法

因此，我們優(yōu)化TPM，以最大化整個去噪過程后的圖像質(zhì)量，從而實現(xiàn)精確的時間預(yù)測。圖像質(zhì)量通過圖像獎勵模型進行衡量?？紤]到整個推理計算圖過于深度，無法進行梯度反向傳播，我們使用鄰近策略優(yōu)化（PPO）進行訓練，其損失函數(shù)公式為：

將在后文中具體說明PPO中使用的動作和優(yōu)勢。

通常，當模型做出一系列預(yù)測時，PPO將每個單獨的預(yù)測視為一個動作并以批量優(yōu)化。然而，最近的RLOO研究指出，當獎勵信號只出現(xiàn)在序列的末尾，并且環(huán)境動態(tài)完全確定時，可以將整個序列視為一個動作，而不會影響性能。因此，為簡化起見，將整個生成過程，包括調(diào)度中的所有預(yù)測時間，視為一個單獨的動作進行優(yōu)化。

其中，(N) 表示生成步驟的總數(shù)。

圖像獎勵與步驟數(shù)的折扣

實驗

實現(xiàn)細節(jié)

數(shù)據(jù)集收集了用于訓練模型的提示詞。這些提示詞是通過 Florence-2和 Llava-Next生成的，用于生成 Laion-Art和 COYO-700M數(shù)據(jù)集的圖像描述，并利用這些提示詞構(gòu)成我們的訓練集。

訓練配置

主要結(jié)果

不同圖像的動態(tài)調(diào)度在前面圖 4 中，展示了使用不同提示詞生成的圖像及其由 TPDM 預(yù)測的對應(yīng)調(diào)度。當輸入較短且簡單的提示時，生成的圖像包含較少的物體和細節(jié)，因此擴散時間會更快地下降，并在相對較少的步驟中達到 0。相反，當提供較長且更復(fù)雜的提示時，模型需要生成更多的視覺細節(jié)，因此擴散時間下降得較慢，以便生成精細的細節(jié)。在這種情況下，TPDM 在生成過程中需要更多的去噪步驟。

調(diào)整 γ 對不同步驟數(shù)的影響公式（9）中的 γ 控制圖像獎勵在更多生成步驟中的折扣方式，從而影響擴散時間在去噪過程中的衰減速度，進而影響我們模型的平均去噪步驟數(shù)。

如下圖 6 所示，當將 γ 從 0.97 降低到 0.85 時，TPDM 趨向于更快速地減少擴散時間，從而減少了采樣步驟，從 15.0 步減少到 7.5 步。此外，與基準模型（黃色線）相比，TPDM（紫色線）在相同的推理步驟數(shù)下，始終能獲得顯著更高的美學分數(shù)，達到了模型效率和生成性能之間的良好平衡。

視覺對比我們的方法在生成細粒度細節(jié)方面表現(xiàn)出色。TPDM 生成的圖像相比 SD3-Medium（圖 5C 中的圖像）和結(jié)果中的圖像，展示了更為真實的筆記本鍵盤。

定量結(jié)果

在多個最先進的擴散模型上應(yīng)用了 TPM，包括 Stable Diffusion 3 Medium、Stable Diffusion 3.5 Large 和 Flux 1.0 dev，展示了 TPM 如何提升它們的性能。主要評估了兩類指標：第一類是客觀指標，包括 FID、與給定提示的對齊度（CLIP-T）以及人類偏好分數(shù)（美學分數(shù) v2 和 HPSv2.1）。第二類是通過用戶研究直接比較不同模型生成的圖像。

定量指標在表 1 中比較了 TPDM 和上述模型。盡管保持競爭力的性能，所有這些模型平均可以在推薦的步驟數(shù)的一半左右生成圖像。

此外，代表人類偏好的指標提高得最多。例如，通過僅使用平均 15.28 步生成的圖像，TPDM-SD3-Medium 獲得了 29.59 的 HPS 分數(shù)，比使用相似步驟的 Stable Diffusion 3 高出 +1.07，比原始的 28 步結(jié)果高出 +0.47。這可能歸因于我們在優(yōu)化中使用的獎勵模型，從而生成符合人類偏好的美學圖像。

用戶研究為了更好地反映人類對這些模型的態(tài)度，通過邀請志愿者對比不同模型生成的圖像，并選擇他們偏好的圖像，進行了用戶研究。

具體來說，對于每個提示，我們提供了從 SD3-Medium 生成的 15 步和 28 步圖像，以及從 TPDM-SD3-Medium 生成的圖像。我們邀請了 15 名志愿者評估 50 個提示生成的圖像。結(jié)果如表 2 所示，表明我們的模型生成的圖像更符合人類的偏好。

模塊架構(gòu)的消融實驗

在本節(jié)中，我們對 TPM 輸入的選擇進行了消融實驗。如表 3 所示，將第一層和最后一層的特征都輸入到 TPM 中，比僅使用其中任意一層的特征表現(xiàn)更好。

結(jié)論與局限性

本文提出了時間預(yù)測擴散模型（TPDM），這是一個具有靈活去噪調(diào)度器的文本到圖像擴散模型，可以針對不同的提示自動調(diào)整去噪調(diào)度。通過引入時間預(yù)測模塊，通過強化學習和獎勵模型有效地訓練了 TPDM?；诋斍邦I(lǐng)先的擴散模型架構(gòu)（Stable Diffusion 3 Medium），我們訓練了一個強大的 MM-DiT 基礎(chǔ) TPDM，并在多個文本到圖像生成基準測試中表現(xiàn)出了競爭力的定量性能。

盡管 TPDM 展現(xiàn)了良好的性能，但仍有一些局限性。例如，在本文中，我們僅為 TPM 設(shè)計了一個相對簡單的架構(gòu)，如何改進該模塊以獲得更好的性能仍然是一個未解之題。其次，我們凍結(jié)了原始模型的參數(shù)，并采用我們的訓練方法更新了擴散模型的參數(shù)，從而取得了更好的結(jié)果，這一方法還需要進一步探索。

本文轉(zhuǎn)自AI生成未來 ，作者：AI生成未來

原文鏈接:??https://mp.weixin.qq.com/s/3Z5Jhdu3u_z-UF1m-glziw??