Note

This page is a reference documentation. It only explains the class signature, and not how to use it. Please refer to the user guide for the big picture.

nidl.losses.InfoNCE¶

class nidl.losses.InfoNCE(temperature=0.1)[source]¶

Bases: Module

Implementation of the InfoNCE loss [1], [2].

This loss function encourages the model to maximize the similarity between positive pairs while minimizing the similarity between negative pairs.

Given a mini-batch of size $n$ , we obtain two embeddings $z_{i}$ and $z_{j}$ representing two different augmented views of the same sample. The InfoNCE (or NT-Xent) loss used in is defined as:

$\mathcal{L}_i = -\log \frac{ \exp\!\big(\operatorname{sim}(z_i, z_j)/\tau\big) }{ \sum\limits_{k=1}^{2N} \mathbf{1}_{[k \ne i]}\, \exp\!\big(\operatorname{sim}(z_i, z_k)/\tau\big) }$

where $\operatorname{sim}(z_i, z_j)$ denotes the cosine similarity between the normalized embeddings $z_i$ and $z_j$ , and $\tau > 0$ is a temperature parameter controlling the concentration of the distribution.

Parameters:

temperature: float, default=0.1: Scale logits by the inverse of the temperature.

References

[1]

Aaron van den Oord, Yazhe Li, and Oriol Vinyals. “Representation learning with contrastive predictive coding.” arXiv preprint arXiv:1807.03748 (2018).

[2]

Ting Chen, Simon Kornblith, Mohammad Norouzi, and Geoffrey Hinton. “A simple framework for contrastive learning of visual representations.” In ICML 2020.

__init__(temperature=0.1)[source]¶: Initialize internal Module state, shared by both nn.Module and ScriptModule.

forward(z1, z2)[source]¶

Forward implementation.

Parameters:

z1: torch.Tensor of shape (batch_size, n_features): First embedded view.
z2: torch.Tensor of shape (batch_size, n_features): Second embedded view.

Returns:

loss: torch.Tensor: The InfoNCE loss computed between z1 and z2.