Generative Modeling by Estimating Gradients of the Data Distribution
本文介紹宋飏提出的帶噪聲擾動的基于得分的生成模型。首先介紹基本的基于得分的生成模型的訓練方法(得分匹配)和采樣方法(朗之萬動力學)。然后基于流形假設,分析了現有方法面臨的低密度區域得分估計不準、標準朗之萬動力學采樣對于多 pattern 分布混合慢的問題。最后提出解決方法,使用多種等級的噪聲對數據分布進行擾動,從而使得在低密度區域有更多的訓練信號;并提出退火的朗之萬動力學采樣方法,基于一個噪聲條件模型,在不同等級噪聲擾動分布上進行采樣,最終采樣出數據樣本。
一、背景知識:Score-based generative modeling
生成模型就是使用數據集來訓練一個模型從 p data ( x ) p_\text{data}(\mathbf{x}) pdata?(x) 中生成新的樣本。假設有從未知的數據分布 p data ( x ) p_\text{data}(\mathbf{x}) pdata?(x) 中 iid 采樣的數據集 { x i ∈ R D } i = 1 N \{\mathbf{x}_i\in\mathbb{R}^D\}_{i=1}^N {xi?∈RD}i=1N?。我們定義分布 p ( x ) p(\mathbf{x}) p(x) 的分數 score 為其概率密度對數的梯度,即 ? x log ? p ( x ) \nabla_\mathbf{x}\log p(\mathbf{x}) ?x?logp(x)。得分網絡 s θ : R D → R D \mathbf{s_\theta}:\ \mathbb{R}^D\rightarrow\mathbb{R}^D sθ?:?RD→RD 是一個參數化網絡,用于估計 p data ( x ) p_\text{data}(\mathbf{x}) pdata?(x) 的分數。score-based 生成模型有兩個關鍵部分:得分匹配(score matching),和朗之萬動力學(Langevin dynamics),分別對應訓練過程和采樣生成過程。score matching 用來訓練一個神經網絡估計數據分布的得分,而朗之萬動力學則能夠根據分布的得分來進行采樣。以下分別進行介紹。
1 score matching for score estimation
score matching 最初是用于基于未知分布的樣本點來訓練非歸一化統計模型(non-normalized statistical model)。之前已經開始有研究將其用于得分估計,本文也是這樣做。使用 score matching,我們無需知道數據分布 p data ( x ) p_\text{data}(\mathbf{x}) pdata?(x) 本身,而是可以直接訓練一個得分網絡 s θ ( x ) \mathbf{s_\theta(x)} sθ?(x) 來估計數據分布的得分 ? x log ? p data ( x ) \nabla_\mathbf{x}\log p_\text{data}(\mathbf{x}) ?x?logpdata?(x)。
score matching 的目標函數:
1 2 E p data ( x ) [ ∣ ∣ s θ ( x ) ? ? x log ? p data ( x ) ∣ ∣ 2 2 ] \frac{1}{2}\mathbb{E}_{p_\text{data}(\mathbf{x})}[||\mathbf{s_\theta}(\mathbf{x})-\nabla_\mathbf{x}\log p_\text{data}(\mathbf{x})||_2^2] 21?Epdata?(x)?[∣∣sθ?(x)??x?logpdata?(x)∣∣22?]
可以等價(至多差一個常數)為:
E p data ( x ) [ tr ( ? x s θ ( x ) ) + 1 2 ∣ ∣ s θ ( x ) ∣ ∣ 2 2 ] \mathbb{E}_{p_\text{data}(\mathbf{x})}[\text{tr}(\nabla_\mathbf{x}\mathbf{s_\theta}(\mathbf{x}))+\frac{1}{2}||\mathbf{s}_\theta(\mathbf{x})||^2_2] Epdata?(x)?[tr(?x?sθ?(x))+21?∣∣sθ?(x)∣∣22?]
其中 ? x s θ ( x ) \nabla_\mathbf{x}\mathbf{s_\theta}(\mathbf{x}) ?x?sθ?(x) 表示 s θ ( x ) \mathbf{s_\theta}(\mathbf{x}) sθ?(x) 的雅可比矩陣。已有研究工作顯示,在一定條件下,方程(3)的最小化解(記作 s θ ? ( x ) \mathbf{s_{\theta^*}}(\mathbf{x}) sθ??(x))也能滿足 s θ ? ( x ) = ? x log ? p data ( x ) \mathbf{s_{\theta^*}}(\mathbf{x})=\nabla_\mathbf{x}\log p_\text{data}(\mathbf{x}) sθ??(x)=?x?logpdata?(x)。實際中,可以使用數據樣本快速估計上式對 p data ( x ) p_\text{data}(\mathbf{x}) pdata?(x) 的期望,但由于跡 tr ( ? x s θ ( x ) ) \text{tr}(\nabla_\mathbf{x}\mathbf{s_\theta}(\mathbf{x})) tr(?x?sθ?(x)) 的計算復雜度太高,因此得分匹配并不適合直接用深層網絡和高維數據訓練。下面我們將討論兩種常見的大規模得分匹配的計算方法。
denoising score matching
denoising score matching 是 score matching 的一種變體,它完全繞過了對跡的計算。denoising score matching 中,首先使用預定義的噪聲分布 q σ ( x ~ ∣ x ) q_\sigma(\tilde{\mathbf{x}}|\mathbf{x}) qσ?(x~∣x) 對數據點 x \mathbf{x} x 進行擾動,然后使用 score matching 估計擾動后的分布,即 q σ ( x ) = △ ∫ q σ ( x ~ ∣ x ) p data ( x ) d x q_\sigma(\mathbf{x})\stackrel{\triangle}{=}\int q_\sigma(\tilde{\mathbf{x}}|\mathbf{x})p_\text{data}(\mathbf{x})d\mathbf{x} qσ?(x)=△∫qσ?(x~∣x)pdata?(x)dx ,的分數。這個目標函數為:
1 2 E q σ ( x ~ ∣ x ) p data ( x ) [ ∣ ∣ s θ ( x ~ ) ? ? x log ? q σ ( x ~ ∣ x ) ∣ ∣ 2 2 ) ] \frac{1}{2}\mathbb{E}_{q_\sigma(\tilde{\mathbf{x}}|\mathbf{x})p_\text{data}(\mathbf{x})}[||\mathbf{s_\theta}(\tilde{\mathbf{x}})-\nabla_\mathbf{x}\log q_\sigma(\tilde{\mathbf{x}}|\mathbf{x})||_2^2)] 21?Eqσ?(x~∣x)pdata?(x)?[∣∣sθ?(x~)??x?logqσ?(x~∣x)∣∣22?)]
已有研究證明,最小化上式的最優解 s θ ? ( x ) \mathbf{s_{\theta^*}}(\mathbf{x}) sθ??(x) 也能滿足 s θ ? ( x ) = ? x log ? q σ ( x ) \mathbf{s_{\theta^*}}(\mathbf{x})=\nabla_\mathbf{x}\log q_\sigma(\mathbf{x}) sθ??(x)=?x?logqσ?(x)。但是注意,只有在噪聲很小的時候,擾動后的分布才約等于數據分布。在生成時,當 q σ ( x ) ≈ p data ( x ) q_\sigma(\mathbf{x})\approx p_\text{data}(\mathbf{x}) qσ?(x)≈pdata?(x) 才能使用噪聲擾動分布的得分近似當做真實數據分布的得分。
siliced score matching
sliced score matching 使用一個隨機映射來對 score matching 中跡 tr ( ? x s θ ( x ) ) \text{tr}(\nabla_\mathbf{x}\mathbf{s_\theta}(\mathbf{x})) tr(?x?sθ?(x)) 的計算進行近似。目標函數是:
E p v E p data [ v T ? x s θ ( x ) v + 1 2 ∣ ∣ s θ ( x ) ∣ ∣ 2 2 ] \mathbb{E}_{p_\mathbf{v}}\mathbb{E}_{p_\text{data}}[\mathbf{v}^T\nabla_\mathbf{x}\mathbf{s_\theta}(\mathbf{x})\mathbf{v}+\frac{1}{2}||\mathbf{s_\theta}(\mathbf{x})||_2^2] Epv??Epdata??[vT?x?sθ?(x)v+21?∣∣sθ?(x)∣∣22?]
其中 p v p_\mathbf{v} pv? 是一個簡單的隨機向量分布(如多元高斯分布)。 v T ? x s θ ( x ) v \mathbf{v}^T\nabla_\mathbf{x}\mathbf{s_\theta}(\mathbf{x})\mathbf{v} vT?x?sθ?(x)v 這一項可以由自動微分框架的前向模式進行高效地計算。不同于 denoising score matching 估計的是擾動后分布的分數,sliced score matching 估計的是未擾動的原始數據的分布,但這會帶來大概四倍的額外計算量。
2 sampling with Langevin dynamics
朗之萬動力學可以在僅知道分布的得分函數 ? x log ? p ( x ) \nabla_\mathbf{x}\log p(\mathbf{x}) ?x?logp(x) 的情況下,采樣出分布 p ( x ) p(\mathbf{x}) p(x) 的樣本點。具體來說,給定固定的步長 ? > 0 \epsilon>0 ?>0 和一個先驗分布中采樣的初始值 x ~ ~ π ( x ) \tilde{\mathbf{x}}\sim\pi\mathbf({x}) x~~π(x) ,朗之萬迭代進行如下計算:
x ~ t = x ~ t ? 1 + ? 2 ? x log ? p ( x ~ t ? 1 ) + ? z t \tilde{\mathbf{x}}_t=\tilde{\mathbf{x}}_{t-1}+\frac{\epsilon}{2}\nabla_\mathbf{x}\log p(\tilde{\mathbf{x}}_{t-1})+\sqrt{\epsilon}\mathbf{z}_t x~t?=x~t?1?+2???x?logp(x~t?1?)+??zt?
其中 z t ~ N ( 0 , I ) \mathbf{z}_t\sim\mathcal{N}(0,\mathbf{I}) zt?~N(0,I)。當 ? → 0 , T → ∞ \epsilon\rightarrow 0,T\rightarrow\infty ?→0,T→∞ 時, x ~ T \tilde{\mathbf{x}}_T x~T? 的分布等于 p ( x ) p(\mathbf{x}) p(x),此時(在一定條件下) x ~ T \tilde{\mathbf{x}}_T x~T? 就相當于是從分布 p ( x ) p(\mathbf{x}) p(x) 中采樣出的一個樣本點了。嚴格來說,在中間步,即 ? > 0 , T < ∞ \epsilon>0,T<\infty ?>0,T<∞ 時,需要使用 Metropolis-Hastings 更新來修正上式的誤差。但在實際中,這個誤差經常是忽略的,本文也假設在 ? \epsilon ? 足夠小, T T T 足夠大的條件下,該誤差可以忽略。
在上述采樣公式中,可以看到,朗之萬動力學只需要知道得分函數 ? x log ? p ( x ) \nabla_\mathbf{x}\log p(\mathbf{x}) ?x?logp(x) 就可以進行采樣生成了。因此,要采樣得到數據分布 p data ( x ) p_\text{data}(\mathbf{x}) pdata?(x) 的樣本,我們只需有一個得分網絡能估計數據分布的得分 s θ ( x ) ≈ ? x log ? p data ( x ) \mathbf{s_\theta}(\mathbf{x})\approx\nabla_\mathbf{x}\log p_\text{data}(\mathbf{x}) sθ?(x)≈?x?logpdata?(x),然后就可以根據朗之萬動力學,使用 s θ ( x ) \mathbf{s_\theta}(\mathbf{x}) sθ?(x) 進行采樣就行了。這就是基于得分的生成模型(score-based generative modeling)的核心思路。
二、面臨挑戰:challenges of score-based model generative modeling
以上介紹了 score-based generative modeling 的基本思想與做法,但實際中,這種生成模型也存在著幾個問題,導致實際中應用還不廣泛。本節將介紹并分析 score-based generative modeling 所面臨的難點。
1 the manifold hypothesis
流形假設認為,現實世界中的數據大都集中在嵌入在高維空間中的低維流形上(換句話說,盡管數據可能存在于一個高維空間,但數據點實際上只在一個低維的子空間(流形)上)。該假設是流形學習的基石,在許多數據集上都能得到驗證。在流形假設下,score-based generative models 面臨著兩個問題。首先,得分 ? x log ? p data ( x ) \nabla_\mathbf{x}\log p_\text{data}(\mathbf{x}) ?x?logpdata?(x) 是高維空間分布的數據梯度,而當數據點 x \mathbf{x} x? 只分布在低維流形上時,該梯度是未定義的;第二,score matching 的目標函數只有在數據分布充滿整個空間中時,得分估計器才是準確的,當數據點只分布在低維流形上時,是不準確的。
數據分布在低維流形上,這對分數估計的影響可以在下圖 sliced score matching 在 CIFAR10 訓練過程的 ssm loss 曲線中看出。左側是直接在 CIFAR10 原圖上進行訓練,可以看到,loss 曲線明顯是震蕩的,非常不穩定,難以收斂;而右側是在原圖上疊加了很小幅度 N ( 0 , 0.0001 ) \mathcal{N}(0,0.0001) N(0,0.0001) 的高斯噪聲,這個噪聲幅度人眼幾乎看不出來,但是卻顯著穩定了訓練進行,loss 曲線可以收斂。
2 low data density regions
數據點在低密度區域的分布太少,會對 score-based generative models 中 score matching 訓練出的噪聲估計器的精確度和朗之萬動力學的 MCMC 采樣過程噪聲帶來困難。以下分別探討數據低密度區域對 score matching 訓練和朗之萬動力學采樣造成的影響。
2.1 inaccurate score estimation with score matching
在低數據密度區域,由于數據稀缺,score mathcing 沒有足夠的參考依據來準確地估計得分。score matching 的目標是最小化得分估計的平均誤差,即 1 2 E p data [ ∣ ∣ s θ ( x ) ? ? x log ? p data ( x ) ∣ ∣ 2 2 ] \frac{1}{2}\mathbb{E}_{p_\text{data}}[||\mathbf{s_\theta}(\mathbf{x})-\nabla_\mathbf{x}\log p_\text{data}(\mathbf{x})||_2^2] 21?Epdata??[∣∣sθ?(x)??x?logpdata?(x)∣∣22?]。在實際中,關于數據分布的期望是通過對數據分布的獨立同分布采樣 { x i } i = 1 N ~ i.i.d p data ( x ) \{\mathbf{x}_i\}_{i=1}^N\stackrel{\text{i.i.d}}{\sim}p_\text{data}(\mathbf{x}) {xi?}i=1N?~i.i.dpdata?(x)。對任意區域 R ? R D \mathcal{R}\sub \mathbb{R}^{D} R?RD,有 p data ( R ) ≈ 0 p_\text{data}(\mathcal{R})\approx 0 pdata?(R)≈0。在大多數情況下,有 { x i } i = 1 N ∩ R = ? \{\mathbf{x}_{i}\}_{i=1}^N\cap \mathcal{R}=\empty {xi?}i=1N?∩R=?,對于該區域中 score matching 沒有足夠的樣本來準確地估計得分。
下圖通過一個簡易的實驗驗證了這一點,該實驗中,作者用 sliced score matching 來擬合一個簡單的數據分布 p data = 1 5 N ( ( ? 5 , 5 ) , I ) + 4 5 N ( ( 5 , 5 ) , I ) p_\text{data}=\frac{1}{5}\mathcal{N}((-5,5),\mathbf{I})+\frac{4}{5}\mathcal{N}((5,5),\mathbf{I}) pdata?=51?N((?5,5),I)+54?N((5,5),I),下圖左側是數據分布真實的得分 ? x log ? p data ( x ) \nabla_\mathbf{x}\log p_\text{data}(\mathbf{x}) ?x?logpdata?(x),右側是模型估計的得分 s θ \mathbf{s_\theta} sθ?。可以看到(看箭頭),只有在數據密集的區域,得分估計才是準確的。
2.2 slow mixing of Langevin dynamics
當數據分布的兩個 pattern 被低密度區域分隔開時,朗之萬動力學對兩個 pattern 的相對權重的預測很慢,從而很難使得采樣結果收斂到真實分布。
考慮一個混合分布 p data ( x ) = π p 1 ( x ) + ( 1 ? π ) p 2 ( x ) p_{\text{data}}(\mathbf{x}) = \pi p_1(\mathbf{x}) + (1 - \pi) p_2(\mathbf{x}) pdata?(x)=πp1?(x)+(1?π)p2?(x),其中 p 1 ( x ) p_1(\mathbf{x}) p1?(x) 和 p 2 ( x ) p_2(\mathbf{x}) p2?(x) 是分布區域不相交的歸一化分布, π ∈ ( 0 , 1 ) \pi \in (0, 1) π∈(0,1)。在 p 1 ( x ) p_1(\mathbf{x}) p1?(x) 的分布區域內,有 $\nabla_x \log p_{\text{data}}(x) = \nabla_x(\log \pi + \log p_1(x)) = \nabla_x \log p_1(x) $,同理在 p 2 ( x ) p_2(\mathbf{x}) p2?(x) 的分布區域內,有 ? x log ? p data ( x ) = ? x ( log ? ( 1 ? π ) + log ? p 2 ( x ) ) = ? x log ? p 2 ( x ) \nabla_x \log p_{\text{data}}(\mathbf{x}) = \nabla_x(\log(1 - \pi) + \log p_2(\mathbf{x})) = \nabla_x \log p_2(\mathbf{x}) ?x?logpdata?(x)=?x?(log(1?π)+logp2?(x))=?x?logp2?(x)。也就是說,在這兩種情況下,得分 ? x log ? p data ( x ) \nabla_x \log p_{\text{data}}(\mathbf{x}) ?x?logpdata?(x) 都與 π \pi π 無關。朗之萬動力學使用 ? x log ? p data ( x ) \nabla_x \log p_{\text{data}}(\mathbf{x}) ?x?logpdata?(x) 來從 p data ( x ) p_{\text{data}}(\mathbf{x}) pdata?(x) 中采樣,所得到的樣本將不依賴于權重系數 π \pi π。在實際中,雖然不同的 pattern 可能存在一定的相交分布區域,這部分相交區域的數據密度很低,因此結論是類似的,在這種情況下,理論上朗之萬動力學還是可以生成正確的樣本,但可能需要非常小的步長、非常大的步數。
為了驗證這一分析,作者使用了與上節相同的簡單數據分布進行了朗之萬動力學采樣實驗,結果如下圖。圖中 (a)、(b) 分別是對數據分布直接采樣、根據數據分布的真實 score 進行朗之萬動力學采樣(注意這里的實驗設了一個簡單的數據分布,也就是說數據分布的確切解析式我們是已知的),顯然圖 (b) 朗之萬動力學得到的樣本之間兩種 pattern 的相對密度是不正確的,與上述分析相符。
三、提出方法:Noise Conditional Score Networks: learning and inference
如何解決 score-based generative models 的上述問題呢?
作者發現,使用隨機高斯噪聲對數據分布進行擾動,會使得數據分布更適合于 score-based generative modeling。首先,由于高斯分布遍布整個空間,因此使用其對數據分布進行擾動,分布就不會僅局限在低維流形上了,這使得數據分布的得分在整個空間上都是良定義的,消除了流形假設帶來的難點。第二,大的高斯噪聲可以填滿原本未擾動數據分布的低密度區域,從而 score mathcing 的訓練過程就更多的監督信號,模型估計的得分就更準。第三,使用多種噪聲等級,我們可以獲得一系列噪聲擾動分布,這一系列噪聲擾動分布逐漸逼近真實數據分布。通過使用模擬退火(simulated annealing)和退火重要性采樣(annealed importance sampling)的思想,我們可以利用這些中間分布來提高多模態分布上朗之萬動力學采樣的速率。
基于以上思考,作者提出了對 score-based generative modeling 的改進措施:1)使用多種不同等級的噪聲對數據分布進行擾動;2)訓練單個(條件)神經網絡來同時估計所有噪聲等級所對應分布的得分。訓練結束后,使用朗之萬動力學來進行采樣生成時,我們從最大噪聲所對應分布的得分開始,逐步退火降低噪聲等級,這有助于權衡大小噪聲水平的益處,平穩地將大噪聲水平的益處逐漸轉移到低噪聲水平,直到擾動后的數據幾乎與原始數據無法區分。下面,我們將更詳細地闡述我們方法的細節,包括得分網絡的架構、訓練目標和退火的朗之萬動力學采樣。
1 網絡:noise conditional score networks
記 { σ i } i = 1 L \{\sigma_i\}_{i=1}^L {σi?}i=1L? 是一個正的等比序列,滿足: σ 1 σ 2 = ? = σ L ? 1 σ L > 1 \frac{\sigma_1}{\sigma_2}=\dots=\frac{\sigma_{L-1}}{\sigma_{L}}>1 σ2?σ1??=?=σL?σL?1??>1,它是我們的噪聲等級序列。記 q σ ( x ) = △ ∫ p data ( t ) N ( x ∣ t , σ 2 I ) d t q_\sigma(\mathbf{x})\stackrel{\triangle}{=}\int p_\text{data}(\mathbf{t})\mathcal{N}(\mathbf{x}|\mathbf{t},\sigma^2\mathbf{I})d\mathbf{t} qσ?(x)=△∫pdata?(t)N(x∣t,σ2I)dt 表示經噪聲擾動的數據分布。在我們選定的噪聲等級序列中,要滿足 1) σ 1 \sigma_1 σ1? 足夠大,能夠消除流形假設中提到的低密度區域訓練信號不足問題;2) σ L \sigma_L σL? 足夠小,在最后減少對真實數據分布的影響。我們要訓練一個條件得分網絡,來按條件估計不同噪聲擾動等級下分布的得分,即 ? σ ∈ { σ i } i = 1 L , s θ ( x , σ ) ≈ ? x log ? q σ ( x ) \forall\sigma\in\{\sigma_i\}_{i=1}^L,\ \ \mathbf{s_\theta}(\mathbf{x},\sigma)\approx\nabla_\mathbf{x}\log q_\sigma(\mathbf{x}) ?σ∈{σi?}i=1L?,??sθ?(x,σ)≈?x?logqσ?(x)。注意我們的模型輸入輸出尺寸是相同的,即當 x ∈ R D \mathbf{x}\in\mathbb{R}^D x∈RD 時, s θ ( x , σ ) ∈ R D \mathbf{s_\theta}(\mathbf{x,\sigma})\in\mathbb{R}^D sθ?(x,σ)∈RD。作者將網絡 s θ ( x , σ ) \mathbf{s_\theta}(\mathbf{x},\sigma) sθ?(x,σ) 命名為噪聲條件得分網絡(Noise Conditional Score Networks,NCSN)。
具體網絡結構方面,由于 NCSN 的輸入輸出尺寸相同,作者借鑒密集預測類任務(如圖像分割),選擇了 UNet 作為主體網絡結構,加入了空洞卷積、實例歸一化等技術,并使用一種改進的條件實例歸一化(conditional instance normalization)來將條件 σ \sigma σ 注入到模型中。
2 訓練:learning NCSNs via score matching
作者強調,之前介紹的 sliced score matching 和 denoising score matching 都可以用來訓練 NCSN。本文中考慮到 dsm 訓練更快且更適合我們估計噪聲擾動數據分布的得分的任務,故而采用 dsm 進行訓練。本文選用噪聲分布為
q σ ( x ~ ∣ x ) = N ( x ~ ∣ x , σ 2 I ) q_\sigma(\tilde{\mathbf{x}}|\mathbf{x})=\mathcal{N}(\tilde{\mathbf{x}}|\mathbf{x},\sigma^2\mathbf{I}) qσ?(x~∣x)=N(x~∣x,σ2I)
從而有:
KaTeX parse error: Got function '\tilde' with no arguments as subscript at position 9: \nabla_\?t?i?l?d?e?{\mathbf{x}}\lo…
對于給定的 σ \sigma σ,denosing score matching 的目標函數為:
l ( θ ; σ ) = △ 1 2 E p data ( x ) E x ~ ~ N ( x , σ 2 I ) [ ∣ ∣ s θ ( x ~ , σ ) + x ~ ? x σ 2 ∣ ∣ 2 2 ] l(\theta;\sigma)\stackrel{\triangle}{=}\frac{1}{2}\mathbb{E}_{p_\text{data}(\mathbf{x})}\mathbb{E}_{\tilde{\mathbf{x}}\sim\mathcal{N}(\mathbf{x},\sigma^2\mathbf{I})}[||\mathbf{s_\theta}(\mathbf{\tilde{x}},\sigma)+\frac{\tilde{\mathbf{x}}-\mathbf{x}}{\sigma^2}||_2^2] l(θ;σ)=△21?Epdata?(x)?Ex~~N(x,σ2I)?[∣∣sθ?(x~,σ)+σ2x~?x?∣∣22?]
對于所有的 σ ∈ { σ i } i = 1 L \sigma\in\{\sigma_i\}_{i=1}^L σ∈{σi?}i=1L? 結合起來,統一的目標函數為:
L ( θ ; { σ i } i = 1 L ) = △ 1 L ∑ i = 1 L λ ( σ i ) l ( θ ; σ i ) \mathcal{L}(\theta;\{\sigma_i\}_{i=1}^L)\stackrel{\triangle}{=}\frac{1}{L}\sum_{i=1}^L\lambda(\sigma_i)l(\theta;\sigma_i) L(θ;{σi?}i=1L?)=△L1?i=1∑L?λ(σi?)l(θ;σi?)
其中 λ ( σ i ) > 0 \lambda(\sigma_i)>0 λ(σi?)>0 是一個關于 σ i \sigma_i σi? 的加權函數。假設 s θ ( x , σ ) \mathbf{s_\theta}(\mathbf{x},\sigma) sθ?(x,σ) 訓練得足夠好,對所有的 i ∈ [ 1 , L ] i\in[1,L] i∈[1,L] 都能最小化上述目標函數,那么模型就能準確地估計噪聲擾動數據分布的得分,即 s θ ? ( x , σ ) = ? x log ? q σ i ( x ) \mathbf{s_{\theta^*}}(\mathbf{x},\sigma)=\nabla_\mathbf{x}\log q_{\sigma_i}(\mathbf{{x}}) sθ??(x,σ)=?x?logqσi??(x)。
對于 λ ( ? ) \lambda(\cdot) λ(?) 的選擇,可以有很多種。理想情況下,我們希望對所有的 σ i \sigma_i σi?,損失 λ ( σ i ) l ( θ ; σ i ) \lambda(\sigma_i)l(\theta;\sigma_i) λ(σi?)l(θ;σi?) 的值大致在相同的數量級上。實驗中,作者發現當網絡訓練得比較好時,大致有 ∣ ∣ s θ ( x , σ ) ∣ ∣ 2 ∝ 1 / σ ||\mathbf{s_\theta}(\mathbf{x},\sigma)||_2\propto1/\sigma ∣∣sθ?(x,σ)∣∣2?∝1/σ。作者據此選擇 λ ( σ ) = σ 2 \lambda(\sigma)=\sigma^2 λ(σ)=σ2。這是因為在此選擇下,有 λ ( σ ) l ( θ ; σ ) = σ 2 l ( ( θ ; σ ) ) = 1 2 E [ ∣ ∣ σ s θ ( x ; σ ) + x ~ ? x σ ∣ ∣ 2 2 ] \lambda(\sigma)l(\theta;\sigma)=\sigma^2l((\theta;\sigma))=\frac{1}{2}\mathbb{E}[||\sigma\mathbf{s_\theta}(\mathbf{x};\sigma)+\frac{\tilde{\mathbf{x}}-\mathbf{x}}{\sigma}||^2_2] λ(σ)l(θ;σ)=σ2l((θ;σ))=21?E[∣∣σsθ?(x;σ)+σx~?x?∣∣22?] ,而由于 x ~ ? x σ ~ N ( 0 , I ) \frac{\tilde{\mathbf{x}}-\mathbf{x}}{\sigma}\sim\mathcal{N}(0,\mathbf{I}) σx~?x?~N(0,I)、 ∣ ∣ σ s θ ( x ; σ ) ∣ ∣ 2 ∝ 1 ||\sigma\mathbf{s_\theta}(\mathbf{x};\sigma)||_2\propto 1 ∣∣σsθ?(x;σ)∣∣2?∝1,因此可得此時 λ ( σ ) l ( θ ; σ ) \lambda(\sigma)l(\theta;\sigma) λ(σ)l(θ;σ) 值的數量級與 σ \sigma σ 無關。
3 采樣:NCSN inference via annealed Langevin dynamics
在 NCSN s θ ( x , σ ) \mathbf{s_\theta}(\mathbf{x},\sigma) sθ?(x,σ) 訓練完成后,我們就能夠對不同噪聲擾動數據分布 q σ i ( x ) q_{\sigma_i}(\mathbf{x}) qσi??(x) 的得分 ? x q σ 1 ( x ) \nabla_\mathbf{x}q_{\sigma_1}(\mathbf{x}) ?x?qσ1??(x) 進行較為準確的估計。這樣,就能執行朗之萬動力學采樣,從各個分布 q σ i ( x ) q_{\sigma_i}(\mathbf{x}) qσi??(x) 中采樣出樣本。但是,如前所述,流形假設指出了朗之萬動力學采樣對多 pattern 數據分布采樣時混合較慢的問題,并且我們的 NCSN 需要估計的是多個噪聲擾動分布 q σ i ( x ) q_{\sigma_i}(\mathbf{x}) qσi??(x) 的得分,而不是數據分布 p data ( x ) p_\text{data}(\mathbf{x}) pdata?(x) 的得分。因此,本文提出了一種新的采用方法——退火的朗之萬動力學(Annealed Langevin Dynamics),來使用 NCSN 進行采樣生成。
具體算法如下圖所示。我們從一個固定的先驗分布(如均勻分布)中采樣一個初始樣本。首先對第一級的噪聲擾動數據分布 q σ 1 ( x ) q_{\sigma_1}(\mathbf{x}) qσ1??(x) 執行 T T T 步朗之萬動力學采樣,步長為 α 1 \alpha_1 α1?。然后將對 q σ 1 ( x ) q_{\sigma_1}(\mathbf{x}) qσ1??(x) 的最終采樣結果作為初始樣本,再對 q σ 2 ( x ) q_{\sigma_2}(\mathbf{x}) qσ2??(x) 進行朗之萬動力學采樣,步長為 α 2 \alpha_2 α2?,注意 α 2 < α 1 \alpha_2<\alpha_1 α2?<α1?。迭代進行這種方式,每次使用上一級分布 q σ i ? 1 ( x ) q_{\sigma_{i-1}}(\mathbf{x}) qσi?1??(x) 朗之萬動力學采樣迭代的最后一個樣本,作為下一級分布 q σ i ( x ) q_{\sigma_i}(\mathbf{x}) qσi??(x) 采樣的起始樣本,并逐漸調小步長 α i = ? ? σ i 2 / σ L 2 \alpha_i=\epsilon\cdot\sigma_i^2/\sigma_L^2 αi?=??σi2?/σL2?。最終,我們對最后一級噪聲擾動分布 q σ L ( x ) q_{\sigma_L}(\mathbf{x}) qσL??(x) 執行朗之萬動力學采樣,得到最終的生成結果。由于 σ L ≈ 0 \sigma_L\approx0 σL?≈0 足夠小,因此 q σ L ( x ) q_{\sigma_L}(\mathbf{x}) qσL??(x) 已經非常接近真實的數據分布 p data ( x ) p_\text{data}(\mathbf{x}) pdata?(x)。
由于噪聲擾動分布 { q σ i } i = 1 L \{q_{\sigma_i}\}_{i=1}^L {qσi??}i=1L? 都是由高斯噪聲進行擾動,因此這些分布的存在區域都是充滿整個空間的,從而他們的得分都是良定義的,避免了前文所述的流形假設指出的分布得分未定義的問題。當 σ 1 \sigma_1 σ1? 足夠大時,分布 q σ 1 ( x ) q_{\sigma_1}(\mathbf{x}) qσ1??(x) 的低密度區域就很小,并且不同 pattern 之間也沒有那么分離。從而,得分估計就更準,并且朗之萬動力學采樣對不同 pattern 的混合也更快。因此,朗之萬動力學從 q σ 1 ( x ) q_{\sigma_1}(\mathbf{x}) qσ1??(x) 采樣出的結果就是比較好的樣本。這些樣本很大概率是來自 q σ 1 ( x ) q_{\sigma_1}(\mathbf{x}) qσ1??(x) 的高密度區域,而由于 q σ 2 ( x ) q_{\sigma_2}(\mathbf{x}) qσ2??(x) 與 q σ 1 ( x ) q_{\sigma_1}(\mathbf{x}) qσ1??(x) 非常接近,從而這些樣本同樣很大概率落在 q σ 2 ( x ) q_{\sigma_2}(\mathbf{x}) qσ2??(x) 的高密度區域。而得分估計和朗之萬動力學采樣都是在分布的高密度區域工作得更好,也就是說第一級噪聲擾動分布 q σ 1 ( x ) q_{\sigma_1}(\mathbf{x}) qσ1??(x) 中采樣出來的最終樣本為第二級噪聲擾動分布 q σ 2 ( x ) q_{\sigma_2}(\mathbf{x}) qσ2??(x) 的朗之萬動力學采樣提供了很好的初始化。以此類推,每一級分布 q σ i ? 1 ( x ) q_{\sigma_{i-1}}(\mathbf{x}) qσi?1??(x) 的樣本都為下一級分布 q σ i ( x ) q_{\sigma_i}(\mathbf{x}) qσi??(x) 的采樣提供了很好的初始化。最終我們能從分布 q σ L ( x ) q_{\sigma_L}(\mathbf{x}) qσL??(x) 中采樣到高質量的結果。
步長 α i \alpha_i αi? 的設置也有很多可能。本文的選擇是 α i ∝ σ i 2 \alpha_i\propto\sigma_i^2 αi?∝σi2?。這樣設計的動機是想要固定朗之萬動力學采樣中信噪比 (Signal-to-Noise Ratio, SNR) α i s θ ( x , σ i ) 2 α i z \frac{\alpha_i\mathbf{s_\theta(\mathbf{x},\sigma_i)}}{2\sqrt{\alpha_i}\mathbf{z}} 2αi??zαi?sθ?(x,σi?)? 的尺度。注意有 E [ ∣ ∣ α i s θ ( x , σ i ) 2 α i z ∣ ∣ 2 2 ] ≈ E [ α i ∣ ∣ s θ ( x , σ i ) ∣ ∣ 2 2 4 ] ∝ 1 4 E [ ∣ ∣ σ i s θ ( x , σ i ) ∣ ∣ 2 2 ] \mathbb{E}[||\frac{\alpha_i\mathbf{s_\theta(\mathbf{x},\sigma_i)}}{2\sqrt{\alpha_i}\mathbf{z}}||_2^2]\approx\mathbb{E}[\frac{\alpha_i||\mathbf{s_\theta}(\mathbf{x},\sigma_i)||_2^2}{4}]\propto\frac{1}{4}\mathbb{E}[||\sigma_i\mathbf{s_\theta}(\mathbf{x},\sigma_i)||_2^2] E[∣∣2αi??zαi?sθ?(x,σi?)?∣∣22?]≈E[4αi?∣∣sθ?(x,σi?)∣∣22??]∝41?E[∣∣σi?sθ?(x,σi?)∣∣22?]。我們之前提到,當 NCSN 訓練得比較好的時候,有 ∣ ∣ s θ ( x , σ ) ∣ ∣ 2 ∝ 1 / σ ||\mathbf{s_\theta}(\mathbf{x},\sigma)||_2\propto1/\sigma ∣∣sθ?(x,σ)∣∣2?∝1/σ,此時有 E [ ∣ ∣ σ i s θ ( x , σ i ) ∣ ∣ 2 2 ∝ 1 \mathbb{E}[||\sigma_i\mathbf{s_\theta}(\mathbf{x},\sigma_i)||_2^2\propto1 E[∣∣σi?sθ?(x,σi?)∣∣22?∝1,從而信噪比 E [ ∣ ∣ α i s θ ( x , σ i ) 2 α i z ∣ ∣ 2 2 ] ∝ 1 4 E [ ∣ ∣ σ i s θ ( x , σ i ) ∣ ∣ 2 2 ] ∝ 1 4 \mathbb{E}[||\frac{\alpha_i\mathbf{s_\theta(\mathbf{x},\sigma_i)}}{2\sqrt{\alpha_i}\mathbf{z}}||_2^2]\propto\frac{1}{4}\mathbb{E}[||\sigma_i\mathbf{s_\theta}(\mathbf{x},\sigma_i)||_2^2]\propto\frac{1}{4} E[∣∣2αi??zαi?sθ?(x,σi?)?∣∣22?]∝41?E[∣∣σi?sθ?(x,σi?)∣∣22?]∝41?,與 σ i \sigma_i σi? 無關。
為了展示退火朗之萬動力學的有效性,本文在前面介紹標準朗之萬動力學混合慢的簡單實驗中,對比了本文所提方法。回顧一下,實驗的目標是使用得分從一個簡單的兩個高斯混合的數據分布中進行采樣。退火的朗之萬動力學的采樣結果如圖 ? 所示,明顯可以看到,采樣出的數據能夠反映出兩個高斯分布的比重差異,也就是說能夠更真實地從數據分布中采樣,而標準的朗之萬動力學采樣(圖 (b)顯然沒有做到這一點。
結語
這篇 NCSN 的方法提出在 DDPM 之前,并且對 DDPM 的設計也有啟發。更重要的是,筆者認為,基于得分的生成模型似乎更加直覺,更具有可解釋性。從理解的角度來說,NSCN 更加順暢。但是 DDPM 的效果似乎更好?因此一般認為 DDPM 是真正將 Diffusion 這一類方法做 work 的工作。宋飏博士在之后的 score based SDE 中也已經將 NCSN 和 DDPM 統一到了 SDE 的框架下。