Introducing PC-ALM
Researchers at Sakana AI have developed PC-ALM, an alternative to backpropagation designed for training deep neural networks. The core concept involves layer-local updates, where each layer independently adjusts its parameters based on local predictions, mirroring the brain’s learning process. This approach addresses the limitations of backpropagation, which relies on a global, sequential update scheme.
Performance and Benchmarks
The PC-ALM method was evaluated on residual MLPs with varying widths and depths, ranging from 8 to 128, across the Fashion-MNIST and MNIST datasets. Notably, the team successfully trained 1000-layer residual MLPs on MNIST, maintaining an accuracy within approximately 2 percentage points of backpropagation. Furthermore, the research extended to ResNet-18 on CIFAR-10 and Tiny ImageNet, demonstrating the method’s scalability.
Technical Details
PC-ALM utilizes an augmented Lagrangian to minimize the supervised loss, incorporating a Lagrange multiplier for each layer. This multiplier, denoted as λi, converges to the backprop adjoints in linear networks. The training process alternates between two local steps: a primal gradient step on activations and a dual step to accumulate prediction errors. The research team utilized a PI controller per layer, adjusting the frequency of the oscillations through the α parameter.
Implementation and Availability
The PC-ALM implementation is available as MIT-licensed JAX code, allowing for CPU-based reproduction of the research findings. The reference cell (width 32, depth 32, ReLU, Fashion-MNIST) achieved 78.66% test accuracy for backprop, 68.13% for PC, and 77.75% for PC-ALM. The code is available on GitHub, facilitating experimentation and further development.
