Scaling Forward Gradient With Local Losses
This work addresses the scalability issue of biologically plausible learning algorithms for deep learning practitioners, representing a significant but incremental improvement over existing forward gradient methods.
The paper tackles the high variance problem in forward gradient learning for deep neural networks by proposing architectural and algorithmic modifications, achieving performance matching backprop on MNIST and CIFAR-10 and outperforming previous backprop-free methods on ImageNet.
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.