2.1 Computational framework

The design and software implementation of dynamic neural networks for predictive coding have been shaped by a set of requirements. These networks are made of nodes that can store any number of variables. Some variables might be found in other nodes as well, and some might be unique. Nodes are connected with each other through directed edges, and there can exist multiple types of connections in the same graph, denoting different interactions between the nodes. Computation steps in the graph typically occur locally between adjacent nodes for prediction and prediction errors. Multiple types of computation can be defined. The computation steps can be triggered either reactively, by observing events in the surroundings of a node and reacting to them, or they can be scheduled by pre-allocating a sequence of steps that will propagate the information through the graph. Finally, all these components should be transparent to the network itself when performing a given computation, allowing it to meet the demands of self-monitoring and self-organisation principles.

To observe the set of constraints above, computations should follow a strict functional programming framework, meaning that they should be in-place programmatically pure functions operating on the components of the network. Functional programming is natively supported in Rust and also enforced by JAX [35] to leverage just-in-time (JIT) compilation and automatic differentiation, therefore departing from object-oriented programming (the definition of classes populated with attributes and methods) that is a central feature of Python. This comes with limitations in the way toolboxes’ API can be developed (see, for example, how this can be handled in [41]). To fully meet the dynamic aspects mentioned previously, the update functions should ideally receive and return all of the following components defining a network:

1. A list of attributes - the attributes are dictionaries of parameters of a given node
2. A set of lists of edges - the edges are the directed connections between the nodes. All the possible edge types are grouped into a set.
3. A set of update functions. Each function defines a computation and can be parametrised by the index of the triggering node. Possible computations are, for example, prediction, posterior update, or prediction error between two nodes.
4. Unless using a reactive computation scheme, the function should also have access to an ordered sequence of update functions that apply to individual nodes.

By defining these four components, and by creating functions that can receive and return all of them, the user can generate arbitrarily sized and structured dynamic neural networks for predictive coding (see Fig 1). The first two items define what is usually called a graph, with the addition that it can be directed and multilayered. The last two items shape what is central to predictive coding: the schedule or reactive nature of the propagation of information through the network. Because all of these components are transparent during message-passing computations, learning algorithms can be developed to act on them as a way of inference. Acting on the attributes corresponds to standard inference or other learning algorithms like reinforcement learning. Acting on the size of the networks is comparable to structure learning. Acting on the edges can relate to causal inference, and acting on the update functions or their sequence can implement principles from meta-learning (see 5 for more details on possible applications). The hard constraint on transparency of the network component during message passing makes this framework difficult to implement in other graph/neural network libraries (e.g., see [17,18]). Most of these constraints can be met by using a pure JAX implementation [35] while remaining compatible with transformations like JIT and automatic differentiations. However, some advanced use cases of dynamic reshaping and edge manipulation might result in degraded performance or incompatibility with certain transformations. When using Rust [36] as the backend, all constraints can be met with no cost in terms of performance.

Download:

• PNG
  larger image
• TIFF
  original image

Fig 1. The four components of a dynamic network for predictive coding.

pyhgf represents any dynamic network using the combination of four variables: attributes, edges, update functions, and update sequences. This modularity allows dissociating update steps and connectivity structures and makes these variables part of the inference process. The creation of a network is read from left to right: 1. Attributes. Nodes in the network contain parameters (e.g., sufficient statistics about probability distributions and coupling weights). 2. Edges. Nodes can have multiple connection types with each other (e.g., value and volatility coupling). The network’s structure is represented in an m-dimensional adjacency list that encodes the directed connections with other nodes. Here, dotted and filled lines represent different types of connectivity. 3. Updates. Update functions are deterministic transformations operating locally that can access and modify the four sets of variables at run time. 4. Update sequences. The update sequence shapes belief propagation. It defines the order in which nodes should be updated when a new observation is presented to the input node(s). By default, prediction propagates from the leaves to the root of the network, while the interleaved sequence of prediction errors and posterior updates follow the inverted path. Updates can also be triggered reactively in which case the propagation starts with the activation of a proximal node.

https://doi.org/10.1371/journal.pcbi.1014340.g001

A dynamic network, as implemented in pyhgf is thus defined as the combination of four variables (see Fig 1). Let, for example, be a neural network with K nodes. This network handles in a tuple four parameters of interest:

