Lossy Compression for Lossless Prediction

1 Introduction . Progress in important areas requires processing huge amounts of data . For climate prediction , models are still data-limited [ 1 ] , despite the Natl . Center for Computational Sciences storing 32 million gigabytes ( GB ) of climate data [ 2 ] . For autonomous driving , capturing a realistic range of rare events with current methods requires around 3 trillion GB of data.1 At these scales , data are only processed by task-specific algorithms , and storing data in human-readable formats can be prohibitive . We need compressors that retain only the information needed for algorithmic execution of downstream tasks . Existing lossy compressors are not up to the challenge , because they aim to reconstruct the data for human perception [ 5–10 ] . However , much of perceptual information is not needed to perform the tasks that we care about . Consider classifying images , which can require about 1 MB to store . Classification is typically invariant under small image transformations , such as rescalings or rotations , and could instead be performed using a representation that discards such information ( see Fig . 1 ) . The amount of unnecessary perceptual information is likely substantial , as illustrated by the fact that typical image classification can be performed using a detailed caption , which requires only about 1 kB to store ( 1000× fewer bits ) . 1Kalra and Paddock [ 3 ] estimated that autonomous vehicles would have to drive hundreds of billions of miles to demonstrate reliability in rare events . At 1 TB / hour [ 4 ] and 30 miles / hour , this is 3e12 GB . 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) . Our goal is to quantify the bit-rate needed to ensure high performance on a collection of prediction tasks . In the simple case of a single supervised task , the minimum bit-rate is achieved by compressing predicted labels , and essentially corresponds to the Information Bottleneck ( IB ; [ 11 ] ) . Our challenge , instead , is to ensure good performance on any future tasks of interest , which will rarely be completely known at compression time , or might be too large to enumerate . We overcome this challenge by focusing on sets of tasks that are invariant under user-defined transformations ( e.g. , translation , brightness , cropping ) , as is the case for many tasks of interest to humans [ 12 , 13 ] . This structure allows us to characterize a worst-case invariant task , which bounds the relative predictive performance on all invariant tasks . As a result , the bit-rate required to perform well on all invariant tasks is exactly the rate to compress the worst-case labels . At a high level , the worst-case task is to recognize which examples are transformed versions of one another , and rate savings come from discarding information from those transformations . We also provide two unsupervised neural compressors to target the optimal rates . One is similar to a variational autoencoder [ 14 ] that reconstructs canonical examples ( Fig . 1 ) . Our second is a simple modification of contrastive self-supervised learning ( SSL ; [ 15 ] ) , which allows us to convert pre-trained SSL models into powerful , generic compressors . Our contributions are : • We formalize the notion of compression for downstream predictive tasks . • We characterize the bits needed for high performance on any task invariant to augmentations . • We provide unsupervised objectives to train compressors that approximate the optimal rates . • We show that our compressor outperforms JPEG by orders of magnitude on 8 datasets on which it was never trained ( i.e. , zero-shot ) . E.g. , on ImageNet [ 16 ] , it decreases the bit-rate by 1000× . 2 Rate-distortion theory background . The goal of lossy compression theory is to find the number of bits ( bit-rate ) required to store outcomes x of a random variable ( r.v . ) X , so that it can be reconstructed within a certain tolerance . This is accomplished in Shannon ’ s [ 17 ] rate-distortion ( RD ) theory by mapping X into a r.v . Z with low mutual information I [ X ; Z ] . Specifically , given a distortion measure D [ X , Z ] , the RD theory characterizes the minimal achievable bit-rate for a distortion threshold δ by Rate ( δ ) = min p ( Z|X ) I [ X ; Z ] such that D [ X , Z ] ≤ δ . ( 1 ) In lossy compression , Z is usually a reconstruction of X , i.e. , it aims to faithfully approximate X . As a result , typical distortions , e.g. , the mean squared error ( MSE ) , assume that the sample spaces X , Z of both r.v.s are the same . This assumption is not required . Indeed , any distortion d : X × Z → R≥0 of the form D [ X , Z ] = Ep ( X , Z ) [ d ( X , Z ) ] , where there exists a z ∈ Z such that D [ X , z ] is finite , is a valid choice [ 18 ] . This shows that RD theory can be used outside of reconstructions . In the following we refer to Z as a compressed representation of X to distinguish it from a reconstruction . 3 Minimal bit-rate for high predictive performance . In this section , we characterize the bit-rate needed to represent X to ensure high performance on downstream tasks . Our argument has three high-level steps : ( i ) define a distortion that controls downstream performance when predicting from Z instead of X ; ( ii ) simplify and validate this distortion when desired tasks satisfy an invariance condition ; ( iii ) apply RD theory with the valid distortion . For simplicity , our presentation is relatively informal ; formal proofs are in Apps . A and B . 3.1 A distortion for worst-case predictive performance . Suppose X is an image . Potential downstream tasks might include Ydog , whether the image displays a dog ; or Yhd , whether the image is hand-drawn . Formally , these and other downstream tasks are expressed as T = { Ydog , Yhd , . . . } , a set of random variables that are jointly distributed with X . Let R [ Y |X ] denote the Bayes ( best possible ) risk when predicting Y from X . For ease of presentation in the main paper , we consider only classification tasks T and Bayes risk of the standard log loss R [ Y |X ] : = infq Ep ( X , Y ) [ − log q ( Y |X ) ] . We deal with MSE and regression in Appx . B.6 . In this setting , a meaningful distortion DT [ X , Z ] quantifies the difference between predicting any Y ∈ T from the compressed Z , as opposed to using X . This is the worst-case excess risk , DT [ X , Z ] : = sup Y ∈T R [ Y |Z ] − R [ Y |X ] . ( 2 ) If DT [ X , Z ] = 0 , it is possible to achieve lossless prediction : performing as well using Z as using X . More generally , bounding DT by δ ensures that R [ Y |Z ] − R [ Y |X ] ≤ δ for all tasks in T . However , there are two issues that need to be addressed before Eq . ( 2 ) can be used . First , it is not clear whether DT is a valid distortion for RD theory . Second , the worst excess-risk DT assumes access to all downstream tasks of interest T during compression , which is unrealistic in general . 3.2 Invariant tasks . The tasks that we care about are not arbitrary , and often share structure . One such structure is invariance to certain pre-specified transformations of input data . For example , computer vision tasks are often invariant to mild transformations such as brightness changes . Such invariance structure is common in realistic tasks , as seen by the wide-spread use of data augmentations [ 13 ] in machine learning ( ML ) , which encourage predictions to be the same for an unaugmented x and an augmented x+ . Motivated by this we focus on sets of invariant tasks T . We consider a general notion of invariance , namely invariance specified by an equivalence relation ∼ on X .2 The equivalence induces a partition of X into disjoint equivalence classes , and we are interested in tasks whose conditional distributions are constant within these classes . Definition 1 . The set of invariant tasks of interest with respect to an equivalence ( X , ∼ ) , denoted T∼ , is all random variables Y such that x ∼ x+ =⇒ p ( Y |x ) = p ( Y |x+ ) for any x , x+ ∈ X . 3.3 Rate-distortion theory for invariant task prediction . The key to simplifying DT∼ is the existence of a ( non-unique ) worst-case invariant task , denoted M ( X ) . Such task contains all and only information to which tasks Y ∈ T∼ are not invariant ; we call them maximal invariants . A maximal invariant M ( • ) with respect to ∼ is any function satisfying3 x ∼ x+ ⇐⇒ M ( x ) =M ( x+ ) for any x , x+ ∈ X . ( 3 ) 2As a reminder , ∼ is an equivalence relation iff for all x , x′ , x′′ ∈ X : ( reflexivity ) x ∼ x , ( symmetry ) x ∼ x′ ⇐⇒ x′ ∼ x , and ( transitivity ) x ∼ x′ and x′ ∼ x′′ =⇒ x ∼ x′′ . 3This extends the definition of maximal invariants [ 19 ] beyond invariances to group actions . A maximal invariant removes all information that tasks are invariant to , as it maps equivalent inputs to the same output , i.e. , M ( x ) =M ( x+ ) . Yet , it retains the minimal information needed to perform invariant tasks , by mapping non-equivalent inputs x 6∼ x− to different outputs M ( x ) 6= M ( x− ) . In other words , M ( x ) indexes the equivalence classes . For example , the Euclidean norm is a maximal invariant for rotation invariance , as all vectors that are rotated versions of one another can be characterized by their radial coordinate . For data augmentations , the canonical ( unaugmented ) version of the input is a maximal invariant . Other examples are shown in Fig . 2 . We prove in Appx . B.2 that under weak regularity conditions , maximal invariant tasks exist in T∼ , and that they achieve the supremum in Eq . ( 2 ) . This allows us to show that DT∼ reduces to the Bayes risk of predicting M ( X ) from Z and that it is a valid distortion measure . Crucially , this allows us to quantify downstream performance without enumerating invariant tasks . Proposition 1 . Let ( X , ∼ ) be an equivalence relation and M a maximal invariant that takes at most countably many values , with H [ M ( X ) ] < ∞ . Then DT∼ ( 2 ) with log loss is a valid distortion and DT∼ [ X , Z ] = R [ M ( X ) |Z ] . ( 4 ) Here we used R [ M ( X ) |X ] = 0 , as M is a deterministic function . Also , note that the countable requirement holds when tasks are invariant to some rounding of the input , as is typically the case due to floating-point storage . We accommodate the uncountable case for squared-error loss in Appx . B.6 . With a valid distortion in hand , we invoke the RD theorem with DT∼ to obtain our “ Rate-Invariance ” ( RI ) theorem . The RI theorem characterizes the bit-rate needed to store X while ensuring small log-loss on invariant tasks . We obtain analogous results for squared-error loss . Theorem 2 ( Rate-Invariance ) . Assume the conditions of Prop . 1 . Let δ ≥ 0 , and Rate ( δ ) denote the minimum achievable bit-rate for transmitting Z such that for any Y ∈ T∼ we have R [ Y |Z ] − R [ Y |X ] ≤ δ . Then Rate ( δ ) = 0 if δ ≥ H [ M ( X ) ] and otherwise it is finite and Rate ( δ ) = H [ M ( X ) ] ︸ ︷︷ ︸ information needed to predict T∼ − δ︸︷︷︸ acceptable decrease in predictive loss = H [ X ] ︸ ︷︷ ︸ standard compression − H [ X |M ( X ) ] ︸ ︷︷ ︸ gains due to invarainces − δ . ( 5 ) To ensure lossless prediction , i.e. , Rate ( 0 ) , our theorem states that we require a bit-rate of H [ M ( X ) ] . Intuitively , this is because M ( X ) contains the minimal information needed to predict losslessly any Y ∈ T∼.4 Furthermore , the theorem relates compression and prediction by showing that allowing a δ decrease in log-loss performance on all tasks can save exactly δ bits . Intuitively , this is a linear relationship , because expected log-loss is measured in bits . On the right of Eq . ( 5 ) we further decompose H [ M ( X ) ] into two terms to provide another interpretation : ( i ) H [ X ] , which , for discrete X , is the bit-rate required to losslessly compress X , and ( ii ) H [ X |M ( X ) ] , which quantifies the information removed due to the invariance of desired tasks . Importantly , removing this information does not impact the best possible predictive performance . See Fig . 3 . The bit-rate gains can be substantial , depending on the invariances . Consider compressing a sequence of n i.i.d . fair coin flips . Suppose one is only interested in predicting permutation invariant labels . Then instead of compressing the entire sequence in H [ Xn ] = nH [ X ] = n bits , one could compress the number of heads , which is a maximal invariant for permutation invariance , in O ( log n ) bits.5 As more interesting examples , we recover in Appx . B.4 results from ( i ) unlabeled graph compression [ 20 ] ; ( ii ) multiset compression [ 21 ] ; ( iii ) single task compression ( IB ; [ 11 ] ) . The equivalence ∼ can be induced by any transformations , such as transforming an image to its caption . We use this idea in Sec . 5.3 to obtain > 1000× compression on ImageNet without sacrificing predictive performance . 4We prove in Appx . B.5 that Rate ( 0 ) = H [ M ( X ) ] for any losses used in ML . 5The number of heads K follow a binomial distribution so H [ K ] ∈ O ( logn ) . Here K is also a minimal sufficient statistic for Xn . More generally , if P ( X ′ ) is invariant to ∼ and ∼ is the coarsest such relation , then minimal sufficiency coincides with maximal invariance . In practice , however , P ( X ′ ) will rarely be ∼ invariant .

Summary: This paper proposes lossy compression algorithms that only compress task-relevant information in order to obtain a better bit-rate while maintaining the performance of the downstream task. Specifically, the paper proposes to augment the training samples before encoding, and the decoder aims to reconstruct the original image (another proposed model uses contrastive loss to do this). In this way, the model can learn representations that are invariant of the input augmentations, which could lead to smaller bit-rate.