Self-Supervised Learning with Data Augmentations Provably Isolates Content from Style

1 Introduction . Learning good representations of high-dimensional observations from large amounts of unlabelled data is widely recognised as an important step for more capable and data-efficient learning systems [ 10 , 72 ] . Over the last decade , self-supervised learning ( SSL ) has emerged as the dominant paradigm for such unsupervised representation learning [ 1 , 20 , 21 , 34 , 41 , 47 , 48 , 90 , 91 , 115 , 122 , 125 , 126 ] . The main idea behind SSL is to extract a supervisory signal from unlabelled observations by leveraging known structure of the data , which allows for the application of supervised learning techniques . A common approach is to directly predict some part of the observation from another part ( e.g. , future from past , or original from corruption ) , thus forcing the model to learn a meaningful representation in the process . While this technique has shown remarkable success in natural language processing [ 13 , 23 , 30 , 81 , 84 , 86 , 95 , 99 ] and speech recognition [ 5 , 6 , 100 , 104 ] , where a finite dictionary allows one to output a distribution over the missing part , such predictive SSL methods are not easily applied to continuous or high-dimensional domains such as vision . Here , a common approach is to learn a joint embedding of similar observations or views such that their representation is close [ 7 , 12 , 22 , 44 ] . Different views can come , for example , from different modalities ( text & speech ; video & audio ) or time points . As still images lack such multi-modality or temporal structure , recent advances in representation learning have relied on generating similar views by means of data augmentation . ⇤Joint first author . †Joint senior author . Correspondence to : jvk @ tue.mpg.de Code available at : https : //www.github.com/ysharma1126/ssl_identifiability 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) . In order to be useful , data augmentation is thought to require the transformations applied to generate additional views to be generally chosen to preserve the semantic characteristics of an observation , while changing other “ nuisance ” aspects . While this intuitively makes sense and has shown remarkable empirical results , the success of data augmentation techniques in practice is still not very well understood from a theoretical perspective—despite some efforts [ 17 , 19 , 28 ] . In the present work , we seek to better understand the empirical success of SSL with data augmentation by formulating the generative process as a latent variable model ( LVM ) and studying identifiability of the representation , i.e. , under which conditions the ground truth latent factors can provably be inferred from the data [ 77 ] . Related work and its relation to the current . Prior work on unsupervised representation learning from an LVM perspective often postulates mutually independent latent factors : this independence assumption is , for example , at the heart of independent component analysis ( ICA ) [ 24 , 56 ] and disentanglement [ 10 , 14 , 18 , 49 , 65 , 71 ] . Since it is impossible to identify the true latent factors without any supervisory signal in the general nonlinear case [ 57 , 82 ] , recent work has turned to weakly- or self-supervised approaches which leverage additional information in the form of multiple views [ 39 , 83 , 108 , 129 ] , auxiliary variables [ 58 , 63 ] , or temporal structure [ 45 , 54 , 55 , 69 ] . To identify or disentangle the individual independent latent factors , it is typically assumed that there is a chance that each factor changes across views , environments , or time points . Our work—being directly motivated by common practices in SSL with data augmentation—differs from these works in the following two key aspects ( see Fig . 1 for an overview ) . First , we do not assume independence and instead allow for both nontrivial statistical and causal relations between latent variables . This is in line with a recently proposed [ 105 ] shift towards causal representation learning [ 40 , 76 , 85 , 87 , 106 , 107 , 112 , 123 , 127 ] , motivated by the fact that many underlying variables of interest may not be independent but causally related to each other.1 Second , instead of a scenario wherein all latent factors may change as a result of augmentation , we assume a partition of the latent space into two blocks : a content block which is shared or invariant across different augmented views , and a style block that may change . This is aligned with the notion that augmentations leave certain semantic aspects ( i.e. , content ) intact and only affect style , and is thus a more appropriate assumption for studying SSL . In line with earlier work [ 39 , 54 , 57 , 58 , 63 , 69 , 82 , 83 , 129 ] , we focus on the setting of continuous ground-truth latents , though we believe our results to hold more broadly . Structure and contributions . Following a review of SSL with data augmentation and identifiability theory ( § 2 ) , we formalise the process of data generation and augmentation as an LVM with content and style variables ( § 3 ) . We then establish identifiability results of the invariant content partition ( § 4 ) , validate our theoretical insights experimentally ( § 5 ) , and discuss our findings and their limitations in the broader context of SSL with data augmentation ( § 6 ) . We highlight the following contributions : • we prove that SSL with data augmentations identifies the invariant content partition of the representation in generative ( Thm . 4.2 ) and discriminative learning with invertible ( Thm . 4.3 ) and non-invertible encoders with entropy regularisation ( Thm . 4.4 ) ; in particular , Thm . 4.4 provides a theoretical justification for the empirically observed effectiveness of contrastive SSL methods that use data augmentation and InfoNCE [ 91 ] as an objective , such as SimCLR [ 20 ] ; • we show that our theory is consistent with results in simulating statistical dependencies within blocks of content and style variables , as well as with style causally dependent on content ( § 5.1 ) ; • we introduce Causal3DIdent , a dataset of 3D objects which allows for the study of identifiability in a causal representation learning setting , and use it to perform a systematic study of data augmentations used in practice , yielding novel insights on what particular data augmentations are truly isolating as invariant content and discarding as varying style when applied ( § 5.2 ) . 1E.g. , [ 69 ] , Fig . 11 where dependence between latents was demonstrated for multiple natural video data sets . 2 Preliminaries and background . Self-supervised representation learning with data augmentation . Given an unlabelled dataset of observations ( e.g. , images ) x , data augmentation techniques proceed as follows . First , a set of observation-level transformations t 2 T are specified together with a distribution pt over T . Both T and pt are typically designed using human intelligence and domain knowledge with the intention of not changing the semantic characteristics of the data ( which arguably constitutes a form of weak supervision ) .2 For images , for example , a common choice for T are combinations of random crops [ 113 ] , horizontal or vertical flips , blurring , colour distortion [ 52 , 113 ] , or cutouts [ 31 ] ; and pt is a distribution over the parameterisation of these transformations , e.g. , the centre and size of a crop [ 20 , 31 ] . For each observation x , a pair of transformations t , t0 ⇠ pt is sampled and applied separately to x to generate a pair of augmented views ( x̃ , x̃0 ) = ( t ( x ) , t0 ( x ) ) . The joint-embedding approach to SSL then uses a pair of encoder functions ( g , g0 ) , i.e . deep nets , to map the pair ( x̃ , x̃0 ) to a typically lower-dimensional representation ( z̃ , z̃0 ) = ( g ( x̃ ) , g0 ( x̃0 ) ) . Often , the two encoders are either identical , g = g0 , or directly related ( e.g. , via shared parameters or asynchronous updates ) . Then , the encoder ( s ) ( g , g0 ) are trained such that the representations ( z̃ , z̃0 ) are “ close ” , i.e. , such that sim ( z̃ , z̃0 ) is large for some similarity metric sim ( · ) , e.g. , the cosine similarity [ 20 , 129 ] , or negative L2 norm [ 129 ] . The advantage of directly optimising for similarity in representation space over generative alternatives is that reconstruction can be very challenging for high-dimensional data . The disadvantage is the problem of collapsed representations.3 To avoid collapsed representations and force the encoder ( s ) to learn a meaningful representation , two main families of approaches have been used : ( i ) contrastive learning ( CL ) [ 20 , 47 , 48 , 91 , 115 , 126 ] ; and ( ii ) regularisation-based SSL [ 21 , 41 , 128 ] . The idea behind CL is to not only learn similar representations for augmented views ( x̃i , x̃0i ) of the same xi , or positive pairs , but to also use other observations xj ( j 6= i ) to contrast with , i.e. , to enforce a dissimilar representation across negative pairs ( x̃i , x̃0j ) . In other words , CL pulls representations of positive pairs together , and pushes those of negative pairs apart . Since both aims can not be achieved simultaneously with a constant representation , collapse is avoided . A popular CL objective function ( used , e.g. , in SimCLR [ 20 ] ) is InfoNCE [ 91 ] ( based on noise-contrastive estimation [ 42 , 43 ] ) : LInfoNCE ( g ; ⌧ , K ) = E { xi } Ki=1⇠px h PK i=1 log exp { sim ( z̃i , z̃0i ) /⌧ } PK j=1 exp { sim ( z̃i , z̃0j ) /⌧ } i ( 1 ) where z̃ = Et⇠pt [ g ( t ( x ) ) ] , ⌧ is a temperature , and K 1 is the number of negative pairs . InfoNCE ( 1 ) has an interpretation as multi-class logistic regression , and lower bounds the mutual information across similar views ( z̃ , z̃0 ) —a common representation learning objective [ 4 , 9 , 15 , 50 , 75 , 79 , 80 , 97 , 120 ] . Moreover , ( 1 ) can be interpreted as alignment ( numerator ) and uniformity ( denominator ) terms , the latter constituting a nonparametric entropy estimator of the representation as K ! 1 [ 124 ] . CL with InfoNCE can thus be seen as alignment of positive pairs with ( approximate ) entropy regularisation . Instead of using negative pairs , as in CL , a set of recent SSL methods only optimise for alignment and avoid collapsed representations through different forms of regularisation . For example , BYOL [ 41 ] and SimSiam [ 21 ] rely on “ architectural regularisation ” in the form of moving-average updates for a separate “ target ” net g0 ( BYOL only ) or a stop-gradient operation ( both ) . BarlowTwins [ 128 ] , on the other hand , optimises the cross correlation between ( z̃ , z̃0 ) to be close to the identity matrix , thus enforcing redundancy reduction ( zero off-diagonals ) in addition to alignment ( ones on the diagonal ) . Identifiability of learned representations . In this work , we address the question of whether SSL with data augmentation can reveal or uncover properties of the underlying data generating process . Whether a representation learned from observations can be expected to match the true underlying latent factors—up to acceptable ambiguities and subject to suitable assumptions on the generative process and inference model—is captured by the notion of identifiability [ 77 ] . Within representation learning , identifiability has mainly been studied in the framework of ( nonlinear ) ICA which assumes a model of the form x = f ( z ) and aims to recover the independent latents , or sources , z , typically up to permutation or element-wise transformation . A crucial negative result states that , with i.i.d . data and without further assumptions , this is fundamentally impossible [ 57 ] . However , recent breakthroughs have shown that identifiability can be achieved if an auxiliary variable ( e.g. , 2Note that recent work has investigated automatically discovering good augmentations [ 26 , 27 ] . 3If the only goal is to make representations of augmented views similar , a degenerate solution which simply maps any observation to the origin trivially achieves this goal . a time stamp or environment index ) renders the sources conditionally independent [ 45 , 54 , 55 , 58 ] . These methods rely on constructing positive and negative pairs using the auxiliary variable and learning a representation with CL . This development has sparked a renewed interest in identifiability in the context of deep representation learning [ 63 , 64 , 69 , 83 , 102 , 108 , 109 , 129 ] . Most closely related to SSL with data augmentation are works which study identifiability when given a second view x̃ of an observation x , resulting from a modified version z̃ of the underlying latents or sources z [ 39 , 69 , 83 , 101 , 108 , 129 ] . Here , z̃ is either an element-wise corruption of z [ 39 , 69 , 101 , 129 ] or may share a random subset of its components [ 83 , 108 ] . Crucially , all previously mentioned works assume that any of the independent latents ( are allowed to ) change , and aim to identify the individual factors . However , in the context of SSL with data augmentation , where the semantic ( content ) part of the representation is intended to be shared between views , this assumption does not hold .

