DiffusionBlocks: Blockwise Training for Generative Models via Score-Based Diffusion

使用端到端反向传播训练大型神经网络会带来显著的内存瓶颈，限制了前沿人工智能研究的可及性。我们提出了一种新的训练框架——*DiffusionBlocks*，它将神经网络模块解释为在连续时间扩散过程中执行去噪操作。通过将网络划分为可独立训练的模块，并基于等累计概率质量优化噪声水平分配，我们的方法在生成任务中实现了显著的内存效率提升，同时保持了与传统反向传播相当的性能。在图像生成和语言建模任务中的实验表明，该方法能够实现与模块数量成比例的内存减少，同时达到更优的性能表现。*DiffusionBlocks* 为在有限计算资源下普及大规模神经网络训练提供了一条有前景的路径。

论文试图解决大型神经网络训练中的内存瓶颈问题，这一问题限制了研究人员在有限计算资源下进行大规模AI模型的开发和研究。这是一个长期存在的问题，但随着模型规模的扩大变得愈发突出。

论文提出了一种名为DiffusionBlocks的新型训练框架，将神经网络块解释为连续时间扩散过程中的去噪操作，并通过将网络划分为独立可训练块以及优化噪声水平分配来显著降低内存消耗。相比传统端到端反向传播方法，这种方法提供了更高的内存效率，同时保持了竞争力的性能。

实验涵盖了图像生成和语言建模任务，展示了内存使用量随块数增加而减少的同时，性能优于或接近传统方法。论文没有明确提及数据集名称，但可能使用了常见的公开数据集（如CIFAR-10、ImageNet或Penn Treebank）。此外，作者未提及代码开源情况，但提出的工作为进一步探索基于扩散模型的高效训练策略提供了方向，例如更复杂的噪声调度机制或跨任务适用性分析。

最近的相关工作包括：1)《Deep Unfolding for Denoising Autoencoders》探讨了类似解耦思想；2)《Memory-Efficient Training of Diffusion Models》研究了扩散模型训练中的内存优化；3)《Optimizing Large-Scale Neural Networks with Block-Wise Backpropagation》提出了分块反向传播技术。这些研究共同推动了高效训练方法的发展。