Stanford maps two routes past weight transport: local IA at 67.9%, SA matching backprop

Two Routes to Scalable Credit Assignment without Weight Symmetry

Daniel Kunin, Aran Nayebi, Javier Sagastuy-Brena, Surya Ganguli, Jonathan M. Bloom, Daniel L. K. Yamins

ICML 2020 Camera Ready Versi

q-bio.NC, cs.LG, cs.NE, stat.ML

2020-02-29

Local Information Alignment reaches 67.93% ImageNet top-1 on ResNet-18 (backprop 70.06%). Non-local Symmetric Alignment matches backprop through ResNet-152.

What problem this solves

Backprop trains deep nets at ImageNet scale, but it needs the backward weight matrix to be the exact transpose of the forward one. That is the weight transport problem: one synapse would have to read another synapse's strength on the spot. No known neural mechanism does that.

Older workarounds stall at scale. Feedback alignment freezes the backward weights as random matrices; it works on shallow nets and MNIST, then falls apart on deep ImageNet models. Target propagation asks the backward weights to invert the forward map locally, and it fails on ImageNet too. Weight Mirror (Akrout et al., 2019) reported 69.73% top-1 on ResNet-18, nearly matching backprop. This Stanford group reimplemented the published recipe and got 63.5%. A TPE search over 824 settings peaked at 64.07%. Copy those ResNet-18 metaparameters onto deeper or v2 ResNets and the alignment loss becomes NaN within the first thousand steps.

Local rules can run. They do not travel.

Method

Untie the two weight sets. Forward weights θf own the task loss J. Backward weights θb carry the error signal. A regularizer R acts only on the backward weights, so the total objective is J(θf)+R(θb). Each step, θf follows the pseudogradient that θb produces; θb follows ∇R. If R is chosen well, the two sets align over time and the pseudogradient starts looking like a true gradient.

R is a linear mix of five geometric primitives, split by locality. Three are local, using only pre- and post-synaptic activity: decay penalizes the Euclidean norm of the backward weights; amp aligns the reconstruction Bl xl+1 with the layer input xl; null penalizes the norm of that reconstruction. Two are non-local. sparse applies Bl to xl, which is the wrong direction for those synapses. self takes the inner product of Bl and Wl, so a backward synapse has to measure its forward counterpart.

Those pieces rebuild known algorithms and yield new ones. Feedback alignment is R=0. Weight Mirror is amp plus decay. Information Alignment, the new local rule, adds null on top of Mirror. The extra term comes from Oja's normalization of a Hebbian update: ∆Bl=η xl xl+1^T blows up the norm; decay is one brake, and an Oja-style expansion of a unit-norm constraint is approximately the null primitive. When α=γ and xl+1=Wl xl, IA is a quadratically regularized linear autoencoder. Kunin et al. (2019) showed that every critical point of that autoencoder has symmetric encoder and decoder.

The two non-local rules push alignment more explicitly. Symmetric Alignment mixes self and decay; its gradient is proportional to ‖Wl − Bl^T‖². Activation Alignment mixes amp and sparse and aligns activations, ‖Wl xl − Bl^T xl‖². For both, the dynamical matrix is positive semi-definite and the eigenvalue on the symmetric component is zero. Weight Mirror's matrix is indefinite, which is why its symmetric mode either explodes or collapses unless λWM exactly matches xl².

They also tried Adam and parameter-free centering / feature normalization on the backward path. The IA search covered all three levers, 628 settings.

Results

ResNet-18 on ImageNet validation:

Methodtop-1top-5
Weight Mirror (Akrout recipe)63.5%85.16%
Weight Mirror + TPE64.07%85.47%
Mirror + Adam64.40%85.53%
Mirror + Adam + local norms63.41%84.83%
Information Alignment67.93%88.09%
Backprop70.06%89.14%

Adam and the local normalizations barely move ResNet-18 accuracy on their own, but they stop the NaNs when the same metaparameters move to other architectures. IA's ResNet-18 setting, reused as-is on ResNet-50, 50v2, 101v2, and 152v2, beats every other local rule and still trails backprop. The gap grows with depth.

The non-local pair sits next to backprop with almost no tuning:

ModelBackpropSAAA
ResNet-1870.06%69.84%69.98%
ResNet-5076.05%76.29%75.75%
ResNet-50v277.21%77.18%76.67%
ResNet-101v278.64%78.74%78.35%
ResNet-152v279.31%79.15%78.98%

SA does not need Weight Mirror's alternating learning modes. AA aligns less tightly, so minimizing R needs Adam. Add Gaussian noise to the backward updates and SA holds up better than backprop does under the same noise on its gradients. Kolen-Pollack, in the appendix, matches backprop on ResNet-18 and ResNet-50, then lags on v2 and deeper nets. It uses local primitives, but it still demands that each backward update be the exact transpose of the forward one, which is about as suspect as instantaneous weight symmetry.

Fit ResNet-18 activations to time-averaged macaque V4 and IT responses and every rule that actually trains, other than failed feedback alignment, matches backprop's neural predictivity. An adult animal's functional responses do not reveal which learning rule produced them.

Why it matters

This is a map of composable credit assignment without tied weights, not another local rule that only wins on MNIST. If ImageNet accuracy is the goal, use backprop. If the weights have to stay decoupled, SA is close to drop-in, and IA is the most stable fully local option on the table.

The paper is blunt about the two remaining routes. Keep improving local rules until they stop dropping across architectures, or find a spike-level implementation of the weight estimation that SA and AA assume. Guerguiev et al. already approximated the self term on small networks with regression discontinuity design. This work shows the rate-coded version scales to deep ImageNet models.

Limitations

IA still trails backprop by about 2.1 points, and the gap widens with depth. Nothing in the paper closes that local gap. SA and AA avoid instantaneous transposes, but the self term itself fails local plasticity: backward weights have to read forward weights. Weight estimation is only tested as rate code plus Gaussian noise. Training deep spiking nets, the authors say, is out of scope.

The 69.73% Weight Mirror number did not reproduce. That could be an implementation mismatch, or the original figure may have been optimistic. Do not treat "Mirror already matches backprop" as settled. Almost every experiment is ImageNet plus the ResNet family: no language, no RL, no other vision backbones. Adult cortical fits cannot tell these rules apart, so neuroscience will have to measure synapses during learning, not V4/IT after training.

Terms

Source

What people are saying

Related papers

All paper explainers