近日，美國斯坦福大學教授李飛飛等人打造出一種名為“嫁接”（grafting）的新型架構(gòu)編輯方法，它能在有限的算力條件之下重構(gòu)預訓練擴散 Transformer。研究團隊表示這種嫁接法既簡單又輕便，采用 8 塊英偉達 H100 GPU 在 24 小時內(nèi)即可完成單項實驗，同時僅使用不到 2% 的預訓練計算資源。

（來源：arXiv）

研究中，他們使用高效替代方案替換了自注意力機制和多層感知機，借此構(gòu)建了混合模型，該混合模型的弗雷歇初始距離（FID，F(xiàn)rechet Inception Distance）值位于 2.38 - 2.64 之間，高于基線模型的 2.27。（注：FID 是一種用于評估生成圖像與真實圖像相似度的指標。）

PixArt-∑，是由華為諾亞方舟實驗室聯(lián)合香港大學團隊、大連理工大學團隊和香港科技大學團隊研發(fā)的文本到圖像的擴散 Transformer（DiT，Diffusion Transformer）模型。本次研究之中，基于 12k 的合成數(shù)據(jù)，李飛飛等人將嫁接法用于 PixArt-∑，在評估生成質(zhì)量的 GenEval 評分下降不到 2% 的情況下，讓嫁接模型實現(xiàn)了 1.43 倍的內(nèi)容生成加速，這表明嫁接法確實能被擴展至高分辨率的文本到圖像模型之中。

隨后，他們使用嫁接法陣針對 Meta 公司推出的擴散 Transformer 模型 DiT-XL/2 進行架構(gòu)重構(gòu)，將每對連續(xù)的 Transformer 塊轉(zhuǎn)換為并行結(jié)構(gòu)，從而使模型深度減半。在 14 層的擴散 Transformer 模型中，被重構(gòu)之后的模型實現(xiàn)了更好的生成質(zhì)量，弗雷歇初始距離（FID）值為 2.77。這些結(jié)果證明了嫁接法在短上下文和長上下文設置以及架構(gòu)重構(gòu)中的實用性。

這也說明通過嫁接預訓練的擴散 Transformer，可以探索新的擴散模型設計。總的來說，嫁接法是一種能在低計算成本下探索擴散 Transformer 設計的輕量級方法。

圖 | 相關論文（來源：arXiv）

預訓練模型能否作為研究新架構(gòu)的“腳手架”？

研究團隊表示，與數(shù)據(jù)、算法、計算資源和基準測試一樣，模型架構(gòu)設計也在機器學習中起著核心作用。模型架構(gòu)設計定義了一個可學習的函數(shù)，并涉及到一些關鍵設計決策，比如算子和配置的選擇等。

盡管如此，由于從頭開始訓練模型的成本過高，尤其是在基礎模型比較流行的當下，人們依然很難深入了解哪些架構(gòu)有效、哪些架構(gòu)無效。因此，研究新架構(gòu)仍然是一個挑戰(zhàn)，特別是對于生成模型而言。

類比于“新軟件基于現(xiàn)有代碼開發(fā)”的理念，研究團隊提出這一設想：預訓練模型能否作為研究新架構(gòu)的“腳手架”？

基于這一設想，他們通過探索預訓練模型的架構(gòu)編輯方法，來探索新型網(wǎng)絡架構(gòu)的設計。

研究中，他們專注于研究擴散 Transformer，這是一類廣泛用于圖像生成和視頻生成的生成式 Transformer。

預訓練模型通過實現(xiàn)一個計算圖來執(zhí)行圖像生成或視頻生成等任務。因此，本次研究重點探究這一問題：如何在算力有限的條件之下，通過修改模型計算圖實現(xiàn)架構(gòu)方案的可行性驗證？

此前，業(yè)內(nèi)有人認為卷積設計可以取代擴散 Transformer 中的多頭注意力或多層感知機。實現(xiàn)這一想法的方法之一便是使用卷積算子替換多頭注意力或多層感知機算子，同時還能保持模型質(zhì)量。

而這隨之會提出以下兩個問題：

• 第一個是算子初始化問題：在將新算子集成到計算圖中之前，如何對其進行初始化？
• 第二個是錯誤累積問題：當多個算子集成到計算圖中時，如何有效抑制誤差傳播？

為了解決這些問題，研究團隊提出了嫁接法，嫁接過程具體如下：

• 第一個階段是激活蒸餾：通過回歸目標蒸餾原始算子的激活特征，將原算子功能遷移至新算子。
• 第二個階段是輕量化微調(diào)：使用有限數(shù)據(jù)進行微調(diào)，減輕由于集成多個新算子所引起的錯誤傳播。

（來源：arXiv）

對于架構(gòu)編輯來說，它涉及到算子添加、算子刪除和替換算子等多種策略。本次研究聚焦于算子替換這一核心策略：即將一個算子替換為另一個算子，而其他策略可被視為特殊的替換情況。

架構(gòu)編輯的空間十分廣闊，這就會引發(fā)這樣一個現(xiàn)實問題：應該研究什么類型的替換？

