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

亮點(diǎn)直擊

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

總覽全文

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

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

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

動(dòng)機(jī)

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

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

?

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

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

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

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

方法

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

基礎(chǔ)知識

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

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

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

時(shí)間預(yù)測擴(kuò)散模型（TPDM）

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

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

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

訓(xùn)練算法

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

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

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

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

圖像獎(jiǎng)勵(lì)與步驟數(shù)的折扣

實(shí)驗(yàn)

實(shí)現(xiàn)細(xì)節(jié)

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

訓(xùn)練配置

主要結(jié)果

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

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

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

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

定量結(jié)果

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

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

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

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

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

模塊架構(gòu)的消融實(shí)驗(yàn)

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

結(jié)論與局限性

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

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

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

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