摘要：近日，美国斯坦福大学教授李飞飞等人打造出一种名为“嫁接”（grafting）的新型架构编辑方法，它能在有限的算力条件之下重构预训练扩散 Transformer。研究团队表示这种嫁接法既简单又轻便，采用 8 块英伟达 H100 GPU 在 24 小时内即可完成单

近日，美国斯坦福大学教授李飞飞等人打造出一种名为“嫁接”（grafting）的新型架构编辑方法，它能在有限的算力条件之下重构预训练扩散 Transformer。研究团队表示这种嫁接法既简单又轻便，采用 8 块英伟达 H100 GPU 在 24 小时内即可完成单项实验，同时仅使用不到 2% 的预训练计算资源。

研究中，他们使用高效替代方案替换了自注意力机制和多层感知机，借此构建了混合模型，该混合模型的弗雷歇初始距离（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 中的多头注意力或多层感知机。实现这一想法的方法之一便是使用卷积算子替换多头注意力或多层感知机算子，同时还能保持模型质量。

对于架构编辑来说，它涉及到算子添加、算子删除和替换算子等多种策略。本次研究聚焦于算子替换这一核心策略：即将一个算子替换为另一个算子，而其他策略可被视为特殊的替换情况。

架构编辑的空间十分广阔，这就会引发这样一个现实问题：应该研究什么类型的替换？

为此，研究团队建立了一个自嫁接基线，使用随机初始化的权重替换现有算子。他们发现通过两阶段嫁接过程可以恢复接近基线的模型质量，借此验证了嫁接法的有效性。

在此基础之上，他们使用高效算子替代现有算子，以便在保持质量的同时减少模型的每秒浮点运算次数（FLOP，Floating Point Operations Per Second）。

与此同时，他们还使用了增加模型每秒浮点运算次数的替代方案，以便验证更广泛的架构设计的可能性。

为了系统性地研究这一点，他们构建了一个基于 DiT XL/2 模型的测试平台。利用这个测试平台，他们通过嫁接开发了一系列混合设计方案：比如将 Softmax 注意力替换为门控卷积、局部注意力和线性注意力，以及将多层感知机替换为可变扩展率和卷积变体。（注：Softmax 注意力是注意力机制中的核心计算方式之一，已被广泛用于 Transformer 架构及其变体。）

同时，他们设计了一套架构编辑方案，以用于评估不同嫁接策略对于模型质量的影响。期间，他们聚焦于以下几个核心设计维度：更换哪个算子？用什么替换算子？如何选择要编辑的层？是完全替换还是部分替换？

替换多头注意力和多层感知机算子的动机，源于研究团队的以下实证证据和架构考量：对于多头注意力，注意力局部性分析表明它更适合采用局部算子；对于多层感知机，研究团队决定采用已有的架构思想。

在两个颇有挑战性的生成式建模场景中，他们验证了嫁接法的效果。

在第一个场景中，即在类条件图像生成任务中，嫁接法产生了具有良好质量的混合架构设计效果。对于多头注意力即 Softmax 注意力，研究团队探索了几种替代方案：局部门控卷积、局部注意和线性注意力。对于多层感知机，替代方案包括具有可变扩展比的多层感知机和卷积变体。有趣的是，几种交错混合架构设计实现了 2.38-2.64 弗雷歇初始距离（FID），这表明嫁接法可以构建高质量的混合架构。

在第二个场景中，研究团队通过架构嫁接技术，构建了面向高分辨率文生图（T2I，text-to-image）任务的高效混合架构。他们在一个具有挑战性的现实环境中验证了嫁接法：即使用 PixArt-∑ 模型进行 2048×2048 分辨率的文本到图像生成。这种实验设置集中体现了三大核心挑战：第一个挑战是需要进行 16384 tokens 的长序列处理；第二个挑战是需要进行多模态文本条件集成；第三个挑战是训练数据较为匮乏。期间，他们针对多头注意力算子进行嫁接，之所以这样做是因为它们占生成延迟的 62% 以上。通过使用数量为 12k 的合成数据，让嫁接模型实现了 1.43 倍的加速，评估生成质量的 GenEval 评分下降不到 2%，这表明嫁接法可以扩展到高分辨率的文本到图像模型之中。

将把嫁接法推至自回归模型等其他模型家族

总的来说，嫁接技术展现出了广阔的应用前景，包括将模型从低分辨率生成调整到高分辨率生成，将理解能力和生成能力从短视频扩展到长视频，或在图像编辑等交互式应用中改善用户体验。在这些应用中，即使是 10% 的速度提升也很有益。

需要说明的是，本次研究之中 PixArt-Σ 模型的设置使用合成数据进行嫁接，这可能会将伪影和偏差传播到嫁接后的模型之中。虽然本次研究侧重于架构编辑，但是那些在嫁接下表现良好的架构，在从头开始训练时是否也能表现良好？这仍然是一个未知数。

另外，本次工作主要侧重于预训练扩散 Transformer 的架构编辑，特别是针对多头注意力和多层感知机组件。针对其他架构组件比如归一化层和激活函数，研究团队将在后续进行研究。

在本次研究之中，所开展的实验主要集中在扩散 Transformer 上，而将嫁接法推广到其他模型家族比如自回归模型，是他们未来的研究方向之一。此外，针对应变分自编码器（VAEs，Variational Autoencoders）中的组件嫁接也有待在未来进行探索。

参考资料：

来源：DeepTech深科技一点号