Image Compression with Product Quantized Masked Image Modeling

Alaaeldin El-Nouby, Matthew J. Muckley, Karen Ullrich, Ivan Laptev, Jakob Verbeek, Hervé Jégou

Introduction

Image compression played a crucial role in accelerating multiple events in history: it was a major component of the Voyager mission (Ludwig & Taylor, 2016). Efficient image codecs have accelerated the rapid growth of the internet by enabling the transmission of images in a few dozens of kilobytes, thanks to the emergence of effective lossy methods. This democratization was accompanied by standardization efforts to facilitate the interoperability, which led to the emergence of standards such as the Joint Photographic Experts Groups (JPEG). Subsequent formats have leveraged scientific advances on all components of source coding, ranging from transforms (Antonini et al., 1992), and quantization (Gray & Neuhoff, 1998), to entropy coding (Witten et al., 1987; Taubman, 2000), eventually leading to modern video compression codecs enabling streaming and video-conferencing applications.

Neural methods have recently become increasingly popular for image compression as well as other image processing tasks, such as denoising (Tian et al., 2020), super-resolution (Bruna et al., 2016; Dong et al., 2015; Ledig et al., 2017; Wang et al., 2021) or image reconstruction (Wang et al., 2020; Knoll et al., 2020). In typical scenarios, neural image compression is not necessarily mature enough to take over standard techniques like the BPG format inherited from the High-Efficiency Video Coding standard (Sullivan et al., 2012). This is because they do not offer a significant quantitative advantage over prior works that would justify the higher complexity, which depends on the context and operational constraints. A key advantage of neural compression methods is their enhanced qualitative reconstruction when incorporating an adversarial loss or likewise psycho-visual objectives favoring visually appealing reconstructions (Agustsson et al., 2019; Mentzer et al., 2020). From this perspective, neural compression is related to image generation. The two subfields, however, are currently dominated by different approaches, noticeably they employ different discretization procedures. Indeed, while earlier neural compression methods utilized vector quantization (Agustsson et al., 2017, VQ), recent methods mostly employ scalar quantization (SQ). In contrast, the recent literature on image generation (Chang et al., 2022; Yu et al., 2021; Esser et al., 2021; Rombach et al., 2022) relies on Vector Quantization jointly with a distortion criterion akin to those used in compression.

In this work we aim to reduce the methodological gap and to make a step towards unification of neural image compression and image generation, and allowing image compression to more directly benefit from the rapid advances in image generation methods. Patch-based masking methods for self-supervised learning (Bao et al., 2022; He et al., 2021; El-Nouby et al., 2021a) have recently demonstrated their potential for image generation (Chang et al., 2022). Inspired by this work we propose a compression approach built upon Vector Quantized Variational Auto-Encoders (Oord et al., 2017; Razavi et al., 2019). In this context, we focus on two intertwined questions: (1) How to define a vector quantizer offering a range of rate-distortion operating points? (2) How to define an entropy model minimizing the cost of storing the quantization indexes, while avoiding the prohibitive complexity of an auto-regressive model?

To address the challenges above, we revisit vector quantization in image compression, and investigate product quantization Jégou et al. (2010) (PQ) in a compression system derived from VQ-VAE Oord et al. (2016b). We show that PQ offers a strong and scalable rate-distortion trade-off. We then we focus on the spatial entropy modeling and coding of the quantization indexes in the VQ or PQ latent layer, hence we name our method as (Vector/Product)-Quantized Masked Image Modeling (VQ-MIM and PQ-MIM). To this end, we introduce a multi-stage vector-quantized image model: we gradually reduce the conditional entropy of the patch latent codes by increasing the number of observed patches we condition on for each stage. The conditional distribution over patches is estimated by a transformer model and provided to an entropy coder, symmetrically on the emitter and receiver sides.

In summary, we make the following contributions:

We introduce a novel Masked Image Modeling conditional entropy model that significantly reduces the rates by leveraging the spatial inter-dependencies between latent codes.