(1)

The variable represents the nodes’ attributes. Attributes can be used to register local information, like the sufficient statistics of a probability distribution as well as the coupling weights with other nodes. This variable can also be arbitrarily extended to include other fixed parameters or results from other update steps. In a convolutional neural network, would, for example, encode the activation strength. Most standard learning models optimize attributes that belong to this parameter space.

The second key parameter, tightly linked to the first one, is the adjacency list that controls the shape of the network. Each item in this set registers the directed connection between node k and other nodes. Networks that exhibit different connectivity structures propagate information differently. The set of directed connections can be multivariate (a node can have different types of connections with other nodes), such as in multilayer networks [42]. For example, the nodalised Hierarchical Gaussian Filter [29] assumes two kinds of coupling between nodes: value and volatility coupling. Every edge , therefore, contains m = 2 sets of node indices in this case, m being the adjacency dimension. By comparison, in a standard recurrent neural network, this variable would define the shape of the layers and their connectivity. Critically, in the proposed framework, this variable is transparent to the update function and can be subject to inference and updates.

The third central component is the set of n in-place update functions defining a message passing step operating on the network’s parameter set such as:

(2)

In a convolutional neural network, this set of functions would include the linear product of input and weights, as well as the activation functions. In the generalised HGF [29], this includes three kinds of steps: a prediction (based on previous values and any parent nodes), an update step (based on input from child nodes and the prediction), and the computation of a prediction error. The specific computations in each case depend on the type of coupling between parent and child nodes. Note that the function is parametrized by a target node k to which it applies. This allows for defined local computation where only information from a subset of adjacent nodes is used, such as in particular mean field approximations.

Finally, a fourth component is introduced to control the scheduling of these update steps over time as . This ordered list describes a sequence of functions parametrized on individual nodes. The update order shapes belief propagation. This component is rarely expressed in the form of a parameter in most conventional applications of neural networks, as well as in the previous implementation of the HGF. This sequence is instead scripted outside the network’s closure and, therefore, not accessible during inference or optimization. In the case of predictive coding neural networks, however, finer control over belief propagation might be requested by the user, of a kind that also offers flexibility in the modification of belief propagation dynamically. When using the HGF as implemented in pyhgf, this scheduling can be generated at runtime from the network structure , assuming an ideal belief propagation pattern with a cascade of prediction from the leaves to the roots of the network, and another cascade of prediction error / posterior update pair from the roots to the leaves.

The proposed framework is intended to provide the minimal layout required to create dynamic neural networks for predictive coding. It allows users to create and manipulate the scheduling of updates through a network of nodes while keeping the four components of the network available for inference. Contrary to other neural networks that rely on matrix multiplication for learning, our networks implement local computations that are run sequentially to propagate beliefs along connectivity paths. This also offers a clear dissociation between components that can be developed separately. It is, for example, possible to create alternative message-passing algorithms without having to develop an entire library to simulate the networks, and it is possible to implement existing predictive coding frameworks so users can easily apply them to behavioural data. pyhgf natively support the generalised hierarchical Gaussian filter [29], a recent development of the Hierarchical Gaussian Filter [27,28] into a nodalised version for predictive coding. In this framework, for example, every node in the network represents a probability distribution of a belief about a latent space in the environment. Beliefs are updated through precision-weighted prediction errors coming from nodes in a lower level of the hierarchy and propagated to higher-level nodes. The exact update functions have been derived in their closed form and can work with arbitrary network architectures [29], which makes this model an excellent application of dynamic neural networks as described here. The relatively widespread use of the hierarchical Gaussian filter in computational psychiatry, and the need for advanced Bayesian modelling tools around it, are also good opportunities to extend the original Matlab toolbox [34] by enhancing the modularity and extensibility of the library.

2.2 Optimization and inference

While predictive coding itself originates from fields related to signal processing and information theory [43], the use of predictive coding as a framework for hierarchical inference [7,8] in biological neural networks makes it especially well-suited to fields related to experimental neuroscience and computational psychiatry. In this context, the neural networks are components of a cognitive model of the subject on which the experimenter performs inference (i.e., observing the observer [19]). For example, in the context of the generalised hierarchical Gaussian filter, the user might be interested in inferring the posterior distribution of tonic volatility at different levels of the hierarchy from observed behaviours.

