spd_learn.modules.step_dynamic_trivialization#

spd_learn.modules.step_dynamic_trivialization(model: Module, optimizer: Optimizer | None = None, *, threshold: float = 0.0, optimizer_state: Literal['reset', 'transport', 'keep'] = 'reset') → None[source]#

Recenter native BiMap and exponential SPD charts, enabling SPD on first call.

Call after K completed optimizer updates and backward, outside parametrize.cached(). Effective values and Parameter identities remain unchanged, including frozen parameters. Unsupported target chains raise. Nonparametrized BiMaps, softplus and positive scalars stay unchanged.

Recentering is deterministic and preserves the horizontal Stiefel frame without consuming RNG state. Gradients of recentered Parameters are cleared.

Parameters:

  • model (nn.Module) – Model containing parametrized BiMap or SymmetricPositiveDefinite modules.

  • optimizer (torch.optim.Optimizer or None, default=None) – Optimizer whose state should be updated according to optimizer_state.

  • threshold (float, default=0.0) – Minimum tangent displacement norm required to trigger recentering for a given parameter. When 0.0 (default), all eligible parameters are recentered unconditionally.

  • optimizer_state ({"reset", "transport", "keep"}, default="reset") – How to handle optimizer.state[original] when optimizer is provided: "transport" parallel-transports first-moment buffers (exp_avg / momentum_buffer) into the new tangent frame while preserving step and exp_avg_sq; "keep" leaves state untouched; "reset" clears optimizer.state[original].

Notes

  • Single-subject / cross-session EEG (short travel): When weights stay close to initialization (\(\lesssim 25^\circ\)), static trivialization (\(K = \infty\), the default) is already nearly isometric, so dynamic updates are not needed.

  • Cross-subject, continual adaptation, or long-horizon training (large travel): When covariance shifts push weights far from initialization, periodic recentering (\(K \in [10, 50]\) or threshold > 0) prevents differential distortion. Just as SPDBatchNormMeanVar uses parallel transport (\(G^{-1/2} P G^{-1/2}\)) to re-center batch covariances to the identity, each dynamic trivialization step updates the base point (\(B \leftarrow \phi_B(A), A \leftarrow 0\)) and transports the parameter tangent space back to the identity (\(\Gamma_{B \to I}\)), acting as Riemannian metric preconditioning.