擴散模型爆火,這是首篇綜述與Github論文分類彙總

機器之心專欄

作者:楊靈

本綜述(Diffusion Models: A Comprehensive Survey of Methods and Applications)來自加州大學 & Google Research 的 Ming-Hsuan Yang、北京大學崔斌實驗室以及 CMU、UCLA、蒙特利爾 Mila 研究院等衆研究團隊,首次對現有的擴散生成模型(diffusion model)進行了全面的總結分析,從 diffusion model 算法細化分類、和其他五大生成模型的關聯以及在七大領域中的應用等方面展開,最後提出了 diffusion model 的現有 limitation 和未來的發展方向。

文章鏈接:https://arxiv.org/abs/2209.00796

論文分類彙總 github 鏈接:https://github.com/YangLing0818/Diffusion-Models-Papers-Survey-Taxonomy

介紹

擴散模型(diffusion models)是深度生成模型中新的 SOTA。擴散模型在圖片生成任務中超越了原 SOTA:GAN,並且在諸多應用領域都有出色的表現,如計算機視覺,NLP、波形信號處理、多模態建模、分子圖建模、時間序列建模、對抗性淨化等。此外,擴散模型與其他研究領域有着密切的聯繫,如穩健學習、表示學習、強化學習。然而,原始的擴散模型也有缺點,它的採樣速度慢,通常需要數千個評估步驟才能抽取一個樣本;它的最大似然估計無法和基於似然的模型相比;它泛化到各種數據類型的能力較差。如今很多研究已經從實際應用的角度解決上述限制做出了許多努力,或從理論角度對模型能力進行了分析。

然而,現在缺乏對擴散模型從算法到應用的最新進展的系統回顧。爲了反映這一快速發展領域的進展,我們對擴散模型進行了首個全面綜述。我們設想我們的工作將闡明擴散模型的設計考慮和先進方法,展示其在不同領域的應用,並指出未來的研究方向。此綜述的概要如下圖所示:

儘管 diffusion model 在各類任務中都有着優秀的表現,它仍還有自己的缺點,並有諸多研究對 diffusion model 進行了改善。爲了系統地闡明 diffusion model 的研究進展,我們總結了原始擴散模型的三個主要缺點,採樣速度慢,最大化似然差、數據泛化能力弱,並提出將的 diffusion models 改進研究分爲對應的三類:採樣速度提升、最大似然增強和數據泛化增強。我們首先說明改善的動機,再根據方法的特性將每個改進方向的研究進一步細化分類,從而清楚的展現方法之間的聯繫與區別。在此我們僅選取部分重要方法爲例, 我們的工作中對每類方法都做了詳細的介紹,內容如圖所示:

在分析完三類擴散模型後,我們將介紹其他的五種生成模型 GAN,VAE,Autoregressive model, Normalizing flow, Energy-based model。考慮到擴散模型的優良性質,研究者們已經根據其特性將 diffusion model 與其他生成模型結合,所以爲了進一步展現 diffusion model 的特點和改進工作,我們詳細地介紹了 diffusion model 和其他生成模型的結合的工作並闡明瞭在原始生成模型上的改進之處。Diffusion model 在諸多領域都有着優異的表現,並且考慮到不同領域的應用中 diffusion model 產生了不同的變形,我們系統地介紹了 diffusion model 的應用研究,其中包含如下領域:計算機視覺,NLP、波形信號處理、多模態建模、分子圖建模、時間序列建模、對抗性淨化。對於每個任務,我們定義了該任務並介紹利用擴散模型處理任務的工作,我們將本項工作的主要貢獻總結如下:

擴散模型基礎

生成式建模的一個核心問題是模型的靈活性和可計算性之間的權衡。擴散模型的基本思想是正向擴散過程來系統地擾動數據中的分佈,然後通過學習反向擴散過程恢復數據的分佈,這樣就了產生一個高度靈活且易於計算的生成模型。

A.Denoising Diffusion Probabilistic Models(DDPM)

一個 DDPM 由兩個參數化馬爾可夫鏈組成,並使用變分推斷以在有限時間後生成與原始數據分佈一致的樣本。前向鏈的作用是擾動數據,它根據預先設計的噪聲進度向數據逐漸加入高斯噪聲,直到數據的分佈趨於先驗分佈,即標準高斯分佈。反向鏈從給定的先驗開始並使用參數化的高斯轉換核,學習逐步恢復原數據分佈。用

