Multi-objective models#

Multiobjective contrastive learning.

Starting in CEBRA 0.6.0, we have added support for subspace contrastive learning. This is a method for training models that are able to learn multiple subspaces of the feature space simultaneously.

Subspace contrastive learning requires to use specialized models and criterions. This module specifies a test of classes required for training CEBRA models with multiple objectives. The objectives are defined by the wrapper class cebra.models.multicriterions.MultiCriterions.

Two solvers are currently implemented:

cebra.solver.multiobjective.ContrastiveMultiobjectiveSolverxCEBRA
cebra.solver.multiobjective.SupervisedMultiobjectiveSolverxCEBRA

Support for training CEBRA with multiple criteria.

Note

This module was introduced in CEBRA 0.6.0.

class cebra.models.multicriterions.MultiCriterions(*args: Any, **kwargs: Any)#

Bases: Module

A module for handling multiple loss functions with different criteria.

This module allows combining multiple loss functions, each operating on specific slices of the input data. It supports both supervised and contrastive learning modes.

Parameters:

losses – A list of dictionaries containing loss configurations. Each dictionary should have: - ‘indices’: Tuple of (start, end) indices for the data slice - ‘supervised_loss’: Dict with loss config for supervised mode - ‘contrastive_loss’: Dict with loss config for contrastive mode Loss configs should contain: - ‘name’: Name of the loss function - ‘kwargs’: Optional parameters for the loss function
mode – Either “supervised” or “contrastive” to specify the training mode

The loss functions can be from torch.nn or custom implementations from cebra.models.criterions. Each criterion is applied to its corresponding slice of the input data during forward pass.

Example

>>> import torch
>>> from cebra.data.datatypes import Batch
>>> # Define loss configurations for a hybrid model with both contrastive and supervised losses
>>> losses = [
...     {
...         'indices': (0, 10),  # First 10 dimensions
...         'contrastive_loss': {
...             'name': 'InfoNCE',  # Using CEBRA's InfoNCE loss
...             'kwargs': {'temperature': 1.0}
...         },
...         'supervised_loss': {
...             'name': 'nn.MSELoss',  # Using PyTorch's MSE loss
...             'kwargs': {}
...         }
...     },
...     {
...         'indices': (10, 20),  # Next 10 dimensions
...         'contrastive_loss': {
...             'name': 'InfoNCE',  # Using CEBRA's InfoNCE loss
...             'kwargs': {'temperature': 0.5}
...         },
...         'supervised_loss': {
...             'name': 'nn.L1Loss',  # Using PyTorch's L1 loss
...             'kwargs': {}
...         }
...     }
... ]
>>> # Create sample predictions (2 batches of 32 samples each with 10 features)
>>> ref1 = torch.randn(32, 10)
>>> pos1 = torch.randn(32, 10)
>>> neg1 = torch.randn(32, 10)
>>> ref2 = torch.randn(32, 10)
>>> pos2 = torch.randn(32, 10)
>>> neg2 = torch.randn(32, 10)
>>> predictions = (
...     Batch(reference=ref1, positive=pos1, negative=neg1),
...     Batch(reference=ref2, positive=pos2, negative=neg2)
... )
>>> # Create multi-criterion module in contrastive mode
>>> multi_loss = MultiCriterions(losses, mode="contrastive")
>>> # Forward pass with multiple predictions
>>> losses = multi_loss(predictions)  # Returns list of loss values
>>> assert len(losses) == 2  # One loss per criterion

class cebra.models.multiobjective.MultiobjectiveModel(module, dimensions, renormalize=False, output_mode='overlapping', append_last_dimension=False)#

Bases: Module

Wrapper around contrastive learning models to all training with multiple objectives

Multi-objective training splits the last layer’s feature representation into multiple chunks, which are then used for individual training objectives.

Parameters:

module (Module) – The module to wrap
dimensions (Tuple[int]) – A tuple of dimension values to extract from the model’s feature embedding.
renormalize (bool) – If True, the individual feature slices will be re-normalized before getting returned—this option only makes sense in conjunction with a loss based on the cosine distance or dot product.
output_mode (str) – A mode as defined in MultiobjectiveModel.Mode. Overlapping means that when dimensions are set to (x0, x1, …)`, features will be extracted from 0:x0, 0:x1, .... When mode is set to separate, features are extracted from x0:x1, x1:x2, ....
append_last_dimension (bool) – Defaults to True, and will allow to omit the last dimension in the dimensions argument (which should be equal to the output dimension) of the given model.

Note

This model will be deprecated in a future version. Please use the functionality in cebra.models.multiobjective instead, which provides more versatile multi-objective training capabilities. Instantiation of this model will raise a deprecation warning. The new model is cebra.models.multiobjective.SubspaceMultiobjectiveModel which allows for unlimited subspaces and better configuration of the feature ranges.

class Mode#

Bases: object

Mode for slicing and potentially normalizing the output embedding.

The options are:

OVERLAPPING: When dimensions are set to (x0, x1, …)`, features will be extracted from 0:x0, 0:x1, ....
SEPARATE: Features are extracted from x0:x1, x1:x2, ...

is_valid(mode)#

Check if a given string representation is valid.

Parameters:: mode – String representation of the mode.
Returns:: True for a valid representation, False otherwise.

property get_offset#: See cebra.models.model.Model.get_offset().

property num_output#: See cebra.models.model.Model.num_output.

forward(inputs)#

Compute multiple embeddings for a single signal input.

Parameters:: inputs – The input tensor
Returns:: A tuple of tensors which are sliced according to self.feature_ranges if renormalize is set to true, each of the tensors will be normalized across the first (feature) dimension.

class cebra.models.multiobjective.SubspaceMultiobjectiveModel(module, feature_ranges, renormalize, split_outputs=True)#

Bases: Module

Wrapper around contrastive learning models to all training with multiple objectives

Multi-objective training splits the last layer’s feature representation into multiple chunks, which are then used for individual training objectives.

Parameters:

module (Module) – The module to wrap
dimensions – A tuple of dimension values to extract from the model’s feature embedding.
renormalize (bool) – If True, the individual feature slices will be re-normalized before getting returned—this option only makes sense in conjunction with a loss based on the cosine distance or dot product.

property get_offset#: See cebra.models.model.Model.get_offset().

property num_output#: See cebra.models.model.Model.num_output.

forward(inputs)#

Compute multiple embeddings for a single signal input.

Parameters:: inputs – The input tensor
Returns:: A tuple of tensors which are sliced according to self.feature_ranges if renormalize is set to true, each of the tensors will be normalized across the first (feature) dimension.

class cebra.models.multiobjective.SubspaceMultiobjectiveConvolutionalModel(module, feature_ranges, renormalize, split_outputs=True)#: Bases: SubspaceMultiobjectiveModel, ConvolutionalModelMixin