pyhgf.updates.vectorised.learning.learning_weights_vectorised#

pyhgf.updates.vectorised.learning.learning_weights_vectorised(parent_state, child_state, coupling_fn, kind='precision_weighted', parent_has_constant=False, child_kind='continuous', child_evidence=None)[source]#

Per-layer weight-learning factors for the vectorised deep network.

The single entry point of this module. One weight update has two halves and this returns the pieces of both: the descent gradient that moves the weight’s mean, and the importance increment that raises its precision. They are the first and second derivative of one layer-local energy, computed in one pass rather than selected between.

Both are rank-one and share their parent-side factor. The gradient is \(u \otimes h\) and the increment is \(p \otimes h^2\), with the same coupled parent activation \(h_i = g(\mu_i)\), so three vectors carry both quantities. Returning factors rather than assembled matrices lets a batched caller average over samples and contract once (einsum('bi,bj->ij') / batch): the same arithmetic, without materialising one weight-matrix-sized array per sample.

Both are strictly local: what one weight gets reads only the prediction error and precision at its child and the activation at its parent, never a reduction over the other parents or children in the layer.

The child-side gradient factor starts from the value prediction error \(\delta = \mu_\text{child} - \hat{\mu}_\text{child}\), and kind decides what scales it:

  • "standard" takes \(\delta\) alone, the raw gradient of an unweighted squared error. This coincides with the free-energy gradient only where the child precision is one (the unit-precision output or categorical convention).

  • "precision_weighted" scales by the child’s posterior precision, \(\delta\,\pi_a\). This is the backprop-parity mode: a moved interior belief shifts by (routed error) / posterior precision, so weighting by that same posterior precision cancels the division and reproduces the backpropagated gradient node for node at any precision setting.

  • "synaptic_uncertainty" charges the process noise between the weight and its child on top, \(\delta\,\pi_a / (1 + \Omega_a \xi_a)\), so the gradient and the increment become the two derivatives of one energy. This is the kind the weight-belief rule descends, and the only one it can: dividing it by the weight’s own precision is the Gaussian posterior step, so both halves of that update read the same evidence. It is not backprop parity.

A binary child drops the precision factor either way: its precision field holds the Bernoulli variance \(p(1-p)\), which cancels through the sigmoid in the gradient, so keeping it would count the same term twice. That cancellation belongs to the first derivative only, and the curvature below keeps it.

The child-side importance factor is the exact Hessian diagonal of the layer-local variational energy: that energy is exactly quadratic in the weights for any coupling function, because the coupling acts on the activation while the map \(w_{ai} \mapsto w_{ai} h_i\) stays linear, so one observation adds \(\tilde{\xi}_a h_i^2\) with no approximation and no sampling. Where the evidence comes from depends on the child:

  • A binary or categorical child is where the data are clamped, so the evidence is the curvature of its own likelihood, per unit \(p_a(1 - p_a)\) with \(p_a\) the child’s expected mean. For the categorical case the Hessian of the log-likelihood with respect to the logits is \(\operatorname{diag}(\mathbf{p}) - \mathbf{p}\mathbf{p}^{\top}\), whose diagonal is that expression. No label enters it, since averaging the outer product of the residual over labels drawn from the model returns the same matrix, so this is the model’s own expected curvature rather than an estimate built from observed errors.

  • A continuous child passes on the evidence it received from below, softened by the process noise it had to cross, \(\tilde\xi_a = (1/\xi_a + \Omega_a)^{-1}\), with \(\Omega_a = \gamma_a / \tilde\pi_a\) recovered from the effective precision the prediction sweep writes as \(\gamma_a = \Omega_a \tilde\pi_a\). The evidence \(\xi_a\) itself is supplied by the caller through child_evidence, carried up from the clamped layer by evidence_pullback(), and the weight-belief rule always supplies it.

    Without it, \(\xi_a\) falls back to \(\pi_a - \tilde\pi_a\) from the cache, which is the same quantity in exact arithmetic only where the filter’s own precision chain was seeded with the clamped layer’s likelihood curvature. It is not: a clamped categorical layer enters that chain at unit precision, because that convention is what makes the message it routes exact cross-entropy backpropagation. Reading the difference also subtracts two large nearly-equal precisions, which loses every significant digit once a layer’s precision has accumulated. The fallback is kept for direct callers reading the layer-local quantity; it is not what the rule descends.

Parameters:

  • parent_state (LayerState) – Current state of the parent layer.

  • child_state (LayerState) – Current state of the child layer (with observations), after the update sweep has written its posterior.

  • coupling_fn (Callable) – Coupling function applied to parent means.

  • kind (str) – The metric the gradient is expressed in, one of SEPARABLE_KINDS. It also sets the importance convention: "standard" uses unit observation precision on both halves.

  • parent_has_constant (bool) – If True, the parent layer has a constant input node (mean = 1.0, precision = 1.0) appended to its activations after coupling.

  • child_kind (str) – The child layer’s node kind, "binary", "categorical" or anything else for a continuous one.

  • child_evidence (Array | None) – The child’s evidence precision \(\xi_a\), supplied by a caller that carries it up the stack itself (evidence_pullback()) rather than letting it be recovered here as \(\pi_a - \tilde\pi_a\). When given, it replaces that difference in the importance factor and the fallback below never applies, since a carried evidence is non-negative by construction. The gradient factor is unaffected: it reads the child’s posterior precision either way. None (default) recovers the evidence from the cache.

Returns:

The triple (u, h, p). The child-side gradient factor \(u\) and the child-side importance factor \(p\) have shape (n_children,); the shared parent-side factor \(h\) has shape (n_parents[+1],). The gradient is u[:, None] * h[None, :] and the increment p[:, None] * (h ** 2)[None, :]. Non-finite entries are zeroed, so optax never propagates a NaN or an inf through its moment accumulators.

Return type:

factors

Raises:

ValueError – If kind is unrecognised.