We introduce product quantization for VQ-VAE. This simple PQ-VAE variant offers a strong and scalable rate-distortion trade-off.

When trained with adversarial and perpetual losses, PQ-MIM exhibits a strong performance in terms of perceptual metrics like FID and KID, outperforming HiFiC (Mentzer et al., 2020).

We qualitatively show that PQ-MIM is capable of operating in a hybrid mode, between generative and compression, without requiring further training and finetuning. This allows for higher resilience to corrupted or missing signal where our model can fill-in the missing information.

Related work

Neural image compression Early approaches to neural image compression reach back to the late 1980s (Sonehara, 1989; Sicuranza et al., 1990; Bottou et al., 1998). Recent rapid advances in explicit and implicit density modelling (Goodfellow et al., 2014; Kingma & Welling, 2014; Larochelle & Murray, 2011; Oord et al., 2016a; Rezende & Mohamed, 2015; Salimans et al., 2017) have renewed interest in posing image compression as a learning problem. Due to the connection to variational learning (Gregor et al., 2016; Frey, 1998; Alemi et al., 2018), variational auto-encoders have been the primary choice for lossy image compression. In contrast to standard variational models, the evaluation of neural compression models focuses on achievable bitrates, multi-scale applicability and computational complexity. The initial works in the field used fully convolutional architectures for encoding/decoding (Ballé et al., 2017; Theis et al., 2017; Mentzer et al., 2018). The resulting encoded latent image representations are quantized and compressed via an entropy coder with learned explicit density model or “entropy bottleneck”. Initial variational approaches directly modeled a single level of code densities. Ballé extended these models by introducing a second “hyperprior” that yielded improved performance (Ballé et al., 2018). Hyperprior models have been the basis for several subsequent advances with further improvements for density modeling, such as joint autoregressive models (Minnen et al., 2018), Gaussian mixture/attention (Cheng et al., 2020), and channel-wise auto-regressive models (Minnen & Singh, 2020). Another line of work has proposed to use vector quantization with histogram-based probabilities for image compression (Agustsson et al., 2017; Lu et al., 2019). Contrary to VQ-VAE models, these models typically optimize the rate (or a surrogate of the rate) directly and may include a spatial component for the quantized vectors. Yang et al. (Yang et al., 2020) showed multiple ways to improve the encoding process, including the fine-tuning of the discretization process and employing bits-back coding in the entropy bottleneck. Finally, (Santurkar et al., 2018; Mentzer et al., 2020; Rippel & Bourdev, 2017; Agustsson et al., 2019) showed that altering the distortion metric to include an additional adversarial loss can make a large difference for compression rate. Another interesting line of work considers image compression by training image-specific networks, or network adapters, that map image coordinates to RGB values, and compressing the image-specific parameters (Dupont et al., 2021; 2022; Strüpler et al., 2022).

VQ-models for image generation. There has been significant interest in generative models based in discrete image representations, as introduced by VQ-VAE (Oord et al., 2017; Razavi et al., 2019). A discrete representation of reduced spatial resolution is learned by means of an autoencoder which quantizes the latent representations. This discrete representation is coupled with a strong prior, for example implemented as an autoregressive pixel-CNN model. VQ-GAN (Esser et al., 2021) replaces the prior architecture with a transformer model (Vaswani et al., 2017), and introduces an adversarial loss term to learn an autoencoder with more visually pleasing reconstructions and improved sample quality. This approach has been extended to text-based generative image models by extending the prior to model a longer sequence that combines the discrete image representation with a prefix that encoding the conditioning text. This has yielded impressive results by scaling the model capacity and training data to tens or hundreds of million text-image pairs (Ding et al., 2021; Gafni et al., 2022; Ramesh et al., 2021). A fundamental limitation of autoregressive generative models is that they sample data sequentially, requiring separate non-parallel evaluation of the predictive (transformer) model to sample each token. To alleviate this, several data items can be sampled independently in parallel, conditioning on all previously sampled tokens. This has been leveraged to speed-up pixel-CNNs for small images and video by two to three orders of magnitude (Reed et al., 2017). More recently, MaskGIT (Chang et al., 2022) and follow-up work (Lezama et al., 2022) explored this for generative models of VQ-VAE representations, and find that similar or better sample quality is obtained by parallel sampling of image patch subsets in few steps, reducing the generation time significantly. Despite their success for image synthesis we are not aware of the use of such models for image compression in earlier work.

