[论文解读] Scaling Forward Gradient With Local Losses

Mengye Ren, Simon Kornblith+2 more|arXiv (Cornell University)|Oct 7, 2022

Domain Adaptation and Few-Shot Learning被引用 12

一句话总结

该论文提出了一种生物上合理的深度学习方法，采用前向梯度学习结合激活值扰动与局部贪婪损失，以降低梯度方差并提升可扩展性。通过在激活值而非权重上施加扰动，并采用新型的LocalMixer架构与分组化、局部化的损失函数，该方法在MNIST和CIFAR-10上实现了与反向传播相当的性能，在ImageNet上显著优于以往的无反向传播方法。

ABSTRACT

Forward gradient learning computes a noisy directional gradient and is a biologically plausible alternative to backprop for learning deep neural networks. However, the standard forward gradient algorithm, when applied naively, suffers from high variance when the number of parameters to be learned is large. In this paper, we propose a series of architectural and algorithmic modifications that together make forward gradient learning practical for standard deep learning benchmark tasks. We show that it is possible to substantially reduce the variance of the forward gradient estimator by applying perturbations to activations rather than weights. We further improve the scalability of forward gradient by introducing a large number of local greedy loss functions, each of which involves only a small number of learnable parameters, and a new MLPMixer-inspired architecture, LocalMixer, that is more suitable for local learning. Our approach matches backprop on MNIST and CIFAR-10 and significantly outperforms previously proposed backprop-free algorithms on ImageNet.

研究动机与目标

• 解决标准前向梯度学习在大规模深度网络中梯度方差高且可扩展性差的问题。
• 开发一种生物上合理的反向传播替代方法，避免对称的反向权重传输与全局同步。
• 通过激活值扰动与局部损失函数，降低梯度方差，提升训练稳定性和性能。
• 设计一种新型架构LocalMixer，专为局部学习与模型并行计算而优化。
• 在标准视觉基准（MNIST、CIFAR-10、ImageNet）上展示无需依赖反向传播的实用性能。

提出的方法

• 使用前向自动微分，通过在激活值上施加扰动而非权重来计算方向梯度，从而降低梯度方差。
• 引入大量局部贪婪损失函数，每个损失仅涉及少量参数，以隔离学习过程并减少方差。
• 设计受MLPMixer启发的LocalMixer架构，包含线性token混合与分组通道操作，以增强与局部学习的兼容性。
• 使用停止梯度操作将局部预测与主网络解耦，实现独立的局部优化。
• 在关键位置应用归一化层（BatchNorm、LayerNorm、LocalNorm）以稳定训练并提升泛化能力。
• 实施多阶段训练流程，包含局部与全局预测头，以支持局部与全局信号的传播。

实验结果

研究问题

• RQ1 在大规模深度网络中，与权重扰动相比，激活值扰动能否使前向梯度学习获得更低的方差？
• RQ2 具有分组与分块结构的局部贪婪损失函数如何影响无反向传播学习中的训练稳定性和性能？
• RQ3 像LocalMixer这样的新型架构能否在无需反向传播的情况下，实现有效的局部学习并扩展至ImageNet级别的基准？
• RQ4 归一化方案（BN、LN、LocalNorm）对局部前向梯度训练性能有何影响？
• RQ5 激活值扰动、局部损失与专用架构的组合是否在标准视觉基准上优于现有无反向传播方法？

主要发现

• 与权重扰动相比，激活值扰动可显著降低梯度估计的方差，从而在高维设置中实现更稳定的训练。
• 在MNIST和CIFAR-10上，所提出的局部前向梯度（LG-FG-A）方法实现了与反向传播相当的测试准确率（例如，在CIFAR-10上使用LayerNorm时测试误差为33.48%，而反向传播为30.55%）。
• 在ImageNet上，该方法优于所有先前的无反向传播算法，使用LayerNorm时达到37.41%的top-1准确率，使用LocalNorm时达到36.24%，显著优于现有替代方案。
• 在局部损失函数中增加分组数量，可使前向梯度训练的训练误差降低最多10%，表明梯度方差减少且优化性能提升。
• 带有分组化、局部学习组件的LocalMixer架构，可实现有效的模型并行训练，并在全局信号有限的情况下仍取得优异性能。
• 如LocalNorm等归一化方案能显著提升泛化能力，在前向梯度训练中相比无归一化可将测试误差降低最多5%。