vicreg_loss

selfeeg.losses.losses.vicreg_loss(z1: Tensor, z2: Tensor = None, Lambda: float = 25, Mu: float = 25, Nu: float = 1, epsilon: float = 0.0001) → Tensor[source]

Pytorch implementation of the VICReg loss function [VIC] .

Parameters:

  • z1 (torch.tensor) – 2D tensor with projections of one augmented version of the batch.

  • z2 (torch.tensor, optional) –

    2D projections of the other augmented version of the batch. Can be none if z1 and z2 are cat together. In this case internal split is done, but be sure that the first dimension can be divided by 2.

    Default = None

  • Lambda (float, optional) –

    Coefficient applied to the invariant loss.

    Default = 25

  • Mu (float, optional) –

    Coefficient applied to the variance loss .

    Default = 25

  • Nu (float, optional) –

    Coefficient applied to the covariance.

    Default = 1

  • epsilon (float, optional) –

    Value summed to the variance for stability purposes.

    Default = 1e-4

Returns:

loss (torch.Tensor) – The calculated loss.

References

[VIC]

A. Bardes, J. Ponce, and Y. LeCun, “Vicreg: Variance-invariance-covariance regularization for self-supervised learning,” arXiv preprint arXiv:2105.04906, 2021.

Example

>>> import torch
>>> import selfeeg.losses
>>> torch.manual_seed(1234)
>>> z1 = torch.randn(64, 32)
>>> z2 = torch.randn(64, 32)
>>> loss = losses.vicreg_loss(z1,z2)
>>> print(loss) # will return 53.0773