arXiv:2506.05340cs.LGcs.AI2025-06NeurIPS被引 14

用嫁接法低成本探索扩散模型新架构,效果接近原模型。

Exploring Diffusion Transformer Designs via Grafting

论文配图:Exploring Diffusion Transformer Designs via Grafting
图 1 · 摘自论文原文
  • 通过嫁接预训练扩散模型,替换注意力或MLP结构,无需大规模重训。
  • 多种混合设计FID达2.38-2.64,仅需<2%预训练算力。
  • 可实现速度提升1.43倍,且适合快速原型验证与架构优化。

设计模型架构需权衡操作选择(如注意力、卷积)与配置(如深度、宽度)。但评估这些决策对性能的影响通常需高昂的预训练成本,限制了架构探索。受软件开发中复用代码的启发,我们提出“嫁接”方法,通过编辑预训练的扩散变压器(DiTs),在小计算开销下构建新架构。基于对激活行为和注意力局部性的分析,我们以DiT-XL/2为基础搭建测试平台,研究嫁接对模型质量的影响。通过该平台,我们开发出一系列混合设计:将softmax注意力替换为门控卷积、局部注意力与线性注意力,将MLP替换为可变扩展比及卷积变体。值得注意的是,许多混合设计在<2%预训练算力下即达到良好质量(FID: 2.38–2.64,对比原模型DiT-XL/2为2.27)。进一步地,我们将嫁接应用于文本到图像模型PixArt-Sigma,实现1.43倍加速,且生成评估得分下降不足2%。最后,我们开展案例研究,通过嫁接将DiT-XL/2中每对连续的Transformer块重构为并行结构,使模型深度减半,且在同等深度下获得更优质量(FID: 2.77)。结果表明,嫁接可有效探索从操作替换到架构重组的多样化扩散模型设计。代码与嫁接模型详见https://grafting.stanford.edu

原文摘要 · Abstract (English)

Designing model architectures requires decisions such as selecting operators (e.g., attention, convolution) and configurations (e.g., depth, width). However, evaluating the impact of these decisions on model quality requires costly pretraining, limiting architectural investigation. Inspired by how new software is built on existing code, we ask: can new architecture designs be studied using pretrained models? To this end, we present grafting, a simple approach for editing pretrained diffusion transformers (DiTs) to materialize new architectures under small compute budgets. Informed by our analysis of activation behavior and attention locality, we construct a testbed based on the DiT-XL/2 design to study the impact of grafting on model quality. Using this testbed, we develop a family of hybrid designs via grafting: replacing softmax attention with gated convolution, local attention, and linear attention, and replacing MLPs with variable expansion ratio and convolutional variants. Notably, many hybrid designs achieve good quality (FID: 2.38-2.64 vs. 2.27 for DiT-XL/2) using <2% pretraining compute. We then graft a text-to-image model (PixArt-Sigma), achieving a 1.43x speedup with less than a 2% drop in GenEval score. Finally, we present a case study that restructures DiT-XL/2 by converting every pair of sequential transformer blocks into parallel blocks via grafting. This reduces model depth by 2x and yields better quality (FID: 2.77) than other models of comparable depth. Together, we show that new diffusion model designs can be explored by grafting pretrained DiTs, with edits ranging from operator replacement to architecture restructuring. Code and grafted models: https://grafting.stanford.edu

扩散模型架构设计模型嫁接

Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。