Augmented Lagrangian Predictive Coding: training 1000-layer networks without backpropagation
Pangram verdict · v3.3
We believe that this entire text is human-written.
AI likelihood · overall
HumanArticle text · 1,428 words · 1 segments analyzed
We introduce PC-ALM, a local alternative to backpropagation. PC-ALM trains residual MLPs up to 1000 layers, nearly matching backprop's performance despite using only layer-local dynamics. PC-ALM equips each layer with a feedback control dynamical system that distributes and propagates supervision credit throughout a network. Standard deep learning relies on backpropagation. The brain, however, cannot implement backpropagation, at least not exactly[1, 2]. How the brain solves the multilayer credit assignment problem without explicit use of backprop remains one of the fundamental unsolved problems in neuroscience (though not without progress[3, 4, 5]). There are several reasons the brain can't implement exact backpropagation. One is “phase locking”[6, 2]. Backpropagation runs in three phases, in strict order: 1) a forward pass, then 2) a backward pass, then 3) a weight update. A weight update is locked until the forward and backward passes have completed—a neuron in an early layer must hold its activation and wait for the error signal to arrive. The brain has no known mechanism that could enforce such strict timing coordination across an entire network[1]. In this post, we introduce PC-ALM (Augmented Lagrangian Predictive Coding), a method for training networks that replaces the forward and backward passes of backprop with layer-local dynamical systems. Each layer is coupled only to its neighbors. Instead of forward-then-backward, we run each layer forward in time. When run to convergence, the dynamics of the whole system distribute supervision credit signals quickly and accurately across the entire network. PC-ALM is an extension of standard predictive coding (PC)[7, 8, 9, 10]. PC uses diffusive (i.e. energy-based or "heat flow") coupling between layers. Compared to PC, PC-ALM introduces dual neurons (Lagrange multipliers) per layer, making each layer's local recurrence a PI feedback controller. In the limiting case of linear networks, the dual neurons converge to the exact backprop credit signals, despite using only local computation. We compare PC-ALM to PC and backprop in a suite of experiments. Local training methods such as PC have historically been difficult to scale. Following the PC literature, we use simple tasks (Fashion-MNIST, CIFAR-10, etc.) and networks such as residual MLPs. We show that PC-ALM can successfully propagate supervision credit in 1000-layer neural networks, overcoming standard PC's signal decay problem[11] while remaining layer-local. We focus on deep, small-width networks, a regime in which PC tends to perform poorly. Ultimately, our motive is to understand how distributed systems (such as the brain) can implement gradient computations without backpropagation. Scientific motivations aside, this research may inform energy-efficient deep learning on neuromorphic hardware, where dynamical systems simulation is cheaper than on GPU[12]. Paper: arxiv.org/abs/2605.31022 Code: github.com/SakanaAI/pc-alm Predictive coding: each layer as a dynamical system Before explaining PC-ALM, let us first explain PC, interpreting it from a dynamical systems standpoint to emphasize its role as a backprop alternative. Predictive coding Predictive coding has its roots in Helmholtz's theories of unconscious perception[13]. Rao & Ballard (1999) developed a mathematical framework for PC as a model of visual cortex[14]. The idea of PC is that each layer attempts to model its incoming signals, sending upward only the prediction error (the part that the layer failed to model) to the next layer. Mathematically, predictive coding utilizes a general motif: take a state and update it to reduce a prediction error at the next step, statet+1=statet−η(statet−targett)⏟prediction error By applying this update rule to each layer's activation vector (the “state” is the layer's activation hi; its “target” is the prediction σ(Wihi−1) arriving from the layer below)1, the PC framework effectively sidesteps backprop's need for a synchronized forward and backward pass. To explain this in more detail, let's write a feedforward network as a constrained optimization problem: minimizeθ,h12‖y−WLhL−1‖2subject tohi=σ(Wihi−1),i=1,…,L−1. where L is the network depth, h0:=x the input, y the target, θ={Wi} the weights, hi the layer activations, and σ an activation function such as ReLU. Note that each hi is an optimization variable2. We then construct a new loss function that includes the original supervision loss, together with a quadratic penalty for violations of each layer constraint: FPC(h,θ)=12‖y−WLhL−1‖2+12∑i=1L−1‖hi−σ(Wihi−1)‖2. This is a quadratic relaxation of the constrained problem. FPC is known as the “free energy” of the network[15, 9]. To train a neural network, PC alternates between inference and learning steps: Predictive coding inference for t=1,…,T hi←hi−ηh∇hiFPCfor i=1,…,L−1 learning Wi←Wi−ηθ∇WiFPCfor i=1,…,L Per mini-batch, a forward pass initializes the activations, followed by T inference steps and a single weight update. We set T proportional to network depth; the 1000-layer experiments below use T=2L. Each hi-update reduces the prediction errors between layers adjacent to i. This is because ∇hiFPC depends only on hi−1, hi, and hi+1. Inference requires only nearest-neighbor communication ("message passing") between layers. Explicitly, writing ri=hi−σ(Wihi−1) for the prediction error between layers i−1 and i, the inference update reads3: hi←hi−ηh(ri↑error below−Wi+1⊤(σ′⊙ri+1↑error above))i=1,…,L−1 The bottom h0 is “clamped” (fixed) to an input value and the top of the network is clamped to the target y. Running T update steps, the network settles into states hi for each layer, after which a gradient descent step is taken on the same FPC but now with respect to the weights W (given the current remaining prediction errors and the current state activations). Wi←Wi+ηθ(σ′⊙ri)hi−1⊤i=1,…,L−1 Both the inference step and the weight update are layer-local. The weight update is Hebbian-like, in that it multiplies a postsynaptic error by the presynaptic activity (a delta rule), and the dynamics map onto a neural circuit with explicit error neurons[14, 7]. PC trains deep networks, but exhibits signal decay PC inference in a 32-layer residual MLP (width 16, ReLU) at weight initialization. Credit = per-layer norm of the prediction error ri. Dashed reference: norm of the backprop adjoint (the loss gradient with respect to that layer's activations). Since minimizing the free energy with respect to each hi does not enforce the layer-wise constraints to hold exactly, PC results in a different learning trajectory compared to standard backpropagation. Nevertheless, PC has been shown to successfully train networks on simple tasks. For example, MNIST and Fashion-MNIST in 128-layer residual MLPs with wide layer widths (512 neurons per layer)[16]. However, PC struggles on more complex tasks and networks[17]. Further, PC struggles even on simple networks/tasks if the network width is smaller than its depth[18]. Each layer adjusts its activity to reduce prediction errors with its neighbors. Supervision enters at the output, but must work its way through this chain of local compromises to influence earlier layers. In deep, narrow networks, the resulting credit signal becomes weak long before it reaches the input. This leads to a documented signal-decay problem of PC[11], as illustrated in the above figure. Increasing T lets credit propagate farther, but requires more computation for each training update. Our method, PC-ALM, introduces a way to improve signal propagation of PC networks, retaining the layer-local dynamics of PC and keeping inference budget T proportional to network depth. Augmented Lagrangian Predictive Coding We propose Augmented Lagrangian Predictive Coding, a variant of PC that uses the augmented Lagrangian (AL)[19, 20, 21] in place of PC's FPC: L(h,θ,λ)=12‖y−WLhL−1‖2+∑i=1L−1λi⊤(hi−σ(Wihi−1))+12∑i=1L−1‖hi−σ(Wihi−1)‖2 supervised loss + Lagrangian term + PC energy In each layer, the augmented Lagrangian introduces a Lagrange multiplier (or dual variable) λi∈Rni of the same dimension as hi. The augmented Lagrangian is used extensively in distributed optimization[22], and has motivated many distributed methods for training deep networks[23, 24]. LeCun (1988)[25] showed that the Lagrange multipliers of a constrained network encode its backprop credit signals at equilibrium. The augmented Lagrangian combines this classical construction with the quadratic prediction-error penalties already used by PC. This suggests a simple possibility: can PC’s local dynamics recover those credit signals if we add the multipliers? To use the augmented Lagrangian for training, we make a simple modification to PC: Augmented Lagrangian predictive coding inference for t=1,…,T hi←hi−ηh∇hiLfor i=1,…,L−1primal λi←λi+α(hi−σ(Wihi−1))for i=1,…,L−1dual learning Wi←Wi−ηθ∇WiLfor i=1,…,L Here α is the dual step size. PC-ALM is primal descent, dual ascent on the augmented Lagrangian, versus PC's descent on the energy. By accumulating local prediction errors, the dual variables recover the exact backprop credit signals in a deep linear network. We derive this result in the paper. Thus, at least for linear networks, PC-ALM gives a method for computing exact supervised loss gradients and distributing them throughout a network, using only layer-local dynamics. In the experiments below, we test whether this advantage carries over to nonlinear networks. Mechanistic interpretation To understand, mechanistically, how PC-ALM works, consider a simple scalar network, with hidden unit h=w1x and output y^=w2h. We want to propagate the gradient of the supervised loss 12(y−y^)2 to w1. We attach a multiplier λ to the constraint h=w1x, initialize λ=0, and initialize h at its forward-pass value.