  Data augmentation is a widely used technique and an essential ingredient in
the recent advance in self-supervised representation learning. By preserving
the similarity between augmented data, the resulting data representation can
improve various downstream analyses and achieve state-of-the-art performance in
many applications. Despite the empirical effectiveness, most existing methods
lack theoretical understanding under a general nonlinear setting. To fill this
gap, we develop a statistical framework on a low-dimension product manifold to
model the data augmentation transformation. Under this framework, we introduce
a new representation learning method called augmentation invariant manifold
learning and design a computationally efficient algorithm by reformulating it
as a stochastic optimization problem. Compared with existing self-supervised
methods, the new method simultaneously exploits the manifold's geometric
structure and invariant property of augmented data and has an explicit
theoretical guarantee. Our theoretical investigation characterizes the role of
data augmentation in the proposed method and reveals why and how the data
representation learned from augmented data can improve the $k$-nearest neighbor
classifier in the downstream analysis, showing that a more complex data
augmentation leads to more improvement in downstream analysis. Finally,
numerical experiments on simulated and real datasets are presented to
demonstrate the merit of the proposed method.
