$μ$PC: Scaling Predictive Coding to 100+ Layer Networks

TL;DR

This work demonstrates that predictive coding networks can be scaled to 100+ layers by applying a Depth-muP-inspired reparameterisation, termed μPC. It identifies two key obstacles to large-scale PCN training—ill-conditioned inference landscapes and depth-sensitive forward-pass instability—and proposes a minimal set of desiderata to guide scalable parameterisations. Empirically, μPC enables stable training of very deep residual networks on simple classification tasks with minimal tuning and supports zero-shot transfer of learning rates across widths and depths, while revealing that μPC can approximate BP in a BP-like regime for wide, shallow configurations. The study provides both theoretical insights and practical benchmarks, and releases code to facilitate further exploration of scalable, local learning algorithms.

Abstract

The biological implausibility of backpropagation (BP) has motivated many alternative, brain-inspired algorithms that attempt to rely only on local information, such as predictive coding (PC) and equilibrium propagation. However, these algorithms have notoriously struggled to train very deep networks, preventing them from competing with BP in large-scale settings. Indeed, scaling PC networks (PCNs) has recently been posed as a challenge for the community (Pinchetti et al., 2024). Here, we show that 100+ layer PCNs can be trained reliably using a Depth-$μ$P parameterisation (Yang et al., 2023; Bordelon et al., 2023) which we call "$μ$PC". By analysing the scaling behaviour of PCNs, we reveal several pathologies that make standard PCNs difficult to train at large depths. We then show that, despite addressing only some of these instabilities, $μ$PC allows stable training of very deep (up to 128-layer) residual networks on simple classification tasks with competitive performance and little tuning compared to current benchmarks. Moreover, $μ$PC enables zero-shot transfer of both weight and activity learning rates across widths and depths. Our results serve as a first step towards scaling PC to more complex architectures and have implications for other local algorithms. Code for $μ$PC is made available as part of a JAX library for PCNs.

Figures (42)

• Figure 1: $\mu$PC enables stable training of 100+ layer ResNets with zero-shot learning rate transfer. (Right) Test accuracy of ReLU ResNets with depths $H = \{8, 16, 32, 64, 128 \}$ trained to classify MNIST for one epoch with standard PC, $\mu$PC and BP with Depth-$\mu$P (see §\ref{['exp-details']} for details). Solid lines and shaded regions indicate the mean and $\pm 1$ standard deviation across 3 different random seeds. These results hold across other activation functions (see Fig. \ref{['fig:mupc-vs-pc-mnist-accs-all-act-fns']}). See also Figs. \ref{['fig:mupc-mnist-5-epochs']}-\ref{['fig:cifar-20-epochs']} for asymptotic results with 128-layer ReLU networks trained for multiple epochs on MNIST, Fashion-MNIST and CIFAR10. (Left) Example of zero-shot transfer of the weight and activity learning rates from 16- to 128-layer Tanh networks. See Figs. \ref{['fig:mupc-hyperparam-transfer-tanh']} & \ref{['fig:mupc-hyperparam-transfer-linear']}-\ref{['fig:mupc-hyperparam-transfer-relu']} for an explanation and the complete transfer results across widths as well as depths.
• Figure 2: Wider and particulary deeper PCNs have a more ill-conditioned inference landscape. We plot the condition number of the activity Hessian $\kappa(\mathbf{H}_{\mathbf{z}})$ (lower is better) of randomly initialised fully connected networks as a function of the width $N$ and depth $H$ (see §\ref{['exp-details']} for details). Insets show 2D projections of the landscape of selected networks around the linear solution (Eq. \ref{['eq:pc-infer-solution']}) along the maximum and minimum eigenvectors of the Hessian $\mathcal{F}(\mathbf{z}^* + \alpha \hat{\mathbf{v}}_{\text{min}} + \beta \hat{\mathbf{v}}_{\text{max}})$. Note that the ill-conditioning is much more extreme for ResNets (see Fig. \ref{['fig:resnets-cond-nums-init']}). Results were similar across different seeds.
• Figure 3: The inference landscape of PCNs grows increasingly ill-conditioned with training. We plot the condition number of the activity Hessian (Eq. \ref{['eq:activity-hessian']}) (top) as well as test accuracies (bottom) for fully connected networks of depths $H \in \{8, 16, 32\}$ during one epoch of training. All networks had width $N=128$ and were trained to classify MNIST (see §\ref{['exp-details']} for more details). Similar results are observed for ResNets (Fig. \ref{['fig:sp-train-cond-nums-skips-mnist']}) and Fashion-MNIST (Fig. \ref{['fig:sp-train-cond-nums-GD-fashion']}). Solid lines and shaded regions indicate the mean and standard deviation over 3 random seeds.
• Figure 4: Parameterisations with stable forward passes induce highly ill-conditioned inference landscapes with depth. We plot the conditioning of the activity Hessian of randomly initialised networks over width $N$ and depth $H$ for the $\mu$PC and orthogonal parameterisations. Networks with and without residual connections were used for these respective parameterisations. Note that ReLU networks with orthogonal initialisation cannot achieve stable forward passes (see Fig. \ref{['fig:fwd-pass-stability-depth-params']}). Results were similar across different seeds.
• Figure 5: $\mu$PC enables zero-shot transfer of the weight and activity learning rates across widths $N$ and depths $H$. Minimum training loss (log) achieved by ResNets of varying width and depth trained with $\mu$PC on MNIST across different weight and activity learning rates. All networks had Tanh as nonlinearity (see Figs. \ref{['fig:mupc-hyperparam-transfer-linear']}-\ref{['fig:mupc-hyperparam-transfer-relu']} for other activation functions); those with varying width (first row) had 8 hidden layers, and those with varying the depth (second row) had 512 hidden units (see §\ref{['exp-details']} for details). Each contour was averaged over 3 random seeds.