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

（來源：arXiv）

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

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

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

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

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

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

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

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

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

基于這一設想，他們通過探索預訓練模型的架構編輯方法，來探索新型網絡架構的設計。

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

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

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

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

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

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

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

（來源：arXiv）

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

架構編輯的空間十分廣闊，這就會引發這樣一個現實問題：應該研究什么類型的替換？

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

在此基礎之上，他們使用高效算子替代現有算子，以便在保持質量的同時減少模型的每秒浮點運算次數（FLOP，Floating Point Operations Per Second）。

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

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

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

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

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

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

（來源：arXiv）

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

（來源：arXiv）

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

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

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

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

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

參考資料：

https://arxiv.org/pdf/2506.05340

排版：初嘉實