Product Quantized Masked Image Modeling

This work takes a step towards to closing the gap between neural image compression and image generation methodologies. We revisit vector quantization for image compression and propose an entropy model inspired by masked image modelling. Our compression pipeline, depicted in Figure 2, relies on three neural networks:

The Encoder network E^:X→Z\hat{E}:\mathcal{X}\rightarrow\mathcal{Z} maps input images x∈X\mathbf{x}\in\mathcal{X} to a quantized representation z∈Z\mathbf{z}\in\mathcal{Z}.

The Mask Image Model (MIM) compresses the quantized representations without loss of information. This network is involved both on the compression and decompression side.

The Decoder network G:Z→XG:\mathcal{Z}\rightarrow\mathcal{X} produces an estimate x^=G(E^(x))\hat{\mathbf{x}}=G(\hat{E}(\mathbf{x})) of the original image x\mathbf{x}.

We now detail the architecture, in particular our PQ proposal, of the image model that we employ in the statistical lossless coding, as well as the training scheme.

Product Quantization. In VQ-VAE, the quantizer QQ is simply an online k-means quantizer that produces quantization indices from real-valued vectors. We denote by q∈{1,…,V}T\mathbf{q}\in\{1,\dots,V\}^{T} the map of quantization indices indicating for each element of z\mathbf{z} which of the VV centroids is selected. The higher VV the more precise is the approximation z\mathbf{z}, leading to higher bit rates. For instance, assuming that indices are coded with a naive coding scheme (see next section), the bit rate is doubled when moving to V=256V=256 to V=65536V=65536 centroids.

Neural network. Without loss of generality, we choose all neural network models to be identical. This is not a requirement but this offers the property that the encoder and decoder have identical complexities, and that the memory and compute peaks are identical. More specifically we choose a cross-variance transformer (El-Nouby et al., 2021b) (XCiT), whose complexity is linear with respect to image resolution. In contrast, standard vision transformers (Dosovitskiy et al., 2021) (ViT) are quadratic in the image surface, which is prohibitive for high resolution images that can typically require strong compression. We point out that recent work has shown that Swin-Transformers (Zhu et al., 2021) could be a compelling choice as well in the context of image compression. Formally they also have a quadratic complexity, but this is amortized by the hierarchical structure of this architecture.

2 Image entropy model

In this section we present PQ-MIM. The objective is to compress the discrete representations q\bm{q} without loss of information, producing a bitstream of the compressed representation that can be transmitted or stored. During the decoding stage, we invert the aforementioned lossless compression. This model can be regarded as the VQ/PQ-VAE counterpart of adaptive contextual arithmetic coders, like EBCOT (Taubman, 2000) or CABAC (Richardson, 2004), proposed in early compression standards, in that it couples a conditional probabilistic model with an arithmetic coder.

Lossless compression. A naive manner for lossless compression of the discrete image codes q={qt}t=1T\mathbf{q}=\{q_{t}\}_{t=1}^{T} is to use fixed-length codes. In that case, each code word is assigned to a unique binary representation of equal length, resulting in ⌈log⁡2V⌉\lceil{\log_{2}V}\rceil bits per element qtq_{t}. This approach is computationally very efficient as fixed-length codes are not model-based, and as such do not require likelihood estimation, and because all codes are of equal length by construction, the computation is perfectly parallelizable. However, theoretically this coding scheme could be Shannon optimal only if codewords are uniformly and independently distributed. Those assumptions are not met in practice due to the architecture choices we have made previously: k-means does not produce uniformly distributed indices except in singular cases (Gray & Neuhoff, 1998). More details about losseless compression are covered in Appendix A.

