AutoEncoderLoss

class kooplearn.torch.nn.AutoEncoderLoss(alpha_rec: float = 1.0, alpha_lin: float = 1.0, alpha_pred: float = 1.0)[source]

Bases: Module

Single-step Dynamic Autoencoder (DAE) loss introduced by Lusch et al. [1].

This loss combines three objectives to train dynamic autoencoders:

  1. Reconstruction loss — measures how well the autoencoder reconstructs inputs.

  2. Linearity loss — enforces linear evolution in latent space.

  3. Prediction loss — penalizes errors between predicted and actual encoded outputs.

The total loss is a weighted sum:

\[\mathcal{L} = \alpha_\mathrm{rec} \, \|x - \phi^{-1}(\phi(x)) \|^2 + \alpha_\mathrm{lin} \, \|\phi(y) - K\phi(x) \|^2 + \alpha_\mathrm{pred} \, \|y - \phi^{-1}(K\phi(x))\|^2\]

where \(\phi^{-1}(\phi(x))\) is the reconstruction of \(x\), \(\phi(y)\) is the encoded output, \(K\phi(x)\) is the evolved input latent representation, and \(\phi^{-1}(K\phi(x))\) is the predicted decoded output.

Hint

Check out the Ordered MNIST example for a practical use of this loss function.

Methods

forward(x: Tensor, y: Tensor, x_rec: Tensor, y_enc: Tensor, x_evo: Tensor, y_pred: Tensor) → Tensor[source]

Compute the Dynamic Autoencoder loss.

Parameters:

  • x (torch.Tensor) – Input features of shape (N, D), where N is the batch size and D is the feature dimension.

  • y (torch.Tensor) – Output (target) features. Same shape as x.

  • x_rec (torch.Tensor) – Reconstructed version of the input x produced by the decoder. Same shape as x.

  • y_enc (torch.Tensor) – Encoded latent representation of the target y.

  • x_evo (torch.Tensor) – Evolved latent representation obtained by applying the learned linear operator to the latent encoding of x. Same shape as x.

  • y_pred (torch.Tensor) – Predicted decoded output corresponding to the evolved latent state. Same shape as x.

Returns:

A scalar tensor representing the total dynamic autoencoder loss.

Return type:

torch.Tensor

[1]

Bethany Lusch, J. Nathan Kutz, and Steven L. Brunton. Deep learning for universal linear embeddings of nonlinear dynamics. Nature Communications, November 2018. URL: https://doi.org/10.1038/s41467-018-07210-0, doi:10.1038/s41467-018-07210-0.