Local proxies for memory-efficient training
作者: Wu S. Zihan, Brian A. Yuan, Guillaume Bellec
分类: cs.LG, cs.AI
发布日期: 2026-10-16
💡 一句话要点
提出局部代理网络以解决内存高效训练问题
🎯 匹配领域: 支柱二：RL算法与架构 (RL & Architecture) 支柱九：具身大模型 (Embodied Foundation Models)
关键词: 局部学习 代理网络 内存效率 变换器 蒸馏训练 深度学习 梯度估计 模型微调
📋 核心要点
- 现有的端到端反向传播方法在内存使用上存在瓶颈，需要存储所有网络激活，导致资源消耗高。
- 本文提出通过训练局部代理网络来实现局部学习，能够在不牺牲准确度的情况下，减少内存需求。
- 实验结果表明，代理传播在CIFAR-100数据集上表现出更好的准确度与内存的权衡，且在微调大型语言模型时也表现优异。
📝 摘要（中文）
局部学习算法能够重塑人工智能的资源需求。标准的端到端反向传播方法的内存瓶颈在于需要存储所有网络激活，直到损失梯度被反向传播。局部学习避免了这一点，但通常会以准确度为代价。本文研究如何在变换器中实现准确且内存高效的局部学习。我们为每个变换器块训练一个局部代理网络，通过蒸馏预测全局网络输出。梯度可以独立地通过这个局部代理传播，同时良好地估计反向传播梯度。理论上，这种方法在深度线性网络中是精确的；在非线性残差模型中，使用适当的残差代理时，梯度误差可以被界定。我们通过数值验证这一点，并展示优化蒸馏损失改善了梯度估计。在CIFAR-100上评估代理传播的性能时，我们发现其在准确度与内存的权衡上优于其他局部算法。最后，我们表明该方法在小内存预算下微调大型语言模型时也具有相关性。我们的算法训练速度更快，性能高于与之内存成本相当的高效零阶学习算法。
🔬 方法详解
问题定义：本文旨在解决标准反向传播方法在内存使用上的瓶颈，特别是在变换器模型中，存储所有激活导致的高资源消耗。
核心思路：通过训练局部代理网络来预测全局网络输出，从而实现局部学习，允许独立的梯度传播，同时保持对反向传播梯度的良好估计。
技术框架：整体架构包括多个变换器块，每个块都有一个局部代理网络。训练过程中，局部代理通过蒸馏学习全局网络的输出，优化损失函数以提高梯度估计的准确性。
关键创新：最重要的创新在于局部代理网络的引入，使得在非线性残差模型中也能有效界定梯度误差，与传统方法相比，显著降低了内存需求而不牺牲准确度。
关键设计：在设计中，采用适当的残差代理以界定梯度误差，优化蒸馏损失函数以提高性能，确保局部代理网络的结构能够有效捕捉全局网络的特征。
🖼️ 关键图片
📊 实验亮点
实验结果显示，代理传播在CIFAR-100数据集上实现了更优的准确度与内存使用的权衡，相较于其他局部算法，准确度提升显著。此外，该方法在微调大型语言模型时，训练速度更快，性能超过了与之内存成本相当的零阶学习算法。
🎯 应用场景
该研究的潜在应用领域包括大规模深度学习模型的训练，尤其是在内存受限的环境中，如移动设备和边缘计算。通过提高内存效率，研究成果能够推动更广泛的人工智能应用，尤其是在资源有限的情况下。未来，该方法也可能对其他类型的深度学习模型产生影响，促进更高效的训练策略。
📄 摘要（原文）
Local learning algorithms can reshape resource requirements of artificial intelligence. The memory bottleneck of standard end-to-end backpropagation lies in the need to store all network activations until the loss gradient has been propagated back through. Local learning avoids this, but typically at a cost in accuracy. We study how to achieve accurate and memory efficient local learning in transformers. For each transformer block, we train a local proxy network to predict the global network output via distillation. Gradients can be independently propagated through this local proxy while estimating well the backpropagation gradients. In theory, this approach is exact for deep linear networks; in non-linear residual models, the gradient error can be bounded when using the appropriate residual proxy. We verify this numerically and demonstrate that optimizing the distillation loss improves gradient estimation. Evaluating the performance of proxy propagation on CIFAR-100 in vision transformers, we find that it achieves a better accuracy/memory trade-off compared to other local algorithms. Finally, we show that the method is further relevant for fine-tuning large language models on a small memory budget. Our method trains faster and achieves higher performance than memory-efficient zeroth-order learning algorithms with comparable memory cost.