DiffusionBlocks：逐块训练神经网络，内存降低数倍的新方法

核心思想：打破端到端训练的内存瓶颈

十多年来，端到端反向传播（end-to-end backprop）一直被视为训练深度网络的唯一方式。但这种方法要求将整个网络同时保存在内存中，随着模型深度线性增长的内存需求，正在成为AI训练的资源瓶颈。

端到端反向传播的核心要求是在前向传播过程中保存每一层的中间激活值（activations），以便在反向传播时计算梯度。对于一个拥有L层的网络，内存消耗与L成线性关系。以GPT-4级别的模型为例，其可能拥有超过120层Transformer块，每层的激活值在大批量训练时可能占用数GB内存。虽然梯度检查点（gradient checkpointing）等技术可以用计算换内存，但本质上并未改变这一线性依赖关系。

梯度检查点的工作原理是：在前向传播时只保存部分层（称为检查点层）的激活值，丢弃其余层的激活值；在反向传播需要某层激活值时，从最近的检查点层重新执行前向计算来恢复。例如，对于一个100层的网络，如果每隔10层设置一个检查点，内存消耗可以从O(L)降低到O(√L)，但代价是额外增加约33%的前向计算量。这本质上是一种时间-空间权衡，并未从根本上解耦训练内存与网络深度的关系。工业界还采用了其他策略来应对内存瓶颈，包括张量并行（将单层的参数和计算分布到多个GPU上）、流水线并行（将不同层分配到不同设备上形成流水线）、以及ZeRO优化器状态分片等。然而，这些方法主要依赖于增加硬件资源来分摊内存压力，而非从算法层面减少内存需求本身。

这一瓶颈直接限制了研究者能够训练的模型深度和批量大小。

Sakana AI团队在ICLR 2026上发表的论文提出了一种全新思路——DiffusionBlocks，将网络拆分为独立的块（block），每次只训练一个块，从而大幅降低训练所需的GPU内存，同时保持与端到端训练相当的性能。

技术原理：将前向传播重新解释为扩散去噪过程

核心洞察：前向传播即逐步去噪

DiffusionBlocks的核心创新在于一个精妙的类比：将神经网络的前向传播过程重新解释为扩散模型逐步去噪信号的过程。

扩散模型（Diffusion Models）的核心思想源自非平衡热力学：首先通过一个前向过程逐步向数据添加高斯噪声，直到数据变为纯噪声；然后训练一个神经网络学习逆向过程，即从噪声中逐步恢复原始数据。每一步去噪都是一个小的、可学习的变换，将当前状态向目标分布推进一小步。DDPM（Denoising Diffusion Probabilistic Models）在2020年证明了这种方法可以生成高质量图像。

从数学角度更精确地理解：扩散模型的前向过程定义为一个马尔可夫链 q(x_t | x_{t-1})，在每一步添加方差为β_t的高斯噪声，经过T步后数据近似服从标准正态分布。逆向过程则训练一个参数化模型 p_θ(x_{t-1} | x_t) 来逐步去噪。训练目标通常简化为让网络预测每一步添加的噪声（即ε-prediction），或者直接预测干净数据（即x_0-prediction）。关键的数学性质是：每一步去噪都可以独立训练，因为给定当前噪声水平t和含噪数据x_t，去噪目标是明确定义的，不需要依赖其他时间步的梯度信息。

DiffusionBlocks巧妙地借用了这一数学框架：将网络中每个块的作用类比为扩散模型中的一步去噪——每个块将中间表征向最终目标推进一小步。具体来说，如果将网络的输入视为"含噪"的表征（距离最终目标较远），那么每经过一个块的处理，表征就"去噪"一步，变得更接近最终的目标表征。这种类比之所以成立，是因为深度残差网络中每个块的变换通常是一个小的增量更新（residual update），这与扩散模型中每步去噪的小步推进在数学结构上高度相似。

具体来说，研究者为网络中的每个块显式分配一个角色——将表征（representation）向目标方向推进一小步，使其比前一个块的输出更接近最终目标。这恰恰就是扩散模型在每一步所做的事情。

逐块训练机制详解

