Deep learning for NeuroImaging in Python.

Note

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

class nidl.losses.yaware_infonce.YAwareInfoNCE(kernel: str = 'gaussian', bandwidth: float | list[float] | ndarray | KernelMetric = 1.0, temperature: float = 0.1)[source]¶

Bases: Module

Implementation of y-Aware InfoNCE loss [1].

Compute the y-Aware InfoNCE loss, which integrates auxiliary information into contrastive learning by weighting sample pairs.

Given a mini-batch of size N, two embeddings z1 and z2 representing two views of the same samples and a weighting matrix W computed using auxiliary variables y, the loss is:

$\ell = -\frac{1}{N} \sum_{i,j} \frac{W_{i,j}}{\sum_{k=1}^{N} W_{i, k}} \log \frac{\exp(\text{sim}(z_1^{i}, z_2^{j}) / \tau)} {\sum_{k=1}^{N} \exp(\text{sim}(z_1^{i}, z_2^{k}) / \tau)}$

where sim is the cosine similarity and τ is the temperature.

To compute W, a kernel K is specified (e.g. Gaussian) as well as a bandwidth H and the weighting matrix is:

$W_{i,j} = K\left( H^{-\frac{1}{2}} (y_i-y_j) \right)$

Parameters:

kernel : str in {‘gaussian’, ‘epanechnikov’, ‘exponential’, ‘linear’, ‘cosine’}, default=’gaussian’

Kernel to compute the weighting matrix between auxiliary variables. See PhD thesis, Dufumier 2022 page 94-95.

bandwidth : Union[float, int, List[float], array, KernelMetric], default=1.0

The method used to calculate the bandwidth (“sigma” in [1]) between auxiliary variables:

If bandwidth is a scalar (int or float), it sets the bandwidth to a diagnonal matrix with equal values.

If bandwidth is a 1d array, it sets the bandwidth to a diagonal matrix and it must be of size equal to the number of features in y.

If bandwidth is a 2d array, it must be of shape (n_features, n_features) where n_features is the number of features in y.

If bandwidth is KernelMetric, it uses the pairwise method to compute the similarity matrix between auxiliary variables.

temperature : float, default=0.1

Temperature used to scale the dot-product between embedded vectors

References

[1] Contrastive Learning with Continuous Proxy Meta-Data for 3D MRI Classification, MICCAI 2021

Initialize internal Module state, shared by both nn.Module and ScriptModule.

forward(z1: Tensor, z2: Tensor, labels: Tensor | None = None)[source]¶

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.

labels : Optional[torch.Tensor] of shape (batch_size, n_labels)

Auxiliary variables associated to the input data. If None, the standard InfoNCE loss is returned.

Returns:

loss : torch.Tensor

The y-Aware InfoNCE loss computed between z1 and z2.