Theoretical tools¤

JPC provides the following main theoretical tools that can be used to study deep linear networks (DLNs) trained with PC:

jpc.linear_equilib_energy() to compute the theoretical PC energy at the solution of the activities for DLNs;
jpc.compute_linear_activity_hessian() to compute the theoretical Hessian of the energy with respect to the activities of DLNs;
jpc.compute_linear_activity_solution() to compute the analytical PC inference solution for DLNs.

jpc.linear_equilib_energy(params: typing.Tuple[jaxtyping.PyTree[typing.Callable], typing.Optional[jaxtyping.PyTree[typing.Callable]]], x: typing.Union[jax.Array, numpy.ndarray, numpy.bool, numpy.number, bool, int, float, complex], y: typing.Union[jax.Array, numpy.ndarray, numpy.bool, numpy.number, bool, int, float, complex], *, param_type: str = 'sp', gamma: typing.Optional[jaxtyping.Shaped[Array, '']] = None, output_energy_scaling: typing.Optional[jaxtyping.Shaped[Array, '']] = None, return_rescaling: bool = False) -> typing.Union[jax.Array, typing.Tuple[jax.Array, jax.Array]] ¤

Computes the theoretical PC energy at the solution of the activities for a deep linear network (Innocenti et al. 2024):

$$ \mathcal{F}^* = rac{1}{2B}\sum_i^B (\mathbf{y}_i - \hat{\mathbf{y}}_i)^T \mathbf{S}^{-1}(\mathbf{y}_i - \hat{\mathbf{y}}_i) $$

where $\hat{\mathbf{y}}_i$ is the feedforward prediction, $\mathbf{S} = \mathbf{I}_{d_y}/\lambda + \sum_\ell \mathbf{P}_\ell \mathbf{P}_\ell^T$, $\lambda$ is the output_energy_scaling (output precision, defaulting to 1), and $\mathbf{P}_\ell$ are the scaled layer products returned by compute_linear_equilib_rescaling(). For scalar output this can be written as $\mathcal{F}^* = \mathcal{L}/s$ with feedforward MSE loss $\mathcal{L}$ and scalar rescaling $s = S_{00}$. For example, a one-hidden-layer $\mu$PC MLP with output_energy_scaling = γ²NL gives $s = (1 + \|\mathbf{w}\|^2/N)/(\gamma^2 N L)$.