在传统训练中，所有参数需要联合优化，梯度需要从网络末端一路回传到开头。而在DiffusionBlocks框架下：

• 网络被分割为多个独立的块
• 每个块只需优化自己的局部目标函数
• 训练时只需要为单个块分配内存
• 各块可以独立训练，无需同时加载整个网络

使用局部目标函数训练网络的思想并非全新。早在2006年，Hinton等人提出的逐层预训练（layer-wise pretraining）就是通过局部无监督目标（如受限玻尔兹曼机的对比散度算法）逐层初始化深度网络，这一方法开启了深度学习的复兴。后来的辅助损失（auxiliary losses）方法，如GoogLeNet中的中间分类器，也试图缓解梯度消失问题。此后，研究社区探索了多种局部学习方法：对比预测编码（Contrastive Predictive Coding, CPC）通过预测未来表征来提供局部监督信号；Greedy InfoMax将互信息最大化作为逐层训练目标；Local Learning Rules受神经科学启发，试图用赫布学习等生物学上更合理的规则替代反向传播。然而，这些早期方法通常性能不如端到端训练，核心原因在于它们的局部目标与全局任务目标之间缺乏严格的数学联系——每个局部模块优化的目标可能与最终任务目标不一致，导致各层学到的表征无法有效协同。

DiffusionBlocks的关键突破在于，它通过扩散过程的数学框架为局部目标提供了原则性的理论基础，使得每个块的局部优化目标与全局目标保持一致，从而在性能上真正逼近端到端训练。扩散框架天然地定义了每一步应该"去噪"到什么程度（由噪声调度决定），这为每个块提供了明确且与全局一致的优化方向，解决了此前局部训练方法中目标不对齐的根本问题。

这种方法从根本上将训练内存需求从与网络深度线性相关，降低为仅与单个块的大小相关。

实验验证：五种主流架构全面覆盖

研究团队在五种不同的架构上验证了DiffusionBlocks的有效性：

1. ViT（Vision Transformer）：视觉领域的主流架构。ViT由Google在2020年提出，将图像分割为固定大小的patch序列（如16×16像素），每个patch被线性投影为一个token向量，然后用标准Transformer编码器处理，证明了纯Transformer架构在视觉任务上可以匹敌甚至超越CNN。ViT的成功关键在于大规模预训练——在JFT-300M等大数据集上预训练后，ViT在ImageNet等基准上超越了当时最好的CNN模型。其深层堆叠结构（ViT-Large有24层，ViT-Huge有32层）使其成为DiffusionBlocks的理想测试平台。

2. DiT（Diffusion Transformer）：图像生成的前沿模型。DiT将Transformer架构应用于扩散模型的图像生成任务，由Peebles和Xie在2023年提出，用Transformer块替代了传统扩散模型中的U-Net骨干网络，成为Sora等视频生成模型的技术基础。DiT的核心设计包括自适应层归一化（adaLN-Zero）来注入时间步和类别条件信息，以及在潜空间（latent space）中操作以降低计算成本。DiT证明了Transformer的缩放定律（scaling laws）同样适用于图像生成领域——更大的模型和更多的计算量能持续提升生成质量（以FID分数衡量）。

3. Masked Diffusion：掩码扩散模型，将扩散过程应用于离散token空间，通过逐步掩码和去掩码实现文本等离散数据的生成。与连续空间的高斯扩散不同，离散扩散模型面临的核心挑战是如何在离散空间中定义平滑的噪声添加和去噪过程。Masked Diffusion采用了一种直观的方案：前向过程逐步将token替换为特殊的[MASK]标记，逆向过程则学习预测被掩码位置的原始token。这种方法与BERT的掩码语言模型有相似之处，但通过多步迭代去掩码实现了更高质量的生成。近期的MDLM（Masked Diffusion Language Model）和GenBit等工作表明，离散扩散模型在文本生成质量上正在逼近自回归模型，同时具备并行生成的速度优势。

4. 自回归Transformer：大语言模型的核心架构，通过逐token预测下一个token的方式生成序列，是GPT系列模型的基础。