表示原始數據及其分佈, 則前向鏈的分佈是可由下式表達:

這說明前向鏈是馬爾可夫過程,x_t 是加入 t 步噪音後的樣本,β_t 是事先給定的控制噪聲進度的參數。當

趨於 1 時,x_T 可以近似認爲服從標準高斯分佈。當β_t 很小時,逆向過程的轉移核可以近似認爲也是高斯的:

我們可以將變分下界作爲損失函數進行學習:

B.Score-Based Generative Models(SGM)

上述 DDPM 可以視作 SGM 的離散形式。SGM 構造一個隨機微分方程(SDE)來平滑的擾亂數據分佈,將原始數據分佈轉化到已知的先驗分佈:

和一個相應的逆向 SDE,來將先驗分佈變換回原始數據分佈:

因此,要逆轉擴散過程並生成數據,我們需要的唯一信息就是在每個時間點的分數函數。利用 score-matching 的技巧我們可以通過如下損失函數來學習分數函數:

對兩種方法的進一步介紹和兩者關係的介紹請參見我們的文章。

原始擴散模型的三個主要缺點,採樣速度慢,最大化似然差、數據泛化能力弱。最近許多研究都在解決這些缺點,因此我們將改進的擴散模型分爲三類:採樣速度提升、最大似然增強和數據泛化增強。在接下來的三、四、五節我們將對這三類模型進行詳細的介紹。

採樣加速方法

在應用時,爲了讓新樣本的質量達到最佳,擴散模型往往需要進行成千上萬步計算來獲取一個新樣本。這限制了 diffusion model 的實際應用價值,因爲在實際應用時,我們往往需要產生大量的新樣本,來爲下一步處理提供材料。研究者們在提高 diffusion model 採樣速度上進行了大量的研究。我們對這些研究進行了詳細的闡述。我們將其細化分類爲三種方法:Discretization Optimization,Non-Markovian Process,Partial Sampling。

A.Discretization Optimization 方法優化求解 diffusion SDE 的方法。因爲現實中求解複雜 SDE 只能使用離散解來逼近真正的解,所以該類方法試圖優化 SDE 的離散化方法,在保證樣本質量的同時減少離散步數。SGM 提出了一個通用的方法來求解逆向過程,即對前向和後向過程採取相同的離散方法。如果給定了前向 SDE 的離散方式:

那麼我們就可以以相同的方式離散化逆向 SDE:

這種方法比樸素 DDPM 效果略好一點。進一步,SGM 向 SDE 求解器中加入了一個矯正器,從而讓每一步生成的樣本都有正確的分佈。在求解的每一步,求解器給出一個樣本後,矯正器都使用馬爾可夫鏈蒙特卡羅方法來矯正剛生成的樣本的分佈。實驗表明向求解器中加入矯正器比直接增加求解器的步數效率更高。

B.Non-Markovian Process方法突破了原有 Markovian Process 的限制,其逆過程的每一步可以依賴更多以往的樣本來進行預測新樣本,所以在步長較大時也能做出較好的預測,從而加速採樣過程。其中主要的工作 DDIM,不再假設前向過程是馬爾可夫過程,而是服從如下分佈:

DDIM 的採樣過程可以視爲離散化的神經常微分方程,其採樣過程更高效,並且支持樣本的內插。進一步的研究發現 DDIM 可以視作流形上擴散模型 PNDM 的特例。

C.Partial Sampling方法通過在 generation process 中忽略一部分的時間節點,而只使用剩下的時間節點來生成樣本,直接減少了採樣時間。例如,Progressive Distillation 從訓練好的擴散模型中蒸餾出效率更高的擴散模型。對於訓練好的一個擴散模型,Progressive Distillation 會從新訓練一個擴散模型,使新的擴散模型的一步對應於訓練好的擴散模型的兩步,這樣新模型就可以省去老模型一半的採樣過程。具體算法如下:

不斷循環這個蒸餾過程就能讓採樣步驟指數級下降。

最大似然估計加強

擴散模型在最大似然估計的表現差於基於似然函數的生成模型,但最大化似然估計在諸多應用場景都有重要意義,比如圖片壓縮, 半監督學習, 對抗性淨化。由於對數似然難以直接計算,研究主要集中在優化和分析變分下界(VLB)。我們對提高擴散模型最大似然估計的模型進行了詳細的闡述。我們將其細化分類爲三類方法:Objectives Designing,Noise Schedule Optimization,Learnable Reverse Variance。

