PCX Library: Efficient Predictive Coding

• PCX Library is an open-source, JAX-based framework that streamlines predictive coding research by addressing scalability, performance, and reproducibility challenges.
• It leverages JAX’s just-in-time compilation and transformation primitives to achieve up to 20× speed improvements over legacy implementations.
• The framework supports diverse predictive coding variants and provides standardized benchmarks for direct comparison with backpropagation and other learning rules.

PCX is an open-source, JAX-based software library designed to address the scalability, performance, and reproducibility bottlenecks in research on predictive coding networks (PCNs). It provides high-performance implementations, an extensible API, and a comprehensive suite of standard machine learning benchmarks. PCX enables practitioners to efficiently train and evaluate large predictive coding architectures, supporting both established and novel variants of the predictive coding (PC) family, with direct comparison to backpropagation and other biologically-plausible learning rules (Pinchetti et al., 2024).

1. Design Motivations and Core Principles

PCX was developed to overcome three critical deficits in existing PC research: inefficient legacy implementations, fragmented task benchmarks, and a lack of systematic comparison with standard baselines. Central principles guiding PCX are:

• Performance: Leveraging JAX’s just-in-time (JIT) compilation and array programing, PCX achieves wall-clock speeds up to 20× those of prior Python or Torch-based PC libraries. This enables experiments on mid-to-large scale networks and datasets previously unaffordable for PCN research.
• Simplicity: The PCX API mirrors PyTorch/EQUINOX, using PyTree model objects and automatic state tracking. Layers and per-layer “vode” (latent variable) states are managed internally, allowing users to define models at a high level.
• Extensibility: All components interoperate with JAX’s transformation primitives (e.g., vmap, pmap) and with EQUINOX, permitting seamless inclusion of new layers, optimizers, and inference strategies.

2. Predictive Coding Architecture and Abstractions

PCX formalizes PC training as variational free-energy minimization for a hierarchical Gaussian generative model. For a network of depth $L$ with latent variables $h_0, h_1, \dots, h_L$ and parameterized conditional means $\mu_\ell = f_\ell(h_{\ell-1};\theta_\ell)$, the energy functional is

$$\mathcal{F}(h, \theta) = -\ln P_\theta(h) = \sum_{\ell=0}^L \frac{1}{2} \|\epsilon_\ell\|^2, \qquad \epsilon_\ell = h_\ell - \mu_\ell$$

During learning, PCX alternates between:

• Inference (state updates): For $T$ steps,

$$h_\ell^{(t+1)} = h_\ell^{(t)} - \gamma\, \frac{\partial \mathcal{F}}{\partial h_\ell}(h^{(t)}, \theta)$$

• Parameter updates: Once per batch,

$$\theta_\ell = \theta_\ell - \eta\, \frac{\partial \mathcal{F}}{\partial \theta_\ell}(h^*, \theta), \quad h^* = h^{(T)}$$

Both updates are local, depending on adjacent layers. The user interacts only with high-level model objects; all underlying free-energy gradients and update logic are encapsulated within PCX abstractions. Energy evaluation and differentiation are handled automatically via JAX autodiff.

3. Benchmark Suite and Supported Algorithms

PCX ships with standardized benchmarks for discriminative and generative learning:

• Discriminative: Image classification on MNIST, FashionMNIST, CIFAR-10, CIFAR-100, and Tiny ImageNet. Provided architectures include a 3-layer MLP (128 units), VGG-5, VGG-7, AlexNet, and VGG-19 (optionally with skip connections).
• Generative: Autoencoding (e.g., convolutional decoders for MNIST, CelebA), Monte Carlo PC (Langevin sampling), and associative memory (denoising/retrieval on Tiny ImageNet).
• Algorithmic coverage:
  • Standard PC (with squared error and cross-entropy loss)
  • Incremental PC (iPC)
  • Monte Carlo PC (Langevin dynamics for sampling)
  • EqProp-style nudging (positive/negative/centered, as in Equilibrium Propagation)
  • Backpropagation-based baselines (BP-CE, BP-SE)

The benchmark package enables direct, reproducible comparisons across PC variants and with gradient-based approaches.

