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