Energy functions

jpc.pc_energy_fn(params: Tuple[PyTree[Callable], Optional[PyTree[Callable]]], activities: PyTree[ArrayLike], y: ArrayLike, x: Optional[ArrayLike] = None, loss: str = 'MSE', record_layers: bool = False) -> Array | Array

Computes the free energy for a feedforward neural network of the form

\[ \mathcal{F}(\mathbf{z}; θ) = 1/N \sum_i^N \sum_{\ell=1}^L || \mathbf{z}_{i, \ell} - f_\ell(\mathbf{z}_{i, \ell-1}; θ) ||^2 \]

given parameters \(θ\), free activities \(\mathbf{z}\), output \(\mathbf{z}_L = \mathbf{y}\) and optional input \(\mathbf{z}_0 = \mathbf{x}\) for supervised training. The activity of each layer \(\mathbf{z}_\ell\) is some function of the previous layer, e.g. ReLU\((W_\ell \mathbf{z}_{\ell-1} + \mathbf{b}_\ell)\) for a fully connected layer with biases and ReLU as activation.

Note

The input \(x\) and output \(y\) correspond to the prior and observation of the generative model, respectively.

Main arguments:

• params: Tuple with callable model (e.g. neural network) layers and optional skip connections.
• activities: List of activities for each layer free to vary.
• y: Observation or target of the generative model.
• x: Optional prior of the generative model (for supervised training).

Other arguments:

• loss: Loss function to use at the output layer (mean squared error 'MSE' vs cross-entropy 'CE'). ??? cite "Reference"

  @article{tscshantz2023hybrid,
    title={Hybrid predictive coding: Inferring, fast and slow},
    author={Tscshantz, Alexander and Millidge, Beren and Seth, Anil K and Buckley, Christopher L},
    journal={PLoS Computational Biology},
    volume={19},
    number={8},
    pages={e1011280},
    year={2023},
    publisher={Public Library of Science San Francisco, CA USA}
  }
  

  - record_layers: If True, returns energies for each layer.

Returns:

The total or layer-wise energy normalised by the batch size.

---

jpc.hpc_energy_fn(model: PyTree[typing.Callable], equilib_activities: PyTree[ArrayLike], amort_activities: PyTree[ArrayLike], x: ArrayLike, y: Optional[ArrayLike] = None, record_layers: bool = False) -> Array | Array

Computes the free energy of an amortised PC network

\[ \mathcal{F}(\mathbf{z}^*, \hat{\mathbf{z}}; θ) = 1/N \sum_i^N \sum_{\ell=1}^L || \mathbf{z}^*_{i, \ell} - f_\ell(\hat{\mathbf{z}}_{i, \ell-1}; θ) ||^2 \]

given the equilibrated activities of the generator \(\mathbf{z}^*\) (target for the amortiser), the feedforward guesses of the amortiser \(\hat{\mathbf{z}}\), the amortiser's parameters \(θ\), input \(\mathbf{z}_0 = \mathbf{x}\), and optional output \(\mathbf{z}_L = \mathbf{y}\) for supervised training.

Note

The input \(x\) and output \(y\) are reversed compared to pc_energy_fn (\(x\) is the generator's target and \(y\) is its optional input or prior). Just think of \(x\) and \(y\) as the actual input and output of the amortiser, respectively.

Main arguments:

• model: List of callable model (e.g. neural network) layers.
• equilib_activities: List of equilibrated activities reached by the generator and target for the amortiser.
• amort_activities: List of amortiser's feedforward guesses (initialisation) for the network activities.
• x: Input to the amortiser.
• y: Optional target of the amortiser (for supervised training).

Other arguments:

• record_layers: If True, returns energies for each layer.

Returns:

The total or layer-wise energy normalised by batch size.