MINC-SSL: Representation Learning via Non-Contrastive Mutual Information

This repository contains a Python library implementing the core components of the MINC (Mutual Information Non-Contrastive) self-supervised learning method from the paper "Representation Learning via Non-Contrastive Mutual Information" (arXiv:2504.16667).

The library provides:

• The MINCLoss PyTorch module.
• Utility functions for Exponential Moving Average (EMA) updates of the target network and the auxiliary matrix (Λ).

This library is designed to be integrated into a user's existing PyTorch training pipeline. It does not include the full training loop, data loading, augmentation pipeline, or model architecture (encoder + projector), as these can vary depending on the specific use case.

Installation

1. Clone this repository:

   git clone https://github.com/danielgaskins/minc-ssl.git
   cd minc-ssl

2. Install the package using pip:

   pip install .

   (Optional: For development, use pip install -e .)

Usage

Here's a conceptual example of how to use the minc_ssl library in a PyTorch training script:

import torch
import torch.nn as nn
import torch.optim as optim
import copy

from minc_ssl import MINCLoss, update_target_network, update_aux_matrix

# --- 1. Define your Model ---
# Your model should consist of an encoder (e.g., ResNet) followed by a projector MLP.
# The forward method should return the projector output (embeddings).
# Embeddings should ideally be L2 normalized before being passed to the loss/updates.
class MySelfSupervisedModel(nn.Module):
    def __init__(self, encoder, projector):
        super().__init__()
        self.encoder = encoder
        self.projector = projector

    def forward(self, x):
        features = self.encoder(x)
        embeddings = self.projector(features)
        # Normalize embeddings (important for MINC loss and Lambda update)
        embeddings = F.normalize(embeddings, dim=1, p=2)
        return embeddings

# Example placeholders (replace with actual models)
encoder = nn.Sequential(nn.Conv2d(3, 16, 3, 1), nn.AdaptiveAvgPool2d(1), nn.Flatten()) # Dummy encoder
projector = nn.Sequential(nn.Linear(16, 32), nn.ReLU(), nn.Linear(32, 64)) # Dummy projector
embedding_dim = 64 # Output dimension of the projector

online_network = MySelfSupervisedModel(encoder, projector)
target_network = copy.deepcopy(online_network)

# Freeze target network parameters
for param in target_network.parameters():
    param.requires_grad = False

# --- 2. Initialize MINC Components ---
# Auxiliary matrix Lambda (initialized to zero)
Lambda = torch.zeros(embedding_dim, embedding_dim)

# MINC Loss criterion
minc_criterion = MINCLoss(alpha=2.0, scale=10.0) # Use paper's recommended alpha=2 and a chosen scale

# Optimizer (LARS is used in the paper, but Adam or SGD with momentum also work)
optimizer = optim.Adam(online_network.parameters(), lr=1e-3)

# EMA decay rates
target_ema_decay = 0.996 # gamma
aux_ema_decay = 0.8    # beta

# --- 3. Training Loop Sketch ---
# Assuming you have a DataLoader `train_loader` that yields batches of augmented views `(x1, x2)`
# and you are training on a device (e.g., 'cuda' or 'cpu')
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
online_network.to(device)
target_network.to(device)
Lambda = Lambda.to(device)
minc_criterion.to(device)


# Inside your training loop over epochs:
for epoch in range(num_epochs):
    for (x1, x2) in train_loader: # Assuming data loader provides pairs of augmented views
        x1, x2 = x1.to(device), x2.to(device)

        # Forward pass through online network
        z1_online = online_network(x1)
        z2_online = online_network(x2) # Need embeddings from the second view for the loss

        # Forward pass through target network (no gradients)
        with torch.no_grad():
            z1_target = target_network(x1) # Use the same view as z1_online for target

        # Calculate Loss
        # Loss takes z_target (from view 1 target net) and z_current (from view 2 online net)
        loss = minc_criterion(z_target, z2_online, Lambda)

        # Backward pass and optimize
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()

        # Update target network EMA
        update_target_network(online_network, target_network, target_ema_decay)

        # Update auxiliary matrix Lambda EMA
        # As per Algorithm 1 and Eq (22), Lambda is updated using embeddings from the target network
        update_aux_matrix(Lambda, z1_target, aux_ema_decay)

        # ... (Log loss, etc.)

    # ... (Handle learning rate scheduling, validation, etc.)

Paper Reference

This library is based on the paper:

"Representation Learning via Non-Contrastive Mutual Information" Zhaohan Daniel Guo, Bernardo Avila Pires, Khimya Khetarpal, Dale Schuurmans, Bo Dai arXiv preprint arXiv:2504.16667 (2025) https://arxiv.org/abs/2504.16667

License

This project is licensed under the MIT License - see the LICENSE file for details.