To improve the bitrate, we hence rely on an entropy coder, which provides an inverse pair of functions, encp\texttt{enc}_{p} and decp\texttt{dec}_{p}, achieving near optimal compression rates on sequences of symbols for any distribution pp. The better pp matches the (unknown) underlying data distribution, the better the compression rate (Cover, 1991). Generally, more powerful generative models will ensure better compression performance.

Fully autoregressive generative models p(q)=∏t=1Tp(qt∣q<t)p(\mathbf{q})=\prod_{t=1}^{T}p(q_{t}|\mathbf{q}_{<t}) are powerful (Ding et al., 2021; Esser et al., 2021; Gafni et al., 2022; Ramesh et al., 2021; Yu et al., 2021), however, they are inconvenient in that the likelihood estimation for this type of model is not trivially parallelizable: each patch index must be processed sequentially as it is used to condition subsequent patch indices. Thus, similar to prior works (Chang et al., 2022; Reed et al., 2017) we propose a masked image model, which we use to predict the image patch indices in several stages. Specifically, we partition q\mathbf{q} into SS subsets q1,q2,…,qS\mathbf{q}_{1},\mathbf{q}_{2},\dots,\mathbf{q}_{S} of patch indices, that we refer to as tokens by analogy to language modelling:

We model the elements in each subset conditionally independent given all preceding groups:

Since p(q1)p(\mathbf{q}_{1}) is not conditioned on any previous elements, it fully factorizes over qt∈q1q_{t}\in\mathbf{q}_{1}, and we model it as the marginal distribution over the vocabulary observed on the training data. The non-trivial conditional distributions p(qs∣q<s;θs)p\left(\mathbf{q}_{s}|\mathbf{q}_{<s};\theta_{s}\right) for s≥2s\geq 2 are modeled using transformer networks, which have few inductive biases and have been successful across many tasks, including image generation. An overview of the MIM entropy model is illustrated in Figure 2.

Quincunx partitioning. In practice, we typically use S=5S=5 stages. We have explored different patterns to partition the TT tokens over the SS stages. In particular we consider the “quincunx” regular grid pattern, where in each stage we double the number of tokens to predict, see Figure 4 for an illustration. This multi-level refinement was previously explored for image compression in the context of lifting schemes designed with oriented wavelets (Chappelier & Guillemot, 2006). In our experiments we contrast this partitioning with alternative ones with other patterns and subset cardinalities (Figure 9).

3 Training the PQ-MIM

Reconstruction objective and training. The goal of lossy image compression is to match the image and its reconstruction as closely as possible according to some distortion metric. In this paper, we train our model to reduce the image distortion and the quantization loss using the following objective:

Following VQ-VAE, our quantization objective LPQL_{\textrm{PQ}} consists of an embedding and commitment losses, averaged over the MM different PQ sub-vectors. For the distortion loss LRecL_{\textrm{Rec}}, we present two setups where we use different types of distortion measures:

MSE & MS-SSIM. Typical distortion measures used in the majority of the neural compression literature such as mean squared error (MSE) or multi-scale structural similarity (Wang et al., 2003) (MS-SSIM). For this setup, the model is trained solely using one distortion measure at a time.

Perceptual measures. Alternatively, we report a setup where we utilize perceptual objectives such as LPIPS (Zhang et al., 2018) and adversarial training (Goodfellow et al., 2014) to enhance psycho-visual image quality. The distortion loss is defined as:

where α\alpha and γ\gamma are weighing coefficients and the adversarial loss LAdvL_{\textrm{Adv}} is defined as:

