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: This works suggests a novel method on how to harness the representational power of models such as CLIP. In this task the input is distorted with a known function, such as missing pixels or gaussian noise, blurring. This works proposes training a student contrastively by matching the clip representation of the clean image with student's representation of the distorted image. They propose different matching functions and explain why in their setup it will not collapse if one does not compare against every other.  The experimental setup is extensive and shows much better label efficiency when working with distorted images, they also show better noise extrapolation and better dataset tranferability in compare to baseline. Their baseline is a resnet pretrained on imagenet supervised trained on the distorted imagenet-100.