Pith. sign in

REVIEW 3 major objections 5 minor 15 references

JPC: Flexible Inference for Predictive Coding Networks in JAX

T0 review · 3 major / 5 minor · reviewed 2026-08-11 · deepseek-v4-flash

Challenge this review Re-run · record.json Download PDF Read on arXiv ↗ Review JSON

Pith's one-line read JPC shows that Heun's method, a second-order Runge-Kutta integrator, runs predictive-coding inference significantly faster than standard Euler integration while maintaining accuracy across datasets and network depths.

desk verdict A useful, clean JAX library for predictive coding, but the headline speedup claim conflates solver order with adaptive step-size control and needs an ablation. read the letter →

arxiv 2412.03676 v1 pith:UVLFEVNU submitted 2024-12-04 cs.NE cs.AIcs.LG

Francesco Innocenti , Paul Kinghorn , Will Yun-Farmbrough , Miguel De Llanza Varona , Ryan Singh , Christopher L. Buckley This is my paper · ORCID
classification cs.NEcs.AIcs.LG
keywords predictivecodingJAXODEsolversHeunmethodRunge-Kuttainferencedynamicsgradientflowneuralnetworktraining
verification ladder T0 review T1 audit T2 compute T3 formal

The pith

A machine-rendered reading of the paper's core claim, the machinery that carries it, and where it could break.

The reading

This paper introduces JPC, a JAX library for training predictive coding networks (PCNs), and argues that the standard way of simulating PC inference—Euler integration of the gradient-flow dynamics of neural activities—can be profitably replaced by higher-order ODE solvers. The paper's central empirical claim is that Heun's method, a second-order explicit Runge-Kutta scheme with an adaptive step-size controller, reaches comparable test accuracy significantly faster in wall-clock time than Euler on image classification tasks, with larger gains on deeper networks. The library wraps this in a compact functional API supporting discriminative, generative, and hybrid PC models, and includes a closed-form energy formula for deep linear networks that can signal whether the inference phase has run long enough. If the runtime claim holds, JPC offers a practical speedup for anyone training PCNs and a ready-made baseline for further research on biologically plausible learning.

What carries the argument

The central object is the gradient-flow inference dynamics of a predictive coding network, $\dot{z}_\ell=-\partial F/\partial z_\ell$ for layer activities $z_\ell$ under the energy $F=\sum_\ell\|z_\ell-f_\ell(W_\ell z_{\ell-1})\|^2$. Euler integration is the usual discretisation of this flow; the paper replaces it with Diffrax's ODE solvers, defaulting to Heun's method (a second-order explicit Runge-Kutta scheme) combined with a PID step-size controller, so the solver adaptively chooses how many inference steps to take. This machinery is what produces the runtime comparison that supports the central claim.

What would settle it

On the same one-epoch classification benchmark, run Euler with the same PID step-size controller and tolerance used for Heun, and compare wall-clock times at matched final test accuracy; the central runtime claim would collapse if Euler with adaptive step control matches or beats Heun.

Watch

Extended reading notes

Core claim

On its own terms, the paper establishes a new default for PC inference: instead of performing gradient descent on the energy $F=\sum_\ell\|z_\ell-f_\ell(W_\ell z_{\ell-1})\|^2$ with a fixed Euler step, one can integrate the continuous-time gradient flow $\dot{z}_\ell=-\partial F/\partial z_\ell$ with Diffrax's ODE solvers. Benchmarking one-epoch training on image datasets with 3-, 5-, and 10-hidden-layer feedforward networks, the authors find that Heun's method—a second-order explicit Runge-Kutta solver with a proportional-integral-derivative step-size controller—converges to comparable mean test accuracy in substantially less wall-clock time than Euler, with the advantage growing with network depth. The paper presents this as a library contribution rather than a theoretical proof, and notes that other optimiser-specific hyperparameters could lead to different results.

Load-bearing premise

