Contrastive Mutual Information Learning: Toward Robust Representations without Positive-Pair Augmentations

Micha Livne

Keywords: representation learning, contrastive learning, latent variable models, probabilistic models

TL;DR: We introduce cMIM, a contrastive extension of the probabilistic auto-encoder MIM, that unifies generative and discriminative representation learning without contrastive data augmentation and is robust to small batch sizes.

Abstract:

Learning representations that transfer well to diverse downstream tasks remains a central challenge in representation learning. Existing paradigms---contrastive learning, self-supervised masking, and denoising auto-encoders---balance this challenge with different trade-offs. We introduce the {contrastive Mutual Information Machine} (cMIM), a probabilistic framework that extends the Mutual Information Machine (MIM) with a contrastive objective. While MIM maximizes mutual information between inputs and latents and promotes clustering of codes, it falls short on discriminative tasks. cMIM addresses this gap by imposing global discriminative structure while retaining MIM’s generative fidelity.

Our contributions are threefold. First, we propose cMIM, a contrastive extension of MIM that removes the need for positive data augmentation and is substantially less sensitive to batch size than InfoNCE. Second, we introduce {informative embeddings}, a general technique for extracting enriched features from encoder--decoder models that boosts discriminative performance without additional training and applies broadly beyond MIM. Third, we provide empirical evidence across vision and molecular benchmarks showing that cMIM consistently outperforms MIM and InfoNCE on classification and regression tasks while preserving competitive reconstruction quality.

These results position cMIM as a unified framework for representation learning, advancing the goal of models that serve both discriminative and generative applications effectively.

Primary Area: unsupervised, self-supervised, semi-supervised, and supervised representation learning

Submission Number: 6044