Masked prediction tasks: a parameter identifiability view

Abstract: The vast majority of work in self-supervised learning, both theoretical and empirical (though mostly the latter), have largely focused on recovering good features for downstream tasks, with the definition of "good" often being intricately tied to the downstream task itself. This lens is undoubtedly very interesting, but suffers from the problem that there isn't a "canonical" set of downstream tasks to focus on -- in practice, this problem is usually resolved by competing on the benchmark dataset du jour. In this paper, we present an alternative lens: one of parameter identifiability. More precisely, we consider data coming from a parametric probabilistic model, and train a self-supervised learning predictor with a suitably chosen parametric form. Then, we ask whether we can read off the ground truth parameters of the probabilistic model from the optimal predictor. We focus on the widely used self-supervised learning method of predicting masked tokens, which is popular for both natural languages and visual data. While incarnations of this approach have already been successfully used for simpler probabilistic models (e.g. learning fully-observed undirected graphical models), we focus instead on latent-variable models capturing sequential structures -- namely Hidden Markov Models with both discrete and conditionally Gaussian observations. We show that there is a rich landscape of possibilities, out of which some prediction tasks yield identifiability, while others do not. Our results, borne of a theoretical grounding of self-supervised learning, could thus potentially beneficially inform practice. Moreover, we uncover close connections with uniqueness of tensor rank decompositions -- a widely used tool in studying identifiability through the lens of the method of moments.

• The paper introduces a novel framework that shifts self-supervised learning from feature recovery to parameter identifiability in probabilistic models.
• It reveals that pairwise prediction tasks fail for discrete HMMs but enable unique parameter recovery in conditionally-Gaussian HMMs with the aid of tensor decomposition.
• The study’s findings suggest robust design strategies for self-supervised algorithms and pave the way for further research on generative models and recovery methods.

"Masked Prediction Tasks: A Parameter Identifiability View" (2202.09305)

Introduction

This paper revisits the field of self-supervised learning, with a focus on masked prediction tasks, from a parameter identifiability standpoint. Traditionally, masked prediction tasks are leveraged to develop models that predict missing parts of data, such as words in a sentence or segments of an image, based on the remaining visible parts. This approach has gained traction due to its effectiveness in various domains, including natural language processing and computer vision. However, the paper challenges the superficial focus on downstream task performance and aims to provide a foundational understanding in terms of parameter recovery within probabilistic models.

The authors introduce a novel framework for evaluating self-supervised learning through parameter identifiability. This shifts the focus from feature recovery to determining if the ground truth parameters of a data-generating model can be extracted from models trained with masked prediction tasks. The research particularly zeroes in on the identifiability in Hidden Markov Models (HMMs) with discrete and conditionally Gaussian observations, laying out both theoretical underpinnings and practical implications.

Methodology

The paper employs self-supervised learning to a parametric probabilistic model and explores the identifiability of parameters through masked prediction tasks. Specifically, it considers data generated from fully discrete HMMs and conditionally Gaussian HMMs (\GHMMs). The core question addressed is whether the optimal predictors from these tasks can uniquely recover the parameters of the underlying probabilistic model.

Discrete Hidden Markov Models

For discrete HMMs, the study scrutinizes parameter recovery using pairwise prediction tasks, demonstrating non-identifiability. It reveals that masked prediction, when focused on single tokens, fails to yield uniqueness in parameter recovery due to the product nature of stochastic matrices.

Conditionally-Gaussian Hidden Markov Models

In contrast, the paper finds that prediction tasks on \GHMMs have stronger identifiability properties. This is credited to the complexity of these models' posterior functions, which encapsulate richer information from the data, thus leading to potentially unique recovery of parameters.

Tensor Decompositions

The paper innovatively employs tensor decomposition techniques to address the identifiability problem. By leveraging Kruskal's theorem on tensor rank decompositions, the authors bridge theoretical insights with practical parameter recovery strategies, setting the stage for efficient algorithmic design within self-supervised learning.

Results

The paper presents compelling evidence that masked prediction tasks can proffer parameter identifiability in certain latent sequence models but not all. It establishes:

• Non-Identifiability in Discrete HMMs: Pairwise masked prediction tasks fail to guarantee parameter recovery.
• Identifiability in \GHMMs: Complex posterior functions intrinsic to \GHMMs enable unique parameter recovery via pairwise prediction.
• Tensor Methods: Utilization of tensor decomposition aids in circumventing non-uniqueness issues encountered in matrix factorizations.

Implications and Future Work

The exploration opens avenues for reimagining self-supervised learning, emphasizing the significance of model structure and prediction tasks. It suggests a paradigm where masked prediction can be a potent tool for understanding and isolating fundamental properties of the data-generating model.

For future research, the paper points towards investigating a broader class of generative models, examining the robustness and sample complexity of various prediction tasks, and extending the theoretical basis with real-world applications.

Conclusion

In essence, this paper provides a structured and incisive look into parameter identifiability within masked prediction tasks. It contributes to the growing theoretical framework underpinning self-supervised learning and posits that rigorous parameter recovery can indeed guide the practical deployment and enhancement of these models across diverse scientific and engineering domains. By marrying theoretical nuance with empirical pragmatism, the research advances both understanding and application in the AI landscape.