{'fraction_non_alphanumeric': 0.03679739425150755, 'fraction_numerical': 0.01443725516087856, 'mean_word_length': 4.515901917941248, 'pattern_counts': {'":': 0, '<': 0, '<?xml version=': 0, '>': 0, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 2, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
Learning to Learn to Compress

In this paper we present an end-to-end meta-learned system for image compression. Traditional machine learning based approaches to image compression train one or more neural network for generalization performance. However, at inference time, the encoder or the latent tensor output by the encoder can be optimized for each test image. This optimization can be regarded as a form of adaptation or benevolent overfitting to the input content. In order to reduce the gap between training and inference conditions, we propose a new training paradigm for learned image compression, which is based on meta-learning. In a first phase, the neural networks are trained normally. In a second phase, the Model-Agnostic Meta-learning approach is adapted to the specific case of image compression, where the inner-loop performs latent tensor overfitting, and the outer loop updates both encoder and decoder neural networks based on the overfitting performance. Furthermore, after meta-learning, we propose to overfit and cluster the bias terms of the decoder on training image patches, so that at inference time the optimal content-specific bias terms can be selected at encoder-side. Finally, we propose a new probability model for lossless compression, which combines concepts from both multi-scale and super-resolution probability model approaches. We show the benefits of all our proposed ideas via carefully designed experiments.


I. INTRODUCTION
In recent years, neural networks have been applied to image and video compression with promising results. The research community has followed mainly two directions. In one direction, one or more components of a traditional codec pipeline are implemented as neural networks, such as an in-loop filter [1] or a post-processing filter [2]. In another direction, commonly referred to as end-to-end learned compression, neural networks are used as the main components of the codec [3]. End-to-end learned methods have recently outperformed traditional codecs [4]. These systems typically follow the auto-encoder paradigm, where the encoder and decoder networks operate as non-linear transform and inverse transform, respectively, and a quantization step is used on the latent tensor output by the encoder [5]. The quantized latents are then losslessly encoded, typically by using arithmetic coding with a learned probability model.
A desired feature of machine learning tools is to generalize well to unseen content. However, current architectures and training methods still suffer from domain shift, where the content type at test time is different from the content type considered during training. In the context of image compression, domain shift is likely to cause low rate-distortion (RD) performance at test time. Even when there is not significant domain shift between training and testing, the neural networks in the codec are not optimized on every unseen test image and thus content adaptation may still bring RD gains. One possible approach is to adapt the encoder neural network's parameters [6]. Another approach is to adapt the decoder network's parameters [7]. However, this latter approach would incur heavy bitrate overheads as the adapted parameters need to be signaled to the decoder. Yet another approach consists of adapting the latent tensor that is output by the encoder. As such adaptation processes is not meant for generalization, we refer to it also as benevolent overfitting, or simply as overfitting.
However, neural networks in traditional end-to-end learned methods are trained for the only purpose of generalization and not for maximizing the performance of the above overfitting approaches, thus generating a gap between training and inference stages. We argue that networks shall be trained specifically for both generalization and quick overfitting, to reduce the number of overfitting iterations needed to achieve a certain gain. The overfitting may be considered to be a form of few-shots learning, more specifically 1-shot learning, where a model is adapted given a single example of an image.
In this paper, we describe L 2 C, our end-to-end learned image compression system that we submitted to the 2020 JPEG-AI Learning-based Image Coding Challenge. Our main novel contributions are: • A novel training paradigm for learned image compression, based on meta-learning for latent tensor overfitting. • A novel probability model based on multi-scale progressive statistical model. • A novel and efficient strategy for generating an ensemble of content-specific decoders.

II. PRIOR ART
In end-to-end learned image codecs, the latent tensor output by the encoder is losslessly encoded by an entropy encoder given the value distribution function estimated by the probability model. The probability model is also referred to as an entropy bottleneck layer as it constrains the amount of information that can be passed through the latent tensor [8]. Element-wise Gaussian distribution or other parametric distribution functions have been proposed as the statistical model of the latent tensor [8]- [10]. However, the assumption that the elements in the latent tenor are independent of each other does not hold in practice. Experiments have shown that the spatial correlation existing in the latent tensor degrades the compression performance [8] [11]. The authors of [8] [11] address this problem by transferring extra information as a hyper-prior to reduce this correlation. Instead of assuming the elements in the latent tensor are independent and identically distributed, we propose to use a powerful statistical model that is capable of capturing this spatial correlation. The proposed multi-scale progressive statistical model is based on previous multi-scale models used for lossless image compression [12], [13] and lossy image compression [14].
Overfitting the codec's parameters at inference time for compression has already been explored in the past. In [6], the authors propose to overfit the encoder neural network to the test image. A similar technique was later used also in [15], [16] for image and video compression, respectively. In [7], the authors proposed to signal compressed weightupdates to the decoder. In order to encourage weight-updates to be robust to compression, a weight-update compression loss is used during the inference time overfitting. In [17], the authors propose an ensemble learning-based rate-distortion optimization (RDO) method to enhance end-to-end learned image compression. Multiple networks with same structure but different parametrization are used for the probability model in the lossless coding module. At inference time, the encoder selects the optimal probability model for each image block, and an index is signaled to the decoder. Latent tensor overfitting was proposed in [18], where the authors claim to achieve 0.5 dB gain.
As already mentioned in Section I, a specialized training strategy may bring benefits in terms of overfitting performance. We see overfitting on test images as a 1-shot learning problem. Many of the current state-of-the-art few-shots learning algorithms are based on Model-Agnostic Meta-Learning (MAML) [19], where training is split into an inner loop and an outer loop. In the inner loop, for each task in a batch of tasks, a model is adapted by gradient descent using training data (support set in meta-learning terminology). In the outer loop, the adapted models are evaluated on validation data (query set). The model is then updated by gradient descent on the evaluation losses. As the evaluation is performed using adapted models, the outer loop update computes second-order derivatives. In [20], first-order approximations of MAML were proposed, where the inner loop gradients are considered to be constant. Our proposed meta-learning strategy is inspired by MAML methods but it also presents fundamental differences. In particular, a "task" in MAML is for example a classification task, whereas in L 2 C it is an image. For each "task", support set and query set include different images in MAML, whereas they include the same image in L 2 C. This is because our adaptation is in fact an overfitting operation, whereas in MAML the adaptation aims at achieving generalization on the given task. In a variant of MAML, called CAML [21], instead of adapting the parameters of a neural network, an additional context vector is adapted. Similarly, in L 2 C the adaptation is performed on a tensor. However, CAML presents several differences to L 2 C, in addition to those related to vanilla MAML. First, CAML aims at achieving reduced metaoverfitting, easier parallelization, and better interpretability, whereas L 2 C aims at achieving better inference time overfitting. Second, the nature of the tensor is different: L 2 C considers a latent tensor, output by another neural network. Third, the final goal of L 2 C is image compression. In [22], a neural network is meta-learned so that it can be quickly pruned on the first video frame, and then used for tracking objects in subsequent frames. However, the target of compression is a neural network instead of the input content itself, and the final goal is object tracking instead of compression.

III. METHODS
In this section we describe in detail both the training phase and the inference phase of our end-to-end learned image compression codec L 2 C. An overview is provided in Fig. 1.
The main components of L 2 C are a fully-convolutional auto-encoder, an importance-map module for spatially-varying channel allocation, quantization and dequantization modules, a lossless coding module based on arithmetic coding and a learned probability model, and an ensemble of decoder-side clusters of content-specific bias terms.
At encoder side, the input frame x is projected into latent space by an encoder neural network E θ parametrized by weights θ, obtaining the latent tensor y ∈ R H s , W s ,c , where H, W are the height and width of x, s is the down-sampling factor of E θ , c is the number of filters in the last layer of E θ . The output latent tensor is multiplied by a binary importance mask m (see Section III-A) that zeros-out a spatially-varying number of channels: where indicates element-wise multiplication. The masked output is quantized into z by uniform scalar quantization using b bits and then entropy coded by an arithmetic encoder. In order to allow for back-propagating non-zero gradients, we use the straight-through estimator for quantization, as in [5]. A learned multi-scale context model is used by the arithmetic codec to estimate the probability distribution of next symbols to encode/decode. At decoder side, the entropy decoded bitstream is dequantized intoŷ and input to a decoder neural network D φ parametrized by weights φ, thus obtaining the reconstructed output imagex.

A. Learned Spatially-varying Channel Masking
In order to allow the model to allocate a varying number of channels to different spatial areas of the encoded tensor y (see (1)), we use an additional neural network I ω parametrized by weights ω and with similar architecture as E θ . I ω takes as input the image x and outputs an importance map τ ∈ R In order to encourage masked representationsỹ that have low entropy and thus be more easily predictable by our probability model for arithmetic coding, we use the following constraint in our training objective function: whereτ is the mean value of τ and ζ is a constant representing the target average non-zero ratio in m (and thus in y m ).

B. Probability Model for Lossless Coding
Our proposed probability model implements a multi-scale progressive statistical model based on the technologies used on lossless image compression [12], [13]. We first downsample the latent tensor to a series of low-resolution representations in multiple scales. The low-resolution representations are used as a context in the statistical model for the elements in a highresolution representation. The system processes the multi-scale representations backward. The representation in the last scale is sent without compression since there is no context information for it. Other representations are encoded/decoded based on the estimated distribution function using the already processed low-resolution representations as input. This procedure repeats until the input latent tensor is processed.
Let z (0) be the latent representation to be modeled and z (i) be the low resolution representation of z (0) at scale i, where i = 1, 2, · · · , M and M is the number of scales. The joint distribution function of elements in z (0) is defined by where p z (M ) is the distribution function of the lowresolution representation at the last scale, which is assumed to follow a uniform distribution. We choose the nearest neighbor downsampling method to avoid the extra effort of encoding round-off errors as used in [13]. Note that this follows the design principle of the coupling layer in Real-NVP algorithm [23]. With this design, the elements at low-resolution representations are directly matched to part of the elements in a high-resolution representation.
To further improve the compression rate, we partition the elements to be processed at each scale into multiple groups and process the groups one by one. The groups that have been processed are added to the context. Let g (i) j be group j at scale i, B i be the number groups at scale i, and C (i) (4) can be written as Let z k be a element to be processed in group g We assume p z k |C (i) j follows a mixture of logistic distributions parameterized by the mixture weights, logistic distribution means and scales [24]. The parameters are determined by a function modeled by a deep neural network with C (i) j as its input.
The architecture of the proposed probability model is illustrated in Figure 2. The progressive statistical model component takes C Mixer Mixer̂ ( ) 3 Fig. 2. Architecture of the probability model that implements a multi-scale progressive statistical model.
At the training stage, the entropy of the input latent representation z (0) is estimated using p (i) j for all steps in all scales. At the training stage, outputs p (0) , p (1) , · · · p (M −1) are used to calculate the cross-entropy of the output tensors z (0) , z (1) , · · · , z (M −1) and the sum of these cross-entropies is taken as the compression loss L r (z). At the encoding/decoding stage, the value distribution functions for the elements to be encoded are derived from p (i) j . Then the distribution functions are given to the arithmetic encoder/decoder to encode/decode the elements in z (i) j . This process starts from the last scale where the elements are encoded/decoded using a uniform distribution function. The encoding/decoding continues to every scale until all elements in z (0) are processed.

C. Meta-Learning
The training process is organized in two stages. In the first stage, a conventional training session is performed to achieve generalization, by using the following rate-distortion loss: where L d1 is the negative multi-scale structural similarity (MS-SSIM) [25], L d2 is the mean-squared error (MSE), L d3 is a perceptual loss, L r is the rate loss provided by the probability model, M is the constraint on the importance map defined in (3). λ d1 , λ d2 , λ d3 , λ r , λ m are scalar values that are determined empirically. Inspired by [26], the perceptual loss is where |·| 1 is the 1 -norm, F 22 and F 43 are the ReLU 2-2 and ReLU 4-3 layers of VGG16 [27] pretrained on ImageNet. The MSE loss ensures that pixel-wise differences are minimized, whereas MS-SSIM and the perceptual loss aim at improving the human-perceived quality. The second stage is aimed at training the neural networks so that the performance of latent tensor overfitting is maximized at inference time. E θ and I ω need to be trained to output a latent tensor which can overfit quickly and effectively, and P γ and D φ need to be trained to provide effective gradients during the overfitting process and to reconstruct high quality images from overfitted input latent tensors. We frame the problem as a few-shots learning problem, where the neural networks need to be trained to allow for overfitting the latent tensor in few iterations. As there's only one example image or shot, this can be regarded as 1-shot learning. The system is effectively learning to overfit for the task of image compression or, in other words, is learning to learn to compress (L 2 C). In practice, we fine-tune the neural networks trained in the first stage, by leveraging meta-learning and in particular MAML. The nature of the two nested loops used L 2 C is different than in vanilla MAML. In the inner loop, a latent tensor is overfitted for each image in a batch. In the outer loop, the neural networks are updated based on the average performance of the overfitted latent tensors over all images in the batch.
1) Inner-Loop: A batch of images x = [x 1 , x 2 , .., x B ] is sampled from the training set. This is equivalent to a batch of tasks in vanilla MAML. For each image x i , we perform an overfitting session. The initial step is to run one forward pass of the encoder-side networks, thus obtaining the initial latent tensor before quantization,ỹ The latent tensor is then updated by gradient descent for n overfitting iterations and with learning rate α. For each k ∈ {1, .., n}: 2) Outer-Loop: The overfitted latent tensorỹ (n) i of each image x i is used for computing an evaluation loss, L xi (ỹ (n) i , θ, ω, φ, γ). The neural networks' parameters are then updated by gradient descent with learning rate β: Assuming that n = 1, this is equivalent to: i , θ, ω, φ, γ), θ, ω, φ, γ . (10)

