Optimizing the Noise in Self-Supervised Learning: from Importance Sampling to Noise-Contrastive Estimation

TL;DR

This paper investigates the optimal noise distribution in self-supervised learning, revealing that it often differs from the data distribution and that the benefits of using the optimal noise are modest compared to simpler choices.

Contribution

It provides a theoretical and empirical analysis of noise distribution optimality in Noise-Contrastive Estimation, challenging the common assumption that it should match the data distribution.

Findings

Optimal noise differs from data distribution when the energy is unknown.

Using the data distribution as noise is nearly optimal in many cases.

Theoretical insights connect noise choice to estimator efficiency.

Abstract

Self-supervised learning is an increasingly popular approach to unsupervised learning, achieving state-of-the-art results. A prevalent approach consists in contrasting data points and noise points within a classification task: this requires a good noise distribution which is notoriously hard to specify. While a comprehensive theory is missing, it is widely assumed that the optimal noise distribution should in practice be made equal to the data distribution, as in Generative Adversarial Networks (GANs). We here empirically and theoretically challenge this assumption. We turn to Noise-Contrastive Estimation (NCE) which grounds this self-supervised task as an estimation problem of an energy-based model of the data. This ties the optimality of the noise distribution to the sample efficiency of the estimator, which is rigorously defined as its asymptotic variance, or mean-squared error. In the special case where the normalization constant only is unknown, we show that NCE recovers a family of Importance Sampling estimators for which the optimal noise is indeed equal to the data distribution. However, in the general case where the energy is also unknown, we prove that the optimal noise density is the data density multiplied by a correction term based on the Fisher score. In particular, the optimal noise distribution is different from the data distribution, and is even from a different family. Nevertheless, we soberly conclude that the optimal noise may be hard to sample from, and the gain in efficiency can be modest compared to choosing the noise distribution equal to the data's.