On the Asymptotics of Self-Supervised Pre-training: Two-Stage M-Estimation and Representation Symmetry

TL;DR

This paper develops an asymptotic theory for self-supervised pre-training using two-stage M-estimation, addressing symmetry issues with Riemannian geometry, and provides insights into the interaction between pre-training and fine-tuning.

Contribution

It introduces a novel asymptotic framework for pre-training that accounts for representation symmetry and links pre-training to downstream tasks through orbit-invariance.

Findings

Derived the limiting distribution of downstream test risk.

Applied theory to spectral pre-training, factor models, and Gaussian mixtures.

Achieved substantial improvements over prior bounds in specific cases.

Abstract

Self-supervised pre-training, where large corpora of unlabeled data are used to learn representations for downstream fine-tuning, has become a cornerstone of modern machine learning. While a growing body of theoretical work has begun to analyze this paradigm, existing bounds leave open the question of how sharp the current rates are, and whether they accurately capture the complex interaction between pre-training and fine-tuning. In this paper, we address this gap by developing an asymptotic theory of pre-training via two-stage M-estimation. A key challenge is that the pre-training estimator is often identifiable only up to a group symmetry, a feature common in representation learning that requires careful treatment. We address this issue using tools from Riemannian geometry to study the intrinsic parameters of the pre-training representation, which we link with the downstream predictor through a notion of orbit-invariance, precisely characterizing the limiting distribution of the downstream test risk. We apply our main result to several case studies, including spectral pre-training, factor models, and Gaussian mixture models, and obtain substantial improvements in problem-specific factors over prior art when applicable.