為此，研究團隊建立了一個自嫁接基線，使用隨機初始化的權(quán)重替換現(xiàn)有算子。他們發(fā)現(xiàn)通過兩階段嫁接過程可以恢復接近基線的模型質(zhì)量，借此驗證了嫁接法的有效性。

在此基礎之上，他們使用高效算子替代現(xiàn)有算子，以便在保持質(zhì)量的同時減少模型的每秒浮點運算次數(shù)（FLOP，F(xiàn)loating Point Operations Per Second）。

與此同時，他們還使用了增加模型每秒浮點運算次數(shù)的替代方案，以便驗證更廣泛的架構(gòu)設計的可能性。

為了系統(tǒng)性地研究這一點，他們構(gòu)建了一個基于 DiT XL/2 模型的測試平臺。利用這個測試平臺，他們通過嫁接開發(fā)了一系列混合設計方案：比如將 Softmax 注意力替換為門控卷積、局部注意力和線性注意力，以及將多層感知機替換為可變擴展率和卷積變體。（注：Softmax 注意力是注意力機制中的核心計算方式之一，已被廣泛用于 Transformer 架構(gòu)及其變體。）

同時，他們設計了一套架構(gòu)編輯方案，以用于評估不同嫁接策略對于模型質(zhì)量的影響。期間，他們聚焦于以下幾個核心設計維度：更換哪個算子？用什么替換算子？如何選擇要編輯的層？是完全替換還是部分替換？

替換多頭注意力和多層感知機算子的動機，源于研究團隊的以下實證證據(jù)和架構(gòu)考量：對于多頭注意力，注意力局部性分析表明它更適合采用局部算子；對于多層感知機，研究團隊決定采用已有的架構(gòu)思想。

在兩個頗有挑戰(zhàn)性的生成式建模場景中，他們驗證了嫁接法的效果。

在第一個場景中，即在類條件圖像生成任務中，嫁接法產(chǎn)生了具有良好質(zhì)量的混合架構(gòu)設計效果。對于多頭注意力即 Softmax 注意力，研究團隊探索了幾種替代方案：局部門控卷積、局部注意和線性注意力。對于多層感知機，替代方案包括具有可變擴展比的多層感知機和卷積變體。有趣的是，幾種交錯混合架構(gòu)設計實現(xiàn)了 2.38-2.64 弗雷歇初始距離（FID），這表明嫁接法可以構(gòu)建高質(zhì)量的混合架構(gòu)。

（來源：arXiv）

在第二個場景中，研究團隊通過架構(gòu)嫁接技術(shù)，構(gòu)建了面向高分辨率文生圖（T2I，text-to-image）任務的高效混合架構(gòu)。他們在一個具有挑戰(zhàn)性的現(xiàn)實環(huán)境中驗證了嫁接法：即使用 PixArt-∑ 模型進行 2048×2048 分辨率的文本到圖像生成。這種實驗設置集中體現(xiàn)了三大核心挑戰(zhàn)：第一個挑戰(zhàn)是需要進行 16384 tokens 的長序列處理；第二個挑戰(zhàn)是需要進行多模態(tài)文本條件集成；第三個挑戰(zhàn)是訓練數(shù)據(jù)較為匱乏。期間，他們針對多頭注意力算子進行嫁接，之所以這樣做是因為它們占生成延遲的 62% 以上。通過使用數(shù)量為 12k 的合成數(shù)據(jù)，讓嫁接模型實現(xiàn)了 1.43 倍的加速，評估生成質(zhì)量的 GenEval 評分下降不到 2%，這表明嫁接法可以擴展到高分辨率的文本到圖像模型之中。

（來源：arXiv）

將把嫁接法推至自回歸模型等其他模型家族

總的來說，嫁接技術(shù)展現(xiàn)出了廣闊的應用前景，包括將模型從低分辨率生成調(diào)整到高分辨率生成，將理解能力和生成能力從短視頻擴展到長視頻，或在圖像編輯等交互式應用中改善用戶體驗。在這些應用中，即使是 10% 的速度提升也很有益。

需要說明的是，本次研究之中 PixArt-Σ 模型的設置使用合成數(shù)據(jù)進行嫁接，這可能會將偽影和偏差傳播到嫁接后的模型之中。雖然本次研究側(cè)重于架構(gòu)編輯，但是那些在嫁接下表現(xiàn)良好的架構(gòu)，在從頭開始訓練時是否也能表現(xiàn)良好？這仍然是一個未知數(shù)。

另外，本次工作主要側(cè)重于預訓練擴散 Transformer 的架構(gòu)編輯，特別是針對多頭注意力和多層感知機組件。針對其他架構(gòu)組件比如歸一化層和激活函數(shù)，研究團隊將在后續(xù)進行研究。

在本次研究之中，所開展的實驗主要集中在擴散 Transformer 上，而將嫁接法推廣到其他模型家族比如自回歸模型，是他們未來的研究方向之一。此外，針對應變分自編碼器（VAEs，Variational Autoencoders）中的組件嫁接也有待在未來進行探索。

參考資料：

https://arxiv.org/pdf/2506.05340

排版：初嘉實