Learned Image Compression with Discretized Gaussian Mixture Likelihoods and Attention Modules Zhengxue Cheng1, Heming Sun2;3, Masaru Takeuchi2, Jiro Katto1 1 Department of Computer Science and Communication Engineering, Waseda University, Tokyo, Japan 2 Waseda Research Institute for Science and Engineering, Tokyo, Japan 3 JST, PRESTO, 4-1-8 Honcho, Kawaguchi, Saitama, Japan fzxcheng@asagi., hemingsun@aoni., masaru-t@aoni., [email protected] Abstract Image compression is a fundamental research field and many well-known compression standards have been developed for many decades. Recently, learned com- pression methods exhibit a fast development trend with promising results. However, there is still a performance gap between learned compression algorithms and reigning compression standards, especially in terms of widely used PSNR metric. In this paper, we explore the remaining redundancy of recent learned compression algorithms. We have found accurate entropy models for rate estimation largely affect the optimization of network parameters and thus affect the rate-distortion performance. Therefore, in this paper, we propose to use discretized Gaussian Mixture Figure 1: Visualization of reconstructed images Kodim04 Likelihoods to parameterize the distributions of latent from Kodak dataset with approximately 0:1 bpp. codes, which can achieve a more accurate and flexible entropy model. Besides, we take advantage of recent attention modules and incorporate them into network JPEG2000 [2], HEVC/H.265 [3] and ongoing Versatile architecture to enhance the performance. Experimental Video Coding (VVC) [4] which will be the next generation results demonstrate our proposed method achieves a compression standard expected by the end of 2020. Typi- state-of-the-art performance compared to existing learned cally they rely on hand-crafted creativity to present module- compression methods on both Kodak and high-resolution based encoder/decoder (codec) block diagrams. They use datasets. To our knowledge our approach is the first work fixed transform matrix, intra prediction, quantization, con- to achieve comparable performance with latest compres- text adaptive arithmetic coders and various de-blocking or sion standard Versatile Video Coding (VVC) regarding loop filters to reduce spatial redundancy to improve the cod- PSNR. More importantly, our approach generates more ing efficiency. The standardization of a conventional codec arXiv:2001.01568v3 [eess.IV] 30 Mar 2020 visually pleasant results when optimized by MS-SSIM. The has historically spanned several years. Along with the fast project page is at https://github.com/ZhengxueCheng/ development of new image formats and the proliferation Learned-Image-Compression-with-GMM-and-Attention. of high-resolution mobile devices, existing image compres- sion standards are not expected to be an optimal and general solutions for all kinds of image contents. Various approaches has been investigated for end-to- 1. Introduction end image compression. Recently, the notable approaches are context-adaptive entropy models for learned image Image compression is an important and fundamental re- compression [19, 20, 21] to achieve superior performance search topic in the field of signal processing for many among all the learned codecs. The work [19] proposed a hy- decades to achieve efficient image transmission and storage. perprior to add additional bits to model the entropy model. Classical image compression standards include JPEG [1], The work [20] jointly combine an autoregressive mask con- 1 volution and the hyperprior. The work [21] proposed a In the early stage, some works are proposed to deal with quite similar idea by considering two types of contexts, bit- non-differential quantization and rate estimation to make consuming contexts (that is, hyperprior) and bit-free con- end-to-end training possible, such as [6,7,8]. After that, texts (that is, autoregressive model) to realize a context- some works focus on the design of network structure, which adaptive entropy model. Our method are based on the devel- is able to extract more compact and efficient latent repre- opment of these recent entropy model techniques to further sentation and to reconstruct high-quality images from com- improve the performance. pressed features. For instance, some works [9, 10, 11] use In this paper, our main contribution is to present a recurrent neural networks to compress the residual infor- more accurate and flexible entropy model by leveraging mation recursively, but they mainly relied on binary rep- discretized Gaussian mixture likelihoods. We visualize resentation at each iteration to achieve scalable coding for the spatial redundancy of compressed codes from recent image compression. Some approaches [12, 13, 14] use gen- learned compression techniques. Motivated from it, we pro- erative models to learn the distribution of images using ad- pose to use discretized Gaussian mixture likelihoods to pa- versarial training to achieve better subjective quality at ex- rameterize the distributions, which removes remaining re- tremely low bit rate. Some approaches include a content- dundancy to achieve accurate entropy model, and thus di- weighted strategy [15] or de-correlating different channels rectly lead to fewer required encoding bits. Besides, we using principle component analysis [16], or consider en- adopt a simplified version of attention module into our ergy compaction [17], or use deep residual units to enhance network architecture. Attention module can make learned the network architecture [18]. Recently, several studies in- models pay more attention to complex regions to improve vestigate the adaptive context model for entropy estimation our coding performance with moderate training complexity. to navigate the optimization process of neural network pa- Experimental results demonstrate our proposed method rameters to achieve best tradeoff between reconstruction er- leads to the state-of-the-art performance on both PSNR and rors and required bits (entropy), including [19, 20, 21, 22]. MS-SSIM quality metrics, in comparison with classical im- Entropy estimation techniques have greatly improved the age compression standards HEVC, JPEG2000, JPEG and learned compression algorithms, and the most representa- existing deep learning based compression approaches. To tive methods are hyperprior model and joint model. How- our knowledge, we are the first work to reach very close ever, there is still a gap between the estimated distribution performance with ongoing versatile video coding test model and true marginal distribution of latent representation. VTM 5.2 with intra profile in terms of PSNR. Moreover, Parameterized Model Some image generation studies our method produces visually pleasant reconstructed im- have investigated several parameterized distribution mod- ages when optimizing by MS-SSIM as Fig.1. els. For instance, standard PixelCNN [24] uses full 256- way softmax likelihoods, but they are extremely memory- 2. Related Work consuming. To address this problem, PixelCNN++ [25] proposed a discretized logistic mixture likelihoods to Hand-crafted Compression Existing image compression achieve faster training. Learned lossless image compres- standards, such as JPEG [1], JPEG2000 [2], HEVC [3] and sion work L3C [26] followed the PixelCNN++ to use a lo- VVC [4], reply on hand-crafted module design individually. gistic mixture model. These works are basically used to es- Specifically, these modules include intra prediction, discrete timate the likelihoods of 8-bit pixel values with fixed range cosine transform or wavelet transform, quantization, and of [0; 255]. However, in learned image compression task, entropy coder such as Huffman coder or content adaptive few studies explore the effect of parameterized distribution binary arithmetic coder (CABAC). They design each mod- models. ule with multiple modes and conduct the rate-distortion op- timization to determine the best mode. Especially, VVC [4] supports larger coding unit, more prediction modes, more 3. Proposed Method transform types and other coding tools. Besides, along 3.1. Formulation of Learned Compression Models with the development of classical compression algorithms, some hybrid methods have been proposed, by taking advan- In the transform coding approach [27], image compres- tage of both conventional compression algorithms and latest sion can be formulated by (as Fig. 2(a)) learned super resolution approaches, such as [23]. Learned Compression Recently, we have seen a great y = ga(x; φ) surge of deep learning based image compression ap- y^ = Q(y) (1) proaches utilizing autoencoder architecture [5], which have x^ = gs(y^; θ) achieved a great success with promising results. The de- velopment of previous learning based compression models where x, x^, y, and y^ are raw images, reconstructed images, have spanned several years and have many related works. a latent presentation before quantization, and compressed (a) Baseline (b) Hyperprior (c) Joint (d) Proposed Model Figure 2: Operational diagrams of learned compression models (a)(b)(c) and proposed Gaussian Mixture Likelihoods (d). codes, respectively. φ and θ are optimized parameters of based rate-distortion optimization. The loss function is analysis and synthesis transforms. UjQ represents the quan- tization and entropy coding. During the training, the quanti- L =R(y^) + R(z^) + λ ·D(x; x^) 1 1 zation is approximated by a uniform noise U(− 2 ; 2 ) to gen- = E[− log2(py^jz^(y^jz^))] + E[− log2(pz^j (z^j ))] (3) erate noisy codes