Understanding Augmentation-based Self-Supervised Representation Learning via RKHS Approximation and Regression

TL;DR

This paper provides a theoretical framework for understanding how data augmentation influences self-supervised learning, using RKHS approximation and regression to analyze generalization and the effects of different augmentations.

Contribution

It introduces a geometric and statistical analysis of augmentation-based pretraining, deriving bounds that separate model and augmentation effects, and introduces augmentation complexity as a key factor.

Findings

Generalization bounds for augmentation-based learning free of model complexity

Decomposition of prediction error into estimation and approximation errors

Quantitative comparison of augmentations via augmentation complexity

Abstract

Data augmentation is critical to the empirical success of modern self-supervised representation learning, such as contrastive learning and masked language modeling. However, a theoretical understanding of the exact role of augmentation remains limited. Recent work has built the connection between self-supervised learning and the approximation of the top eigenspace of a graph Laplacian operator, suggesting that learning a linear probe atop such representation can be connected to RKHS regression. Building on this insight, this work delves into a statistical analysis of augmentation-based pretraining. Starting from the isometry property, a geometric characterization of the target function given by the augmentation, we disentangle the effects of the model and the augmentation, and prove two generalization bounds that are free of model complexity. Our first bound works for an arbitrary encoder, where the prediction error is decomposed as the sum of an estimation error incurred by fitting a linear probe with RKHS regression, and an approximation error entailed by RKHS approximation. Our second bound specifically addresses the case where the encoder is near-optimal, that is it approximates the top-d eigenspace of the RKHS induced by the augmentation. A key ingredient in our analysis is the augmentation complexity, which we use to quantitatively compare different augmentations and analyze their impact on downstream performance.