Building a Predictive Coding Library
A JAX-based library that trains the same computational graph with predictive coding or backpropagation, enabling direct comparison of where and how the two learning methods diverge.
Picking up the thread
My earlier work showed that a transformer could learn without backpropagation, using local learning rules. But the comparison had an important limitation: the predictive coding model and the standard baseline were separate implementations. Any difference could therefore come from an implementation detail rather than the learning rule itself.
FabricPC addresses this by providing a JAX framework for predictive coding on general computation graphs. Most existing deep learning libraries are built around backpropagation, so supporting predictive coding required implementing not only the training and inference mechanics, but also the supporting infrastructure for distributed computation. A model is defined once as a graph of nodes, and the same graph can then be trained using either predictive coding or backpropagation. I contributed the transformer components and the infrastructure needed to run and benchmark them at scale.
With the architecture, initialization, data, training schedule, and random seeds held fixed, the learning rule is the only thing that changes. This makes it possible to directly study where predictive coding and backpropagation actually diverge.
The transformer as a graph
The model keeps the familiar transformer shape, but each block is broken into three predictive coding nodes rather than treated as one opaque computation: an attention-residual node, a feed-forward expansion node, and a projection-residual node. Every node holds a latent state, predicts what the node above it should be, and carries a local energy measuring how wrong that prediction was. The vocabulary projection at the top uses cross-entropy against the next character; everything below it uses a simple Gaussian error.
Training a batch runs a short settling loop before any weight moves:
- Clamp the input and the next-character target; initialise every latent by one ordinary forward pass.
- Sweep the latents a few times to reduce the total energy, letting the clamped target pull the states into agreement.
- Once settled, update each node's weights from its own local error alone.
E_i = ½ · ‖ z_i − μ_i ‖² (local energy at each node)
ΔW_i ∝ ε_i · z_parentᵀ (update from local error only)
Two details make this work on a real transformer. Residual connections mix an identity path
with a learned one, so the graph scales only the learned inputs, using a depth-aware factor
borrowed from µPC. And dropout masks are frozen for the whole settling loop, resampling them
mid-inference would keep changing the target the latents are chasing. The number of settling
sweeps, T, is again the knob that matters most.
A matched comparison
The test bed is deliberately small and clean: character-level Tiny Shakespeare, a two-block model with 413,249 parameters, five seeds. Predictive coding uses four settling steps. The backpropagation control is the identical graph with the settling loop switched off and a single backward pass in its place.
Across the five seeds the predictive coding transformer reached a test perplexity of 7.12 ± 0.05, against 6.00 ± 0.05 for backpropagation, a small, consistent gap of about 1.1, at roughly 1.46× the training time because of the settling sweeps. On the same graph, local learning lands within touching distance of the global backward pass. That is a very different picture from the order of magnitude gaps you see when the two models are only loosely matched.
Where the two methods diverge
The aggregate gap is only part of the story. Because both conditions run on the same graph, I could line up their parameter updates module by module and measure how closely each predictive coding update points in the same direction as the backpropagation update it replaces.
With four settling steps the agreement is almost total at the output and absent at the input. The vocabulary projection lines up with backpropagation almost perfectly (≈ 0.99), while the embedding shows no agreement at all (≈ −0.08). The learning signal arrives at the top of the network and does not reach the bottom. Adding more settling steps pushes agreement down into the earlier layers, at thirty steps the embedding climbs to about 0.89, which confirms this is finite-step error propagation working its way inward.
The catch: more agreement with backpropagation did not mean better results. Four steps gave the lowest test perplexity; fourteen and thirty steps matched backpropagation more closely everywhere and generalised worse. Truncated settling seems to act as a mild regulariser.
So closeness to backpropagation is not the target to aim for. A predictive coding update can differ from the backward pass and still be the better update. That reframes the question from "how do we make it match?" to "where does the local signal actually reach, and why?"
Inference matters, depth does not
The same sweeps confirmed the pattern from the earlier work. Settling depth is the axis that moves performance: one step is far too few, four is the sweet spot, many more overfit. Stacking more transformer blocks, by contrast, did nothing measurable at this scale, every depth I tried performed the same within seed to seed noise. Spend compute on letting layers settle, not on making the model taller.
Where it stands
This is ongoing research, so I would rather be precise than promotional. At this scale the predictive coding transformer trains stably and reproducibly, responds cleanly to its hyperparameters, and trails a well-tuned backpropagation baseline by a small, steady margin. Matching that baseline was never the near-term goal; understanding the learning rule was.
The current direction follows the alignment finding one step further: rather than measuring agreement only at the end of training, I track it across training. The early picture is that the backpropagation-aligned "credit front" reaches deep into the network at initialisation and then retreats toward the output as training proceeds, leaving the earliest layers increasingly under-driven. If that holds up, the interesting story is not that predictive coding fails to match backpropagation, but where and when its learning signal thins out, which is a more useful thing to know.