D. Adapting Decoders' Parameters
We propose to have multiple sets of overfitted decoder's parameters, from which the encoder can choose at inference time in order to adapt the decoding process to the test image. To this end, after training has completed, we overfit a subset of decoder's parameters to each 256x256 patch of the training set. In [28], the encoder adapts a neural network that is used as a post-processing filter within a conventional decoder. The authors found that updating only the bias terms is a good trade-off between gain in reconstruction quality and bitrate overhead incurred by signaling the updated weights. Thus, we overfit only the bias terms of the convolutional layers of D φ . The obtained sets of overfitted bias terms are clustered by k-means into 255 clusters. These clusters represent contentspecific bias terms. At inference time, the optimal cluster can be signaled by using only 8 bits. One index is reserved for signaling that no adaptation is needed and the default bias terms shall be used.

E. Rate Control and Inference Pipeline
The JPEG-AI Challenge defines the following 8 target bitrates: {r T 0 = 2.0, r T 1 = 1.5, r T 2 = 1.0, r T 3 = 0.75, r T 4 = 0.5, r T 5 = 0.25, r T 6 = 0.12, r T 7 = 0.06}, measured as bits-perpixel (BPP). We achieve multiple bitrates by several strategies: For each image x i and target bitrate r T j , with j ∈ {0, 1, 2, 3, 4, 5, 6, 7}, the codec with closest target bitrate is selected. For example, for r T 0 , codec C 0 is selected. Initially, the codec is run using b = 8 quantization bits. If the achieved bitrate is higher than r T 0 , b is decreased. Selection of the codec and of the number of bits is repeated until the bitrate is within the allowed margin for the target bitrate. Then, latent tensor overfitting is performed for further adjusting the bitrate and for optimizing the reconstruction quality. After overfitting, the optimal content-specific bias terms for D φ are selected for each image patch.

