預(yù)訓(xùn)練模型能否作為探索新架構(gòu)設(shè)計(jì)的“底座” ？

最新答案是：yes！

簡(jiǎn)單來(lái)說(shuō)，按照研究人員設(shè)計(jì)好的架構(gòu)方案從頭訓(xùn)練模型，往往是檢驗(yàn)一個(gè)架構(gòu)是否有效的重要手段。

但問(wèn)題在于，從頭訓(xùn)練模型的成本也太高了！

對(duì)此，包括李飛飛團(tuán)隊(duì)在內(nèi)的研究人員提出了一種被稱為“Grafting（嫁接）”的新思路——

直接將預(yù)訓(xùn)練好的模型作為“底座”，通過(guò)修改其組件來(lái)研究新架構(gòu)。

這就好比軟件開(kāi)發(fā)中，程序員常基于現(xiàn)有代碼修改而非重寫(xiě)，以此省時(shí)省力。

基于這一思路，他們重點(diǎn)關(guān)注了DiTs這一廣泛用于圖像和視頻生成的Transformer模型。

具體而言，這群人先是構(gòu)建了一個(gè)基于DiT-XL/2設(shè)計(jì)的測(cè)試平臺(tái)，以方便后續(xù)研究“嫁接”對(duì)模型質(zhì)量的影響，然后實(shí)際使用“嫁接”技術(shù)開(kāi)發(fā)了一系列混合設(shè)計(jì)。

結(jié)果發(fā)現(xiàn)，許多混合設(shè)計(jì)在使用不到2%的預(yù)訓(xùn)練計(jì)算量的情況下，獲得了和原來(lái)大差不差的模型性能。

將這一方法應(yīng)用于文生圖模型PixArt-Σ，其生成速度提高了1.43倍，但生成圖像的質(zhì)量只下降了不到2%。

以上說(shuō)明，“嫁接”確實(shí)能成為一種輕量級(jí)、高效的架構(gòu)探索工具，可以讓研究者在缺少計(jì)算資源的情況下測(cè)試新想法。

下面詳細(xì)揭秘團(tuán)隊(duì)提出的這種新方法——

兩階段架構(gòu)編輯法

眾所周知，模型架構(gòu)設(shè)計(jì)通常涉及測(cè)試不同的組件（如注意力機(jī)制、卷積層）和配置（如模型深度、寬度）。

而作為一種架構(gòu)編輯方法，“嫁接”主要通過(guò)修改預(yù)訓(xùn)練DiTs的計(jì)算圖來(lái)實(shí)現(xiàn)新架構(gòu)的驗(yàn)證，具體則主要通過(guò)激活蒸餾和輕量級(jí)微調(diào)這兩個(gè)關(guān)鍵階段來(lái)實(shí)現(xiàn)。

所謂計(jì)算圖，是指模型內(nèi)部的運(yùn)算邏輯結(jié)構(gòu)——由多層Transformer塊組成，每個(gè)塊包含自注意力（MHA）、多層感知器（MLP）等算子（Operator），這些算子按特定順序連接，形成執(zhí)行生成任務(wù)的“數(shù)據(jù)流路徑”。

要實(shí)現(xiàn)這種修改替換，關(guān)鍵要解決兩個(gè)問(wèn)題：

問(wèn)題1：在將新算子整合到計(jì)算圖之前，應(yīng)該如何初始化新算子？

如果簡(jiǎn)單地把新算子的權(quán)重隨機(jī)初始化，它可能一開(kāi)始就會(huì)和模型的其他部分不協(xié)調(diào)，導(dǎo)致模型性能下降。

問(wèn)題2：如何減輕因替換多個(gè)算子而導(dǎo)致的錯(cuò)誤累積？

當(dāng)替換多個(gè)算子時(shí)，每個(gè)替換都可能引入一些誤差。一旦誤差逐漸積累，最終可能會(huì)導(dǎo)致模型性能大幅下降。

對(duì)此，新方法采用了以下兩階段架構(gòu)編輯法：

• 激活蒸餾（Activation Distillation）：新算子（如卷積）初始化時(shí)，通過(guò)回歸任務(wù)學(xué)習(xí)原算子的“行為”，即用少量數(shù)據(jù)訓(xùn)練新算子，使其輸出與原算子的激活值盡可能接近。
• 輕量級(jí)微調(diào)（Lightweight Fine-tuning）：替換多個(gè)組件后，用有限數(shù)據(jù)進(jìn)行端到端微調(diào)，減少誤差累積，恢復(fù)模型性能。

