Compressive Visual Representations

1 Introduction . Individuals develop mental representations of the surrounding world that generalize over different views of a shared context . For instance , a shared context could be the identity of an object , as it does not change when viewed from different perspectives or lighting conditions . This ability to represent views by distilling information about the shared context has motivated a rich body of self-supervised learning work [ 54 , 4 , 12 , 30 , 33 , 47 ] . For a concrete example , we could consider an image from the ImageNet training set [ 60 ] as a shared context , and generate different views by repeatedly applying different data augmentations . Finding stable representations of a shared context corresponds to learning a minimal high-level description since not all information is relevant or persistent . This explicit requirement of learning a concise representation leads us to prefer objectives that are compressive and only retain the relevant information . Recent contrastive approaches to self-supervised visual representation learning aim to learn representations that maximally capture the mutual information between two transformed views of an image [ 54 , 4 , 12 , 33 , 40 ] . The primary idea of these approaches is that this mutual information †Main contributors 1Code available at https : //github.com/google-research/compressive-visual-representations 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) . corresponds to a general shared context that is invariant to various transformations of the input , and it is assumed that such invariant features will be effective for various downstream higher-level tasks . However , although existing contrastive approaches maximize mutual information between augmented views of the same input , they do not necessarily compress away the irrelevant information from these views [ 12 , 33 ] . As shown in [ 26 , 27 ] , retaining irrelevant information often leads to less stable representations and to failures in robustness and generalization , hampering the efficacy of the learned representations . An alternative state-of-the-art self-supervised learning approach is BYOL [ 30 ] , which uses a slow-moving average network to learn consistent , view-invariant representations of the inputs . However , it also does not explicitly capture relevant compression in its objective . In this work , we modify SimCLR [ 12 ] , a state-of-the-art contrastive representation method , by adding information compression using the Conditional Entropy Bottleneck ( CEB ) [ 27 ] . Similarly , we show how BYOL [ 30 ] representations can also be compressed using CEB . By using CEB we are able to measure and control the amount of information compression in the learned representation [ 26 ] , and observe its impact on downstream tasks . We empirically demonstrate that our compressive variants of SimCLR and BYOL , which we name C-SimCLR and C-BYOL , significantly improve accuracy and robustness to domain shifts across a number of scenarios . Our primary contributions are : • Reformulations of SimCLR and BYOL such that they are compatible with information-theoretic compression using the Conditional Entropy Bottleneck [ 26 ] . • An exploration of the relationship between Lipschitz continuity , SimCLR , and CEB compression , as well as a simple , tractable lower bound on the Lipschitz constant . This provides an alternative explanation , in addition to the information-theoretic view [ 26 , 27 , 1 , 2 ] , for why CEB compression improves SimCLR model robustness . • Extensive experiments supporting our hypothesis that adding compression to the state-of-theart self-supervised representation methods like SimCLR and BYOL can significantly improve their performance and robustness to domain shifts across multiple datasets . In particular , linear evaluation accuracies of C-BYOL are even competitive with the supervised baselines considered by SimCLR [ 12 ] and BYOL [ 30 ] . C-BYOL reaches 76.0 % and 78.8 % with ResNet-50 and ResNet-50 2x respectively , whereas the corresponding supervised baselines are 76.5 % and 77.8 % respectively . 2 Methods . In this section , we describe the components that allow us to make distributional , compressible versions of SimCLR and BYOL . This involves switching to the Conditional Entropy Bottleneck ( CEB ) objective , noting that the von Mises-Fisher distribution is the exponential family distribution that corresponds to the cosine similarity loss function used by SimCLR and BYOL , and carefully identifying the random variables and the variational distributions needed for CEB in SimCLR and BYOL . We also note that SimCLR and CEB together encourage learning models with a smaller Lipschitz constant , although they do not explicitly enforce that the Lipschitz constant be small . 2.1 The Conditional Entropy Bottleneck . In order to test our hypothesis that compression can improve visual representation quality , we need to be able to measure and control the amount of compression in our visual representations . To achieve this , we use the Conditional Entropy Bottleneck ( CEB ) [ 26 ] , an objective function in the Information Bottleneck ( IB ) [ 66 ] family . Given an observation X , a target Y , and a learned representation Z of X , CEB can be written as : CEB ≡ min Z βI ( X ; Z|Y ) − I ( Y ; Z ) ( 1 ) = min Z β ( H ( Z ) −H ( Z|X ) −H ( Z ) +H ( Z|Y ) ) −H ( Y ) +H ( Y |Z ) ( 2 ) = min Z β ( −H ( Z|X ) +H ( Z|Y ) ) +H ( Y |Z ) ( 3 ) where H ( · ) and H ( ·|· ) denote entropy and conditional entropy respectively . We can drop the H ( Y ) term because it is constant with respect to Z. I ( Y ; Z ) is the useful information relevant to the task , or the prediction target Y . I ( X ; Z|Y ) is the residual information Z captures about X when we already know Y , which we aim to minimize . Compression strength increases as β increases . We define e ( z|x ) as the true encoder distribution , where z is sampled from ; b ( z|y ) , a variational approximation conditioned on y ; d ( y|z ) , the decoder distribution ( also a variational approximation ) which predicts y conditioned on z . As shown in [ 26 ] , CEB can be variationally upper-bounded : vCEB ≡ min e ( z|x ) , b ( z|y ) , d ( y|z ) Ex , y∼p ( x , y ) , z∼e ( z|x ) β ( log e ( z|x ) − log b ( z|y ) ) − log d ( y|z ) ( 4 ) There is no requirement that all three distributions have learned parameters . At one limit , a model ’ s parameters can be restricted to any one of the three distributions ; at the other limit , all three distributions could have learned parameters . If e ( · ) has learned parameters , its distributional form may be restricted , as we must be able to take gradients through the z samples.2 The only requirement on the b ( · ) and d ( · ) distributions is that we be able to take gradients through their log probability functions . InfoNCE . As shown in [ 26 ] , besides parameterizing d ( y|z ) , it is possible to reuse b ( z|y ) to make a variational bound on the H ( Y |Z ) term . As I ( Y ; Z ) = H ( Y ) −H ( Y |Z ) and H ( Y ) is a constant with respect to Z : H ( Y |Z ) ≤ Ex , y∼p ( x , y ) , z∼e ( z|x ) log b ( z|y ) ∑K k=1 b ( z|yk ) ( 5 ) where K is the number of examples in a minibatch . Eq . ( 5 ) is also known as the contrastive InfoNCE bound [ 54 , 56 ] . The inner term , d ( y|z ) ≡ b ( z|y ) ∑K k=1 b ( z|yk ) , ( 6 ) is a valid variational approximation of the true but unknown p ( y|z ) . Fischer [ 26 ] calls Eq . ( 6 ) the CatGen decoder because it is a categorical distribution over the minibatch that approximates the generative decoder distribution . 2.2 C-SimCLR : Compressed SimCLR . The InfoNCE bound [ 54 ] enables many contrastive visual representation methods to use it to capture shared context between different views of an image as a self-supervised objective [ 12 , 13 , 33 , 15 , 40 ] . In this work , we show how to compress the SimCLR [ 12 ] model , but the method we discuss is generally applicable to other InfoNCE-based models . SimCLR applies randomized augmentations to an image to create two different views , x and y ( which we also refer to as x′ ) , and encodes both of them with a shared encoder , producing representations rx and ry . Both rx and ry are l2-normalized . The SimCLR version of the InfoNCE objective is : LNCE ( rx , ry ) = − log e 1 τ r T y rx∑K k=1 e 1 τ r T yk rx ( 7 ) where τ is a temperature term and K is the number of views in a minibatch . SimCLR further makes its InfoNCE objective bidirectional , such that the final objective becomes LNCE ( rx , ry ) + LNCE ( ry , rx ) = − log e 1 τ r T y rx∑K k=1 e 1 τ r T yk rx − log e 1 τ r T x ry∑K k=1 e 1 τ r T xk ry ( 8 ) We can observe the following : exp ( 1τ r T y rx ) in Eq . ( 7 ) corresponds to the unnormalized b ( z|y ) in Eq . ( 5 ) . e ( ·|x ) generates z = rx , whilst ry and ryk are distribution parameters of b ( ·|y ) and b ( ·|yk ) respectively . e ( ·|x ) and b ( ·|y ) share model parameters . von Mises-Fisher Distributional Representations . The cosine-similarity-based loss ( Eq . ( 7 ) ) is commonly used in contrastive learning and can be connected to choosing the von Mises-Fisher ( vMF ) distribution for e ( ·|x ) and b ( ·|y ) [ 32 , 70 ] . vMF is a distribution on the ( n− 1 ) -dimensional hyper-sphere . The probability density function is given by fn ( z , µ , κ ) = Cn ( κ ) eκµ T z , where µ and κ are denoted as the mean direction and concentration parameter respectively . We assume κ is a 2For example , e ( z|x ) could not generally be a mixture distribution , as sampling the mixture distribution has a discrete component , and we can not easily take gradients through discrete samples . constant . The normalization term Cn ( κ ) is a function of κ and equal to κ n/2−1 ( 2π ) n/2In/2−1 ( κ ) , where Iv denotes the modified Bessel function of the first kind at order v. By setting the mean direction µ to ry , concentration κb of b ( ·|y ) to 1/τ , and rx to z , we can connect the SimCLR objective ( Eq . ( 7 ) ) to the distributional form of InfoNCE ( Eq . ( 5 ) ) e 1 τ r T y rx∑K k=1 e 1 τ r T yk rx = Cn ( κb ) e κbr T y rx∑K k=1 Cn ( κb ) e κbrTyk rx = fn ( rx , ry , κb ) ∑K k=1 fn ( rx , ryk , κb ) = b ( rx|y ) ∑K k=1 b ( rx|yk ) ( 9 ) z = rx is a deterministic unit-length vector , so we can view e ( ·|x ) as a spherical delta distribution , which is equivalent to a vMF with rx as the mean direction and κe → ∞ . We can further extend the forward encoder to have non-infinite κe , which results in a stochastic z . These allow us to have SimCLR in a distributional form with explicit distributions e ( ·|x ) and b ( ·|y ) and satisfy the requirements of CEB discussed in Sec . 2.1 . Compressing SimCLR with Bidirectional CEB . Figure 1 illustrates the Compressed SimCLR ( C-SimCLR ) model . The model learns a compressed representation of an view X that only preserves information relevant to predicting a different view Y by switching to CEB . As can be seen in Eq . ( 3 ) , the CEB objective treats X and Y asymmetrically . However , as shown in [ 26 ] , it is possible to learn a single representation Z of both X and Y by having the forward and backward encoders act as variational approximations of each other : CEBbidir ≡ min Z βXI ( X ; Z|Y ) − I ( Y ; Z ) + βY I ( Y ; Z|X ) − I ( X ; Z ) ( 10 ) ≡ min Z βX ( −H ( Z|X ) +H ( Z|Y ) ) +H ( Y |Z ) ( 11 ) + βY ( −H ( Z|Y ) +H ( Z|X ) ) +H ( X|Z ) ≤ min e ( ·|· ) , b ( ·|· ) , c ( ·|· ) , d ( ·|· ) Ex , y∼p ( x , y ) [ ( 12 ) Ezx∼e ( zx|x ) [ βX ( log e ( zx|x ) − log b ( zx|y ) ) − log d ( y|zx ) ] + Ezy∼e ( zy|y ) [ βY ( log ( e ( zy|y ) − log ( b ( zy|x ) ) − log c ( x|zy ) ] ] where d ( ·|· ) and c ( ·|· ) are the InfoNCE variational distributions of b ( ·|· ) and e ( ·|· ) respectively . e and b use the same encoder to parameterize mean direction in SimCLR setting . Since SimCLR is trained with a bidirectional InfoNCE objective , Eq . ( 12 ) gives an easy way to compress its learned representation . As in SimCLR , the deterministic hx ( in Fig . 1 ) is still the representation used on downstream classification tasks .

Summary: The paper posits that compressing the representation during learning improves the robustness of said representation. To demonstrate their argument, they derive an objective for SimCLR and BYOL that explicitly compresses the representation using the Conditional Entropy Bottleneck. The addition of this compression improves the performance for both BYOL and SimCLR on ImageNET. They also demonstrate that their method improves the robustness of the encoder and its representation to perturbation of the input data.