A.Objectives Designing方法利用擴散 SDE 推倒出生成數據的對數似然與分數函數匹配的損失函數的關係。這樣通過適當設計損失函數,就可以最大化 VLB 和對數似然。Song et al. 證明了可以設計損失函數的權重函數,使得 plug-in reverse SDE 生成樣本的似然函數值小於等於損失函數值,即損失函數是似然函數的上界。分數函數擬合的損失函數如下:

我們只需將權重函數λ(t)設爲擴散係數 g(t)即可讓損失函數成爲似然函數的 VLB,即:

B.Noise Schedule Optimization通過設計或學習前向過程的噪聲進度來增大 VLB。VDM 證明了當離散步數接近無窮時,損失函數完全由信噪比函數 SNR(t)的端點決定:

那麼在離散步數接近無窮時,可以通過學習信噪比函數 SNR(t)的端點最優化 VLB,而通過學習信噪比函數中間部分的函數值來實現模型其他方面的改進。

C.Learnable Reverse Variance方法學習反向過程的方差,從而較少擬合誤差,可以有效地最大化 VLB。Analytic-DPM 證明,在 DDPM 和 DDIM 中存在反向過程中的最優期望和方差:

使用上述公式和訓練好的分數函數,在給定前向過程的條件下,最優的 VLB 可以近似達到。

數據泛化增強

擴散模型假設數據存在於歐幾里得空間,即具有平面幾何形狀的流形,並添加高斯噪聲將不可避免地將數據轉換爲連續狀態空間,所以擴散模型最初只能處理圖片等連續性數據,直接應用離散數據或其他數據類型的效果較差。這限制了擴散模型的應用場景。數個研究工作將擴散模型推廣到適用於其他數據類型的模型,我們對這些方法進行了詳細地闡釋。我們將其細化分類爲兩類方法:Feature Space Unification,Data-Dependent Transition Kernels。

A.Feature Space Unification方法將數據轉化到統一形式的 latent space,然後再 latent space 上進行擴散。LSGM 提出將數據通過 VAE 框架先轉換到連續的 latent space 上後再在其上進行擴散。這個方法的難點在於如何同時訓練 VAE 和擴散模型。LSGM 表明由於潛在先驗是 intractable 的,分數匹配損失不再適用。LSGM 直接使用 VAE 中傳統的損失函數 ELBO 作爲損失函數,並導出了 ELBO 和分數匹配的關係:

該式在忽略常數的意義下成立。通過參數化擴散過程中樣本的分數函數,LSGM 可以高效的學習和優化 ELBO。

B.Data-Dependent Transition Kernels方法根據數據類型的特點設計 diffusion process 中的 transition kernels,使擴散模型可以直接應用於特定的數據類型。D3PM 爲離散型數據設計了 transition kernel,可以設爲 lazy random-walk,absorbing state 等。GEODIFF 爲 3D 分子圖數據設計了平移 - 旋轉不變的圖神經網絡,並且證明了具有不變性的初分佈和 transition kernel 可以導出具有不變性的邊緣分佈。假設 T 是一個平移 - 旋轉變換,如:

那麼生成的樣本分佈也有平移 - 旋轉不變性:

和其他生成模型的聯繫

在下面的每個小節中,我們首先介紹其他五類重要的生成模型,並分析它們的優勢和侷限性。然後我們介紹了擴散模型是如何與它們聯繫起來的,並說明通過結合擴散模型來改進這些生成模型。VAE,GAN,Autoregressive model, Normalizing flow, Energy-based model 和擴散模型的聯繫如下圖所示:

擴散模型的應用

在本節中,我們分別介紹了擴散模型在計算機視覺、自然語言處理、波形信號處理、多模態學習、分子圖生成、時間序列以及對抗學習等七大應用方向中的應用,並對每類應用中的方法進行了細分並解析。例如在計算機視覺中可以用 diffusion model 進行圖像補全修復(RePaint):

在多模態任務中可以用 diffusion model 進行文本到圖像的生成(GLIDE):

還可以在分子圖生成中用 diffusion model 進行藥物分子和蛋白質分子的生成(GeoDiff):

應用分類彙總見表:

未來研究方向