Can Local Learning Match Self-Supervised Backpropagation?

TL;DR

The best local-SSL rule with the CLAPP loss function matches the performance of a comparable global BP-SSL with InfoNCE or CPC-like loss functions, and improves upon state-of-the-art for local SSL on these benchmarks.

Abstract

While end-to-end self-supervised learning with backpropagation (global BP-SSL) has become central for training modern AI systems, theories of local self-supervised learning (local-SSL) have struggled to build functional representations in deep neural networks. To establish a link between global and local rules, we first develop a theory for deep linear networks: we identify conditions for local-SSL algorithms (like Forward-forward or CLAPP) to implement exactly the same weight update as a global BP-SSL. Starting from the theoretical insights, we then develop novel variants of local-SSL algorithms to approximate global BP-SSL in deep non-linear convolutional neural networks. Variants that improve the similarity between gradient updates of local-SSL with those of global BP-SSL also show better performance on image datasets (CIFAR-10, STL-10, and Tiny ImageNet). The best local-SSL rule with the CLAPP loss function matches the performance of a comparable global BP-SSL with InfoNCE or CPC-like loss functions, and improves upon state-of-the-art for local SSL on these benchmarks.

Figures (7)

• Figure 1: Overview of BP and local-SSL(A) To train the target synapse (blue), BP gradients need to be propagated down a one-to-one matching error network (red) that is gated (green) by the feedforward activations. (B): In Local-SSL, plasticity is modulated by predictive signals from same-layer lateral projections (red arrow) and global scalar values (green). (C) We show that local-SSL with feedback from the top layer better approximates BP. (D) Notations for theoretical analysis: In Section 3, we compare gradient updates of local-SSL and global BP-SSL. Both algorithms share the same feedforward weight $W^l$ and projection $B_*^L$ in the last layer. Trainable projections $B^l$ are assumed to reach optimum $B_*^l$ instantaneously. (Details in Section 3). (E) Performance summary: On STL-10, our theory-guided improvements (orange bars) improve local-SSL algorithms. More details in Section \ref{['sec:empirical_results']}.
• Figure 2: Numerical verification of Theorem \ref{['th1']}. Cosine similarity of the gradient update between BP-SSL training and local-SSL rules across layers (error bar indicating 95% confidence interval computed through different batches of input). Simulations are for theorem \ref{['th1']} as well as the cases when conditions are dropped. Random fixed B means that the $B^l$ are not optimized to be $B_*^l$. ReLU MLP means adding ReLU activation at each layer. Non-orthogonal W means that each $W^l$ is randomly initialized with the default uniform distribution in PyTorch.
• Figure 3: Numerical verification that direct feedback improves backpropagation approximation. Cosine similarity of the gradient update between local-SSL and BP-SSL across layers. A: Simulation for theorem \ref{['th:low_dim']} on a 6-layer deep linear network with shrinking widths, both for linear $f(x)=-x$ and non-linear (softplus) $f(x) = \log(1+\exp(-x))$. B: Comparison of local-SSL and theoretical optimum of BP approximation. Simulations are conducted on a 6-layer ReLU MLP trained with MNIST input. Gray: top-down feedback projection $B^l$ is fixed with random initialization. Purple: train top-down feedback explicitly with definition \ref{['def:optimal_fb']}.
• Figure 4: Numerical verification that spatial dependence improves BP approximation.(A): Illustration of spatial dependence: without spatial dependence (left), the same $B^l$ is used to project $c^l$ (blue) onto neurons across the feature map, so they share the same color; with spatial dependence (right), lateral projections $B^l$ depend on the locations $(x_1, y_1) \text{ and } (x_2, y_2)$, of $z^l$ and $c^l$ in their layer's feature map. Only nearby neurons with the same color share the same projection $B^l$. (B) and (C): Cosine similarity of the gradient update between BP-SSL and local-SSL across layers. For B, values are obtained from 4-layer linear convenets with optimal $B_*^l$. Kernel and stride have length 2 for simplicity. For C, values are obtained from VGG models trained on STL10. Purple: theoretical optimal update $\Delta_\circ$ obtained from training following the definition \ref{['def:optimal_fb']}. (Details in Appendix \ref{['appendix:train']}). The plotted measurements are performed after 100 epochs of training.
• Figure 5: Comparison between CLAPP illing_2021 and CLAPP++. Orange bars are theory guided algorithmic changes.

Theorems & Definitions (7)

• Theorem 3.1
• proof
• Corollary 3.2
• Theorem 3.3
• Definition 3.4
• proof
• proof