Speedrunning ImageNet Diffusion

📄 arXiv: 2512.12386v1 📥 PDF

作者: Swayam Bhanded

分类: cs.CV

发布日期: 2025-12-13


💡 一句话要点

提出SR-DiT,结合多种优化策略加速ImageNet扩散模型训练,显著提升效率。

🎯 匹配领域: 支柱四:生成式动作 (Generative Motion)

关键词: 扩散模型 Transformer 图像生成 ImageNet 训练加速 Token路由 表征对齐

📋 核心要点

  1. 现有扩散模型训练效率提升方法研究分散,缺乏系统整合与协同优化。
  2. SR-DiT框架整合token路由、架构改进和训练修改,并以表征对齐为基础,实现高效训练。
  3. 实验表明,SR-DiT在ImageNet-256上以更小模型和更少迭代次数达到媲美更大模型的性能。

📝 摘要(中文)

最近的进展显著提高了扩散Transformer的训练效率。然而,这些技术大多是孤立地研究的,忽略了结合多种方法所产生的潜在协同效应。我们提出了SR-DiT(Speedrun Diffusion Transformer),一个系统地整合了token路由、架构改进和训练修改的框架,并以表征对齐为基础。我们的方法在ImageNet-256上仅使用一个1.4亿参数的模型,在40万次迭代中,无需分类器引导,就实现了FID 3.49和KDD 0.319,与参数量为6.85亿的模型经过更长时间训练的结果相当。据我们所知,这是该模型尺寸下的最先进结果。通过广泛的消融研究,我们确定了哪些技术组合最有效,并记录了协同效应和不兼容性。我们将我们的框架作为一个计算可访问的基线发布,以供未来研究使用。

🔬 方法详解

问题定义:论文旨在解决扩散模型在ImageNet等大型数据集上训练效率低下的问题。现有方法通常孤立地研究各种优化策略,未能充分挖掘它们之间的协同效应,导致训练成本高昂,阻碍了扩散模型在实际应用中的推广。

核心思路:论文的核心思路是将多种已有的优化策略(包括token路由、架构改进和训练修改)进行系统性地整合,并通过表征对齐作为基础,从而充分发挥它们之间的协同效应,实现训练效率的最大化。这种整合的思路旨在打破各个优化策略之间的壁垒,使得它们能够相互促进,共同提升训练效率。

技术框架:SR-DiT框架的核心是一个扩散Transformer模型,其训练流程包括以下几个主要阶段:首先,使用token路由机制来减少计算量,加速模型训练;其次,采用架构改进策略,例如更高效的网络结构设计,以提升模型的表达能力;然后,通过训练修改,例如调整学习率策略或引入正则化项,来优化模型的训练过程;最后,以表征对齐作为基础,确保不同优化策略之间能够协同工作,避免出现冲突或抵消效应。

关键创新:论文的关键创新在于系统性地整合了多种优化策略,并以表征对齐作为基础,从而实现了训练效率的显著提升。与现有方法相比,SR-DiT不是简单地堆砌各种优化策略,而是通过精心设计,使得它们能够相互协同,共同提升性能。此外,论文还通过消融实验,深入分析了不同优化策略之间的协同效应和不兼容性,为未来的研究提供了重要的参考。

关键设计:在token路由方面,论文采用了动态稀疏化的方法,根据token的重要性自适应地调整计算量。在架构改进方面,论文采用了更高效的Transformer结构,例如使用线性注意力机制或减少参数量。在训练修改方面,论文采用了自适应学习率调整策略和权重衰减正则化方法。表征对齐通过最小化不同优化策略下模型输出表征之间的差异来实现。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

SR-DiT在ImageNet-256上仅使用1.4亿参数的模型,经过40万次迭代,无需分类器引导,就实现了FID 3.49和KDD 0.319。这一结果与参数量为6.85亿的模型经过更长时间训练的结果相当,表明SR-DiT在模型效率方面取得了显著的提升,代表了当前模型尺寸下的领先水平。

🎯 应用场景

该研究成果可应用于图像生成、图像编辑、超分辨率等领域,尤其是在计算资源受限的情况下,能够以更低的成本训练出高性能的扩散模型。这有助于推动扩散模型在移动设备、嵌入式系统等平台的应用,并促进相关技术的普及。

📄 摘要(原文)

Recent advances have significantly improved the training efficiency of diffusion transformers. However, these techniques have largely been studied in isolation, leaving unexplored the potential synergies from combining multiple approaches. We present SR-DiT (Speedrun Diffusion Transformer), a framework that systematically integrates token routing, architectural improvements, and training modifications on top of representation alignment. Our approach achieves FID 3.49 and KDD 0.319 on ImageNet-256 using only a 140M parameter model at 400K iterations without classifier-free guidance - comparable to results from 685M parameter models trained significantly longer. To our knowledge, this is a state-of the-art result at this model size. Through extensive ablation studies, we identify which technique combinations are most effective and document both synergies and incompatibilities. We release our framework as a computationally accessible baseline for future research.