The load-bearing premise is that the runtime comparison is fair: Euler with a few fixed step sizes and Heun with an adaptive controller are matched by picking the run with the highest mean test accuracy on each hyperparameter grid, so a different matching rule—say equal integration error or equal step count—could change the speedup.

Editorial extensions

If this is right

  • Any PCN training loop that currently uses fixed-step Euler inference can switch to a second-order solver and expect faster wall-clock training at similar accuracy, with the largest gains appearing in 10-hidden-layer networks.
  • The same ODE-solver interface works across discriminative, generative, and hybrid PC models, so solver improvements carry over to generative and amortised settings without reimplementation.
  • The closed-form energy for deep linear PCNs acts as a convergence diagnostic: more inference steps correlate with better test accuracy on MNIST and Fashion-MNIST, offering an empirical stopping criterion.
  • Because the library is jitted and under 1000 lines of code, the runtime gains come with a reproducible and extensible codebase, not a black-box implementation.

Reading between the lines

Editorial extensions of the paper, not claims the author makes directly.

  • The runtime comparison pairs fixed-step Euler with adaptively stepped Heun and selects the highest-accuracy run on each grid; my inference, not the paper's, is that an adaptive Euler baseline or a tolerance-matched comparison could narrow the reported speedup.
  • If the speedup transfers to generative and hybrid PC models, which the paper does not benchmark for runtime, ODE-solver inference could make predictive coding competitive with backpropagation in wall-clock time on structured tasks.
  • The correlation between the closed-form energy and test accuracy points toward an automatic early-stopping rule for inference, an extension the paper leaves implicit.
  • The paper's closing suggestion of stochastic differential solvers raises the prospect of combining the library's wrapper with Langevin-type sampling for better generative performance at lower cost, but this is untested.
Share X Bluesky LinkedIn Reddit HN

Signed reviews

No signed human review yet.

Editorial analysis

A structured set of objections, weighed in public.

Desk editor's note, referee report, and a circularity audit.

Referee Report

3 major / 5 minor

Summary. The paper introduces JPC, a JAX library for training predictive coding networks (PCNs). JPC provides a simple high-level API and an advanced API built on Equinox, Diffrax, and Optax, and claims to be less than 1000 lines of code. The main empirical claim, developed in Section 3, is that using a second-order ODE solver (Heun) to integrate the PC inference gradient flow yields significantly faster wall-clock runtimes than standard Euler integration, with comparable test accuracy across several datasets and network depths. The paper also presents a theoretical tool from prior work (the closed-form energy of deep linear PCNs, Eq. 4) and a demonstration that the theoretical energy can indicate whether enough inference steps have been performed. The library code is publicly available.

Significance. If the runtime claim is substantiated, the contribution is practically useful: PC training is iterative, and a faster inference integrator could make PC more competitive. The library itself is a valuable community resource, being concise, functional, and built on well-maintained JAX libraries. The explicit use of Diffrax ODE solvers and the inclusion of theoretical diagnostics are also useful. However, the central speedup claim is currently supported by an experimental comparison that confounds solver order with adaptive step-size control and lacks essential statistical and reproducibility details. These issues need to be resolved before the paper's headline result can be accepted.

