Stochastic Layer-wise Learning: Scalable and Efficient Alternative to Backpropagation

TL;DR

Stochastic Layer-wise Learning (SLL) targets the memory and update-locking bottlenecks of backpropagation by decomposing the global objective into layer-wise ELBOs under a Markov assumption. It uses fixed stochastic projections to generate auxiliary posteriors and a Bhattacharyya surrogate for the layer-wise KL term, with optional dropout, enabling strictly local updates that remain aligned with the global objective. The authors prove that the mean of layer-wise ELBOs lower-bounds the network ELBO and demonstrate that SLL achieves competitive accuracy across MLPs, CNNs, and ViTs from MNIST to ImageNet with substantial memory savings, including up to 4x in some settings. This work provides a practical, scalable path to modular local learning that preserves global representational coherence without full backpropagation.

Abstract

Backpropagation underpins modern deep learning, yet its reliance on global gradient synchronization limits scalability and incurs high memory costs. In contrast, fully local learning rules are more efficient but often struggle to maintain the cross-layer coordination needed for coherent global learning. Building on this tension, we introduce Stochastic Layer-wise Learning (SLL), a layer-wise training algorithm that decomposes the global objective into coordinated layer-local updates while preserving global representational coherence. The method is ELBO-inspired under a Markov assumption on the network, where the network-level objective decomposes into layer-wise terms and each layer optimizes a local objective via a deterministic encoder. The intractable KL in ELBO is replaced by a Bhattacharyya surrogate computed on auxiliary categorical posteriors obtained via fixed geometry-preserving random projections, with optional multiplicative dropout providing stochastic regularization. SLL optimizes locally, aligns globally, thereby eliminating cross-layer backpropagation. Experiments on MLPs, CNNs, and Vision Transformers from MNIST to ImageNet show that the approach surpasses recent local methods and matches global BP performance while memory usage invariant with depth. The results demonstrate a practical and principled path to modular and scalable local learning that couples purely local computation with globally coherent representations.

Figures (8)

• Figure 1: Overview of Stochastic Layer-wise Learning (SLL). (a) SLL treats each hidden activation $h_l$ as a latent variable and projects it to $v_l$ via a random matrix. The local ELBO comprises a log-likelihood term and a KL surrogate that promotes inter-layer consistency. Optimizing this loss yields the improved projection $\tilde{v}_l$ and its corresponding activation $\tilde{h}_l$. (b) SLL optimizes each layer independently using a prediction loss $\mathcal{L}_{\text{pred}}(v_l, y)$ from the log-likelihood and a feature alignment loss $\mathcal{L}_{\text{BC}}^{per}(v_l, v_{l-1}^{\text{sg}})$ approximating the KL term. Arrows denote forward computation (black), local updates (red), and stop-gradient paths (slashed).
• Figure 2: Peak training memory on (a) MLPs (1024 neurons/layer) as a function of depth. BP memory scales linearly, while SLL remains constant; (b) CNNs on the Imagenette without pooling layers. Each convolution layer uses a kernel size of 3 and 64 output channels; (c) ViTs on Imagenette. For fair comparison, we are using SGD as the optimizer in training.
• Figure 3: t-SNE visualization of activations and random projections on MNIST, colored by class.
• Figure 4: (a) Training curves of a 3-layer MLP on CIFAR-10 via SLL. Ablation study: (b) random projection dimension in a 3×1000 MLP trained on CIFAR-100; activations are downsampled to $d$-dim via adaptive pooling before projection. (c) network width in SLL on CIFAR-100, showing that wider layers significantly enhance performance and stability.
• Figure 5: Activation and weight distributions from VGG-11 trained with BP and SLL on Imagenette.