PC-ALM Trains 1,000-Layer Networks Within 2 Points of Backprop

Sakana AI has published Augmented Lagrangian Predictive Coding (PC-ALM), a training algorithm that eliminates backpropagation's requirement for global phase-locking while recovering credit signals that match backprop's gradients in the linear case. The result: residual MLPs 1,000 layers deep trained within roughly 2 percentage points of backpropagation on MNIST — a depth regime where standard predictive coding collapses entirely. The MIT-licensed JAX reference implementation runs on CPU, making the method immediately reproducible without specialised hardware.

The motivation is well-grounded. As architectural specificity increasingly outperforms raw GPU scaling, training algorithms that escape backprop's sequential dependency become a legitimate systems concern, not just a neuroscience curiosity.

Why standard predictive coding breaks at depth

Predictive coding treats every hidden activation as an optimisation variable and penalises squared mismatch between each layer's activation and the prediction from the layer below. Supervision enters at the output and must diffuse inward through a chain of local quadratic penalties. Innocenti et al. characterised this PC-BP gap as worst when width is smaller than depth: in deep, narrow networks the credit signal fades before reaching early layers.

The augmented Lagrangian fix

PC is the quadratic-penalty relaxation of a constrained optimisation problem — minimise the supervised loss subject to hᵢ = σ(Wᵢhᵢ₋₁) at every layer. PC-ALM replaces that relaxation with the full augmented Lagrangian, attaching a Lagrange multiplier λᵢ ∈ ℝᵈⁱ (same dimension as hᵢ) to each layer's constraint while retaining PC's quadratic penalty. Setting λ = 0 recovers PC exactly.

Each inference step alternates between two local operations: a primal gradient step on activations, and a dual accumulation step λᵢ ← λᵢ + αrᵢ that integrates the layer's prediction error rᵢ. Completing the square reveals that every primal step is a standard PC step with the prediction target shifted by −λᵢ/ρ. After T steps the weight update acts on the composite signal λᵢ + ρrᵢ. The authors interpret this as a PI controller per layer: the prediction error is the proportional term, the Lagrange multiplier is the integral term. Setting α = 0 yields PC; setting α = ρ with the inner problem solved exactly yields the classical method of multipliers.

In linear networks the team proves that under a spectral-radius stability condition — ηₕσᵢ²(2ρ + α) < 4 per mode — PC-ALM converges to the KKT point where each λᵢ equals the exact backprop adjoint, a relationship LeCun identified in 1988. Unlike PC's monotone gradient flow, PC-ALM's iteration matrix has complex eigenvalues that produce damped oscillations; α controls oscillation frequency but not the decay rate.

Benchmark results

The sweep covers residual MLPs with widths and depths ranging from 8 to 128, tested on both MNIST and Fashion-MNIST under the mean-field parameterisation of Innocenti et al., trained for 1 epoch with an inference budget of T = 2L. PC-ALM matches backprop across every width, depth, and activation function (identity, tanh, ReLU) tested, while PC degrades sharply in deep, narrow configurations. Extended experiments include 1,000-layer residual MLPs on MNIST (width 32, ReLU, 5 epochs) staying within approximately 2 percentage points of BP, and ResNet-18 results on CIFAR-10 and Tiny ImageNet where PC-ALM outperforms PC on every task attempted.

AI Mastery analysis

The proof that PC-ALM converges to exact backprop adjoints holds only for linear networks under a stability bound. Nonlinear networks — every practical case — carry no such guarantee, making the empirical width-depth sweep the actual evidence of correctness rather than the theoretical result. The T = 2L inference budget is non-trivial: at 1,000 layers that is 2,000 inner steps per outer weight update, substantially increasing compute per gradient step even if each step is parallelisable across layers.

The memory cost is concrete and stated: PC-ALM requires twice the activation memory of standard PC because both activations and Lagrange multipliers must be stored per layer. For a 1,000-layer network that is a meaningful overhead that practitioners targeting memory-constrained hardware will need to budget explicitly. This is the kind of production architecture tradeoff that determines deployment viability more than benchmark accuracy alone.

The PI-controller framing is analytically clean but flags a practical tuning burden. The dual rate α interacts with the activity step ηₕ and the penalty coefficient ρ through the per-mode stability bound. Getting all three into a stable regime across varied architectures will require either principled initialisation rules or hyperparameter search — neither of which the current reference implementation addresses beyond the paper's specific configurations.

PC-ALM is a credible step toward biologically plausible, hardware-friendly training that does not sacrifice accuracy at scale. Whether the 2× memory overhead and multi-step inference budget prove acceptable will depend on whether future work can reduce T without reopening the PC-BP gap. The CPU-runnable, MIT-licensed JAX codebase lowers the barrier for the community to probe exactly that question.

Primary source

Sakana AI Researchers Introduce PC-ALM, a Layer-Local Alternative to Backpropagation That Trains 1000-Layer Networks — MarkTechPost