arXiv:2506.14202cs.LGcs.AI2025-06被引 7

将Transformer分块训练转化为可独立优化的去噪过程,显著降低内存占用。

DiffusionBlocks: Block-wise Neural Network Training via Diffusion Interpretation

  • 利用残差连接对应动力系统更新,转为去噪过程实现分块独立训练
  • 在多种Transformer架构上性能媲美端到端训练,内存需求随块数线性下降
  • 适合大规模生成模型训练,尤其适用于视觉与扩散类任务

端到端反向传播需存储所有层的激活值,导致内存瓶颈限制模型扩展。现有分块训练方法依赖经验式局部目标,且多局限于分类任务。本文提出DiffusionBlocks,一种将基于Transformer的网络转化为真正独立可训练块的理论框架,在保持与端到端训练相当性能的同时,支持可扩展的分块训练。核心思想是残差连接天然对应动力系统更新,通过最小修改将其转化为去噪过程,各块可借助分数匹配目标独立学习。该独立性使每次仅需梯度更新一个块,内存开销随块数成比例降低。在多种Transformer架构(视觉、扩散、自回归、递归深度、掩码扩散)上的实验表明,DiffusionBlocks训练性能与端到端相当,并可在实际任务中实现可扩展分块训练。该方法为跨多样化架构的现代生成任务提供了理论扎实的可扩展训练路径。代码见https://github.com/SakanaAI/DiffusionBlocks。

原文摘要 · Abstract (English)

End-to-end backpropagation requires storing activations throughout all layers, creating memory bottlenecks that limit model scalability. Existing block-wise training methods offer means to alleviate this problem, but they rely on ad-hoc local objectives and remain largely unexplored beyond classification tasks. We propose $\textit{DiffusionBlocks}$, a principled framework for transforming transformer-based networks into genuinely independent trainable blocks that maintain competitive performance with end-to-end training. Our key insight leverages the fact that residual connections naturally correspond to updates in a dynamical system. With minimal modifications to this system, we can convert the updates to those of a denoising process, where each block can be learned independently by leveraging the score matching objective. This independence enables training with gradients for only one block at a time, thereby reducing memory requirements in proportion to the number of blocks. Our experiments on a range of transformer architectures (vision, diffusion, autoregressive, recurrent-depth, and masked diffusion) demonstrate that DiffusionBlocks training matches the performance of end-to-end training while enabling scalable block-wise training on practical tasks beyond small-scale classification. DiffusionBlocks provides a theoretically grounded approach that successfully scales to modern generative tasks across diverse architectures. Code is available at https://github.com/SakanaAI/DiffusionBlocks .

Transformer分块训练扩散模型内存优化

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