Training the entropy model. Our MIM module is an XCiT transformer that accepts TT tokens as input representing the image patches. During training, we randomly mask a set of tokens by sampling from a uniform distribution U(0,1)U(0,1). The masked tokens are replaced with a mask embedding vector, while the observed token indices are mapped to their corresponding continuous representation using an embedding look-up table. The MIM module outputs a context vector for every masked token which in turn is passed to MM linear heads, representing the different PQ sub-vector indices, followed by a softmax to yield a distribution p(qs∣q<s)p({\bm{q}}_{s}|{\bm{q}}_{<s}). The module is trained using a standard cross-entropy objective. While we train the autoencoder and the MIM modules simultaneously, we do not backpropagate gradients from the MIM module to the encoder EE or the quantization parameters, so the two components can be trained separately in sequence.

Experiments

We first present our experimental setup in Section 4.1 and the present our results in Section 4.2. We provide ablation studies in Section 4.3. We will share our code and models.

Rate-distortion control. For all our experiments we fix the codebook size V=256V=256 and only vary the number of sub-vectors M∈{2,4,6}M\in\{2,4,6\} for two different down-sampling factors f∈{8,16}f\in\{8,16\}.

For the results reported in Figure 5, the models are trained using solely MSE (η=0.5\eta=0.5) or MS-SSIM (η=10.0\eta=10.0) distortion losses for their corresponding plots. As for the models trained with perceptual objectives (Figure 6 and Table LABEL:tab:discriminator), they are trained with a weighted sum of MSE, LPIPS (α=1\alpha=1) and adversarial loss (γ=0.1\gamma=0.1). We use a ProjectedGAN Discriminator (Sauer et al., 2021) architecture. The perceptual training is initialized with an MSE only trained checkpoint and trained for 50 epochs using ImageNet with a learning rate of 10−410^{-4} and weight decay of 5.10−55.10^{-5}. Similar to HiFiC (Mentzer et al., 2020), we freeze the encoder during the perceptual training. Additionally, we find that clipping the gradient norm to a maximum value of 4.04.0 improves the training stability. Our discriminator takes only the decoded image as an input and does not rely on any other conditional signal.

MIM implementation details. Our MIM module is an XCiT-L12 with 12 layers and embedding dimension of 768. MIM is trained simultaneously with the PQ-VAE, but the gradients are not backpropagated to the encoder or the quantizer parameters. The PQ indicies qq are split into inputs and targets for the MIM model as defined by the quincunx partitioning pattern. By default we use S=5S=5 stages. All stages are processed with the same MIM model. The masked patches are replaced by a learnable “mask” token embedding. The loss is computed only for the masked patches. Since every token is assigned MM PQ indices, the output of the MIM transformer is passed to MM separate linear heads to predict MM softmax normalized distributions over their corresponding codebooks. The marginal distribution of the codebook is computed as a normalized histogram over the ImageNet training set. For entropy coding, we use the implementation of the torchachttps://github.com/fab-jul/torchac arithmetic coder.

Datasets. We train our models using ImageNet (Deng et al., 2009). For data augmentation, we apply random resized cropping to 256 ⁣× ⁣256256\!\times\!256 images and horizontal flipping. For evaluation and comparison to prior work, we use Kodak (Kodak, 1993) and Tecnick (Asuni & Giachetti, 2014) datasets for PSNR and MS-SSIM. Moreover, we compute the perceptual metrics (FID (Heusel et al., 2017), KID (Bińkowski et al., 2018)) for perceptually trained models using the CLIC 2020 test-set (Toderici et al., 2020) (428 images) using the same patch cropping scheme detailed by Mentzer et al. (2020).

Baselines. We compare to several existing neural compression baselines: the scale hyperprior model (Ballé et al., 2018), mean & scale hyperprior (Minnen et al., 2018), GMM hyperprior (Cheng et al., 2020), and HiFiC (Mentzer et al., 2020). Among the non-neural codecs, we compare to the popular BPG (Bellard, ), WebP, and JPEG2000.

2 Main experimental results

