DiffusionBlocks：逐个模块训练神经网络
DiffusionBlocks 是一种区块式神经网络训练方法，它将网络划分为多个区块，并利用扩散框架独立训练每个区块，从而将内存消耗降低到单个区块的水平，同时保持与端到端训练相当的性能。该方法在 ICLR 2026 上发表。
现代人工智能在日常生活、数学、编程等任务中展现出卓越性能。然而，前沿模型通常拥有数百亿甚至更多参数，需要数千个 GPU 进行训练，只有少数组织具备开发资源。这一资源需求的一大原因在于训练方法本身：当今的神经网络（包括 Transformer）通过端到端优化进行训练，所有参数联合学习，导致内存需求随模型规模增长。理想情况下，我们希望能不受硬件限制地训练任意大小的模型，这正是 DiffusionBlocks 的动机——让 AI 开发更加可及。
DiffusionBlocks 是一种将网络划分为多个区块并依次独立训练的方法。在训练过程中，只需为一个区块分配内存，而不是整个网络，从而显著降低内存消耗。同时，通过扩散框架，其性能与端到端优化相当。该研究由 Makoto Shing、Masanori Koyama 和 Takuya Akiba 共同完成，并在 ICLR 2026 上展示。
背景：深度学习中的内存瓶颈
现代 AI 的发展受益于规模化——模型越大、数据越多，性能越好。但规模化也带来了内存成本的增加。端到端反向传播需要保留所有中间状态（激活值），因此内存消耗与网络深度成正比。由于 Transformer 的扩展主要依靠增加层数，深度直接驱动了内存成本的增长。
区块式训练的实践
区块式训练是一种自然解决方案：将网络划分为较小区块并独立训练，使训练内存降至单个区块的需求。先前的工作在图像分类上有所探索，但通常落后于端到端训练，且难以扩展到生成任务和现代 Transformer 架构。关键限制在于：局部目标的设计缺乏理论基础，且适用范围有限。
DiffusionBlocks 方法
DiffusionBlocks 明确为每个区块分配角色，使网络整体完成从输入到目标的任务。每个区块只需优化自身目标，可独立于其他区块进行训练。具体地，每个区块负责“逐步接近目标”的动力学过程，这与扩散模型的本质一致。基于残差连接对应于微分方程离散步骤的观察，区块动力学可被视为扩散模型的反向（去噪）过程，从而实现原理性独立训练。
将现有网络转换为 DiffusionBlocks 只需三步：
- 分区：将 L 层划分为 B 个区块。
- 分配噪声范围：每个区块负责一定范围的“接近目标程度”。
- 添加条件模块：使区块识别其分配范围。
训练时，每次迭代随机采样一个区块，内存消耗降至约 1/B。
验证结果
DiffusionBlocks 在三个任务域的五种架构上进行了验证：图像分类（ViT）、图像生成（DiT）和文本生成（掩码扩散、自回归 Transformer、循环深度 Transformer）。在所有情况下，DiffusionBlocks 均实现了与端到端训练相当的性能，同时内存使用大幅降低。
扩展到循环深度模型
DiffusionBlocks 自然扩展到循环深度模型（循环 Transformer），这类模型将同一网络重复应用 K 次。传统的时间反向传播（BPTT）计算成本高，而 DiffusionBlocks 视角允许在训练时仅进行一次前向传播，显著降低计算开销而不牺牲性能。
未来工作
两个重要方向：一是理论分析为何内存效率与性能均可提升（可能源于自然的课程学习效应）；二是扩展到预训练大模型，通过微调将现有模型转换为 DiffusionBlocks，使大规模模型的训练和下游使用对个人研究者、学生和小型实验室更为可行。
DiffusionBlocks 是 Sakana AI 提高 AI 效率更广泛努力的一部分，为训练阶段的内存消耗提供了新的优化方向。有关更多详情，请参阅完整论文。