On Masked Pre-training and the Marginal Likelihood

TL;DR

This paper reveals that masked pre-training optimizes the model's marginal likelihood, providing a theoretical basis for its success and suggesting new directions for Bayesian self-supervised learning.

Contribution

It establishes a theoretical link between masked pre-training and marginal likelihood maximization, enhancing understanding of its effectiveness.

Findings

Masked pre-training corresponds to maximizing marginal likelihood.

Theoretical insights explain the success of masked pre-training.

Empirical results confirm the theory in large language models.

Abstract

Masked pre-training removes random input dimensions and learns a model that can predict the missing values. Empirical results indicate that this intuitive form of self-supervised learning yields models that generalize very well to new domains. A theoretical understanding is, however, lacking. This paper shows that masked pre-training with a suitable cumulative scoring function corresponds to maximizing the model's marginal likelihood, which is de facto the Bayesian model selection measure of generalization. Beyond shedding light on the success of masked pre-training, this insight also suggests that Bayesian models can be trained with appropriately designed self-supervision. Empirically, we confirm the developed theory and explore the main learning principles of masked pre-training in large language models.