Comparison to existing (neural) codecs. We compare PQ-MIM to other approaches across a wide range of bitrates in Figure 5All results for the baselines are reported using the authors official repositories. Note that in our evaluation we consider bitrates that are an order of magnitude lower than what is typically studied in the literature: most previous studies were limited to 0.1 bpp and above, see e.g (Ballé et al., 2018; Cheng et al., 2020; Minnen et al., 2018). The extremely low bitrates we consider make it possible to transmit a 256 ⁣× ⁣256256\!\times\!256 image in an SMS or a tweet (280 characters)For example, a bitrate of 0.03 yields 256×256×0.03/8≈246256\times 256\times 0.03/8\approx 246 bytes for a 256 ⁣× ⁣256256\!\times\!256 image. as shown in Figure 7. PQ-MIM achieves a strong and competitive performance for both Kodak and Tecnick datasets, outperforming all prior neural and standard codecs with the exception of GMM hyperprior (Cheng et al., 2020). We observe that PQ-MIM is particularly strong for low rates, making it a good fit for extreme compression scenarios. Moreover, PQ-MIM exhibits a particularly strong MS-SSIM performance which was designed to model the human visual contrast perception (Wang et al., 2004; 2003).

Perceptual metrics comparison. In Figure 6, we compare PQ-MIM to HiFiC (Mentzer et al., 2020) in perceptual quality measures like FID and KID as well as MS-SSIM. HiFiC is based on the mean & scale hyperprior model (Minnen et al., 2018), but adds adversarially trained discriminator model to improve the perceptual quality of the image reconstructions. PQ-MIM, with perceptual training, achieves a strong performance for all reported metrics, outperforming HiFiC for all operating points.

3 Analysis and Ablations

Model size and architecture. In Figure 9, we analyze the effect of using XCiT of different capacities for the autoencoder with respect to rate-distortion trade-off. We observe that the performance improves with higher capacity autoencoders, but their is a diminishing return with further increase in capacity. For all our experiments we use an XCiT-L6 since it achieves the best performance.

Masking patterns. In Table 2, we compare to predicting tokens one-by-one autoregressively in a raster-scan order, the same pattern used in VQ-VAE based generative image models such as DALL-E (Ramesh et al., 2021) and VQ-GAN (Esser et al., 2021). In contrast to PQ-MIM, raster-scan models require causal attention, which makes XCiT not a good fit. We use a standard ViT model instead. However, due to the quadratic complexity of ViT and the high resolution of images typically used for evaluation of compression method (e.g. Tecnick), our autoregressive variant consistently exceeded the memory limits, even when using A100 GPUs with 40GB memory. Moreover, raster scan fully-autoregressive models results in extremely expensive FLOP count since it needs TT separate forward passes per image. On the other hand, our stage-wise MIM with quincunx pattern requires 4 evaluations, and does not scale with the image resolution as is the case for raster, making it a more practical solution.

Number of prediction stages. We compare the quincunx masking pattern with masking based on confidence score following MaskGiT (Chang et al., 2022) patch selection procedure for image generation. In the latter case, at a given step, we pick the patches to transmit dynamically based on their confidence score. The confidence is defined as the maximum across the probabilities assigned over the vocabulary by the model. Note that for decoding the same confidence score can be used to identify the group of tokens to decode. For the quincunx and confidence based masking policy we use the same 5-stage scheme, in which the number of tokens in subsequent groups doubles in size. We ablate the number of prediction stages SS for the quincunx and confidence-based sampling. For the latter we consider two options: (i) a linear scheme where each group of tokens contains (approximately) the same number of tokens; and (ii) a doubling scheme where each subsequent group of tokens is double the size of the previous group, as is also used for the quincunx pattern. Every point on the curves in Figure 9 corresponds to the bitrate when encoding/decoding with a given number of steps. For example, S=2S=2 steps means we have only two steps each encoding/decoding 50% of the patches. For the 3-steps doubling schedule the groups have sizes of 1/4, 1/4, 1/2, and so on. For all three patterns the bitrate monotonically decreases with the number of steps, as more tokens can be predicted from larger contexts.