major comments (3)
  1. [Section 3, Figure 1] The comparison between Euler and Heun does not isolate the effect of solver order. The text states 'Unlike Euler, Heun employed a standard Proportional–Integral–Derivative step size controller.' This means the observed speedup could be entirely due to adaptive step-size selection rather than Heun's second-order accuracy. The authors should provide ablations, such as Euler with the same PID controller, Heun without a controller, or a comparison at matched integration error or matched step counts, before claiming that the second-order solver is faster.
  2. [Section 3, Figures 3–6] The hyperparameter selection rule, retaining the run with the highest mean test accuracy across dt and T grids, does not ensure that the two solvers operate at comparable inference accuracy or integration error. Heun could reach the same test accuracy with looser tolerances and fewer, less-accurate steps. The paper does not report step counts or integration error, and the runs are only one epoch long. The phrase 'not significantly different' is used without any statistical test. Without this information, the wall-clock comparison is difficult to interpret and may be biased.
  3. [Section 3, Figure 1 and experimental setup] Essential experimental details are missing: the hardware used, the ODE solver tolerance settings (e.g., rtol/atol for Heun's PID controller), the number of training iterations (only 'one epoch' is stated), and how the wall-clock time was measured (e.g., whether jit compilation time was fully excluded per iteration). The shaded regions in Figure 1 show variation across three initializations, but no confidence interval or statistical test is provided for the runtime difference. For a paper whose central claim is about runtime, these details are necessary for reproducibility and for assessing the magnitude of the reported speedup.
minor comments (5)
  1. [Figure 2] The caption contains a typo: 'MINST' should be 'MNIST'.
  2. [References] Reference [13] appears to have a typo: 'Tscshantz' should be 'Tschantz'.
  3. [Section 3] The caveat at the end of Section 3 acknowledges that optimiser-specific hyperparameters could change results, which is helpful, but it does not address the confound between solver order and adaptive step-size control that affects the main claim.
  4. [Figures 4–6] The accuracy curves in Figures 3–6 do not include error bars or confidence intervals, making it hard to verify the statement that accuracy differences are 'not significantly different'.
  5. [Section 4.2] In the code sketch, the variable name 'op t_ st at e' appears to be a typo for 'opt_state'.

Circularity Check

0 steps flagged · score 0.0 of 10

No significant circularity: the runtime speedup is a direct empirical comparison and the borrowed theoretical energy is an independent closed-form result, not a fitted input.

full rationale

The main claimed result — that Heun inference is faster than Euler on PCNs — is an experimental measurement, not a derivation from fitted parameters. The paper sweeps dt and T and selects the run with highest mean test accuracy; this is a benchmark selection rule, not a parameter fitted to the runtime quantity being reported, so it does not make the speedup true by construction. The only self-citation with substantive content is Eq. (4), attributed to the authors' prior work [3]. That expression is a parameter-free closed-form formula for the inference-equilibrium energy of deep linear networks under stated assumptions (identity activations, ∇zF = 0), and the paper reports numerical agreement (Figure 2) rather than defining the energy in terms of the formula. Thus the citation is real mathematical support and not a load-bearing self-citation. The paper's own caveat that hyperparameter choices could change results, and the confound that Heun used a PID step-size controller while Euler did not, weaken the attribution of speed to solver order, but that is an experimental fairness/correctness concern rather than circular reasoning. No equation or claim in the paper reduces to its own inputs by construction, so no circular step is established.

Assumptions & free parameters 4 free parameters · 4 assumptions · 0 invented entities

The paper's central speedup claim is empirical and rests on benchmark protocol choices (dt, T selection) and on standard PC assumptions (energy function, gradient-flow inference). The theoretical tool additionally imports a closed-form result from the authors' prior paper [3]. No new physical or mathematical entities are introduced.

free parameters (4)
  • Euler step size dt (selected per dataset and depth) = one of {0.5, 0.1, 0.05}, chosen by highest mean test accuracy (Figs 4-6)
    The runtime and accuracy of Euler inference depend on dt; the comparison retains the best-accuracy dt for each task, making the speedup a best-case result.
  • Euler integration horizon T (number of steps, selected per dataset and depth) = one of {5, 10, 20, 50, 100, 200, 500}, chosen by highest mean test accuracy
    The number of Euler steps sets the per-step cost; T is swept and the best-accuracy T is retained, with ties broken toward smaller T.
  • Heun initial step size dt (selected per dataset and depth) = one of {0.5, 0.1, 0.05}, chosen by highest mean test accuracy
    Heun's adaptive PID controller starts from this dt; the swept value that gives the best accuracy is retained.
  • Heun integration upper limit T (selected per dataset and depth) = one of {5, 10, 20, 50, 100, 200, 500}, chosen by highest mean test accuracy
    The integration horizon bounds the inference time for Heun; the best-accuracy T is used for the runtime comparison.
assumptions (4)
  • domain assumption The energy function F (Eq. 1), a sum of squared prediction errors, defines the PCN training objective.
    Section 2 defines PCNs via this energy following standard predictive coding literature; the paper does not derive it.
  • domain assumption Inference is the gradient flow z_dot = -dF/dz, and Euler gradient descent on the energy is its discretization.
    Section 2, after Eq. 3: 'standard GD on the energy... can be seen as the Euler discretisation of the gradient system.' The equivalence is standard but assumes the continuous-time limit is the correct object.
  • domain assumption The closed-form inference-equilibrium energy for deep linear networks (Eq. 4) from reference [3] is correct and applies to the trained linear networks in Figure 2.
    Section 4.2 relies on [3]; the paper does not re-derive Eq. 4 and assumes it holds for the networks used.
  • domain assumption The Diffrax ODE solvers and their default step-size controller produce accurate solutions of the gradient flow at a cost well represented by wall-clock time.
    Section 4 builds on Diffrax, Equinox, and Optax; the paper assumes these libraries are correctly implemented and that wall-clock time is a valid measure of solver cost under the chosen controller and tolerances.

how reviews work

0 comments
Cite this review

Pith. "Pith review of JPC: Flexible Inference for Predictive Coding Networks in JAX." pith.science (2026). https://pith.science/paper/UVLFEVNU

@misc{pith2026241203676,
  author       = {Pith},
  title        = {Pith review of: JPC: Flexible Inference for Predictive Coding Networks in JAX},
  year         = {2026},
  howpublished = {\url{https://pith.science/paper/UVLFEVNU}},
  note         = {Machine review of arXiv:2412.03676}
}
read the original abstract

We introduce JPC, a JAX library for training neural networks with Predictive Coding. JPC provides a simple, fast and flexible interface to train a variety of PC networks (PCNs) including discriminative, generative and hybrid models. Unlike existing libraries, JPC leverages ordinary differential equation solvers to integrate the gradient flow inference dynamics of PCNs. We find that a second-order solver achieves significantly faster runtimes compared to standard Euler integration, with comparable performance on a range of tasks and network depths. JPC also provides some theoretical tools that can be used to study PCNs. We hope that JPC will facilitate future research of PC. The code is available at https://github.com/thebuckleylab/jpc.

Figures

Figures reproduced from arXiv: 2412.03676 by the authors.

Figure 2
Figure 2. Theoretical PC energy for deep linear networks (Eq. 4) can help predict whether more inference could lead to better performance. We compare the theoretical energy with the numerical energy for different upper limits t of inference integration, as well as test accuracies, for a 10-hidden-layer, 300-width linear network trained to classify MINST with learning rate 1e −3 and batch size 64. Results were consistent acros… view at source ↗
Figure 4
Figure 4. Maximum mean test accuracy on MNIST achieved with Euler and Heun as a function of different step sizes dt and upper integration limits T. For the results in [PITH_FULL_IMAGE:figures/full_fig_p008_4.png] view at source ↗
Figure 5
Figure 5. Same results as [PITH_FULL_IMAGE:figures/full_fig_p009_5.png] view at source ↗
Figures from the paper (2 more)
Figure 6
Figure 6. Figure 6: Same results as [PITH_FULL_IMAGE:figures/full_fig_p009_6.png]
Figure 7
Figure 7. Figure 7: Same results as [PITH_FULL_IMAGE:figures/full_fig_p009_7.png]

Discussion (0). Continue with ORCID to comment.

Reference graph

Works this paper leans on

15 extracted references · 6 canonical work pages

  1. [3]

    Innocenti, E

    F. Innocenti, E. M. Achour, R. Singh, and C. L. Buckley. Only strict saddles in the energy landscape of predictive coding networks? arXiv preprint arXiv:2408.11979, 2024

    arXiv 2024

  2. [1]

    C. L. Buckley, C. S. Kim, S. McGregor, and A. K. Seth. The free energy principle for action and perception: A mathematical review. Journal of Mathematical Psychology, 81:55–79, 2017

    work page 2017

  3. [2]

    Babuschkin, K

    DeepMind, I. Babuschkin, K. Baumli, A. Bell, S. Bhupatiraju, J. Bruce, P. Buchlovsky, D. Bud- den, T. Cai, A. Clark, I. Danihelka, A. Dedieu, C. Fantacci, J. Godwin, C. Jones, R. Hemsley, T. Hennigan, M. Hessel, S. Hou, S. Kapturowski, T. Keck, I. Kemaev, M. King, M. Kunesch, L. Martens, H. Merzic, V . Mikulik, T. Norman, G. Papamakarios, J. Quan, R. Ring...

    work page 2020

  4. [4]

    P. Kidger. On neural differential equations. arXiv preprint arXiv:2202.02435, 2022. 6

    arXiv 2022

  5. [5]

    Kidger and C

    P. Kidger and C. Garcia. Equinox: neural networks in jax via callable pytrees and filtered transformations. arXiv preprint arXiv:2111.00254, 2021

    arXiv 2021

  6. [6]

    Legrand, L

    N. Legrand, L. Weber, P. T. Waade, A. H. M. Daugaard, M. Khodadadi, N. Mikuš, and C. Mathys. pyhgf: A neural network library for predictive coding. arXiv preprint arXiv:2410.09206, 2024

    arXiv 2024

  7. [7]

    Millidge, T

    B. Millidge, T. Salvatori, Y . Song, R. Bogacz, and T. Lukasiewicz. Predictive coding: towards a future of deep learning beyond backpropagation? arXiv preprint arXiv:2202.09467, 2022

    arXiv 2022

  8. [8]

    Millidge, A

    B. Millidge, A. Seth, and C. L. Buckley. Predictive coding: a theoretical and experimental review. arXiv preprint arXiv:2107.12979, 2021

    arXiv 2021

Show all 15 references
  1. [9]

    Oliviers, R

    G. Oliviers, R. Bogacz, and A. Meulemans. Learning probability distributions of sensory inputs with monte carlo predictive coding. PLOS Computational Biology, 20(10):e1012532, 2024

    2024

  2. [10]

    Pinchetti, C

    L. Pinchetti, C. Qi, O. Lokshyn, G. Olivers, C. Emde, M. Tang, A. M’Charrak, S. Frieder, B. Menzat, R. Bogacz, et al. Benchmarking predictive coding networks–made simple. arXiv preprint arXiv:2407.01163, 2024

    2024 arXiv

  3. [11]

    Rosenbaum

    R. Rosenbaum. On the relationship between predictive coding and backpropagation. Plos one, 17(3):e0266102, 2022

    2022

  4. [12]

    Salvatori, A

    T. Salvatori, A. Mali, C. L. Buckley, T. Lukasiewicz, R. P. Rao, K. Friston, and A. Ororbia. Brain- inspired computational intelligence via predictive coding. arXiv preprint arXiv:2308.07870, 2023

    2023

  5. [13]

    Tscshantz, B

    A. Tscshantz, B. Millidge, A. K. Seth, and C. L. Buckley. Hybrid predictive coding: Inferring, fast and slow. PLoS Computational Biology, 19(8):e1011280, 2023

    2023

  6. [14]

    van Zwol, R

    B. van Zwol, R. Jefferson, and E. L. Broek. Predictive coding networks and inference learning: Tutorial and survey. arXiv preprint arXiv:2407.04117, 2024

    2024

  7. [15]

    Zahid, Q

    U. Zahid, Q. Guo, and Z. Fountas. Sample as you infer: Predictive coding with langevin dynamics. arXiv preprint arXiv:2311.13664, 2023. 7 A Supplementary figures 100 400 700 10 15 20 25 Euler Heun Training iteration Test accuracy (%) Loading [MathJax]/extensions/MathMenu.js 10...

    2023 arXiv

Pith tools

Reviewed August 11, 2026 · model on record in the stance chip above.