A transformer that never runs backpropagation
Every layer predicts what the next one will do, measures how wrong it was, and corrects itself using only that local error.
The question
Backpropagation requires every layer to wait for an error signal computed at the output and passed back through the entire network. That creates a strict sequential dependency, needs the whole computation graph held in memory, and has no clear counterpart in biological neural systems. Real neurons do not appear to have access to a global error term.
Predictive coding proposes something local instead. Each layer holds a latent state, emits a prediction of the layer above, and receives back only the discrepancy between that prediction and what actually arrived. It updates itself from that. Everything a layer needs is available at the layer.
The theory is well established in computational neuroscience. Whether you can build a working transformer language model on it was much less clear. That was the question I took on.
How the architecture works
The model keeps the familiar transformer shape, with token and position embeddings, multi-head attention, an MLP block and an output projection to vocabulary logits. What changes is the learning rule inside every one of those components.
Each layer runs an iterative refinement loop over T inference steps:
- Initialise the layer's latent state
x. - Predict the activity of the next layer, giving
μ. - Compare against what the next layer actually shows:
ε = target − μ. - Correct both the state and the weights from that error alone.
x(t+1) = x(t) + η · ε(t)
ΔW = η · ε(t)ᵀ · x(t)
Both update rules read only quantities available at that layer. T, the number of
times this settling loop runs before weights are committed, turns out to be the most
consequential knob in the system.
Adapting attention
Attention was the hardest part to port. An early version updated the query, key and value
projections through a shared pathway, which pushed the heads toward redundancy and stopped
them specialising. I reworked it so each head's Q, K and
V update independently from their own error signals, restoring the head
diversity that makes multi-head attention worth having. I also added optional lateral
connections as a separate module, so the inhibitory dynamics predictive coding theory calls
for could be tested without entangling them with the core update rule.
Building the apparatus
A research claim is only as good as the infrastructure that tests it. Much of this work was making the experiments fast enough and repeatable enough to be worth believing.
- Multi-GPU training with DistributedDataParallel, including a dtype mismatch that surfaced only under mixed precision and barrier placement that was silently desynchronising ranks.
- FlashAttention behind a flag, with a documented CUDA and driver matrix.
- KV caching, so sampling from a model with an iterative inference loop stayed tractable.
- A BPE tokeniser replacing the borrowed GPT-2 one, giving a vocabulary matched to the corpus.
- Two-phase Bayesian search with Optuna over both the architecture and the predictive coding hyperparameters, optimising a combined energy and perplexity objective.
- Seeded runs, benchmark configurations committed to the repository, and structured logs.
- Three corpora behind one config switch: Tiny Shakespeare, Penn Treebank, and an OpenWebText subset.
Why a combined objective? Energy, the model's total prediction error, is what predictive coding actually minimises. Perplexity is what tells you whether it is any good at language. The experiments below are why you cannot tune on energy alone.
What the ablations showed
With the infrastructure in place I ran controlled sweeps varying exactly one thing at a time, holding every other hyperparameter at its tuned value, five epochs per configuration.
Iterative inference is where the gains are
T from 2 to 10. Perplexity falls steadily; energy climbs, then spikes.
Validation perplexity improves monotonically as T rises, from roughly
835 at T=2 down to about 100 at T=10. That is an
order-of-magnitude gain from nothing but letting each layer settle longer before committing an
update. Depth of inference substitutes for a good deal of what backpropagation buys you.
Energy tells a different story. It rises gently through T=8 and then jumps sharply
at T=10, to roughly ten times its previous value, at exactly the point where
perplexity is at its best. At the useful end of the range the two metrics move in opposite
directions.
Stacking more blocks does not help
Across 2 to 6 blocks, validation perplexity stays inside a narrow band of roughly 834 to 851, with no trend, just noise. Energy meanwhile increases with every block added. More depth costs more compute and more prediction error, and returns nothing.
The finding: for this architecture at this scale, iterative inference is the axis that matters and architectural depth is not. Spend compute on letting layers settle, not on stacking more of them.
Energy taken alone would have pointed the wrong way on both sweeps. It prefers shallow models and few iterations, which is the opposite of what the language modelling metric wants.
Where it stands
This is live research and I would rather be accurate than promotional about it. A predictive coding transformer at this scale does not match a well-tuned backpropagation baseline on perplexity. That was never the near-term bar. The question was whether the architecture trains stably, responds coherently to its hyperparameters, and where its leverage sits. To all three the answer is now yes, with evidence.
The work continues in FabricPC, a JAX library from SingularityNET where I contributed the transformer components. It trains the same architecture with either predictive coding or backpropagation, which makes the comparison a much more direct one.
Two things I took from it. The most interesting result was a negative one, depth not mattering, and it only became visible because the sweeps were controlled and the configs pinned. And choosing the objective is a research decision rather than a detail: optimising the theoretically principled quantity alone would have quietly selected the worst models.