Summary: This paper studies the effect of data augmentation in self-supervised representation learning. They take a latent variable perspective where there is a content component and style component of the latent variable, and consider settings where the content is invariant to the augmentation whereas the style component changes a lot under the data augmentation. In the latent variable model, arbitrary dependencies between the content and style latent variables are allowed.    The theoretical contribution of this paper is to show that, under certain assumptions on how the augmentation influences the latent variables, 1) a generative model for the data would perfectly identify the true content variables (up to an invertible function) and 2) minimizing the (entropy-regularized) squared difference between encoder outputs on two different augmentations of the data would also give an encoder which perfectly recovers the content variables (also up to an invertible function). The assumptions for these results are that the augmentations don't change the content latent variable and only the style latent variable, and that the augmentations can map the coordinates of the style latent variables to any value.   This paper also provides experiments for verifying aspects of the theory, first on synthetic Gaussian data, and also in a semi-synthetic setting of rendered images of 3D objects where different latents control different aspects of the image. They compare the effect of using InfoNCE objective on pairs of examples obtained via data augmentation against the effect of using InfoNCE objective on pairs of examples obtained via directly modifying the style latent variables. They show that training InfoNCE with modifying the latent variables does indeed preserve a lot of content information while discarding style information. Data augmentation is observed to behave in a similar way in some instances, though the equivalence is not exact.