An Analysis of Self-supervised Pre-training with Dependent Samples

Abstract: Self-supervised learning relies on so-called data augmentations $φ(x)$ of unlabeled datapoints $x$ --- for example, masking random pixels in an image $x$ --- that should leave the label of $x$ invariant and are often used to learn a lower-complexity invariant subspace $\cal V$ for downstream tasks. In practice, such augmentations ${ φ_l(x_i) }$ are pooled together to learn $\cal V$, despite obvious inter-dependencies between different augmentations $φ_l(x), φ_k(x)$ of the same datapoint $x$. However, theoretical works on the subject typically consider procedures that avoid such dependencies, and are therefore limited to operate on smaller subsets of independent data. We show in this work that pooling augmentations together, despite inter-dependencies, is a better alternative than the baseline of partitioning the data into subsets of independent data. More precisely, in the context of estimating $\cal V$, the statistical estimation error bounds for pooling are never worse than the partitioning baseline, and in some cases --- such as masking or noise injection-based augmentations over a shallow neural network --- naive pooling leads to faster rates in terms of the number of augmentations. The benefits of pooling are particularly prominent when the correlations between different augmentations $φ_l(x), φ_k(x)$ have mild effects on estimation or help decrease the estimation variance. The analysis, therefore, yields new insights into the success of pooling augmented samples in self-supervised pre-training, and provides an intuition behind the practical preference towards using many augmentations.