Overfitting for Fun and Profit: Instance-Adaptive Data Compression

1 INTRODUCTION . The most common approach to neural lossy compression is to train a variational autoencoder ( VAE ) like model on a training dataset to minimize the expected RD cost D + βR ( Theis et al. , 2017 ; Kingma & Welling , 2013 ) . Although this approach has proven to be very successful ( Ballé et al. , 2018 ) , a model trained to minimize expectedRD cost over a full dataset is unlikely to be optimal for every test instance because the model has limited capacity , and both optimization and generalization will be imperfect . The problem of generalization will be especially significant when the testing distribution is different from the training distribution , as is likely to be the case in practice . Suboptimality of the encoder has been studied extensively under the term inference suboptimality ( Cremer et al. , 2018 ) , and it has been shown that finetuning the encoder or latents for a particular instance can lead to improved compression performance ( Lu et al. , 2020 ; Campos et al. , 2019 ; Yang et al. , 2020b ; Guo et al. , 2020 ) . This approach is appealing as no additional information needs to be added to the bitstream , and nothing changes on the receiver side . Performance gains however are limited , because the prior and decoder can not be adapted . ∗Equal contribution †Qualcomm AI Research is an initiative of Qualcomm Technologies , Inc. ‡Work done during internship at Qualcomm AI Research In this paper we present a method for full-model instance-adaptive compression , i.e . adapting the entire model to a single data instance . Unlike previous work , our method takes into account the costs for sending not only the latent prior , but also the decoder model updates , as well as quantization of these updates . This is achieved by extending the typical RD loss with an additional model rate term M that measures the number of bits required to send the model updates under a newly introduced model prior , resulting in a combined RDM loss . As an initial proof of concept , we show that this approach can lead to very substantial gains in RD performance ( ∼ 1 dB PSNR gain at the same bitrate ) on the problem of I-frame video coding , where a set of key frames , sampled from a video at 2 fps , are independently coded using an I-frame ( image compression ) model . Additionally , we show how the model rate bits are distributed across the model , and ( by means of an ablation study ) quantify the individual gains achieved by including a model-rate loss and using quantization-aware finetuning . The rest of this paper is structured as follows . Section 2 discusses the basics of neural compression and related work on adaptive compression . Section 3 presents our method , including details on the RDM loss , the choice of the model prior , its quantization , and the ( de ) coding procedure . In Sections 4 and 5 we present our experiments and results , followed by a discussion in Section 6 . 2 PRELIMINARIES AND RELATED WORK . 2.1 NEURAL DATA COMPRESSION . The standard approach to neural compression can be understood as a particular kind of VAE ( Kingma & Welling , 2013 ) . In the compression literature the encoder qφ ( z|x ) is typically defined by a neural network parameterized by φ , with either deterministic output ( so qφ ( z|x ) is one-hot ) ( Habibian et al. , 2019 ) or with fixed uniform [ 0 , 1 ] noise on the outputs ( Ballé et al. , 2018 ) . In both cases , sampling z ∼ qφ ( z|x ) is used during training while quantization is used at test time . The latent z is encoded to the bitstream using entropy coding in conjunction with a latent prior pθ ( z ) , so that coding z takes about − log pθ ( z ) bits ( up to discretization ) . On the receiving side , the entropy decoder is used with the same prior pθ ( z ) to decode z and then reconstruct x̂ using the decoder network pθ ( x|z ) ( note that we use the same symbol θ to denote the parameters of the prior and decoder jointly , as in our method both will have to be coded and added to the bitstream ) . From these considerations it is clear that the rate R and distortion D can be measured by the two terms in the following loss : LRD ( φ , θ ) = β Eqφ ( z|x ) [ − log pθ ( z ) ] ︸ ︷︷ ︸ R +Eqφ ( z|x ) [ − log pθ ( x|z ) ] ︸ ︷︷ ︸ D . ( 1 ) This loss is equal ( up to the tradeoff parameter β and an additive constant ) to the standard negative evidence lower bound ( ELBO ) used in VAE training . The rate term of ELBO is written as a KL divergence between encoder and prior , but since DKL ( q , p ) = R − H [ q ] , and the encoder entropy H [ q ] is constant in our case , minimizing the KL loss is equivalent to minimizing the rate loss . Neural video compression is typically decomposed into the problem of independently compressing a set of key frames ( i.e . I-frames ) and conditionally compressing the remaining frames ( Lu et al. , 2019 ; Liu et al. , 2020 ; Wu et al. , 2018 ; Djelouah et al. , 2019 ; Yang et al. , 2020a ) . In this work , we specifically focus on improving I-frame compression . 2.2 ADAPTIVE COMPRESSION . A compression model is trained on a dataset D with the aim of achieving optimal RD performance on test data . However , because of limited model capacity , optimization difficulties , or insufficient data ( resulting in poor generalization ) , the model will in general not achieve this goal . When the test data distribution differs from that of the training data , generalization will not be guaranteed even in the limit of infinite data and model capacity , and perfect optimization . A convenient feature of neural compression however is that a model can easily be finetuned on new data or data from a specific domain . A model can for instance ( further ) be trained after deployment , and Habibian et al . ( 2019 ) showed improved RD gains after finetuning a video compression model on footage from a dash cam , an approach dubbed adaptive compression ( Habibian et al. , 2019 ) . In adaptive compression , decoding requires access to the adapted prior and decoder models . These models ( or their delta relative to a pretrained shared model ) thus need to be signaled . When the amount of data coded with the adapted model is large , the cost of signaling the model update will be negligible as it is amortized . However , a tradeoff exists , the more restricted the domain of adaptation , the more we can expect to gain from adaptation ( e.g . an image compared to a video or collection of videos ) . In this paper we consider the case where the domain of adaptation is a set of I-frames from a single video , resulting in costs for sending model updates which become very relevant . 2.3 CLOSING THE AMORTIZATION GAP . Coding model updates can easily become prohibitively expensive when the model is adapted for every instance . However , if we only adapt the encoder or latents , no model update needs to be added to the bitstream , since the encoder is not needed for decoding as the latents are sent anyway . We can thus close , or at least reduce , the amortization gap ( the difference between qφ ( z|x ) and the optimal encoder ; Cremer et al . ( 2018 ) ) without paying any bits for model updates . Various authors have investigated this approach : Aytekin et al . ( 2018 ) ; Lu et al . ( 2020 ) adapt the encoder , while Campos et al . ( 2019 ) ; Yang et al . ( 2020b ) ; Guo et al . ( 2020 ) adapt the latents directly . This simple approach was shown to provide a modest boost in RD performance . 2.4 ENCODING MODEL UPDATES . As mentioned , when adapting ( parts of ) the decoder or prior to an instance , model updates have to be added to the bitstream in order to enable decoding . Recent works have proposed ways to finetune parts of the model , while keeping the resulting bitrate overhead small . For instance Klopp et al . ( 2020 ) train a reconstruction error predicting network at encoding time , quantize its parameters , and add them to the bitstream . Similarly ( Lam et al. , 2019 ; 2020 ) propose to finetune all parameters or only the convolutional biases , respectively , of an artifact removal filter that operates after decoding . A sparsity-enforcing and magnitude-suppressing penalty is leveraged , and additional thresholding is applied to even more strictly enforce sparsity . The update vector is thereafter quantized using k-means clustering . Finally , Zou et al . ( 2020 ) finetune the latents in a meta-learning framework , in addition to updating the decoder convolutional biases , which are quantized by k-means and thereafter transmitted . All these methods perform quantization post-training , leading to a potentially unbounded reduction in performance . Also , albeit the use of regularizing loss terms , no valid proxy for the actual cost of sending model updates is adopted . Finally , none of these methods performs adaptation of the full model . The field of model compression is related to our work as the main question to be answered is how to most efficiently compress a neural network without compromising on downstream task performance ( Han et al. , 2016 ; Kuzmin et al. , 2019 ) . Bayesian compression is closely related , where the model weights are sent under a model prior ( Louizos et al. , 2017 ; Havasi et al. , 2018 ) as is the case in our method . Instead of modeling uncertainty in parameters , we however assume a deterministic posterior ( i.e . point estimate ) . Another key difference with these works is that we send the model parameter updates relative to an existing baseline model , which enables extreme compression rates ( 0.02-0.2 bits/param ) . This concept of compressing updates has been used before in the context of federated learning ( McMahan et al. , 2017 ; Alistarh et al. , 2017 ) as well . We distinguish ourselves from that context , as there the model has to be compressed during every iteration , allowing for error corrections in later iterations . We only transmit the model updates once for every data instance that we finetune on . 3 FULL-MODEL INSTANCE-ADAPTIVE COMPRESSION . In this section we present full-model finetuning on one instance , while ( during finetuning ) taking into account both model quantization and the costs for sending model updates . The main idea is described in Section 3.1 , after which Section 3.2 and 3.3 respectively provide details regarding the model prior and its quantization . The algorithm is described in Section 3.4 . 3.1 FINETUNING AT INFERENCE TIME . Full-model instance-adaptive compression entails finetuning of a set of global model parameters { φD , θD } ( obtained by training on dataset D ) on a single instance x . This results in updated parameters φ , θ , of which only θ has to be signaled in the bitstream . In practice we only learn the changes with respect to the global model , and encode the model updates δ = θ − θD of the decoding model . In order to encode δ , we introduce a continuous model prior p ( δ ) to regularize these updates , and use the quantized counterpart p [ δ̄ ] for entropy ( de ) coding them ( more on quantization in Section 3.3 ) . The overhead for sending quantized model update δ̄ is given by model rate ĎM = − log p [ δ̄ ] , and can be approximated by its continuous counterpart M = − log p ( δ ) ( see Appendix A.1 for justification ) . Adding this term to the standard RD loss using the same tradeoff parameter β , we obtain the instance-adaptive compression objective : LRDM ( φ , δ ) = LRD ( φ , θD + δ̄ ) + β ( − log p ( δ ) ) ︸ ︷︷ ︸ M . ( 2 ) At inference time , this objective can be minimized directly to find the optimal model parameters for transmitting datapoint x . It takes into account the additional costs for encoding model updates ( M term ) , and incorporates model quantization during finetuning ( LRD evaluated at θ̄ = θD + δ̄ ) .

Summary: This paper investigates how to improve the test time performance of learned image compression models through finetuning of the full model. The authors finetune the model (both the model parameters and the prior on the latent space) for every test-time instance, appending the model updates to the bitstream. The model updates are coded according to a discretised, mean-zero Gaussian distribution with a single learned variance. They demonstrate that this approach yields a superior rate-distortion curve than the non-finetuned model on a set of I-frame video data.