5. 循环深度Transformer（Looped Transformer）：迭代式应用同一网络的架构，通过权重共享实现参数高效的深度扩展。这一思想可以追溯到Universal Transformer（2018年），它将同一个Transformer块反复应用于输入序列，使模型具备"思考更多步"的能力。深度均衡模型（Deep Equilibrium Models, DEQ）将这一思想推向极致，将无限次迭代建模为寻找不动点的过程。循环深度Transformer的优势在于：参数量不随有效深度增长（因为权重共享），且可以根据任务难度动态调整迭代次数。然而，其训练难度远高于标准前馈网络，因为需要通过整个迭代过程进行反向传播。

这两类Transformer架构（ViT和DiT）都具有深层堆叠的特点，非常适合DiffusionBlocks的逐块训练策略。在所有测试场景中，DiffusionBlocks都达到了与端到端训练竞争性的性能，同时仅使用了一小部分内存。

对循环深度Transformer的特殊价值

值得特别关注的是DiffusionBlocks在循环深度（Looped）Transformer上的应用。这类模型通过迭代地应用同一个网络来增加有效深度，但传统训练需要昂贵的时间反向传播（BPTT）。

时间反向传播（Backpropagation Through Time, BPTT）是训练循环网络的标准方法，需要将循环展开为等效的深层前馈网络，然后对整个展开序列进行反向传播。对于循环深度Transformer，如果同一个块被迭代应用T次，BPTT需要保存T步的所有中间激活值，内存消耗为O(T)，计算复杂度也随T线性增长。当T很大时（如数十到数百次迭代），这使得训练变得极其昂贵。

实践中，研究者常采用截断BPTT（Truncated BPTT）作为折中方案：只展开最近K步（K << T）进行反向传播，忽略更早时间步的梯度贡献。这虽然降低了内存和计算成本，但引入了梯度偏差——模型无法学习到需要跨越K步以上的长程依赖关系。另一种替代方案是隐式微分（implicit differentiation），如DEQ模型所采用的方法，它通过在不动点处直接求解雅可比方程来计算梯度，避免了展开过程，但需要额外的线性方程求解步骤，且对不动点的存在性和稳定性有严格要求。

通过DiffusionBlocks的视角，研究者可以用训练时的单次前向传播替代多次迭代的BPTT，这不仅节省了内存，还大幅简化了训练流程。DiffusionBlocks通过将每次迭代视为独立的去噪步骤，只需单次前向传播即可计算局部损失，完全避免了BPTT的展开过程。这意味着无论循环迭代次数T有多大，训练时的内存消耗始终为O(1)——仅需保存单次迭代的激活值。这一特性使得研究者可以在推理时自由增加迭代次数以提升模型能力，而无需担心训练阶段的内存限制，为循环深度架构的实际部署扫清了关键障碍。

意义与展望：重新定义深度网络训练范式

DiffusionBlocks的意义远不止于内存优化：

• 降低训练门槛：使得在有限硬件资源上训练更深的模型成为可能。当前训练一个大型语言模型通常需要数千块高端GPU组成的集群，成本高达数千万美元。DiffusionBlocks通过将内存需求从O(L)降低到O(1)（相对于网络深度），有望让学术实验室和中小型企业也能参与到大模型训练的前沿研究中。

• 理论贡献：建立了深度网络训练与扩散过程之间的数学联系。这一联系不仅具有实用价值，还为理解深度网络的工作机制提供了新的理论视角——将深度网络的逐层变换理解为一个连续动力系统的离散化（类似于Neural ODE的思想），而扩散框架则为这个动力系统的每一步提供了明确的优化目标。

• 实用性强：适用于当前主流的多种架构，包括LLM和视觉模型

• 可扩展性：为未来更大规模模型的训练提供了新的技术路径。DiffusionBlocks的逐块独立训练特性天然支持异步并行——不同的块可以在不同的设备上同时训练，无需像传统流水线并行那样进行复杂的调度和通信。这种特性与当前分布式训练系统的发展方向高度契合。

随着模型规模持续增长，训练资源的瓶颈只会越来越突出。DiffusionBlocks提供了一种原则性的解决方案，有望改变我们训练深度神经网络的基本范式。该工作已在GitHub上开源，研究者可以直接复现和扩展这一方法。

核心要点

核心要点