4. Empirical Results and Scalability

PCX demonstrates state-of-the-art performance within the PC paradigm on all included datasets. Table 1 exemplifies PC, iPC, and nudging performance versus backpropagation. For example, on MNIST, iPC achieves $98.45\pm0.09\%$ test accuracy, exceeding standard BP for the same setting. On CIFAR-100, the centered nudging (CN) variant reaches $67.19\pm0.24\%$ (top-1), closely tracking BP baselines.

For generative tasks, PC and iPC outperform or match backpropagation in mean squared error on MNIST, CIFAR-10, and CelebA autoencoding. On MCPC (Langevin) for MNIST, the Fréchet Inception Distance (FID) is $2.53\pm0.17$, surpassing a standard VAE baseline.

Scalability is validated by wall-clock profiling. On CIFAR-100 (VGG-5), PCX requires 5.5 s/epoch (A100 GPU) compared to 3.2 s for BP and 110 s for a legacy PC implementation. For AlexNet (161M parameters), PCX trains in 29 s/epoch (vs. 10 s for BP). Training time scales linearly with both layer count $L$ and inference steps $T$. The framework achieves efficient utilization of JAX’s parallelization where possible, though across-layer vmap is currently limited.

Dataset	PC-CE	iPC	BP-CE
MNIST	98.11	98.45	98.07
CIFAR-100 (Top-1)	60.00	56.07	60.82
Tiny ImageNet	41.29	29.94	43.72

Excerpts from Table 1; see (Pinchetti et al., 2024) for full details.

5. Usage Model and Extensibility

PCX installation is via PyPI:

pip install pcx

A minimal classifier is built as follows:

import pcx.nn as pxnn import pcx.predictive_coding as pxc class MLPClassifier(pxc.EnergyModule): def __init__(self, in_dim, h_dim, out_dim): self.layers = [ pxnn.Linear(in_dim, h_dim), pxnn.Linear(h_dim, out_dim), ] self.vodes = [pxc.Vode((h_dim,)), pxc.Vode((out_dim,))] def __call__(self, x, y=None): h = x for layer, vode in zip(self.layers, self.vodes): u = jnn.relu(layer(h)) h = vode(u) if y is not None: self.vodes[-1].set("h", y) return u

Benchmark runs (e.g., MNIST MLP) invoke:

python examples/s4_1_discriminative_mode/run_mnist.py \ --model MLPClassifier --h_dim 128 --T 7 \ --optim_h sgd --lr_h 5e-2 \ --optim_w adamw --lr_w 1e-4 --weight_decay 1e-4

PCX utilizes JAX transformations for batching (vmap) and distributed execution (pmap) where feasible, and supports modular extensions via its high-level abstractions.

6. Limitations and Open Challenges

Despite substantial advances, several issues remain:

• Inference speed: The requirement for $T$ iterative inference steps per batch induces a compute overhead relative to single-pass backpropagation; eliminating or mitigating this gap is open.
• Optimization stability: AdamW optimizer is unstable for wide layers and large state rates $\gamma$, while SGD is robust but slower.
• Deep network scaling: PC suffers on very deep networks without skip connections; for VGG-19 on CIFAR-10, performance drops from $\sim74\%$ to $\sim25\%$ without skips.
• Parallelism: Layerwise parallelism (vmap over layers) is limited by the sequential nature of inference; future JAX developments may enable improved parallel execution.

Suggested future research directions include new inference schemes that reduce $T$, advanced initialization and regularization methods, optimizer designs tailored to PC’s local energy landscape, and architectural modifications (skip-connections, attention) facilitating stable deep learning.

7. Significance and Impact

PCX is the first “batteries-included” high-performance framework for predictive coding research. By consolidating model abstractions, energy-based learning dynamics, and a comprehensive, reproducible benchmark suite, PCX sets rigorous new baselines for biological credit-assignment algorithms. Its empirical analysis exposes enduring bottlenecks in PC—most notably, inference cost, top-layer error concentration, and optimization sensitivity—establishing a foundation for methodologically robust progress in scalable predictive coding (Pinchetti et al., 2024).