并且，為了評(píng)估“嫁接”本身的效果，研究正式開(kāi)始前還引入了自嫁接（self-grafting）作為對(duì)照實(shí)驗(yàn)。

所謂自嫁接，是指將現(xiàn)有MHA、MLP等替換為相同類型但權(quán)重隨機(jī)初始化的算子 。

其作用主要有三個(gè)：

• 評(píng)估在不改變架構(gòu)的情況下，“嫁接”過(guò)程本身對(duì)模型的影響；
• 為后續(xù)比較不同的替換方案提供一個(gè)基準(zhǔn)性能，便于判斷新方案的優(yōu)劣；
• 研究影響模型性能的各種因素，比如數(shù)據(jù)規(guī)模大小、回歸目標(biāo)的選擇以及超參數(shù)設(shè)置等。

結(jié)果發(fā)現(xiàn)，在實(shí)際操作中，僅需8k樣本就能實(shí)現(xiàn)較好的初始化。

此外，即便替換DiT-XL/2中所有的多頭注意力（MHA）或多層感知器（MLP）層，僅使用10%的訓(xùn)練數(shù)據(jù)進(jìn)行微調(diào)，模型也能正常恢復(fù)。

實(shí)驗(yàn)結(jié)果

研究人員進(jìn)行了三項(xiàng)實(shí)驗(yàn)，并得出以下主要結(jié)論：

實(shí)驗(yàn)1：混合架構(gòu)實(shí)驗(yàn)，驗(yàn)證替換的可行性。

通過(guò)將DiT-XL/2中的注意力層MHA替換為滑動(dòng)窗口注意力（SWA）或門(mén)控卷積（Hyena-X），在50%替換比例下，F(xiàn)ID僅比基線高0.4（FID值越低，說(shuō)明越接近原始性能）。

而100%全替換會(huì)導(dǎo)致FID驟降（數(shù)值>75），生成質(zhì)量崩潰，這說(shuō)明并非所有層都能被局部算子替代，即模型中存在“必須依賴全局信息” 的層，而另一部分層可接受局部計(jì)算。

團(tuán)隊(duì)還嘗試將DiT-XL/2中的感知器層MLP也進(jìn)行了替換，結(jié)果在將MLP的擴(kuò)展比改成r=3或r=6的情況下，就算全換掉，模型效果也挺好，這說(shuō)明MLP寬度改起來(lái)不容易出問(wèn)題。

一言以蔽之，多種混合設(shè)計(jì)的生成質(zhì)量均接近原模型，且計(jì)算成本不到預(yù)訓(xùn)練的2%。

實(shí)驗(yàn)2：文本到圖像生成實(shí)驗(yàn)，驗(yàn)證新架構(gòu)的有效性。

接下來(lái)，研究人員對(duì)文生圖模型PixArt-Σ進(jìn)行了“嫁接”，將MHA替換為Hyena-X，結(jié)果使用12k合成數(shù)據(jù)微調(diào)后，實(shí)現(xiàn)了1.43倍速度提升（從235ms→164ms），GenEval分?jǐn)?shù)從49.75→47.78（下降小于2%）。

實(shí)驗(yàn)3：并行化改造實(shí)驗(yàn)，驗(yàn)證架構(gòu)重組的有效性。

通過(guò)將DiT-XL/2的28層順序塊轉(zhuǎn)為14層并行塊（每對(duì)順序塊并行執(zhí)行），在深度減半的情況下，模型生成質(zhì)量?jī)?yōu)于同類深度模型。

這驗(yàn)證了，并行架構(gòu)在減少深度的同時(shí)可提升質(zhì)量，可用作模型輕量化的思路。

不過(guò)最后，團(tuán)隊(duì)也提到了研究的局限性。一是僅在DiT-XL/2模型上進(jìn)行了驗(yàn)證，二是僅測(cè)試了替換成Hyena-X和SWA的效果，結(jié)論的普適性受限。

但不管怎樣，團(tuán)隊(duì)認(rèn)為“嫁接”這種方法在探索新的模型架構(gòu)方面顯示出很大的潛力，尤其是在需要高效利用計(jì)算資源的場(chǎng)景中。

BTW，目前研究所涉及的22種“嫁接”模型均已開(kāi)源。

論文：
https://grafting.stanford.edu/
博客：
https://www.liquid.ai/research/exploring-diffusion-transformer-designs-via-grafting
開(kāi)源地址：
https://huggingface.co/grafting
https://github.com/keshik6/grafting