李飞飞团队新作：DiT不训练直接改架构，模型深度减半，质量还提高了
仅用于站内搜索，没有排版格式，具体信息请跳转上方微信公众号内链接
机器之心报道
编辑：欣东、陈陈
本文介绍了一种名为「嫁接」的技术，用于在小计算预算下通过编辑预训练DiffusionTransformers（简称DiTs）来探索新的模型架构设计。这种方法允许研究者在不从头开始训练模型的情况下，通过替换模型中的某些算子（如MLP）来创建新的混合架构，从而在保持模型质量的同时减少计算量。
模型架构设计在机器学习中扮演着核心角色，与数据、算法、算力和基准测试一样重要。它定义了模型函数、算子选择（如注意力机制、卷积）和配置设定（如模型深度、宽度）等等模型要素。
尽管如此，由于从头训练模型的成本过高——尤其人们难以获得关于架构设计的深刻洞见（即哪些方案有效、哪些无效）。因此，研究新架构仍是一项挑战，对生成模型而言尤为如此。
在本文中，来自斯坦福大学、LiquidAI等机构的研究者探索了这一问题，即对预训练模型进行架构编辑来研究新架构。
论文链接：https ://arxiv. org/pdf/2506. 05340v1
论文主页：https ://grafting. stanford.edu/
论文标题：ExploringDiffusionTransformerDesignsviaGrafting
具体而言，该研究提出了一种编辑预训练扩散transformer（DiT）的简单方法，即Grafting（嫁接），该方法可以在较小的计算预算下实现新的架构。
嫁接过程如下：
（i）激活蒸馏：此阶段通过回归目标（regressionobjective）蒸馏原始算子的激活特征，将其功能迁移至新算子。该阶段核心在于实现算子间的功能传递。
（ii）轻量级调优：此阶段通过使用有限的数据进行调优，减轻了由于集成多个新算子而导致的误差传播。
此外，架构编辑还涵盖多种策略，如添加、删除和替换算子。
本文还基于DiT-XL/2构建了一个测试平台，以研究嫁接对模型质量的影响。
利用该测试平台，本文通过嫁接技术开发了一系列混合设计：用门控卷积、局部注意力和线性注意力取代Softmax注意力，用可变扩展率和卷积变体取代MLP。
值得注意的是，许多混合设计使用不到2%的预训练计算资源就实现了良好的质量（FID：2. 38–2. 64，而DiT-XL/2为2. 27）。然后，本文嫁接了一个文本转图像模型(PixArt-Σ)，实现了1. 43倍的加速，而GenEval分数下降不到2%。
最后，本文展示了一个案例研究，该研究通过嫁接技术将每对序列Transformer模块转换为并行模块，从而重构了DiT-XL/2。这将模型深度减少到原来一半，并获得了比其他同等深度模型更高的质量（FID：2. 77）。
总而言之，该研究展示了可以通过预训练DiT来探索新的扩散模型设计，其修改范围涵盖从算子替换到架构重构。
嫁接扩散Transformer
两阶段嫁接方法
嫁接旨在通过编辑预训练模型的计算图来实现新架构。由于该研究专注于用替代方案替换现有算子，这引出了两个问题：
问题1：在将新算子集成到计算图之前，应该如何初始化？
对应第一阶段：通过激活蒸馏进行初始化。由于DiT的激活是连续且平滑的，这可以被视为一个回归问题：
问题2：当多个算子集成到计算图时，如何减轻误差传播？
对应第二阶段：轻量级调优。随着更多算子被替换，初始化误差会不断传播，导致与预训练模型的行为出现偏差。
本文采用端到端微调来缓解阶段1的累积误差。微调目标函数如公式1所示。
实践中，本文发现，即使替换DiT-XL/2中的所有MHA或MLP层，仅使用10%的训练数据也能恢复竞争性能。
自嫁接基准
在研究新的架构设计之前，该研究引入了自嫁接（self-grafting），这是一种简单的对照设置：将现有算子（如MHA、MLP）替换为相同类型但权重随机初始化的算子。这样可以保持计算图的结构——包括算子类型和参数数量——但改变了具体的计算过程。自嫁接有三方面作用：（1）评估在不改变架构的情况下嫁接流程本身的效果；（2）为比较不同的替换方案提供一个性能基准；（3）研究影响性能的因素，如数据规模、回归目标和超参数。
激活行为分析以及自嫁接结果
本文首先分析了DiT-XL/2层中的MHA和MLP算子激活行为。在这两种情况下，本文观察到激活值存在较大差异，尤其是在较深的层中（表1(i,ii)）。
经过分析，本文得出通过选择特定于算子的回归目标，可以实现高质量的初始化。
如表1(iii,iv)所示，回归目标的选择会影响性能。对于MHA，L1实现了最佳FID（2. 51），其次是Huber（2. 55）和L2（2. 58）。对于MLP，L2表现最佳（2. 33），而L1表现不佳（2. 83）；值得注意的是，MLP的参数量是MHA的2倍。
这表明高质量的初始化需要量身定制的、激活感知的策略。
研究还发现，使用10%的数据进行完全自嫁接可实现接近基线的性能。表明在适度的数据和计算预算下完全自嫁接是可行的。
实验
实验I：通过嫁接实现混合架构
本节实验围绕这个问题进行：当现有算子被高效的替代方案取代时，我们能否保持模型质量？
为了探究这个问题，本文研究了以下嫁接过程：
1. 待替换算子的类型——MHA或MLP；
2. 替换算子的类型——例如卷积；
3. 层选择策略——替换所有层中的算子或使用启发式选择；
4. 替换率——全部替换或部分替换。
为了实验，该研究构建了一个测试平台，并提出两种层选择策略：完全替换和交错替换。测试平台详见表3。
此外，该研究还引入了Hyena-X和Hyena-Y两种新的高效门控卷积算子，并设计为MHA的直接替代品。Figure3展示了它们的结构。
MHA结果。通过嫁接替换DiT-XL/2中的MHA算子，获得了良好的质量-效率权衡。主要发现如下：
在交错嫁接下，较小的感受野表现出惊人的效果。实验发现，在50%交错替换比例下，滑动窗口注意力（SWA）、Hyena-X/Y和Mamba-2等替代方案均能保持FID分数与基线（2. 27）差距在0. 5以内。尤其值得注意的是，尽管SWA和Hyena变体的感受野有限（卷积核K=4/窗口w=4），其FID下降幅度却极小。
替换策略：交错替换vs. 完全替换。将交错替换比例从50%提升至75%时，性能通常下降，但SWA在75%交错替换下仍有效（FID=3. 09）。100%替换时，性能急剧恶化（所有FID>75），这与局部性分析一致，表明只有部分层是局部且适合嫁接的。
数据规模和层选择的消融实验结果。
MLP结果显示通过嫁接的方式替换MLP算子是有效的。
经过实验，得出要点1：嫁接对于在较小的计算预算下构建具有良好生成质量的高效混合架构非常有效。交错设计尤其有效。
实验II：通过嫁接改进文本到图像的扩散Transformers
结果。嫁接模型在实时计算速度（wall-clocktime）上实现了1. 43倍的提升，同时生成评估分数（GenEval）仅出现小幅下降（47. 78vs.49. 75）。特定属性的指标（Attribute-specificmetrics）基本保持可比，并且定性样本也展现出良好的对齐度和质量。在一些纹理区域观察到了局部性的失真（artifacts），这可能是由于LoRA的适应能力以及所使用的合成数据质量不高所致（失败案例详见图D. 3，D. 4）
要点2：在文生图DiTs中成功应用嫁接技术，构建的混合架构在实现显著加速的同时，生成质量损失极小。
了解更多内容，请参考原论文。
©THEEND
转载请联系本公众号获得授权
投稿或寻求报道：liyazhou@jiqizhixin. com
