用擴(kuò)散模型生成網(wǎng)絡(luò)參數(shù),LeCun點(diǎn)贊尤洋團(tuán)隊(duì)新研究
如果你有被 Sora 生成的視頻震撼到,那你就已經(jīng)見識(shí)過擴(kuò)散模型在視覺生成方面的巨大潛力。當(dāng)然,擴(kuò)散模型的潛力并不止步于此,它在許多其它不同領(lǐng)域也有著讓人期待的應(yīng)用前景,更多案例可參閱機(jī)器之心不久前的報(bào)道《爆火Sora背后的技術(shù),一文綜述擴(kuò)散模型的最新發(fā)展方向》。
近日,來自新加坡國立大學(xué)的尤洋團(tuán)隊(duì)、加州大學(xué)伯克利分校、Meta AI Research 的一項(xiàng)新研究找到了擴(kuò)散模型的一個(gè)新應(yīng)用方向:用來生成神經(jīng)網(wǎng)絡(luò)的模型參數(shù)!
- 論文地址:https://arxiv.org/pdf/2402.13144.pdf
- 項(xiàng)目地址:https://github.com/NUS-HPC-AI-Lab/Neural-Network-Diffusion
- 論文標(biāo)題:Neural Network Diffusion
這樣一來,似乎就可以使用現(xiàn)有的神經(jīng)網(wǎng)絡(luò)來輕松生成新的神經(jīng)網(wǎng)絡(luò)了!這一點(diǎn)讓 Yann LeCun 也忍不住點(diǎn)贊分享。不僅如此,生成的模型的性能表現(xiàn)甚至還能超越原始模型。
擴(kuò)散模型的起源可以追溯到非平衡熱力學(xué)。2015 年時(shí),Jascha Sohl-Dickstein 等在論文《Deep Unsupervised Learning using Nonequilibrium Thermodynamics》中最早使用擴(kuò)散過程來漸進(jìn)式地移除輸入中的噪聲,進(jìn)而生成清晰的圖像。
之后的 DDPM 和 DDIM 等研究工作優(yōu)化了擴(kuò)散模型,使其訓(xùn)練范式有了前向和反向過程的鮮明特點(diǎn)。
當(dāng)時(shí),擴(kuò)散模型生成的圖像的質(zhì)量還未達(dá)到理想水平。
GuidedDiffusion 這項(xiàng)工作進(jìn)行了充分的消融研究并發(fā)現(xiàn)了一個(gè)更好的架構(gòu);這項(xiàng)開創(chuàng)性的工作開始讓擴(kuò)散模型在圖像質(zhì)量上超越基于 GAN 的方法。之后出現(xiàn)的 GLIDE、Imagen、DALL·E 2 和 Stable Diffusion 等模型已經(jīng)可以生成照片級(jí)真實(shí)感的圖像。
盡管擴(kuò)散模型在視覺生成領(lǐng)域已經(jīng)取得了巨大成功,但它們?cè)谄渌I(lǐng)域的潛力還相對(duì)欠開發(fā)。
新加坡國立大學(xué)、加州大學(xué)伯克利分校、Meta AI Research 近日的這項(xiàng)研究則發(fā)掘出了擴(kuò)散模型的一個(gè)驚人能力:生成高性能的模型參數(shù)。
要知道,這項(xiàng)任務(wù)與傳統(tǒng)的視覺生成任務(wù)存在根本性的差異!參數(shù)生成任務(wù)的重心是創(chuàng)造能在給定任務(wù)上表現(xiàn)良好的神經(jīng)網(wǎng)絡(luò)參數(shù)。之前已有研究者從先驗(yàn)和概率建模方面探索過這一任務(wù),比如隨機(jī)神經(jīng)網(wǎng)絡(luò)和貝葉斯神經(jīng)網(wǎng)絡(luò)。但是,之前還沒有人研究使用擴(kuò)散模型來生成參數(shù)。
如圖 1 所示,仔細(xì)觀察神經(jīng)網(wǎng)絡(luò)的訓(xùn)練過程與擴(kuò)散模型,可以發(fā)現(xiàn)基于擴(kuò)散的圖像生成方法與隨機(jī)梯度下降(SGD)學(xué)習(xí)過程有一些共同點(diǎn):1)神經(jīng)網(wǎng)絡(luò)的訓(xùn)練過程和擴(kuò)散模型的反向過程都可以被看作是從隨機(jī)噪聲/初始化轉(zhuǎn)換成特定分布的過程;2)通過多次添加噪聲,可以將高質(zhì)量圖像和高性能參數(shù)降級(jí)為簡單分布,比如高斯分布。
該團(tuán)隊(duì)基于上述觀察提出了一種用于參數(shù)生成的新方法:neural network diffusion,即神經(jīng)網(wǎng)絡(luò)擴(kuò)散,縮寫為 p-diff,其中的 p 是指參數(shù)(parameter)。
該方法的思路很直接,就是使用標(biāo)準(zhǔn)的隱擴(kuò)散模型來合成神經(jīng)網(wǎng)絡(luò)的參數(shù)集,因?yàn)閿U(kuò)散模型能夠?qū)⒔o定的隨機(jī)分布轉(zhuǎn)換為一個(gè)特定的分布。
他們的方法很簡單:組合使用一個(gè)自動(dòng)編碼器和一個(gè)標(biāo)準(zhǔn)隱擴(kuò)散模型來學(xué)習(xí)高性能參數(shù)的分布。
首先,對(duì)于一個(gè)使用 SGD 優(yōu)化器訓(xùn)練的模型參數(shù)子集,訓(xùn)練一個(gè)自動(dòng)編碼器來提取這些參數(shù)的隱含表征。然后,使用一個(gè)標(biāo)準(zhǔn)隱擴(kuò)散模型從噪聲開始合成隱含表征。最后,用經(jīng)過訓(xùn)練的自動(dòng)編碼器來處理合成的隱含表征,得到新的高性能模型參數(shù)。
這種新方法表現(xiàn)出了這兩個(gè)特點(diǎn):1)在多個(gè)數(shù)據(jù)集和架構(gòu)上,其性能表現(xiàn)能在數(shù)秒時(shí)間內(nèi)與其訓(xùn)練數(shù)據(jù)(即 SGD 優(yōu)化器訓(xùn)練的模型)媲美,甚至還能有所超越;2)生成的模型與訓(xùn)練得到的模型差異較大,這說明新方法能夠合成新參數(shù),而不是記憶訓(xùn)練樣本。
神經(jīng)網(wǎng)絡(luò)擴(kuò)散
介紹擴(kuò)散模型
擴(kuò)散模型通常由前向和反向過程構(gòu)成,這些過程組成一個(gè)多步驟的鏈?zhǔn)竭^程并且可通過時(shí)間步驟索引。
前向過程。給定一個(gè)樣本 x_0 ~ q(x),前向過程是在 T 個(gè)步驟中逐漸添加高斯噪聲,得到 x_1、x_2……x_T。
反向過程。不同于前向過程,反向過程的目標(biāo)是訓(xùn)練一個(gè)能遞歸地移除 x_t 中的噪聲的去噪網(wǎng)絡(luò)。該過程是多個(gè)步驟的反向過程,此時(shí) t 從 T 一路降至 0。
神經(jīng)網(wǎng)絡(luò)擴(kuò)散方法概述
神經(jīng)網(wǎng)絡(luò)擴(kuò)散(p-diff)這種新方法的目標(biāo)是基于隨機(jī)噪聲生成高性能參數(shù)。如圖 2 所示,該方法包含兩個(gè)過程:參數(shù)自動(dòng)編碼器和參數(shù)生成。
給定一組已經(jīng)過訓(xùn)練的高性能模型,首先選取其參數(shù)的一個(gè)子集并將其展平為一維向量。
之后,使用一個(gè)編碼器提取這些向量的隱含表征,同時(shí)還有一個(gè)解碼器負(fù)責(zé)基于這些隱含表征重建出參數(shù)。
然后,訓(xùn)練一個(gè)標(biāo)準(zhǔn)的隱擴(kuò)散模型來基于隨機(jī)噪聲合成這種隱含表征。
訓(xùn)練之后,就可使用 p-diff 通過這樣的鏈?zhǔn)竭^程來生成新參數(shù):隨機(jī)噪聲 → 反向過程 → 已訓(xùn)練的解碼器 → 生成的參數(shù)。
實(shí)驗(yàn)
該團(tuán)隊(duì)在論文中給出了詳細(xì)的實(shí)驗(yàn)設(shè)置,可幫助其他研究者復(fù)現(xiàn)其結(jié)果,詳見原論文,我們這里更關(guān)注其結(jié)果和消融研究。
結(jié)果
表 1 是在 8 個(gè)數(shù)據(jù)集和 6 種架構(gòu)上與兩種基準(zhǔn)方法的結(jié)果比較。
基于這些結(jié)果,可以得到以下觀察:1)在大多數(shù)實(shí)驗(yàn)案例中,新方法能取得與兩種基準(zhǔn)方法媲美或更優(yōu)的結(jié)果。這表明新提出的方法可以高效地學(xué)習(xí)高性能參數(shù)的分布,并能基于隨機(jī)噪聲生成更優(yōu)的模型。2)新方法在多個(gè)不同數(shù)據(jù)集上的表現(xiàn)都很好,這說明這種方法具有很好的泛化性能。
消融研究和分析
表 2(a) 展示了不同訓(xùn)練數(shù)據(jù)規(guī)模(即原始模型的數(shù)量)的影響??梢钥吹?,不同數(shù)量的原始模型的最佳結(jié)果之間的性能差異其實(shí)不大。
為了研究 p-diff 在其它歸一化層深度上的有效性,該團(tuán)隊(duì)還探索了新方法合成其它淺層參數(shù)的性能。為了保證 BN 參數(shù)的數(shù)量相等,該團(tuán)隊(duì)為三組 BN 層(它們位于不同深度的層之間)實(shí)現(xiàn)了新提出的方法。實(shí)驗(yàn)結(jié)果見表 2(b),可以看到在所有深度的 BN 層設(shè)置上,新方法的表現(xiàn)(最佳準(zhǔn)確度)都優(yōu)于原始模型。
噪聲增強(qiáng)的目的是提升訓(xùn)練自動(dòng)編碼器的穩(wěn)健性和泛化能力。該團(tuán)隊(duì)對(duì)噪聲增強(qiáng)在輸入?yún)?shù)和隱含表征方面的應(yīng)用進(jìn)行了消融研究。結(jié)果見表 2(c)。
此前,實(shí)驗(yàn)評(píng)估的都是新方法在合成模型參數(shù)子集(即批歸一化參數(shù))方面的效果。那么我們不禁要問:能否使用此方法合成模型的整體參數(shù)?
為了解答這個(gè)問題,該團(tuán)隊(duì)使用兩個(gè)小型架構(gòu)進(jìn)行了實(shí)驗(yàn):MLP-3 和 ConvNet-3。其中 MLP-3 包含三個(gè)線性層和 ReLU 激活函數(shù),ConvNet-3 則包含三個(gè)卷積層和一個(gè)線性層。不同于之前提到的訓(xùn)練數(shù)據(jù)收集策略,該團(tuán)隊(duì)基于 200 個(gè)不同的隨機(jī)種子從頭開始訓(xùn)練了這些架構(gòu)。
表 3 給出了實(shí)驗(yàn)結(jié)果,其中將新方法與兩種基準(zhǔn)方法(原始方法和集成方法)進(jìn)行了比較。其中報(bào)告了 ConvNet-3 在 CIFAR-10/100 以及 MLP-3 在 CIFAR-10 和 MNIST 上的結(jié)果比較和參數(shù)數(shù)量。
這些實(shí)驗(yàn)表明新方法在合成整體模型參數(shù)方面的有效性和泛化能力,也就是說新方法實(shí)現(xiàn)了與基準(zhǔn)方法相當(dāng)或更優(yōu)的性能。這些結(jié)果也能體現(xiàn)新方法的實(shí)際應(yīng)用潛力。
但該團(tuán)隊(duì)也在論文中表明目前還無法合成 ResNet、ViT 和 ConvNeXt 等大型架構(gòu)的整體參數(shù)。這主要是受限于 GPU 內(nèi)存的極限。
至于為什么這種新方法能夠有效地生成神經(jīng)網(wǎng)絡(luò)參數(shù),該團(tuán)隊(duì)也嘗試探索分析了原因。他們使用 3 個(gè)隨機(jī)種子從頭開始訓(xùn)練了 ResNet-18 并對(duì)其參數(shù)進(jìn)行了可視化,如圖 3 所示。
他們通過 min-max 歸一化方法,分別得到了不同層的參數(shù)分布的熱力圖?;诰矸e層(Conv.-layer2)和全連接層(FC-layer18)的可視化結(jié)果,可以看到這些層中確實(shí)存在一定的參數(shù)模式。通過學(xué)習(xí)這些模式,新方法就能生成高性能的神經(jīng)網(wǎng)絡(luò)參數(shù)。
p-diff 是單純靠記憶嗎?
p-diff 看起來能生成神經(jīng)網(wǎng)絡(luò)參數(shù),但它究竟是生成參數(shù)還是僅僅記住了參數(shù)呢?該團(tuán)隊(duì)就此做了一番研究,比較了原始模型和生成模型的差異。
為了進(jìn)行量化比較,他們提出了一個(gè)相似度指標(biāo)。簡單來說,這個(gè)指標(biāo)就是通過計(jì)算兩個(gè)模型在錯(cuò)誤預(yù)測(cè)結(jié)果上的交并比(IoU)來確定它們的相似度。然后他們基于此進(jìn)行了一些比較研究和可視化。比較結(jié)果見圖 4。
圖 4(a) 報(bào)告了原始模型和 p-diff 模型之間的相似度比較,其中涉及 4 種比較方案。
可以看到,生成的模型之間的差異比原始模型之間的差異大得多。另外,原始模型和生成的模型之間的最大相似度也低于原始模型之間的相似度。這足以表明,p-diff 可以生成與其訓(xùn)練數(shù)據(jù)(即原始模型)不同的新參數(shù)。
該團(tuán)隊(duì)也將新方法與微調(diào)模型和添加噪聲的模型進(jìn)行了比較。結(jié)果見圖 4(b)。
可以看到,微調(diào)模型和添加噪聲的模型很難超過原始模型。此外,微調(diào)模型或添加噪聲的模型與原始模型之間的相似度非常高,這表明這兩種操作方法無法獲得全新且高性能的模型。但是,新方法生成的模型則表現(xiàn)出了多樣的相似度以及優(yōu)于原始模型的性能。
該團(tuán)隊(duì)也比較了隱含表征。結(jié)果見圖 4(c)??梢钥吹?,p-diff 可以生成全新的隱含表征,而添加噪聲方法只會(huì)在原始模型的隱含表征周圍進(jìn)行插值。
該團(tuán)隊(duì)也可視化了 p-diff 過程的軌跡。具體而言,他們繪出了在推理階段的不同時(shí)間步驟生成的參數(shù)軌跡。圖 5(a) 給出了 5 條軌跡(使用了 5 種不同的隨機(jī)噪聲初始化)。圖中紅心是原始模型的平均參數(shù),灰色區(qū)域是其標(biāo)準(zhǔn)差(std)。
隨著時(shí)間步驟增加,生成的參數(shù)整體上會(huì)更接近原始模型。但是也能看出,這些軌跡的終點(diǎn)(橙色三角形)與平均參數(shù)仍然有些距離。另外,這五條軌跡的形狀也很多樣化。
最后,該團(tuán)隊(duì)研究了原始模型的數(shù)量(K)對(duì)生成的模型的多樣性的影響。圖 5(b) 可視化地展示了不同 K 時(shí)原始模型與生成的模型之間的最大相似度。具體來說,他們的做法是生成 50 個(gè)模型,通過持續(xù)生成參數(shù),直到生成的 50 個(gè)模型在所有情況下的表現(xiàn)均優(yōu)于 76.5%。
可以看到,當(dāng) K=1 時(shí),相似度很高且范圍窄,說明這時(shí)候生成的模型基本是記憶了原始模型的參數(shù)。隨著 K 增大,相似度范圍也變大了,這表明新方法可以生成與原始模型不同的參數(shù)。