Towards Scaling Deep Neural Networks with Predictive Coding: Theory and Practice

TL;DR

This work investigates predictive coding (PC) as a brain-inspired alternative to backpropagation (BP) for training deep networks. It develops a theory that PC operates as a trust-region optimiser using curvature information, and shows that the equilibrated PC energy is a weight-dependent, rescaled mean-squared error, with saddles that are often easier to escape than in BP. The thesis then introduces μPC, a Depth-μP-inspired parameterisation enabling stable training of 100+ layer PCNs, and presents JPC, an open-source JAX library for flexible PC training and analysis. While PC currently faces hardware- and scalability challenges that limit practical advantages over BP on standard GPUs, the work provides deep insights into PC inference and learning dynamics, with implications for neuroscience and future energy-based AI systems, and highlights hardware co-design as a key direction for scaling PC methods. The combined theory, empirical demonstrations on linear and nonlinear networks, and toolchain (μPC, JPC) lay a foundation for further exploration of brain-inspired learning and its potential future hardware realizations.

Abstract

Backpropagation (BP) is the standard algorithm for training the deep neural networks that power modern artificial intelligence including large language models. However, BP is energy inefficient and unlikely to be implemented by the brain. This thesis studies an alternative, potentially more efficient brain-inspired algorithm called predictive coding (PC). Unlike BP, PC networks (PCNs) perform inference by iterative equilibration of neuron activities before learning or weight updates. Recent work has suggested that this iterative inference procedure provides a range of benefits over BP, such as faster training. However, these advantages have not been consistently observed, the inference and learning dynamics of PCNs are still poorly understood, and deep PCNs remain practically untrainable. Here, we make significant progress towards scaling PCNs by taking a theoretical approach grounded in optimisation theory. First, we show that the learning dynamics of PC can be understood as an approximate trust-region method using second-order information, despite explicitly using only first-order local updates. Second, going beyond this approximation, we show that PC can in principle make use of arbitrarily higher-order information, such that for feedforward networks the effective landscape on which PC learns is far more benign and robust to vanishing gradients than the (mean squared error) loss landscape. Third, motivated by a study of the inference dynamics of PCNs, we propose a new parameterisation called "$μ$PC", which for the first time allows stable training of 100+ layer networks with little tuning and competitive performance on simple tasks. Overall, this thesis significantly advances our fundamental understanding of the inference and learning dynamics of PCNs, while highlighting the need for future research to focus on hardware co-design if PC is to compete with BP at scale.

Figures (71)

• Figure 1: Landscape geometry and gradient descent dynamics of BP vs PC on a toy network. Training loss and energy landscapes of an example 1MLP trained with BP (left) and PC (right), plotted both as surfaces (top) and contours with superimposed gradient fields (bottom). Surfaces are plotted at the same scale for comparison, and vector fields are standardised for visualisation (see § \ref{['ch3:toy-exp']} for more details). The energy landscape of PC is plotted at the (approximate) inference equilibrium $\mathcal{F}|_{\nabla_z\mathcal{F} \approx 0}$ (see also Figure \ref{['ch3:fig:energy-land-infer-dynamics']} for a visualisation of the landscape inference dynamics). Note that this is essentially the same plot as the left column of Figure \ref{['ch4:fig:toy-examples']} in the next chapter.
• Figure 2: The PC weight update direction is significantly closer to optimal than BP on 1MLPs. For the first $5$ training batches, we plot the mean cosine similarity between the optimal weight direction $\Delta \mathbf{w}^*$ and the update $\Delta \mathbf{w}$ computed by (i) PC, $- \nabla_{\mathbf{w}} \mathcal{F}|_{\nabla_z \mathcal{F} \approx 0}$; (ii) BP with SGD, $- \nabla_{\mathbf{w}} \mathcal{L}$; and (iii) a trust-region Newton (TRN) method, $- (\mathbf{H} + \lambda \mathbf{I})^{-1}\nabla_{\mathbf{w}} \mathcal{L}$ with $\lambda = 2$. Shaded regions indicate the standard error of the mean (SEM) across 10 random weight initialisations.
• Figure 3: PC can train deeper chains significantly faster than BP. Mean training loss of 1D networks (deep chains) trained with BP and PC (see \ref{['ch3:chain-exps']} for details). Rows and columns indicate different activation functions (Linear, Tanh and ReLU) and number of hidden layers $H = \{1, 5, 10\}$, respectively. Each network type was optimised for learning rate, and training was terminated when the loss stopped decreasing. Shaded regions represent the SEM across 3 different initialisations.
• Figure 4: Faster convergence of PC in deep and wide networks trained on MNIST. Mean training loss of deep ($H = 10$) and wide ($N = 500$) networks trained to classify MNIST for 3 random initialisation (see \ref{['ch3:chain-exps']} for details). As for Figure Figure \ref{['ch3:fig:chain-test-losses']}, training was terminated whenever the loss stopped decreasing. SEMs are not visible.
• Figure 5: Empirical verification of the theoretical equilibrated energy of deep linear networks (Theorem \ref{['ch4:thm1']}). For different datasets, we plot the energy (Eq. \ref{['ch4:eq:pc-energy']}) at the numerical inference equilibrium $\mathcal{F}|_{\nabla_{\mathbf{z}}\mathcal{F} \approx 0}$ for DLNs with different number of hidden layers $H \in \{2, 5, 10\}$ (see §\ref{['ch4:exp-details']} for more details), observing an excellent match with the theoretical prediction (Eq. \ref{['ch4:eq:dln-equilib-energy']}).

Theorems & Definitions (5)

• Definition 1
• Definition A.1
• Definition A.2
• Theorem A.1
• Theorem A.2