!!! note

  This expression assumes no biases. It could also be generalised to 
  other network architectures (e.g. ResNets) and parameterisations
  (see [Innocenti et al. 2025](https://openreview.net/forum?id=lSLSzYuyfX&referrer=%5Bthe%20profile%20of%20Francesco%20Innocenti%5D(%2Fprofile%3Fid%3D~Francesco_Innocenti1))). 
  However, note that the equilibrated energy for ResNets and other
  parameterisations can still be computed by getting the activity solution
  with [`jpc.compute_linear_activity_solution()`](https://thebuckleylab.github.io/jpc/api/Theoretical%20tools/#jpc.compute_linear_activity_solution) 
  and then plugging this into the standard PC energy 
  [jpc.pc_energy_fn()](https://thebuckleylab.github.io/jpc/api/Energy%20functions/#jpc.pc_energy_fn).

!!! example

  In practice, this means that if you run, at any point in training, the 
  inference dynamics of any PC linear network to equilibrium, then 
  [`jpc.pc_energy_fn()`](https://thebuckleylab.github.io/jpc/api/Energy%20functions/#jpc.pc_energy_fn) 
  will return the same energy value as this function. For a demonstration, see
  [this example notebook](https://thebuckleylab.github.io/jpc/examples/linear_net_theoretical_energy/).

  ![](linear_net_theoretical_energy.png)

??? cite "Reference"

  ```bibtex
  @article{innocenti2025only,
      title={Only Strict Saddles in the Energy Landscape of Predictive Coding Networks?},
      author={Innocenti, Francesco and Achour, El Mehdi and Singh, Ryan and Buckley, Christopher L},
      journal={Advances in Neural Information Processing Systems},
      volume={37},
      pages={53649--53683},
      year={2025}
  }
  ```

Main arguments:

params: Tuple with callable network layers and optional skip connections.
x: Network input.
y: Network output.

Other arguments:

param_type: Determines the parameterisation. Options are "sp" (standard parameterisation), "mupc" (μPC), or "ntp" (neural tangent parameterisation). See _get_param_scalings() for the specific scalings of these different parameterisations. Defaults to "sp".
gamma: Optional scaling factor for the output layer. If provided, the output layer parameter scaling is multiplied by 1/gamma. Defaults to None (no additional scaling).
output_energy_scaling: Optional multiplier for the output-layer energy term. Note that this equals the precision (inverse covariance) of the generative distribution at the output layer. Defaults to None (equivalent to a scaling of 1).
return_rescaling: If True, also returns the rescaling matrix S. Defaults to False.

Returns:

Mean total theoretical energy over a data batch and optionally the rescaling matrix S.

jpc.compute_linear_equilib_energy_grads(params: typing.Tuple[jaxtyping.PyTree[typing.Callable], typing.Optional[jaxtyping.PyTree[typing.Callable]]], x: typing.Union[jax.Array, numpy.ndarray, numpy.bool, numpy.number, bool, int, float, complex], y: typing.Union[jax.Array, numpy.ndarray, numpy.bool, numpy.number, bool, int, float, complex], *, param_type: str = 'sp', gamma: typing.Optional[jaxtyping.Shaped[Array, '']] = None, output_energy_scaling: typing.Optional[jaxtyping.Shaped[Array, '']] = None) -> PyTree[jax.Array] ¤

Computes the gradient of the linear equilibrium energy with respect to model parameters $∇_θ \mathcal{F}^*$.

Main arguments:

params: Tuple with callable model layers and optional skip connections.
x: Network input.
y: Network output.

Other arguments:

param_type: Determines the parameterisation. Options are "sp" (standard parameterisation), "mupc" (μPC), or "ntp" (neural tangent parameterisation). See _get_param_scalings() for the specific scalings of these different parameterisations. Defaults to "sp".
gamma: Optional scaling factor for the output layer. If provided, the output layer parameter scaling is multiplied by 1/gamma. Defaults to None (no additional scaling).
output_energy_scaling: Optional multiplier for the output-layer energy term. Note that this equals the precision (inverse covariance) of the generative distribution at the output layer. Defaults to None (equivalent to a scaling of 1).

Returns:

Parameter gradients for the network.

jpc.update_linear_equilib_energy_params(params: typing.Tuple[jaxtyping.PyTree[typing.Callable], typing.Optional[jaxtyping.PyTree[typing.Callable]]], optim: optax._src.base.GradientTransformation | optax._src.base.GradientTransformationExtraArgs, opt_state: typing.Union[jax.Array, numpy.ndarray, numpy.bool, numpy.number, typing.Iterable[ForwardRef(ArrayTree)], typing.Mapping[typing.Any, ForwardRef(ArrayTree)]], x: typing.Union[jax.Array, numpy.ndarray, numpy.bool, numpy.number, bool, int, float, complex], y: typing.Union[jax.Array, numpy.ndarray, numpy.bool, numpy.number, bool, int, float, complex], *, param_type: str = 'sp', gamma: typing.Optional[jaxtyping.Shaped[Array, '']] = None, output_energy_scaling: typing.Optional[jaxtyping.Shaped[Array, '']] = None) -> typing.Dict ¤

Updates parameters of a linear network by taking gradients of the linear equilibrium energy with a given optax optimiser.

Main arguments:

params: Tuple with callable model layers and optional skip connections.
optim: optax optimiser, e.g. optax.sgd().
opt_state: State of optax optimiser.
x: Network input.
y: Network output.

Other arguments:

param_type: Determines the parameterisation. Options are "sp" (standard parameterisation), "mupc" (μPC), or "ntp" (neural tangent parameterisation). See _get_param_scalings() for the specific scalings of these different parameterisations. Defaults to "sp".
gamma: Optional scaling factor for the output layer. If provided, the output layer parameter scaling is multiplied by 1/gamma. Defaults to None (no additional scaling).
output_energy_scaling: Optional multiplier for the output-layer energy term. Note that this equals the precision (inverse covariance) of the generative distribution at the output layer. Defaults to None (equivalent to a scaling of 1).

Returns:

Dictionary with updated model, parameter gradients, and optimiser state.

`jpc.compute_linear_activity_hessian(Ws: PyTree[jax.Array], *, use_skips: bool = False, param_type: str = 'sp', activity_decay: bool = False, diag: bool = True, off_diag: bool = True, gamma: typing.Optional[jaxtyping.Shaped[Array, '']] = None, output_energy_scaling: typing.Optional[jaxtyping.Shaped[Array, '']] = None) -> Array` ¤

Computes the theoretical Hessian matrix of the PC energy with respect to the activities for a linear network, $(\mathbf{H}_{\mathbf{z}})_{\ell k} := \partial^2 \mathcal{F} / \partial \mathbf{z}_\ell \partial \mathbf{z}_k \in \mathbb{R}^{(NH)×(NH)}$ where $N$ and $H$ are the width and number of hidden layers, respectively (Innocenti et al., 2025).

Info

This function can be used (i) to study the inference landscape of linear PC networks and (ii) to compute the analytical solution with jpc.compute_linear_activity_solution().

Warning

This was highly hard-coded for quick experimental iteration with different models and parameterisations. The computation of the blocks could be implemented much more elegantly by fetching the transformation for each layer.

Reference

@article{innocenti2025mu,
    title={$$ackslash$mu $ PC: Scaling Predictive Coding to 100+ Layer Networks},
    author={Innocenti, Francesco and Achour, El Mehdi and Buckley, Christopher L},
    journal={arXiv preprint arXiv:2505.13124},
    year={2025}
}

Main arguments:

Ws: List of all the network weight matrices.

Other arguments:

use_skips: Whether to assume one-layer skip connections at every layer except from the input and to the output. False by default.
param_type: Determines the parameterisation. Options are "sp" (standard parameterisation), "mupc" (μPC), or "ntp" (neural tangent parameterisation). See _get_param_scalings() for the specific scalings of these different parameterisations. Defaults to "sp".
activity_decay: $\ell^2$ regulariser for the activities.
diag: Whether to compute the diagonal blocks of the Hessian.
off-diag: Whether to compute the off-diagonal blocks of the Hessian.
gamma: Optional scaling factor for the output layer. If provided, the output layer parameter scaling is multiplied by 1/gamma. Defaults to None (no additional scaling).
output_energy_scaling: Optional multiplier for the output-layer energy term (output precision λ). Defaults to None (equivalent to 1).

Returns:

The activity Hessian matrix of size $NH×NH$ where $N$ is the width and $H$ is the number of hidden layers.

jpc.compute_linear_activity_solution(model: PyTree[equinox.nn._linear.Linear], x: typing.Union[jax.Array, numpy.ndarray, numpy.bool, numpy.number, bool, int, float, complex], y: typing.Union[jax.Array, numpy.ndarray, numpy.bool, numpy.number, bool, int, float, complex], *, use_skips: bool = False, param_type: str = 'sp', activity_decay: bool = False, gamma: typing.Optional[jaxtyping.Shaped[Array, '']] = None, output_energy_scaling: typing.Optional[jaxtyping.Shaped[Array, '']] = None, epsilon: Shaped[Array, ''] = 0.0, hessian: typing.Optional[jax.Array] = None) -> PyTree[jax.Array] ¤

Computes the theoretical solution for the PC activities of a linear network (Innocenti et al., 2025).

\[ \mathbf{z}^* = \mathbf{H}_{\mathbf{z}}^{-1}\mathbf{b} \]

where $(\mathbf{H}_{\mathbf{z}})_{\ell k} := \partial^2 \mathcal{F} / \partial \mathbf{z}_\ell \partial \mathbf{z}_k \in \mathbb{R}^{(NH)×(NH)}$ is the Hessian of the energy with respect to the activities, and $\mathbf{b} \in \mathbb{R}^{NH}$ is a sparse vector depending only on the data and associated weights. The activity Hessian is computed analytically using jpc.compute_linear_activity_hessian().

Info

This can be used to study how linear PC networks learn when they perform perfect inference.

Reference

@article{innocenti2025mu,
    title={$$ackslash$mu $ PC: Scaling Predictive Coding to 100+ Layer Networks},
    author={Innocenti, Francesco and Achour, El Mehdi and Buckley, Christopher L},
    journal={arXiv preprint arXiv:2505.13124},
    year={2025}
}

Main arguments:

model: Linear network defined as a list of Equinox Linear layers.
x: Network input.
y: Network output.

Other arguments:

use_skips: Whether to assume one-layer skip connections at every layer except from the input and to the output. False by default.
param_type: Determines the parameterisation. Options are "sp" (standard parameterisation), "mupc" (μPC), or "ntp" (neural tangent parameterisation). See _get_param_scalings() for the specific scalings of these different parameterisations. Defaults to "sp".
activity_decay: $\ell^2$ regulariser for the activities.
gamma: Optional scaling factor for the output layer. If provided, the output layer parameter scaling is multiplied by 1/gamma. Defaults to None (no additional scaling).
output_energy_scaling: Optional multiplier for the output-layer energy term. Note that this equals the precision (inverse covariance) of the generative distribution at the output layer. Defaults to None (equivalent to a scaling of 1).
epsilon: Small regularization value added to the diagonal of the Hessian matrix before inversion to improve numerical stability. Defaults to 0..
hessian: Optional Hessian matrix to use instead of computing it analytically using jpc.compute_linear_activity_hessian(). Defaults to None.

Returns:

List of theoretical activities for each layer.