擴(kuò)散模型=流匹配?谷歌DeepMind博客深度詳解這種驚人的等價(jià)性
機(jī)器之心報(bào)道機(jī)器之心編輯部擴(kuò)散模型和流匹配實(shí)際上是同一個(gè)概念的兩種不同表達(dá)方式嗎?從表面上看,這兩種方法似乎各有側(cè)重:擴(kuò)散模型專注于通過迭代的方式逐步去除噪聲,將數(shù)據(jù)還原成清晰的樣本。而流匹配則側(cè)重于構(gòu)建可逆變換系統(tǒng),目標(biāo)是學(xué)習(xí)如何將簡單的基礎(chǔ)分布精確地映射到真實(shí)數(shù)據(jù)分布。因?yàn)榱髌ヅ涞墓胶芎唵危⑶疑蓸颖镜穆窂胶苤苯樱罱絹碓绞苎芯空邆兊臍g迎,于是很多人都在問:「到底是擴(kuò)散模型好呢?還是流匹配好?」現(xiàn)在,這個(gè)困擾已得到解答。Google DeepMind 的研究團(tuán)隊(duì)發(fā)現(xiàn),原來擴(kuò)散模型和流匹配就像一枚硬幣的兩面,本質(zhì)上是等價(jià)的 (尤其是在流匹配采用高斯分布作為基礎(chǔ)分布時(shí)),只是不同的模型設(shè)定會導(dǎo)致不同的網(wǎng)絡(luò)輸出和采樣方案。這無疑是個(gè)好消息,意味著這兩種框架下的方法可以靈活搭配,發(fā)揮組合技了。比如在訓(xùn)練完一個(gè)流匹配模型后,不必再局限于傳統(tǒng)的確定性采樣方法,完全可以引入隨機(jī)采樣策略。鏈接:https://diffusionflow.github.io在這篇博客的開頭,作者們寫道:「我們的目標(biāo)是幫助大家能夠自如地交替使用這兩種方法,同時(shí)在調(diào)整算法時(shí)擁有真正的度 —— 方法的名稱并不重要,重要的是理解其本質(zhì)。」擴(kuò)散模型與流匹配擴(kuò)散模型主要分為前向過程和反向兩個(gè)階段。前向過程用公式表示為:其中 z_t 是在時(shí)間點(diǎn) t 時(shí)的帶噪聲數(shù)據(jù),x 代表原始數(shù)據(jù),ε 代表隨機(jī)噪聲,a_t 和 σ_t 是控制噪聲添加程度的參數(shù)。若滿足,稱為「方差保持」,意味著在每個(gè)時(shí)間步驟中,噪聲的方差保持不變或接近不變。DDIM 采樣器的反向過程用公式表示為:其中,而在流匹配中,前向過程視為數(shù)據(jù) x 和噪聲項(xiàng) ε 之間的線性插值:采樣人們普遍認(rèn)為,這兩個(gè)框架在生成樣本的方式上有所不同:流匹配采樣是確定性的,具有直線路徑,而擴(kuò)散模型采樣是隨機(jī)性的,具有曲線路徑。下面文章將澄清這一誤解:首先關(guān)注更簡單的確定性采樣,稍后再討論隨機(jī)情況。假設(shè)你想使用訓(xùn)練好的降噪器模型將隨機(jī)噪聲轉(zhuǎn)換為數(shù)據(jù)點(diǎn)。可以先回想一下 DDIM 的更新,有趣的是,重新排列項(xiàng)可以用以下公式來表達(dá),這里涉及幾組網(wǎng)絡(luò)輸出和重新參數(shù)化:我們再回到公式(4)中的流匹配更新,和上述方程看起來很相似。如果在最后一行將網(wǎng)絡(luò)輸出設(shè)為,并令,可以得到、這樣我們就恢復(fù)了流匹配更新!更準(zhǔn)確地說,流匹配更新可以被視為重參數(shù)化采樣常微分方程(ODE)的歐拉積分:對于 DDIM 采樣器而言,普遍存在以下結(jié)論:DDIM 采樣器對于應(yīng)用于噪聲調(diào)度α_t、σ_t的線性縮放是不變的,因?yàn)榭s放不會影響和,這對于其他采樣器來說并不成立,例如概率流 ODE 的歐拉采樣器。為了驗(yàn)證上述結(jié)論,本文展示了使用幾種不同的噪聲調(diào)度得到的結(jié)果,每種調(diào)度都遵循流匹配調(diào)度,并具有不同的縮放因子。如下圖,隨意調(diào)整滑塊,在最左側(cè),縮放因子是 1,這正是流匹配調(diào)度,而在最右側(cè),縮放因子是。可以觀察到 DDIM(以及流匹配采樣器)總是給出相同的最終數(shù)據(jù)樣本,無論調(diào)度的縮放如何。對于概率流 ODE 的歐拉采樣器,縮放確實(shí)會產(chǎn)生真正的差異:可以看到路徑和最終樣本都發(fā)生了變化。看到這里,需要思考一下。人們常說流匹配會產(chǎn)生直線路徑,但在上圖中,其采樣軌跡看起來是彎曲的。在下面的交互式圖表中,我們可以通過滑塊更改右側(cè)數(shù)據(jù)分布的方差。不過,在像圖像這樣的真實(shí)數(shù)據(jù)集上找到這樣的直線路徑要復(fù)雜得多。但結(jié)論仍然是相同的:最優(yōu)的積分方法取決于數(shù)據(jù)分布。我們可以從確定性采樣中得到的兩個(gè)重要結(jié)論:采樣器的等價(jià)性:DDIM 與流匹配采樣器等價(jià),并且對噪聲調(diào)度的線性縮放不變。對直線性的誤解:流匹配調(diào)度僅在模型預(yù)測單個(gè)點(diǎn)時(shí)才是直線。訓(xùn)練對于擴(kuò)散模型,學(xué)習(xí)模型是通過最小化加權(quán)均方誤差(MSE)損失來完成的:流匹配也符合上述訓(xùn)練目標(biāo):網(wǎng)絡(luò)應(yīng)該輸出什么下面總結(jié)了文獻(xiàn)中提出的幾個(gè)網(wǎng)絡(luò)輸出,包括擴(kuò)散模型使用的幾個(gè)版本和流匹配使用的其中一個(gè)版本。然而,在實(shí)踐中,模型的輸出可能會產(chǎn)生非常大的影響。例如,基于相似的原因,在低噪聲水平下是有問題的,因?yàn)闆]有信息量,并且錯(cuò)誤在中被放大了。因此,一種啟發(fā)式方法是選擇一個(gè)網(wǎng)絡(luò)輸出,它是、的組合,這適用于和流匹配矢量場如何選擇加權(quán)函數(shù)加權(quán)函數(shù)是損失函數(shù)中最重要的部分,它平衡了圖像、視頻和音頻等數(shù)據(jù)中高頻和低頻分量的重要性。這一點(diǎn)至關(guān)重要,因?yàn)檫@些信號中的某些高頻分量是人類無法感知的。如果通過加權(quán)情況來查看損失函數(shù),可以得出以下結(jié)果:即公式 (7) 中的條件流匹配目標(biāo)與擴(kuò)散模型中常用的設(shè)置相同。下面繪制了文獻(xiàn)中常用的幾個(gè)加權(quán)函數(shù)。流匹配加權(quán)(也稱為 v-MSE + 余弦調(diào)度加權(quán))會隨著 λ 的增加而呈指數(shù)下降。該團(tuán)隊(duì)在實(shí)驗(yàn)中發(fā)現(xiàn)了另一個(gè)有趣的聯(lián)系:Stable Diffusion 3 加權(quán) [9](這是流匹配的一種重新加權(quán)版本)與擴(kuò)散模型中流行的 EDM 加權(quán) [10] 非常相似。如何選擇訓(xùn)練噪聲調(diào)度?最后討論訓(xùn)練噪聲調(diào)度,因?yàn)樵谝韵乱饬x上,它對訓(xùn)練的重要程度最低:1. 訓(xùn)練損失不會隨訓(xùn)練噪聲調(diào)度變化。具體來說,損失函數(shù)可以重寫為它只與端點(diǎn)(λ_max, λ_min)有關(guān),但與中間的調(diào)度 λ_t 無關(guān)。在實(shí)踐中,應(yīng)該選擇合適的 λ_max, λ_min,使得兩端分別足夠接近干凈數(shù)據(jù)和高斯噪聲。λ_t 可能仍然會影響訓(xùn)練損失的蒙特卡洛估計(jì)量的方差。一些文獻(xiàn)中提出了一些啟發(fā)式方法來在訓(xùn)練過程中自動調(diào)整噪聲調(diào)度。這篇博文有一個(gè)很好的總結(jié):https://sander.ai/2024/06/14/noise-schedules.html#adaptive2. 類似于采樣噪聲調(diào)度,訓(xùn)練噪聲調(diào)度不會隨線性擴(kuò)展(linear scaling)而變化,因?yàn)槿藗兛梢暂p松地將線性擴(kuò)展應(yīng)用于 z_t,并在網(wǎng)絡(luò)輸入處進(jìn)行 unscaling 以獲得等價(jià)性。噪聲調(diào)度的關(guān)鍵定義屬性是對數(shù)信噪比 λ_t。3. 人們可以根據(jù)不同的啟發(fā)式方法為訓(xùn)練和采樣選擇完全不同的噪聲調(diào)度:對于訓(xùn)練,最好有一個(gè)噪聲調(diào)度來最小化蒙特卡洛估計(jì)量的方差;而對于采樣,噪聲調(diào)度與 ODE / SDE 采樣軌跡的離散化誤差和模型曲率更相關(guān)。總結(jié)下面給出了訓(xùn)練擴(kuò)散模型 / 流匹配的一些要點(diǎn):加權(quán)中的等價(jià)性:加權(quán)函數(shù)對于訓(xùn)練很重要,它平衡了感知數(shù)據(jù)不同頻率分量的重要性。流匹配加權(quán)與常用的擴(kuò)散訓(xùn)練加權(quán)方法相同。訓(xùn)練噪聲調(diào)度的不重要性:噪聲調(diào)度對訓(xùn)練目標(biāo)的重要性要小得多,但會影響訓(xùn)練效率。網(wǎng)絡(luò)輸出的差異:流匹配提出的網(wǎng)絡(luò)輸出是新的,它很好地平衡了更深入地理解采樣器這一節(jié)將更詳細(xì)地介紹各種不同的采樣器。回流算子流匹配中的回流(Reflow)運(yùn)算是使用直線將噪聲與數(shù)據(jù)點(diǎn)連接起來。通過基于噪聲運(yùn)行一個(gè)確定性的采樣器,可以得到這些 (數(shù)據(jù),噪聲) 對。然后,可以訓(xùn)練模型,使之可以根據(jù)給定噪聲直接預(yù)測數(shù)據(jù),而無需采樣。在擴(kuò)散技術(shù)的相關(guān)文獻(xiàn)中,這同樣的方法是最早的蒸餾技術(shù)之一。確定性采樣器與隨機(jī)采樣器此前已經(jīng)討論了擴(kuò)散模型或流匹配的確定性采樣器。另一種方法是使用隨機(jī)采樣器,例如 DDPM 采樣器。執(zhí)行一個(gè)從 λ_t 到 λ_t+Δλ 的 DDPM 采樣步驟完全等價(jià)于執(zhí)行一個(gè)到 λ_t+2Δλ 的 DDIM 采樣步驟,然后通過執(zhí)行前向擴(kuò)散重新噪聲化到 λ_t+Δλ。也就是說,通過前向擴(kuò)散重新噪聲化恰好逆轉(zhuǎn)了 DDIM 所取得的一半進(jìn)展。為了理解這一點(diǎn),讓我們看一個(gè) 2D 示例。從相同的高斯分布混合開始,我們可以執(zhí)行一個(gè)小的 DDIM 采樣步驟,左圖帶有更新反轉(zhuǎn)的符號,右圖則是一個(gè)小的前向擴(kuò)散步驟:對于單個(gè)樣本而言,這些更新的行為完全不同:反轉(zhuǎn)的 DDIM 更新始終將每個(gè)樣本推離分布模式,而擴(kuò)散更新完全是隨機(jī)的。但是,在匯總所有樣本時(shí),更新后得到的分布是相同的。因此,如果執(zhí)行 DDIM 采樣步驟(不反轉(zhuǎn)符號),然后執(zhí)行前向擴(kuò)散步驟,則整體分布與更新之前的分布保持不變。通過重新加噪來撤消的 DDIM 步驟的比例是一個(gè)超參數(shù),并且可以選擇(即不必一定是 DDIM 步驟的一半)。這個(gè)超參數(shù)在《Elucidating the design space of diffusion-based generative models》中被稱為 level of churn,可譯為「攪動水平」。有趣的是,將攪動添加到采樣器的效果是:減少采樣過程早期做出的模型預(yù)測對最終樣本的影響,并增加對后續(xù)預(yù)測的權(quán)重。如下圖所示:在這里,我們使用余弦噪聲調(diào)度以及預(yù)測將不同采樣器都運(yùn)行了 100 個(gè)采樣步驟。忽略非線性相互作用,采樣器產(chǎn)生的最終樣本可以寫成采樣過程中做出的預(yù)測和高斯噪聲 e 的加權(quán)和:這些預(yù)測的權(quán)重 h_t 顯示在 y 軸上,而 x 軸上顯示不同的擴(kuò)散時(shí)間 t。DDIM 會在此設(shè)置下對預(yù)測賦予相等的權(quán)重,而 DDPM 則更注重在采樣結(jié)束時(shí)所做的預(yù)測。另請參閱《Dpm-solver++: Fast solver for guided sampling of diffusion probabilistic models》以了解中這些權(quán)重的解析表達(dá)式。SDE 和 ODE 視角前面,我們已經(jīng)觀察到擴(kuò)散模型和流匹配算法之間的等價(jià)性。下面將使用 ODE 和 SDE 來形式化地描述正向過程和采樣的等價(jià)性,以實(shí)現(xiàn)理論上的完整性。擴(kuò)散模型擴(kuò)散模型的前向過程涉及到隨時(shí)間推移逐漸破壞一個(gè)數(shù)據(jù),而該過程可使用以下隨機(jī)微分方程(SDE)來描述:其中 dz 是無窮小的高斯(即布朗)。f_t 和 g_t 決定了噪聲調(diào)度。其生成過程由前向過程的逆過程給出,其公式為:其中 ?log p_t 是前向過程的分?jǐn)?shù)。請注意,這里引入了一個(gè)附加參數(shù) η_t,它控制的是推理時(shí)的隨機(jī)性。這與之前介紹的攪動(churn)參數(shù)有關(guān)。當(dāng)離散化后向過程時(shí),如果 η_t=0,則是恢復(fù) DDIM;如果 η_t=1,則是恢復(fù) DDPM。流匹配流匹配中 x 和 ε 之間的插值可以用以下常微分方程(ODE)描述:假設(shè)該插值為其生成過程只是在時(shí)間上反轉(zhuǎn)這個(gè) ODE,并將 u_t 替換為其對 z_t 的條件期望。這是隨機(jī)插值(stochastic interpolants)的一個(gè)特例 —— 在這種情況下,它可以泛化成 SDE:其中 ε_t 控制著推理時(shí)的隨機(jī)性。兩個(gè)框架的等價(jià)性這兩個(gè)框架都分別由三個(gè)超參數(shù)定義:擴(kuò)散的三個(gè)參數(shù)是 f_t、g_t、η_t,而流匹配的三個(gè)參數(shù)是 α_t、σ_t、ε_t。通過從一組超參數(shù)推導(dǎo)得到另一組超參數(shù),可以顯示這兩組超參數(shù)的等價(jià)性。從擴(kuò)散到流匹配:從流匹配到擴(kuò)散:總之,除了訓(xùn)練考慮和采樣器選擇之外,擴(kuò)散和高斯流匹配沒有根本區(qū)別。結(jié)語讀到這里,想必你已經(jīng)理解了擴(kuò)散模型和高斯流匹配的等價(jià)性。不過,文中重點(diǎn)介紹的是流匹配為該領(lǐng)域帶來的兩個(gè)新模型規(guī)范:網(wǎng)絡(luò)輸出:流匹配提出了一種網(wǎng)絡(luò)輸出的向量場參數(shù)化方案,并且其不同于擴(kuò)散文獻(xiàn)中使用的方案。當(dāng)使用高階采樣器時(shí),網(wǎng)絡(luò)輸出可能會有所不同。它也可能影響訓(xùn)練動態(tài)。采樣噪聲調(diào)度:流匹配利用了簡單的采樣噪聲調(diào)度 α_t = 1-t 和 σ_t = t,并且更新規(guī)則與 DDIM 相同。該團(tuán)隊(duì)最后表示:「如果能通過實(shí)證方式研究這兩個(gè)模型規(guī)范在不同的真實(shí)應(yīng)用中的重要性,那一定會很有趣。我們將此留給了未來的工作。」參考文獻(xiàn)Flow matching for generative modelingLipman, Y., Chen, R.T., Ben-Hamu, H., Nickel, M. and Le, M., 2022. arXiv preprint arXiv:2210.02747.Flow straight and fast: Learning to generate and transfer data with rectified flowLiu, X., Gong, C. and Liu, Q., 2022. arXiv preprint arXiv:2209.03003.Building normalizing flows with stochastic interpolantsAlbergo, M.S. and Vanden-Eijnden, E., 2022. arXiv preprint arXiv:2209.15571.Stochastic interpolants: A unifying framework for flows and diffusionsAlbergo, M.S., Boffi, N.M. and Vanden-Eijnden, E., 2023. arXiv preprint arXiv:2303.08797.Denoising diffusion implicit modelsSong, J., Meng, C. and Ermon, S., 2020. arXiv preprint arXiv:2010.02502.Score-based generative modeling through stochastic differential equationsSong, Y., Sohl-Dickstein, J., Kingma, D.P., Kumar, A., Ermon, S. and Poole, B., 2020. arXiv preprint arXiv:2011.13456.Understanding diffusion objectives as the elbo with simple data augmentationKingma, D. and Gao, R., 2024. Advances in Neural Information Processing Systems, Vol 36.Diffusion is spectral autoregression [HTML]Dieleman, S., 2024.Scaling rectified flow transformers for high-resolution image synthesisEsser, P., Kulal, S., Blattmann, A., Entezari, R., Muller, J., Saini, H., Levi, Y., Lorenz, D., Sauer, A., Boesel, F. and others,, 2024. Forty-first International Conference on Machine Learning.Elucidating the design space of diffusion-based generative modelsKarras, T., Aittala, M., Aila, T. and Laine, S., 2022. Advances in neural information processing systems, Vol 35, pp. 26565—26577.Knowledge distillation in iterative generative models for improved sampling speed [PDF]Luhman, E. and Luhman, T., 2021. arXiv preprint arXiv:2101.02388.Denoising diffusion probabilistic modelsHo, J., Jain, A. and Abbeel, P., 2020. Advances in neural information processing systems, Vol 33, pp. 6840—6851.Progressive Distillation for Fast Sampling of Diffusion ModelsSalimans, T. and Ho, J., 2022. International Conference on Learning Representations.Dpm-solver++: Fast solver for guided sampling of diffusion probabilistic modelsLu, C., Zhou, Y., Bao, F., Chen, J., Li, C. and Zhu, J., 2022. arXiv preprint arXiv:2211.01095.?THE END轉(zhuǎn)載請聯(lián)系本公眾號獲得授權(quán)投稿或?qū)で髨?bào)道:liyazhou@jiqizhixin.com

粵公網(wǎng)安備 44011502001135號