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 shows how to use the Conditional Entropy Bottleneck (CEB)[1] in a self-supervised fashion.  Thus resulting in a new method for training approximately minimal self-supervised  (SSL) representations, i.e., a new member of the SSL information bottleneck family (although extremely similar to [2]). Specifically, they show how to use CEB in conjunction with SimCLR and BYOL. Furthermore, they provide good arguments to suggest that their proposed objective will learn smoother encoders than standard self-supervised models, which can help robustness. Through extensive experiments,  they show significant performance gains of the proposed methods compared to SimCLR and BYOL.  - [1] Fischer, Ian. "The conditional entropy bottleneck." Entropy 22.9 (2020): 999. - [2] Federici, Marco, et al. "Learning robust representations via multi-view information bottleneck." arXiv preprint arXiv:2002.07017 (2020).