面向内存高效训练的局部代理网络

arXiv:2610.23145v1 发布类型：新 摘要：局部学习算法能够重塑人工智能的资源需求。标准端到端反向传播的内存瓶颈在于：在损失梯度反向传播完成之前，需要存储网络的全部激活值。局部学习可以避免这一点，但通常会以牺牲准确率为代价。我们研究了如何在 Transformer 中实现准确且内存高效的局部学习。对于每个 Transformer 模块，我们通过蒸馏训练一个局部代理网络，使其预测全局网络的输出。梯度可以独立地通过该局部代理网络传播，同时较为准确地估计反向传播梯度。从理论上看，对于深度线性网络，这种方法是精确的；对于非线性残差模型，在使用适当的残差代理网络时，梯度误差可以得到界定。我们通过数值实验验证了这一点，并证明优化蒸馏损失能够改善梯度估计。在视觉 Transformer 上使用 CIFAR-100 评估代理传播的性能时，我们发现，与其他局部算法相比，该方法实现了更优的准确率与内存权衡。最后，我们展示了该方法对于在有限内存预算下微调大型语言模型同样具有价值。与内存成本相当的内存高效零阶学习算法相比，我们的方法训练速度更快且性能更高。