Sakana AI 提出 DiffusionBlocks：一种将残差网络转换为可独立训练去噪模块的块状训练框架

Sakana AI 与东京大学的研究人员联合提出了一种名为 DiffusionBlocks 的新型训练框架，旨在解决深度神经网络训练中的内存瓶颈问题。该框架将 Transformer 网络划分为多个独立的块（block），每个块均可独立训练，从而将训练内存消耗降低至原来的 1/B（B 为块数），同时保持与端到端训练相当的性能。

在传统的端到端反向传播训练中，需要存储每一层的中间激活值，内存消耗随网络深度线性增长。尽管激活检查点（activation checkpointing）可以减少激活内存，但无法减少参数、梯度和优化器状态所占用的内存。例如，使用 Adam 优化器时，每一层需要存储参数、梯度以及动量与方差两个状态，总内存约为参数量的 4 倍。块状训练通过将网络分为 B 块并独立训练，将内存需求降低至约 1/B。然而，如何为每个块设计合理的局部目标一直是个挑战。此前的方法如 Hinton 的正向-正向算法（Forward-Forward）和贪婪逐层训练依赖于特定的局部目标，性能往往不如端到端训练，且主要局限于分类任务。

DiffusionBlocks 的核心见解在于利用残差网络与扩散模型之间的深层联系。残差网络的更新形式为 z_ℓ = z_{ℓ-1} + f_θℓ(z_{ℓ-1})，这相当于常微分方程的欧拉离散化。研究者证明，这一更新形式恰好对应于基于分数的扩散模型中的概率流 ODE（在方差爆炸（VE）框架下）。通过反向扩散过程的欧拉离散化，可以得到与残差连接结构完全匹配的更新规则。因此，一叠残差块可以被视为离散化的去噪步骤。在基于分数的扩散模型中，分数匹配目标可以在每个噪声水平上独立优化，这意味着每个块可以使用自己的局部目标独立训练，无需块间通信。

将标准残差网络转换为 DiffusionBlocks 需要三步：首先，将 L 层网络划分为 B 个连续块；其次，定义噪声分布并分配噪声区间给每个块，推荐使用对数正态分布；最后，通过向每个块输入添加目标的噪声版本和条件归一化（AdaLN）来实现噪声条件化。训练时每次迭代只采样一个块，其他块不进行计算，内存消耗仅为 L/B 层。

在划分策略上，DiffusionBlocks 采用等概率划分而非均匀划分。均匀划分忽略了不同噪声水平的去噪难度差异，而等概率划分使每个块处理相同概率质量下的区间，从而在 CIFAR-10 上实现了更低的 FID（38.03 vs 43.53）。

实验评估覆盖了五种架构和多个数据集。在 CIFAR-100 上使用 ViT，DiffusionBlocks 达到了 59.30% 的准确率，而端到端基线为 60.25%；在 CIFAR-10 上使用 DiT-S/2，FID 为 37.20 vs 39.83；在 ImageNet 256×256 上使用 DiT-L/2，FID 为 10.63 vs 12.09；在 text8 上使用掩码扩散模型（MDM），BPC 为 1.45 vs 1.56；在 LM1B 上使用自回归 Transformer，MAUVE 为 0.71 vs 0.50；在 OpenWebText 上，MAUVE 为 0.82 vs 0.85。此外，对于循环深度模型 Huginn，DiffusionBlocks 以约 10 倍的计算缩减实现了 MAUVE 0.70 vs 0.49。这些结果表明，DiffusionBlocks 在内存大幅减少的同时，性能与端到端训练相当，甚至在某些情况下更好。

与同期工作 NoProp 相比，DiffusionBlocks 是唯一结合连续时间公式和块状训练的方法，在 CIFAR-100 上准确率达到 46.88%，仅比端到端反向传播低不到 1 个百分点。

DiffusionBlocks 的优势包括：基于分数匹配的原则性理论基础，无需任务特定修改即可适用于多种架构，B 倍的内存减少，以及扩散模型推理时 B 倍的计算减少。等概率划分显著优于均匀划分，且对于循环深度模型，它用单次前向传播代替了多次 BPTT。此外，块可以在 GPU 上并行训练，通信开销为零。

局限性包括：需要输入和输出维度匹配，目前无法应用于 U-Net 类架构；仅在从头训练的模型上验证过；缺乏选择最优块数的原则性方法；噪声条件化带来轻微的时间开销；在 OpenWebText 上某些指标略差。

总体而言，DiffusionBlocks 为训练深层 Transformer 网络提供了一种内存高效且理论上合理的替代方案，尤其适用于扩散模型和循环深度模型。