Efficient Distribution Matching of Representations via Noise-Injected Deep InfoMax

2024-10-09

arXiv.org

AI Summary

In self-supervised representation learning, aligning encoder output distributions with predefined priors (e.g., uniform or Gaussian) remains challenging. To address this, we propose the Noise-Injected Normalized Encoder (NINE) within the Deep InfoMax framework: after normalizing representations, independent random noise is injected—requiring no modification to the mutual information maximization objective, no additional discriminator, and no extra loss terms—yet enabling exact matching to any absolutely continuous prior distribution. This work is the first to integrate controllable distribution modeling into the InfoMax paradigm while preserving model simplicity and training stability. We validate NINE on downstream tasks including generative modeling, disentangled representation learning, and anomaly detection. Experimental results demonstrate substantial improvements in distribution matching fidelity, with only negligible computational overhead.

Technology Category

• Machine Learning: Representation Learning
• Computer Vision: Representation Learning for Vision
• Search and Optimization: Learning to Search

Application Category

• Graph Algorithms and Modeling for the Web: Graph embeddings and representation learning for Web-related graphs
• Search and Retrieval-Augmented AI: Web learning to rank, online learning, and counterfactual learning for ranking
• User Modeling, Personalization and Recommendation: Fairness-aware retrieval and ranking

Abstract

Deep InfoMax (DIM) is a well-established method for self-supervised representation learning (SSRL) based on maximization of the mutual information between the input and the output of a deep neural network encoder. Despite the DIM and contrastive SSRL in general being well-explored, the task of learning representations conforming to a specific distribution (i.e., distribution matching, DM) is still under-addressed. Motivated by the importance of DM to several downstream tasks (including generative modeling, disentanglement, outliers detection and other), we enhance DIM to enable automatic matching of learned representations to a selected prior distribution. To achieve this, we propose injecting an independent noise into the normalized outputs of the encoder, while keeping the same InfoMax training objective. We show that such modification allows for learning uniformly and normally distributed representations, as well as representations of other absolutely continuous distributions. Our approach is tested on various downstream tasks. The results indicate a moderate trade-off between the performance on the downstream tasks and quality of DM.

Problem

• Self-supervised Learning
• Distribution Alignment
• Performance Enhancement

Innovation

• Improved Depth Information Maximization
• Noise Introduction
• Pattern Matching Enhancement