Inverse Problems Leveraging Pre-trained Contrastive Representations

1 Introduction . Modern representation learning networks like CLIP [ 35 ] are showing incredible performance for image classification , even for zero-shot problems with labels not seen during training . Training these encoders comes at a staggering cost and requires datasets and computing resources only available to very few organizations . In this paper we show how to leverage this pretrained power for a new family of problems in the presence of image corruptions or other types of measurements . Inverse problems involve reconstructing an unknown vector x from measurements y = A ( x ) . Typically , the forward operator A corrupts the unknown vector x and reduces its dimension , i.e . the observations y live in a lower-dimensional space compared to x . In the special case of linear inverse problems , the forward operator is simply a matrix and the measurements are in the form y = Ax + noise . Special cases of linear inverse problems include image denoising , inpainting , super-resolution , compressed sensing used in medical tomography , seismic geological imaging and many others , see e.g . [ 31 ] for a recent overview . ∗Equal contribution . Code available at https : //github.com/Sriram-Ravula/Contrastive-Inversion . 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) . In this paper , we introduce the study of a new family of inverse problems : reconstructing the representation of an image given a corrupted or measured input . Formally , if a ( clean ) image is x and its CLIP representation is R ( x ) , we would like to obtain that representation by only observing a highly corrupted input A ( x ) . This is impossible if the forward process A removes information needed to obtain the representation . Surprisingly , we show that we can recover representations that are useful for downstream tasks , even from extremely corrupted versions of the image . We introduce a robust encoder S that is trained to imitate the behavior of the pretrained CLIP encoder acting on clean images x . However , the input to the robust encoder is only corrupted images A ( x ) that are created by applying the forward operator on x . Our approach is illustrated in Figure 2 . The teacher encoder is the pretrained CLIP , and the student encoder is our robust encoder operating on corrupted images . Formally , the robust encoder S ( A ( x ) ) is trained to approximate R ( x ) using a contrastive loss . Many applications , such as object recognition with low-cost cameras , remote sensing and aerial imaging rely on noisy or blurry data and can face occlusions or sensor corruptions . As we demonstrate in our experiments in the Appendix , normal CLIP fails on highly corrupted images . Our procedure allows us to transfer the power of CLIP to heavily corrupted images in downstream tasks , with relatively little extra training . 1.1 Results . We show that our method is able to obtain useful representations even under extreme corruptions such as removing 90 % of the pixels as shown in the top panel of Figure 1 . The highly corrupted images enter our robust encoder and the obtained representation is used in a linear classifier to produce ImageNet-100 labels . Our main result is that our method outperforms a pretrained ResNet ( of the same size as our robust encoder ) fine-tuned end-to-end on labeled distorted images . Using less labeled data For some corruption levels , we are able to outperform end-to-end fine-tuned ResNets using as little as 10 % of labeled samples . This is even when the fine-tuned baseline uses 100 % of ImageNet-100 labels for training . The primary advantage of our model is that our robust encoder observes representations from the pretrained CLIP encoder , which was trained with a much bigger dataset compared to ImageNet . Still , the fact that this implicit advantage in the representations and 10 % of labeled data is sufficient to outperform a supervised trained ResNet fine-tuned with 10 times more labeled data is very surprising and illustrates the power and versatility of pretrained representation learners . Robustness to noise and data shifts Our method is very robust to changes in the forward operators , data statistics and label shifts . We experiment with three classes of forward operators : random pixel masking , additive Gaussian noise , and Gaussian blurring distortions . For each , we train and test on a wide-range of distortion levels , from slight to severe corruption . We show that our robust encoder produces useful representations even when the level of corruption is outside its training domain . We show that our representations are useful for a wide range of tasks , without requiring knowledge of the task when the robust encoder is trained . We illustrate excellent classification accuracy across five datasets , frequently outperforming end-to-end supervised baselines trained with knowledge of the target task . Our experiments include a chest X-ray COVID pneumonia task which has very different morphology compared to ImageNet . Surprisingly , the same universal representations , combined with a custom linear probe are very successful across all tasks . Contrastive versus MSE training We formulate the contrastive student training method as a regularization upon the simple mean squared error loss between student embeddings of a distorted image and teacher embeddings of the clean version of that image . We analyze the effects of this regularization on the training dynamics . Our results empirically show that simple MSE is worse in most cases , further strengthening our argument for the usefulness of contrastive learning in this setting . 2 Related Work . Robust Image Recognition It has been shown that classifiers that perform well on clean image tasks are not robust to common image distortions [ 21 ] . Several datasets have been proposed specifically to benchmark generalization of classifier performance to natural distortions [ 18 , 38 ] . Related to our work , [ 45 ] and [ 12 ] fine-tune pretrained classifiers on distorted data , which seems to yield better performance than end-to-end training . However even under these training processes , modern classifiers exhibit inferior performance to human vision on distorted data [ 12 ] . Inverse problems There is significant recent literature on solving inverse problems including denoising , inpainting and deconvolution for deblurring . While classical techniques rely on sparsity based priors e.g . [ 30 , 37 ] , recent techniques include data-driven deep-learning methods [ 1 , 10 , 13 , 33 ] as well as combinations of sparsity and generative methods [ 11 ] and untrained deep nets [ 16 ] . Our work focuses on recovering the representation of an image , as opposed to an image itself , so it is also related to task-aware sensing [ 23 ] . However , our approach is fundamentally different from all previous inverse problems since instead of trying to reconstruct the image x , we aim to reconstruct a representation R ( x ) , which lives in a different space . The other important distinction is that even though our corruption processes are linear in the pixel space , they are non-linear with respect to the representation vector we try to recover . In effect , we are solving a non-linear inverse problem in a supervised way using a contrastive loss . In the Appendix , we compare to methods that attempt to solve the inverse problem in pixel space , then apply a classifier on the recovered image . Contrastive Representation Learning Self-supervised representation learning has recently exploded in popularity , largely due to the success of constrastive losses in learning representations from unlabeled data [ 4 , 6 , 7 , 15 , 17 , 43 ] . These techniques are able to generate highly general embeddings of images that are effective for many types of downstream tasks , even on domains that were not explicitly considered in training [ 27 ] . Contrastive losses generally operate based on a simple push-pull principle : images desired to be close in embedding space are pushed together , while unrelated images are pulled apart . One particularly popular choice of contrastive loss is the InfoNCE loss , derived from techniques for noise contrastive estimation [ 14 ] , and popularized in the self-supervised setting in [ 32 ] . Several works have considered adversarial training frameworks to yield representations that are more robust to adversarial attacks [ 20 , 22 , 25 ] . However these works only consider adversarial robustness and not robustness to common corruptions . Our approach is similar to that of [ 24 ] , which employs a variant of the InfoNCE loss for a supervised setting . This differs from our work in that we exclusively focus on robustness to natural corruptions . Moreover , our contrastive training step can be performed without any task-specific labeled data . This is made possible by the use of the powerful embeddings provided by CLIP , and allows the embeddings to be used for downstream tasks on multiple datasets . Knowledge Transfer Methods Our work is closely related to prior works which aim to distill , reduce , or transfer the knowledge from one network to another for a specific task [ 3 , 19 , 29 , 44 ] . Of note is [ 40 ] , where the authors use a contrastive objective to transfer representations from a teacher network to a student network . Our work diverges in that we do not transfer from a larger , more powerful teacher to a smaller student , but rather transfer between a teacher and student of the same architecture initialized from the same weights . In addition , although the authors test on cross-modal transfer tasks such as transferring between color channels , transferring representations between clean and distorted images is a different task : we try to extract the same high-level information from less data as opposed to different , but related data . 3 Method . Our problem is to recover the embeddings of clean images when we only have access to highly corrupted versions of the images . The encoder R ( · ) is assumed to yield high-quality representations for a variety of domains . The distortion process A ( · ) is assumed to be a known forward operator that greatly distorts images . We assume A ( · ) is sufficiently severe as to inhibit the performance of the encoder R ( · ) , but not so severe that recovery is impossible . From a collection of input images { xi } Ni=1 , we are only given access to the distorted images and the representations of the clean images { A ( xi ) , R ( xi ) } Ni=1 . Our approach is to learn a student function S so that S ( A ( · ) ) is equally useful as the teacher representation R ( · ) . We measure the utility of a representation by the performance on an unspecified downstream supervised learning task . Least squares Loss One potential approach for the task of recovering the teacher ’ s embeddings from the corrupted inputs is to minimize the expected ` 2 distance in embedding space between the clean teacher embeddings R ( · ) and the predicted student representations S ( A ( · ) ) . If both R ( · ) and S ( · ) are constrained to have ` 2-normalized outputs , then the empirical least squares loss becomes : L̂MSE ( S ; R , A ) : = −1 N N∑ i=1 〈S ( A ( xi ) ) , R ( xi ) 〉 , ( 1 ) where we have expanded and subtracted out constant terms . However , this process may not yield the most effective embedding of the corrupted data . Due to either limitations in the training process , the severity of the distortion process , or the generalization properties of the approximate minimizers of L̂MSE , the learned embeddings S ( A ( · ) ) may not be as useful for downstream tasks as embeddings learned using the other losses we consider . Contrastive Loss Inspired by recent advances in self-supervised representation learning , we learn S by minimizing a contrastive loss . We consider the following variant on the popular InfoNCE loss : L̂contr ( S ; τ , R , A ) : = −1 N N∑ i=1 log exp ( K ( i , i ) /τ ) ∑N j=1 exp ( K ( i , j ) /τ ) , ( 2 ) where K ( i , j ) : = 〈S ( A ( xi ) ) , R ( xj ) 〉 measures the similarity between the learned embedding of A ( xi ) and the clean embedding of xj , and τ is a temperature hyperparameter . We follow [ 5 , 42 ] and rewrite L̂contr in terms of explicit ‘ pull ‘ and ‘ push ‘ terms as : L̂contr ( S ; τ , R , A ) = 1 τ L̂MSE ( S ; R , A ) + L̂unif ( S ; τ , R , A ) , ( 3 ) where the second term , referred to as the uniformity term , is defined as : L̂unif ( S ; τ , R , A ) : = 1 N N∑ i=1 log N∑ j=1 eK ( i , j ) /τ . ( 4 ) The first term of L̂contr is simply L̂MSE which encourages alignment of S ( A ( xi ) ) with R ( xi ) , and the L̂unif term encourages the learned representations for corrupted data to be dissimilar from all other representations of clean data . The primary difference between Equation 3 and the InfoNCE loss commonly used in self-supervised learning is the choice of the similarity measureK ( · , · ) . Without access to specified target embeddings , K ( i , j ) is typically chosen to be the inner product between projections of the embeddings of xi and xj , ( or between xi and a positive example in the case of K ( i , i ) ) . Here , the uniformity term is necessary to prevent representation collapse , where in our setting L̂MSE alone suffices to prevent this degenerate case . We note that alternative choices for K ( · , · ) in L̂unif may be employed . For example , we could also contrast the student embeddings S ( A ( · ) ) across two different images . We ablate against other choices in the experimental section and find that our choice of K ( · , · ) is one of several that exhibits comparable performance . Effect of uniformity term To examine the effect of the uniformity term of the loss on the training dynamics , we consider the gradients of L̂contr with respect to the parameters of the encoder S. We first decompose L̂unif to consider the contributions of each individual data point : L̂unif ( S ; τ , R , A ) = 1 N N∑ i=1 L̂unifi ( S ; τ , R , A ) , L̂unifi ( S ; τ , R , A ) : = log N∑ j=1 e K ( i , j ) τ . Then the gradient of L̂contr with respect to the parameters of S may be written as ∇SL̂contr ( S ; τ , R , A ) = −1 τ ∇SL̂MSE ( S ; R , A ) + 1 τN N∑ i=1 N∑ j=1 wi ( j ) ∇SK ( i , j ) , where wi ( j ) : = exp ( K ( i , j ) /τ − L̂unifi ( S ; T , R , A ) ) . This can be interpreted as follows . The weighting of the first term by τ−1 balances the gradients ∇SK ( i , i ) and ∑ j ∇SK ( i , j ) , as is common with other choices of contrastive losses ( c.f . Theorem 2 in [ 41 ] ) . Noting that for each i , the weights wi ( j ) sum to 1 , the gradient of each individual uniformity term L̂unifi is a convex combination of the gradients of the similarity terms K ( i , · ) . The individual similarity terms are weighted exponentially proportionally to each K ( i , j ) ’ s contribution to the total uniformity loss . As τ decreases , the weights place greater emphasis on terms that are most similar . As τ is commonly chosen to be less than 1 , the dynamics automatically reflect the influence of the ‘ hardest ’ negative examples . When is perfect recovery possible ? Finally , we describe conditions which allow for perfect recovery of the clean embeddings from the corrupted images in the training set . The first condition is that the corruption process should not be too destructive : it suffices to assume that the implication A ( xi ) = A ( xj ) =⇒ R ( xi ) = R ( xj ) holds for all xi , xj in the training set , i.e. , there exists a function which attains exact recovery . In this case , any S∗ that minimizes L̂MSE has S∗ ( A ( xi ) ) = R ( xi ) for all xi in the training set . To argue for the same recovery guarantees when optimizing L̂contr , we need to assume that the teacher R ( · ) provides a well-separated embedding of the training data . Proposition 1 . If R is a minimizer of the uniformity term R ∈ argmin f N∑ i=1 log N∑ j=1 exp 〈f ( xi ) , f ( xj ) 〉 τ , then any encoder S∗ ∈ argminS L̂contr ( S ; τ , R , A ) exactly recovers the target embedding , S∗ ( A ( xi ) ) = R ( xi ) for all xi in the training set . Training procedure Our final method consists of ( 1 ) the contrastive step , in which the student learns representations for the distorted images and ( 2 ) a fine-tuning step , where we train a linear classifier on top of the learned representations . During the second step , the student encoder is kept frozen .

Summary: The paper consider the following problem: Let $x$ be an image, let $R$ be a feature map or representation obtained through the CLIP network, and let $A$ be a distortion process. The representation takes as input an image and yields a feature representation. The distortion process also takes an image as input, and yields a distorted image (e.g., it deletes pixels).  The paper assumes we are given the representations and corresponding distorted images, i.e., a collection $\{A(x_i), R(x_i)\}$, and the goal is to learn a student function $S(A( \cdot ))$ that is equally useful as the original representation $R(x)$ for a classification task. The paper learns the student $S$ by minimizing a contrastive loss, and the utility of the representation obtained by the student $S$ is measured by the performance on a supervised learning task, specifically on ImageNet-100.   The paper provides simulation results demonstrating that the method learns useful representations. The paper's main result is to show that the method 'outperforms a pretrained ResNet, of the same size as the robust encoder, fine-tuned end-to-end on labeled distorted images'. In addition to this result, the paper provides a number of ablations studies and compares to two reasonable baselines.