spectral_contrastive_loss¶

Spectral contrastive loss for self-supervised learning.

Originally introduced by HaoChen et al. [1] and adapted for evolution operators in Turri et al. [2] and Jeong et al. [3].

The loss encourages alignment of paired samples (x_i, y_i) while discouraging alignment of unpaired samples:

\[\mathcal{L}(x, y) = \frac{1}{N(N-1)}\sum_{i \neq j}\langle x_{i}, y_{j} \rangle^2 - \frac{2}{N}\sum_{i=1}^N\langle x_{i}, y_{i} \rangle\]

Hint

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

Parameters:

x (ArrayLike) – Input features of shape (N, D), where N is the number of samples and D is the feature dimension.
y (ArrayLike) – Output features of shape (N, D). Must have the same shape as x.

Returns:

Scalar loss value.

Return type:

jax.Array

Raises:

ValueError – If x and y do not have the same shape or if x is not 2-dimensional.

References

[1]

Jeff Z HaoChen, Colin Wei, Adrien Gaidon, and Tengyu Ma. Provable guarantees for self-supervised deep learning with spectral contrastive loss. Advances in neural information processing systems, 34:5000–5011, 2021.

[2]

Giacomo Turri, Luigi Bonati, Kai Zhu, Massimiliano Pontil, and Pietro Novelli. Self-supervised evolution operator learning for high-dimensional dynamical systems. arXiv preprint arXiv:2505.18671, 2025.

[3]

Minchan Jeong, J Jon Ryu, Se-Young Yun, and Gregory W Wornell. Efficient parametric svd of koopman operator for stochastic dynamical systems. arXiv preprint arXiv:2507.07222, 2025.