On the one hand, using confidence-based masking patterns, doubling and linear schemes lead to mostly comparable bitrates for the same number of steps, with further improvement for linear with more steps however with diminishing return. On the other hand, quincunx provides a stronger reduction ratio with strictly 5 steps, even when compared with confidence-based masking with higher number of steps.

Predicting missing patches. When encoding/decoding the discrete image representation q\mathbf{q}, we reduce the bitrate with our models by their ability to predict the remaining tokens given those of preceding stages. To illustrate the predictive ability of our model, we consider an experiment where we remove a subset of the tokens (sampled randomly), and use our model to fill in the missing patches by conditioning on the observed ones. We then use the PQ-VAE decoder to decode the discrete latent codes including the filled in ones. In Figure 10 we show the results obtained when removing 10%, 20%, 30% and 50% of the patches. Our MIM is able to model redundancies among the patches, which corroborates our findings of its compression ability.

4 Limitations

Image compression, and compression of visual data in general, is an important technology to scale the distribution of visual data. This is ever more important to cope with the growing quantity of visual data that is streamed in the form of video and for augmented and virtual reality applications. Compression is also critical to allow users with low-bandwidth connections to benefit from applications relying on sharing of image, video, or virtual reality data.

Caveats of learned neural compression models: biases and performance. As with any machine learning model, potential biases in the training data may be transferred into the model via training. In our setting, this can affect the autoencoding reconstruction abilities for content under-represented in the training data, as well as the compression abilities of the model for such data. Such biases should be assessed before deployment of the model. Beyond rate-distortion trade-offs, important evaluation dimensions include the energy and latency performance of compression models. Current neural compression methods, including ours, need to be further optimized to be competitive with existing codecs on these criteria.

Specific limitations of our approach. In our work, we specifically focus on a high compression regime, with compression rate lower than 0.6 bit per pixel. This is a favorable case for our approach: it benefits from the generative model capability inherited from the VQ-VAE. Compared to VQ, PQ allows for higher rates, but it becomes comparatively less effective when increasing the number of subquantizers, as one can deduce from the slopes of the rate-distortion curves in Fig. 5. It is also computationally more expensive.

Conclusion

In this paper, we have revisited vector quantization for neural image compression. We introduced a product-quantization variants of VQ-VAE and shown that it has a better scaling properties in terms of bit-rate. Additionally, we introduced a novel conditional entropy model based on masked image modeling. We have shown that combined with the qunicunx partitioning pattern, PQ-MIM provides strong reduction in bit-rate. PQ-MIM exhibits a competitive performance for PSNR and MS-SSIM metric compared to strong neural compression baselines. Furthermore, when we train PQ-MIM using perceptual losses, it provides a strong performance on multiple metrics (e.g. FID, KID and MS-SSIM) compared the strong baseline of HiFiC. Finally, we have shown that PQ-MIM can operate in a hybrid compression/generation mode where it can fill the gaps for non-transmitted patches.

References

Appendix A Lossless Compression

Block codes assign a unique binary code of equal length to each event in X\mathcal{X}. The expected length per symbol of such codes is l(m)=⌈log⁡2∣X∣⌉l(m)=\lceil\log_{2}|\mathcal{X}|\rceil bits.

We can improve this result if we consider to give shorter code words to more frequently occurring symbols and longer codes to less frequently occurring ones. In fact building on this idea, it can be shown that the smallest expected achievable average code length per point using an approximation to the true data distribution p(x)p(x), is given by the cross-entropy:

This quantity is bounded from below (seen via Jensen’s inequality) by the entropy of the “true” probability distribution pdp_{d}, i.e. when p=pdp=p_{d} (MacKay, 2003; Cover, 1991).

Entropy coders reach this bound up to a small constant ϵ\epsilon. An entropy codec is a tuple of an inverse pair of functions, encpX\texttt{enc}_{p_{X}} and decpX\texttt{dec}_{p_{X}} that achieve near optimal compression rates on sequences of symbols.

Appendix B Qualitative Samples