Skip to content

Analytical tools¤

jpc.linear_equilib_energy(network: PyTree[equinox.nn._linear.Linear], x: ArrayLike, y: ArrayLike) -> Array ¤

Computes the theoretical equilibrated PC energy for a deep linear network (DLN).

\[ \mathcal{F}^* = 1/N \sum_i^N (\mathbf{y}_i - W_{L:1}\mathbf{x}_i)^T S^{-1}(\mathbf{y}_i - W_{L:1}\mathbf{x}_i) \]

where the rescaling is \(S = I_{d_y} + \sum_{\ell=2}^L (W_{L:\ell})(W_{L:\ell})^T\), and we use the shorthand \(W_{k:\ell} = W_k \dots W_\ell\) for \(\ell, k \in 1,\dots, L\). See the reference below for more details.

Note

This expression assumes no biases.

Reference 
@article{innocenti2024only,
  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={arXiv preprint arXiv:2408.11979},
  year={2024}
}

Main arguments:

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

Returns:

Mean total analytical energy over a batch or dataset.

jpc.linear_activities_solution(network: PyTree[equinox.nn._linear.Linear], x: ArrayLike, y: ArrayLike) -> PyTree[Array] ¤

Computes the theoretical solution for the PC activities of a deep linear network (DLN).

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

where \(A\) is a sparse block diagonal matrix depending only on the weights, and \(\mathbf{b} = [W_1 \mathbf{x}, \mathbf{0}, \dots, W_L^T \mathbf{y}]^T\). In particular, \(A_{\ell,k} = I + W_\ell^T W_\ell\) if \(\ell = k\), \(A_{\ell,k} = -W_\ell\) if \(\ell = k+1\), \(A_{\ell,k} = -W_\ell^T\) if \(\ell = k-1\), and \(\mathbf{0}\) otherwise, for \(\ell, k \in [2, \dots, L]\).

Note

This expression assumes no biases.

Main arguments:

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

Returns:

List of theoretical activities for each layer.