IV. EXPERIMENTS
This section presents the experimental setup and results. The number of channels in the convolutional layers of E θ , I ω and D φ is 192 except for the last layer, which is one 1 for I ω and 3 for D φ . Training was performed on patches of size 256x256 extracted from the JPEG-AI dataset, using a batch-size of 12. Adam optimizer was used for the first training stage and for the outer loop updates of the meta-learning fine-tuning stage. The initial training was performed for 200 epochs, whereas the meta-learning fine-tuning was performed for 5 epochs. In the inner-loop of the meta-learning fine-tuning, 4 overfitting iterations were used. The learning rates for the initial training, the outer loop updates and the inner loop overfitting were 0.0001, 0.0001, 0.1, respectively. Table I reports results for 4 images in the JPEG-AI test set. The number of parameters in E θ , I ω , P γ , D φ are 1.6M, 1.6M, 1.9M, 2.7M, respectively. We performed measurements for the inference time on a representative image, jpegai10. For each patch, the encoding used 10 overfitting iterations and then searched for the best biases among 255 sets. Encoding and decoding took 290.2 and 14.3 seconds, respectively, on a workstation equipped with Intel ® Core ™ i7-7820X CPU @ 3.60GHz, 8 cores, 64GB RAM, Nvidia RTX 2080 Ti GPU with 12 GB memory, PyTorch 1.3.1, Python 3.7.
We now compare our meta-learning fine-tuning to a baseline fine-tuning, in terms of latent tensor overfitting performance. For each number of iterations in {4, 10, 100}, each image and each fine-tuning method, we search for the best learning rate and report the results averaged over the test set in Table II, in terms of percentage increase in average loss drop due to overfitting. As a concrete example, for image jpegai02 and 4 overfitting iterations, the baseline model increases MS-SSIM from 0.9721 to 0.9727 (0.0006 MS-SSIM gain) and BPP from 0.7666 to 0.7694 (0.0028 BPP increase), whereas L 2 C model increases MS-SSIM from 0.9701 to 0.9712 (0.0011 MS-SSIM gain) and BPP from 0.6787 to 0.6812 (0.0025 BPP increase). In Table III we show the effectiveness of the latent tensor overfitting with our meta-learned model in decreasing the bitrate and in improving the reconstruction quality. For each of these two experiments, we randomly selected three images. Finally, we evaluated our content-specific bias terms. The mean MS-SSIM gain is 0.0029, averaged over all test images and bitrates.

V. CONCLUSIONS
We presented our end-to-end learned solution for the JPEG-AI challenge, which introduces a new training paradigm for learned image compression, based on meta-learning. The models are trained to allow for a more effective latent tensor overfitting at inference stage. Also, we proposed a novel probability model for the lossless coding module, and a mechanism for generating and using content-specific decoder-side parameters. In our experiments, we showed how these proposed idea bring benefits in terms of rate-distortion performance.