Understanding Diffusion Models: A Unified Perspective(六)(完結)
- 文章概括
- 指導(Guidance)
- 分類器指導
- 無分類器引導(Classifier-Free Guidance)
- 總結
文章概括
引用:
@article{luo2022understanding,title={Understanding diffusion models: A unified perspective},author={Luo, Calvin},journal={arXiv preprint arXiv:2208.11970},year={2022}
}
Luo, C., 2022. Understanding diffusion models: A unified perspective. arXiv preprint arXiv:2208.11970.
原文: https://arxiv.org/abs/2208.11970
代碼、數據和視頻:https://arxiv.org/abs/2208.11970
系列文章:
請在 《 《 《文章 》 》 》 專欄中查找
指導(Guidance)
到目前為止,我們專注于僅對數據分布 p ( x ) p(x) p(x)進行建模。然而,我們通常還對學習條件分布 p ( x ∣ y ) p(x|y) p(x∣y)感興趣,這將使我們能夠通過條件信息 y y y明確控制生成的數據。這構成了圖像超分辨率模型(如級聯擴散模型[Cascaded Diffusion Models][18])的基礎,同時也是圖像-文本模型(如DALL-E 2[19]和Imagen[7])的前沿技術。
一種自然的方式是簡單地在每次迭代中將條件信息與時間步信息一起添加。回憶公式32中的聯合分布:
p ( x 0 : T ) = p ( x T ) ∏ t = 1 T p θ ( x t ? 1 ∣ x t ) p(x_{0:T}) = p(x_T) \prod_{t=1}^T p_\theta(x_{t-1}|x_t) p(x0:T?)=p(xT?)t=1∏T?pθ?(xt?1?∣xt?)
然后,為了將其轉化為一個條件擴散模型,我們可以在每個轉換步驟中簡單地添加任意條件信息 y y y,表示為:
p ( x 0 : T ∣ y ) = p ( x T ) ∏ t = 1 T p θ ( x t ? 1 ∣ x t , y ) p(x_{0:T}|y) = p(x_T) \prod_{t=1}^T p_\theta(x_{t-1}|x_t, y) p(x0:T?∣y)=p(xT?)t=1∏T?pθ?(xt?1?∣xt?,y)
例如, y y y可以是在圖像-文本生成中的文本編碼,也可以是用于執行超分辨率的低分辨率圖像。因此,我們可以像之前一樣,通過預測 x ^ θ ( x t , t , y ) ≈ x 0 \hat{x}_\theta(x_t,t,y) \approx x_0 x^θ?(xt?,t,y)≈x0?、 ? ^ θ ( x t , t , y ) ≈ ? 0 \hat{\epsilon}_\theta(x_t,t,y) \approx \epsilon_0 ?^θ?(xt?,t,y)≈?0?或 s θ ( x t , t , y ) ≈ ? log ? p ( x t ∣ y ) s_\theta(x_t,t,y) \approx \nabla \log p(x_t|y) sθ?(xt?,t,y)≈?logp(xt?∣y),學習VDM的核心神經網絡,以實現每種所需的解釋和實現。
這種基礎公式的一個警告是,通過這種方式訓練的條件擴散模型可能會傾向于忽略或弱化任何給定的條件信息。因此,提出了指導(Guidance)作為一種更明確控制模型對條件信息權重的方法,但代價是樣本多樣性的降低。兩種最流行的指導形式被稱為分類器指導(Classifier Guidance)[10, 20]和無分類器指導(Classifier-Free Guidance)[21]。
這段內容主要講解了如何將擴散模型從無條件生成(直接建模數據分布 p ( x ) p(x) p(x))擴展到條件生成(建模條件分布 p ( x ∣ y ) p(x|y) p(x∣y)),并提到了一種潛在的問題,即條件信息 y y y在訓練過程中可能被模型忽略或弱化。
1. 從無條件生成到條件生成
1.1 無條件生成建模 擴散模型在無條件生成建模中專注于直接建模數據分布 p ( x ) p(x) p(x)。聯合分布形式(回憶公式 32): p ( x 0 : T ) = p ( x T ) ∏ t = 1 T p θ ( x t ? 1 ∣ x t ) . p(x_{0:T}) = p(x_T) \prod_{t=1}^T p_\theta(x_{t-1} | x_t). p(x0:T?)=p(xT?)t=1∏T?pθ?(xt?1?∣xt?).
- x 0 : T x_{0:T} x0:T?:表示數據從初始噪聲 x T x_T xT? 到去噪還原 x 0 x_0 x0? 的一系列中間步驟。
- p θ ( x t ? 1 ∣ x t ) p_\theta(x_{t-1} | x_t) pθ?(xt?1?∣xt?):表示從時間 t t t 的噪聲分布 x t x_t xt? 逐步還原到 t ? 1 t-1 t?1 的噪聲分布 x t ? 1 x_{t-1} xt?1?。
- p ( x T ) p(x_T) p(xT?):初始噪聲的分布,通常為標準高斯分布。
1.2 條件生成建模 我們希望生成樣本不僅符合 p ( x ) p(x) p(x),還能夠通過條件信息 y y y 控制生成過程,即建模條件分布 p ( x ∣ y ) p(x|y) p(x∣y)。 為了實現條件生成,只需在無條件模型的每個步驟中添加條件信息 y y y:
p ( x 0 : T ∣ y ) = p ( x T ) ∏ t = 1 T p θ ( x t ? 1 ∣ x t , y ) . p(x_{0:T} | y) = p(x_T) \prod_{t=1}^T p_\theta(x_{t-1} | x_t, y). p(x0:T?∣y)=p(xT?)t=1∏T?pθ?(xt?1?∣xt?,y).公式含義
- p ( x T ) p(x_T) p(xT?):與無條件模型相同,初始噪聲的分布仍然為標準高斯分布。
- p θ ( x t ? 1 ∣ x t , y ) p_\theta(x_{t-1} | x_t, y) pθ?(xt?1?∣xt?,y):條件擴散模型的核心,用于學習在給定條件 y y y 下從 x t x_t xt? 轉換到 x t ? 1 x_{t-1} xt?1? 的概率。
如何引入條件信息 y y y 條件信息 y y y 的形式是任意的,比如:
- 在圖像-文本生成任務中, y y y 是文本編碼。
- 在圖像超分辨率任務中, y y y 是低分辨率圖像。
模型通過神經網絡參數化 p θ ( x t ? 1 ∣ x t , y ) p_\theta(x_{t-1} | x_t, y) pθ?(xt?1?∣xt?,y),可以在每次迭代中引入 y y y,如:
- 預測 x ^ θ ( x t , t , y ) ≈ x 0 \hat{x}_\theta(x_t, t, y) \approx x_0 x^θ?(xt?,t,y)≈x0?:從條件 y y y 中還原出原始數據 x 0 x_0 x0?。
- 預測 ? ^ θ ( x t , t , y ) ≈ ? 0 \hat{\epsilon}_\theta(x_t, t, y) \approx \epsilon_0 ?^θ?(xt?,t,y)≈?0?:預測初始噪聲。
- 預測 s θ ( x t , t , y ) ≈ ? log ? p ( x t ∣ y ) s_\theta(x_t, t, y) \approx \nabla \log p(x_t | y) sθ?(xt?,t,y)≈?logp(xt?∣y):學習條件得分函數。
2. 條件擴散模型忽略條件信息的問題
2.1 問題來源 條件擴散模型的訓練目標類似于無條件模型,只是多了條件 y y y: E p ( x 0 : T , y ) [ ∥ s θ ( x t , t , y ) ? ? log ? p ( x t ∣ y ) ∥ 2 2 ] . \mathbb{E}_{p(x_{0:T}, y)} \left[ \| s_\theta(x_t, t, y) - \nabla \log p(x_t | y) \|_2^2 \right]. Ep(x0:T?,y)?[∥sθ?(xt?,t,y)??logp(xt?∣y)∥22?].
- 損失函數關注的是得分函數 s θ ( x t , t , y ) s_\theta(x_t, t, y) sθ?(xt?,t,y) 對于數據 x t x_t xt? 的預測精度。
- 目標是讓模型通過學習條件分布 p ( x t ∣ y ) p(x_t | y) p(xt?∣y) 更好地指導采樣。
問題: 在很多情況下,模型可以選擇忽略 y y y,只關注數據分布 p ( x ) p(x) p(x),也能在損失函數上取得較低值。這是因為:
- 在 y y y 的影響很弱時,模型可以退化為無條件擴散模型,只預測 p ( x t ) p(x_t) p(xt?) 的得分函數,而不使用 y y y。
- 如果 y y y 與 x t x_t xt? 的關系較復雜,而簡單的無條件預測已經足夠降低損失,模型傾向于忽略 y y y。
2.2 為什么模型可能忽略 y y y 這個現象稱為“條件退化問題”(Conditional Degradation Problem),其主要原因包括:
目標函數對條件信息的依賴較弱:
- 損失函數中沒有顯式地對條件信息 y y y 強化約束,模型更容易忽略 y y y。
- 在高維數據空間中, p ( x t ) p(x_t) p(xt?) 的特征可能已經足夠顯著,使得條件信息 y y y 顯得無足輕重。
訓練數據中條件信息的質量和多樣性不足:
- 如果 y y y 的分布較簡單或信息量不足,模型可能會學習到“忽略 y y y”更容易優化。
無條件預測較為簡單:
- 如果無條件生成已經能夠很好地降低損失,模型可能會放棄更復雜的條件生成。
3. 示例:條件擴散中的問題
3.1 圖像超分辨率中的條件擴散 任務:給定低分辨率圖像 y y y,生成對應的高分辨率圖像 x x x。
條件擴散公式: p ( x 0 : T ∣ y ) = p ( x T ) ∏ t = 1 T p θ ( x t ? 1 ∣ x t , y ) . p(x_{0:T} | y) = p(x_T) \prod_{t=1}^T p_\theta(x_{t-1} | x_t, y). p(x0:T?∣y)=p(xT?)t=1∏T?pθ?(xt?1?∣xt?,y). 模型可能的問題:
- 如果模型發現低分辨率圖像 y y y 信息不足以顯著影響生成結果,它可能退化為無條件生成,僅根據 p ( x ) p(x) p(x) 生成高分辨率圖像。
- 結果是生成的圖像可能與 y y y 無關,僅符合高分辨率圖像的分布。
3.2 圖像-文本生成中的條件擴散 任務:給定文本描述 y y y,生成對應的圖像 x x x。
條件擴散公式: p ( x 0 : T ∣ y ) = p ( x T ) ∏ t = 1 T p θ ( x t ? 1 ∣ x t , y ) . p(x_{0:T} | y) = p(x_T) \prod_{t=1}^T p_\theta(x_{t-1} | x_t, y). p(x0:T?∣y)=p(xT?)t=1∏T?pθ?(xt?1?∣xt?,y). 模型可能的問題:
- 如果文本 y y y 的語義信息較模糊,模型可能傾向于生成一個與 p ( x ) p(x) p(x) 相符的“通用圖像”。
- 比如,輸入“貓在沙發上”,但生成的結果可能是一只貓,而無關沙發。
4. 解決條件信息弱化問題
為了解決模型傾向于忽略條件信息的問題,可以采取以下方法:
加強條件信息的權重:
- 在訓練中通過損失函數顯式強調 ( y ) 的影響,比如引入條件正則化項:
E p ( x 0 : T , y ) [ ∥ s θ ( x t , t , y ) ? ? log ? p ( x t ∣ y ) ∥ 2 2 + λ ∥ f ( y ) ? g ( x t ) ∥ 2 2 ] , \mathbb{E}_{p(x_{0:T}, y)} \left[ \| s_\theta(x_t, t, y) - \nabla \log p(x_t | y) \|_2^2 + \lambda \| f(y) - g(x_t) \|_2^2 \right], Ep(x0:T?,y)?[∥sθ?(xt?,t,y)??logp(xt?∣y)∥22?+λ∥f(y)?g(xt?)∥22?],
其中 f ( y ) f(y) f(y) 和 g ( x t ) g(x_t) g(xt?) 表示條件和數據的特征提取函數。改進條件編碼方式:
- 增加條件編碼的質量,比如在圖像-文本任務中引入更強大的文本編碼器(如預訓練的 CLIP 模型)。
增加條件數據多樣性:
- 確保訓練數據中條件 y y y 的分布足夠廣泛,覆蓋各種可能的情況。
使用交叉注意力機制:
- 在神經網絡中引入交叉注意力層,使條件信息 y y y 更加顯著地影響生成過程。
總結
- 條件擴散模型通過引入條件信息 y y y,在每次迭代中將無條件生成擴展為條件生成。
- 條件信息可能被忽略的原因包括目標函數弱化對 y y y 的依賴、數據分布的復雜性以及模型訓練中的簡化偏好。
- 可以通過加強條件編碼、調整損失函數和改進訓練數據等方法解決這一問題。
分類器指導
讓我們從擴散模型的基于分數(score-based)形式開始,其中我們的目標是在任意噪聲水平 t t t下學習條件模型的分數 ? log ? p ( x t ∣ y ) \nabla\log p(x_t|y) ?logp(xt?∣y)。回顧一下, ? \nabla ?是 ? x t \nabla_{x_t} ?xt??的簡寫,為了簡潔起見采用此形式。根據貝葉斯公式,我們可以推導出以下等價形式:
? log ? p ( x t ∣ y ) = ? log ? ( p ( x t ) p ( y ∣ x t ) p ( y ) ) ( 163 ) = ? log ? p ( x t ) + ? log ? p ( y ∣ x t ) ? ? log ? p ( y ) ( 164 ) = ? log ? p ( x t ) ? unconditional?score + ? log ? p ( y ∣ x t ) ? adversarial?gradient ( 165 ) \begin{aligned} \nabla \log p(x_t|y) &= \nabla \log \left(\frac{p(x_t)p(y|x_t)}{p(y)}\right) \quad & (163) \\ &= \nabla \log p(x_t) + \nabla \log p(y|x_t) - \nabla \log p(y) \quad & (164) \\ &= \underbrace{\nabla \log p(x_t)}_{\text{unconditional score}} + \underbrace{\nabla \log p(y|x_t)}_{\text{adversarial gradient}} \quad & (165) \end{aligned} ?logp(xt?∣y)?=?log(p(y)p(xt?)p(y∣xt?)?)=?logp(xt?)+?logp(y∣xt?)??logp(y)=unconditional?score ?logp(xt?)??+adversarial?gradient ?logp(y∣xt?)???(163)(164)(165)?
我們利用了一個事實,即 log ? p ( y ) \log p(y) logp(y)關于 x t x_t xt?的梯度為零。
我們最終推導的結果可以理解為學習一個無條件的分數函數,并將其與分類器 p ( y ∣ x t ) p(y|x_t) p(y∣xt?)的對抗梯度相結合。因此,在分類器引導(Classifier Guidance)中,無條件擴散模型的分數函數按照之前的推導進行學習,同時訓練一個分類器,該分類器接收任意帶噪的 x t x_t xt?并嘗試預測條件信息 y y y。然后,在采樣過程中,用于退火Langevin動態的整體條件分數函數被計算為無條件分數函數與帶噪分類器的對抗梯度之和。
為了引入細粒度控制,以鼓勵或阻止模型考慮條件信息,分類器引導通過一個超參數 γ \gamma γ對帶噪分類器的對抗梯度進行縮放。在分類器引導下學習的分數函數可以總結如下:
? log ? p ( x t ∣ y ) = ? log ? p ( x t ) + γ ? log ? p ( y ∣ x t ) ( 166 ) \nabla\log p(x_t|y) = \nabla\log p(x_t) + \gamma\nabla\log p(y|x_t) \quad (166) ?logp(xt?∣y)=?logp(xt?)+γ?logp(y∣xt?)(166)
直觀來看,當 γ = 0 \gamma=0 γ=0時,條件擴散模型完全忽略條件信息的學習,而當 γ \gamma γ值較大時,條件擴散模型會學習生成嚴格符合條件信息的樣本。這種情況以犧牲樣本多樣性為代價,因為模型僅會生成那些即使在有噪聲的情況下也能輕松再現所提供條件信息的數據。
條件引導(Classifier Guidance)的一個顯著缺點是它依賴于單獨訓練的分類器。由于分類器必須處理任意噪聲輸入,而大多數現有的預訓練分類模型并未針對這種情況進行優化,因此分類器必須與擴散模型同時進行特定的訓練。
無分類器引導(Classifier-Free Guidance)
無分類器引導(Classifier-Free Guidance)中,作者摒棄了單獨訓練分類器模型的方法,而是選擇了無條件擴散模型和條件擴散模型。為了推導無分類器引導下的得分函數,我們可以首先重排公式165,得到:
? log ? p ( y ∣ x t ) = ? log ? p ( x t ∣ y ) ? ? log ? p ( x t ) ( 167 ) \nabla \log p(y|x_t) = \nabla \log p(x_t|y) - \nabla \log p(x_t) \quad (167) ?logp(y∣xt?)=?logp(xt?∣y)??logp(xt?)(167)
然后,將其代入公式166,我們得到:
? log ? p ( x t ∣ y ) = ? log ? p ( x t ) + γ ( ? log ? p ( x t ∣ y ) ? ? log ? p ( x t ) ) (168) = ? log ? p ( x t ) + γ ? log ? p ( x t ∣ y ) ? γ ? log ? p ( x t ) (169) = γ ? log ? p ( x t ∣ y ) ? conditional?score + ( 1 ? γ ) ? log ? p ( x t ) ? unconditional?score (170) \begin{aligned} \nabla \log p(x_t|y) &= \nabla \log p(x_t) + \gamma \left(\nabla \log p(x_t|y) - \nabla \log p(x_t)\right) \quad &\text{(168)} \\ &= \nabla \log p(x_t) + \gamma \nabla \log p(x_t|y) - \gamma \nabla \log p(x_t) \quad &\text{(169)} \\ &= \underbrace{\gamma \nabla \log p(x_t|y)}_\text{conditional score} + \underbrace{(1 - \gamma)\nabla \log p(x_t)}_\text{unconditional score} \quad &\text{(170)} \end{aligned} ?logp(xt?∣y)?=?logp(xt?)+γ(?logp(xt?∣y)??logp(xt?))=?logp(xt?)+γ?logp(xt?∣y)?γ?logp(xt?)=conditional?score γ?logp(xt?∣y)??+unconditional?score (1?γ)?logp(xt?)???(168)(169)(170)?
再一次, γ \gamma γ是一個控制學習的條件模型對條件信息關注程度的參數。當 γ = 0 \gamma=0 γ=0時,學習的條件模型完全忽略條件信息并學習一個無條件擴散模型。當 γ = 1 \gamma=1 γ=1時,模型明確地學習不帶引導的普通條件分布。當 γ > 1 \gamma>1 γ>1時,擴散模型不僅優先考慮條件得分函數,還會朝遠離無條件得分函數的方向移動。換句話說,它降低了生成不使用條件信息的樣本的概率,更傾向于生成明確使用條件信息的樣本。這也會導致樣本多樣性下降,但以生成準確匹配條件信息的樣本為代價。
由于學習兩個獨立的擴散模型代價較高,我們可以將條件擴散模型和無條件擴散模型合并為一個單一的條件模型;通過將條件信息替換為固定的常量值(例如零),可以查詢無條件擴散模型。這本質上是在條件信息上隨機執行dropout操作。無分類器引導(Classifier-Free Guidance)的優雅之處在于,它讓我們能夠更靈活地控制條件生成過程,同時不需要額外訓練多個擴散模型,僅需訓練一個單一的擴散模型即可。
總結
讓我們回顧一下在研究過程中的發現。首先,我們將變分擴散模型(VDM)推導為一種馬爾科夫分層變分自編碼器(HVAE)的特殊情況,其中三個關鍵假設使得證據下界(ELBO)的計算和優化變得可行且可擴展。接著,我們證明了優化VDM歸結為學習一個神經網絡來預測以下三種潛在目標之一:從任意噪聲化的圖像中恢復原始源圖像,從任意噪聲化圖像中預測原始噪聲,或在任意噪聲水平下的噪聲化圖像的得分函數。然后,我們深入探討了學習得分函數的意義,并將其與基于得分的生成建模(Score-based Generative Modeling)的視角明確關聯起來。最后,我們探討了如何使用擴散模型來學習條件分布。
總之,擴散模型作為生成模型表現出了令人難以置信的能力;實際上,它們為當前最先進的基于文本條件的圖像生成模型(如Imagen和DALL-E 2)提供了支持。此外,這些模型背后的數學理論極其優雅。然而,仍有一些需要考慮的缺點:
- 這很可能并非我們人類自然建模和生成數據的方式;我們并不是通過將樣本生成為隨機噪聲后迭代去噪的方式進行數據建模。
- VDM無法生成可解釋的潛變量(Latents)。而VAE通過對編碼器的優化,有望學習一個具有結構化的潛變量空間,而在VDM中,每個時間步的編碼器已經是一個線性高斯模型,無法靈活優化。因此,中間的潛變量僅限于原始輸入的噪聲版本。
- 潛變量被限制為與原始輸入相同的維度,這進一步增加了學習有意義、壓縮的潛在結構的難度。
- 采樣是一個代價高昂的過程,因為在這兩種方法下都需要運行多個去噪步驟。請記住,其中一個限制是必須選擇足夠大的時間步數 T T T,以確保最終的潛變量完全是高斯噪聲;在采樣過程中,我們必須遍歷所有這些時間步以生成樣本。
最后,擴散模型的成功凸顯了分層變分自編碼器(HVAE)作為生成模型的強大能力。我們已經表明,當我們將其推廣到無限的潛變量層次時,即使編碼器是簡單的,潛變量維度是固定的,并且假設了馬爾科夫轉換,我們仍然能夠學習到強大的數據模型。這表明,在一般的、深層的HVAE中,通過學習復雜的編碼器和語義上有意義的潛變量空間,可能會實現進一步的性能提升。
致謝:我想感謝Josh Dillon、Yang Song、Durk Kingma、Ben Poole、Jonathan Ho、Yiding Jiang、Ting Chen、Jeremy Cohen和Chen Sun,他們審閱了本文的草稿并提供了許多有幫助的修改建議和評論。非常感謝!