This kind of reverse inference requires the use of techniques like MCMC sampling or gradient descent, which involves the evaluation of several instances of a network, as well as the gradient at evaluation, to find parameters maximizing likelihood. In the Matlab HGF toolbox [27,28], the inference step is implemented using a variant of the BFGS algorithm, which can be difficult to apply in the context of multilevel models, where there is a particularly pressing need for both high performance and the benefits of automatic differentiation. The pyhgf codebase is entirely written in Python and, as of version 0.2.0, can use JAX [35] as a computational backend which can easily deploy code on CPU, GPU and TPU. JAX offers a rapidly growing ecosystem for machine learning [37] and artificial intelligence that already includes toolboxes that are conceptually related to predictive coding and Hierarchical Gaussian Filters, such as state-space modelling (https://github.com/probml/dynamax), reinforcement learning [44], neural networks [38,41] or graph neural networks [18]. We leverage the automatic differentiation and just-in-time compilation offered by JAX [35] to let the networks interface smoothly with other optimization and inference libraries like PyMC [39] that support a large range of sampling or variational methods, including Hamiltonian Monte-Carlo methods such as the No-U-Turn Sampler (NUTS) [45], an approach that has proved to be highly efficient when scaling to high-dimensional problems. While dimensionality was not a major concern for individual model fittings, this can become problematic if we want to model group-level parameters, and therefore estimate a large number of networks together with hyperpriors (multilevel modelling). Assessing group-level estimates is a crucial step for studies in computational psychiatry, where gaining insights into computational parameters at the population level can inform further diagnosis and classification. In pyhgf, it is possible to apply multilevel modelling to any dynamic neural network handled by the toolbox. In the next section, we will focus on the development workflow using the standard three-level Hierarchical Gaussian Filter as an example.

3.1 Generative model, forward fitting and parameter inference

The standard workflow begins by constructing a Bayesian network representing the agent’s generative model of the environment (Fig 2A left panel). In pyhgf, this structure is defined using the Network class, which stores nodes, parameters, edges, and belief update sequences. The resulting graph can be interpreted as a Bayesian network for forward inference (prediction) and used to simulate belief dynamics.

Listing 1. Creating a network

from pyhgf.model import Network

# create a new network --------------------------------

# This structure is known as the three-level binary HGF

binary_hgf = (

  Network()

  .add_nodes(kind="binary-state")

  .add_nodes(kind="continuous-state", value_children=0)

  .add_nodes(kind="continuous-state", volatility_children=1)

)

# visualisation utility to inspect the network’s structure

binary_hgf.plot_network()

In most applications, however, the model is fitted to observations. Model inversion proceeds through iterative belief propagation at each time step: (1) predictions are propagated from leaves to the root of the network, (2) a new observation is incorporated, and (3) prediction errors propagate upward to update posterior beliefs (Fig 2A. right panel). The resulting belief trajectories depend on node attributes such as the tonic volatility at the second level (), which modulates the precision of the implied normal distribution and thus the effective learning rate.

Observations are inserted at the first node (x₁), whose inferred mean () corresponds to the probability of observing category 1. This value is predicted by a continuous node (x₂) representing the logit-transformed probability. A third node (x₃) controls volatility through volatility coupling, allowing fluctuations in x₂. If volatility is assumed constant, the third node can be omitted, resulting in a two-level binary HGF. The internal plotting utilities allow convenient inspection of network structure and debugging, and will produce a figure similar to Fig 2C (right panel).

Listing 2. Belief trajectories

from pyhgf import load_data

u, y = load_data("binary")

# provide a vector of binary observations to the network

binary_hgf.input_data(input_data=u)

# visualisation utility to inspect beliefs trajectories

binary_hgf.plot_trajectories()

We provide a sequence of observations using the same class. Fitting the model automatically triggers iterative calls of the belief propagation update. We log the network’s state at each time point in so-called belief trajectories, which describe the agent’s inferred states over time. To evaluate model fit to behaviour, we compute the Bayesian surprise of observed responses using a response function that maps beliefs to actions.

Listing 3. Creating a network

from pyhgf.math import binary_surprise

import jax.numpy as jnp

def binary_softmax(

  hgf,

  response_function_inputs,

  response_function_parameters,

):

  “”“Surprise under the binary sofmax model.”””

  # the expected values at the first level of the HGF

  beliefs = hgf.node_trajectories[0]["expected_mean"]

  # the binary surprises

  surprise = binary_surprise(x=response_function_inputs, expected_mean=beliefs)

  # ensure that inf is returned if the model cannot fit

  surprise = jnp.where(jnp.isnan(surprise), jnp.inf, surprise)

  return surprise

# compute the sum of the binary surprise

# lower values indicate that the model is a better fit to the participant’s behavior

binary_hgf.surprise(

  response_function = binary_softmax,

  response_function_inputs = y

).sum()

The response function specifies the likelihood of observed decisions given the inferred beliefs, and the resulting surprise corresponds to the negative log-probability of the data under the model [19]. We now have all the ingredients to infer parameters governing belief dynamics or decision noise by providing an MCMC sampling algorithm with information about our custom distribution. This interface can be handled automatically using the HGFDistribution class, which will create the relevant structure for sampling inside a PyMC model [39].

Listing 4. Creating a network

from pyhgf.distribution import HGFDistribution

import pymc as pm

# create a custom distribution that can be plugged into a PyMC model

hgf_logp_op = HGFDistribution(

  n_levels = 2,

  model_type = “binary,”

  input_data = u[jnp.newaxis,:],

  response_function = binary_softmax,

  response_function_inputs=y[jnp.newaxis,:]

)

# create a PyMC model with one uniform prior over the tonic volatility

with pm.Model() as three_levels_binary_hgf:

  # Set a prior over the evolution rate at the second level.

  tonic_volatility_2 = pm.Uniform("tonic_volatility_2", -3.5, 0.0)

  # Call the pre-parametrized HGF distribution here.

  # All parameters are set to their default value, except omega_2.

  pm.Potential(“hgf_loglike,” hgf_logp_op(tonic_volatility_2 = tonic_volatility_2))

# visualisation utility to inspect the PyMC computational graph

# note that while this is also a network, this is only related to the sampling procedurepm.model_to_graphviz(three_levels_binary_hgf)

# sample

with three_levels_binary_hgf:

  three_level_hgf_idata = pm.sample(chains=2, cores=1)

The procedure above describes the perceptual model and explains how beliefs evolve in the network as new observations are made. We then assume that an agent uses available beliefs at time k to inform decisions and actions. How to convert beliefs into actions depends on the problem we try to solve. We assume that decisions are generated from the inferred probability through a logistic sigmoid response function parameterised by an inverse temperature parameter t. By estimating parameters such as and t, we obtain the posterior density through MCMC sampling (Fig 2B). The posterior means can then be used to refit the model and recover the belief trajectories most consistent with the observed behaviour (Fig 2C). Because belief updates rely on closed-form variational updates, model inversion is deterministic given fixed inputs and parameters.

3.2 Bayesian multilevel modelling, parameter recovery and model comparison

Experiments typically involve multiple participants, requiring joint inference over multiple parameter sets. To illustrate this scenario, we simulated responses for 50 participants using the same observation sequence u but different parameter values for and t (Fig 3A). For each participant, we generated a response vector from the response model and fitted the HGF using the procedure described above.

Since participant fits are independent, the inference can be performed either iteratively or jointly in a single-level Bayesian model (Fig 3B). Parameter recovery was assessed by comparing simulated and inferred values of and t. As shown in Fig 3C, recovered parameters closely followed the identity line, indicating reliable recovery from the behavioural data.

Alternative generative models can also be compared on the same dataset, for instance, by modifying network structure or response functions. Here we compared two models differing only in the response function: one with a fixed inverse temperature (t = 1) and one in which t was free and inferred. Model comparison was performed using leave-one-out cross-validation (LOO) implemented in ArviZ [47]. As expected, the model estimating t achieved a higher expected log pointwise predictive density (ELPD) (Fig 3 D.). Such comparisons should primarily be guided by theoretical considerations and complemented by prior predictive checks and posterior predictive validation.

Finally, population-level inference can be performed using multilevel models in which individual parameters are drawn from group distributions (Fig 3B). This hierarchical approach enables estimation of population parameters and increases statistical power for group comparisons. Panel E. in Fig 3 shows posterior densities of the group means for and t, together with empirical group means. In both cases, the empirical means fall within the 94% highest density interval, indicating reliable population-level estimation.

4.1 Generalised Bayesian filtering in static networks

Predictive coding models can be viewed as dynamic implementations of Bayesian networks in which belief updates occur through prediction-error-driven message passing [21,22,29]. Even when the network structure is fixed, this formulation provides a flexible way to approximate variational inference in complex generative models [9,50]. One advantage is that arbitrary-sized neural networks can represent arbitrary complex generative models without having to rethink an entire optimization algorithm. While nodes are implicitly tracking one parameter value through unidimensional normal distributions, more complex probabilistic models can be constructed by combining nodes that track the sufficient statistics of any exponential-family distributions [5], including multivariate cases. This approach in itself offers considerable modularity in real-time probabilistic modelling, as the development of complex variational update algorithms can be replaced by the manipulation of nodes in a large network. Such flexibility may extend the use of the HGF and related predictive coding filters to a broader range of real-time probabilistic modelling problems [51,52].

One straightforward consequence is that networks can be expanded along multiple dimensions: horizontally (by adding more input nodes), vertically (adding more parent nodes), or using multivariate dependencies in which nodes can influence multiple descendants or receive input from multiple parents (see Fig 4A.). These extensions allow the modelling of richer generative structures, including mixtures of distributions or custom coupling functions between variables. One straightforward application is the training of deep neural networks to solve prediction and classification tasks (e.g., MNIST or CIFAR), such as what has been implemented in the machine learning literature [53–55]. While work still needs to be done to port the existing predictive coding framework to such an application, this direction is now enabled by our framework. But another potential application is the integration of multimodal data streams frequently encountered in cognitive neuroscience. For example, separate branches could model physiological signals (e.g., respiration, heart rate, EEG, or fMRI) and behavioural outcomes while sharing higher-level volatility estimates. Such models provide a principled way to capture interactions between behavioural and physiological processes within a single generative framework, rather than analysing these modalities separately.

4.2 Structure learning, causal inference, and meta-learning in dynamic networks

Beyond flexible connectivity, pyhgf also allows structural variables to be modified during belief propagation. Predictive coding updates are driven by precision-weighted prediction errors, and very large errors may destabilise inference if the model structure is overly restrictive. Dynamic reconfiguration of the network offers an alternative strategy: instead of forcing existing beliefs to accommodate unexpected observations, the model can adapt its structure to better capture the underlying generative process.

This mechanism can naturally express several forms of adaptive learning. For example, nodes can be added or removed during inference, enabling non-parametric model growth in response to unexpected observations (Fig 4B). Similar ideas have recently emerged in research on self-growing neural networks and lifelong learning in reinforcement learning [56,57]. In predictive coding networks, such structural adaptation may support branching, splitting, or merging of model components as environmental complexity changes. Our framework now supports the creation of such networks and behaviours.

Structural adaptation can also target the connectivity between nodes. Because the prediction step of a predictive coding network corresponds to a Bayesian generative graph, modifying edges effectively changes the assumed causal structure between variables. Dynamic adjustment of edges, therefore, provides a natural mechanism for causal discovery from observational data (Fig 4B). Causal inference is increasingly recognised as a core component of biological and artificial learning systems [58–60]. Within the pyhgf framework, causal relations can be inferred in real time, enabling the study of time-varying causal dependencies or volatility coupling on causal edges.

Finally, structural flexibility can extend to the belief propagation dynamics themselves. Since propagation consists of a sequence of update functions, both the sequence and the functions can, in principle, be modified during inference. In pyhgf, update functions are stored as part of the network representation and can therefore evolve as learning progresses. This enables forms of meta-learning in which the system adapts its own inference algorithm to improve predictive accuracy (Fig 4B.). Conceptually, this approach combines ideas from Bayesian non-parametric modelling, where distributions over functions are inferred, with meta-learning approaches in reinforcement learning [61].

In addition to these core computational internal changes, and closer to active inference models, the package also supports the modular extension of the response functions used by the agent (i.e., the beliefs-to-action function). This framework is central for reinforcement learning applications and often requires custom response functions tailored to specific tasks. Examples of how to use custom response functions can be found in the online documentation.