Transformer-based Transform Coding

1 INTRODUCTION . Transform coding ( Goyal , 2001 ) is the dominant paradigm for compression of multi-media signals , and serves as the technical foundation for many successful coding standards such as JPEG , AAC , and HEVC/VVC . Codecs based on transform coding divide the task of lossy compression into three modularized components : transform , quantization , and entropy coding . With the advancement of deep learning , data compression is entering into a new era of learning-based schemes ( Ballé et al. , 2021 ) , in which all three components can be enhanced by deep neural networks : autoencoder networks are adopted as flexible nonlinear transforms , deep generative models are used as powerful learnable entropy models , and various differentiable quantization schemes are proposed to aid endto-end training . Thanks to these advancements , we have seen rapid progress in the domain of image and video compression . Particularly , the hyperprior line of work ( Ballé et al. , 2018 ; Minnen et al. , 2018 ; Lee et al. , 2019 ; Agustsson et al. , 2020 ; Minnen & Singh , 2020 ) has led to steady progress of neural compression performance over the past two years , reaching or even surpassing state-of-theart traditional codecs . For example , in image compression , BPG444 was surpassed by a neural codec in 2018 ( Minnen et al. , 2018 ) , and ( Cheng et al. , 2020 ; Xie et al. , 2021 ; Ma et al. , 2021 ; Guo et al. , 2021 ; Wu et al. , 2020 ) have claimed on-par or better performance than VTM ( a test model of the state-of-the-art non-learned VVC standard ) . One general trend in the advancement of neural image compression schemes is to develop ever more expressive yet expensive prior models based on spatial context ( Minnen et al. , 2018 ; Lee et al. , 2019 ; Ma et al. , 2021 ; Guo et al. , 2021 ; Wu et al. , 2020 ) . However , the rate-distortion improvement from context based prior modeling often comes with a hefty price tag1 in terms of decoding complexity . Noteably , all existing works that claimed on-par or better performance than VTM ( Cheng et al. , 1In the extreme case when a latent-pixel-level spatial autoregressive prior is used , decoding of a single 512x768 image requires no less than 1536 interleaved executions of prior model inference and entropy decoding ( assuming the latent is downsampled by a factor of 16x16 ) . 2020 ; Xie et al. , 2021 ; Ma et al. , 2021 ; Guo et al. , 2021 ; Wu et al. , 2020 ) rely on slow and expensive spatial context based prior models . The development of nonlinear transforms , on the other hand , are largely overlooked . This leads us to the following questions : can we achieve the same performance as that of expensive prior models by designing a more expressive transform together with simple prior models ? And if so , how much more complexity in the transform is required ? Interestingly , we show that by leveraging and adapting the recent development of vision transformers , not only can we build neural codecs with simple prior models that can outperform ones built on expensive spatial autoregressive priors , but do so with smaller transform complexity compared to its convolutional counterparts , attaining a strictly better ratedistortion-complexity trade-off . As can be seen in Figure 1 , our proposed neural image codec SwinT-ChARM can outperform VTM-12.1 at comparable decoding time , which , to the best of our knowledge , is a first in the neural compression literature . As main contributions , we 1 ) extend Swin-Transformer ( Liu et al. , 2021 ) to a decoder setting and build Swin-transformer based neural image codecs that attain better rate-distortion performance with lower complexity compared with existing solutions , 2 ) verify its effectiveness in video compression by enhancing scale-space-flow , a popular neural P-frame codec , and 3 ) conduct extensive analysis and ablation study to explore differences between convolution and transformers , and investigate potential source of coding gain . 2 BACKGROUND & RELATED WORK . Conv-Hyperprior The seminal hyperprior architecture ( Ballé et al. , 2018 ; Minnen et al. , 2018 ) is a two-level hierarchical variational autoencoder , consisting of a pair of encoder/decoder ga , gs , and a pair of hyper-encoder/hyper-decoder ha , hs . Given an input image x , a pair of latent y = ga ( x ) and hyper-latent z = ha ( y ) is computed . The quantized hyper-latent ẑ = Q ( z ) is modeled and entropycoded with a learned factorized prior . The latent y is modeled with a factorized Gaussian distribution p ( y|ẑ ) = N ( µ , diag ( σ ) ) whose parameter is given by the hyper-decoder ( µ , σ ) = hs ( ẑ ) . The quantized version of the latent ŷ = Q ( y−µ ) +µ is then entropy coded and passed through decoder gs to derive reconstructed image x̂ = gs ( ŷ ) . The tranforms ga , gs , ha , hs are all parameterized as ConvNets ( for details , see Appendix A.1 ) . Conv-ChARM ( Minnen & Singh , 2020 ) extends the baseline hyperprior architecture with a channel-wise auto-regressive model ( ChARM ) 2 , in which latent y is split along channel dimension into S groups ( denoted as y1 , . . . , yS ) , and the Gaussian prior p ( ys|ẑ , ŷ < s ) is made autoregressive across groups where the mean/scale of ys depends on quantized latent in the previous groups ŷ < s . In practice , S = 10 provides a good balance of performance and complexity and is adopted here . Spatial AR models Most of recent performance advancements of neural image compression is driven by the use of spatial auto-regressive/context models . Variants include causal global prediction ( Guo et al. , 2021 ) , 3D context ( Ma et al. , 2021 ) , block-level context ( Wu et al. , 2020 ) , nonlocal context ( Li et al. , 2020 ; Qian et al. , 2021 ) . One common issue with these designs is that decoding 2For details refer to Figure 11 and 12 in the Appendix . can not be parallelized along spatial dimensions , leading to impractical3decoding latency , especially for large resolution images . ConvNet-based transforms While the design space of prior models is extensively explored , nonlinear transforms , as an important component , have received less attention . A standard convolution encoder-decoder with GDN ( Ballé et al. , 2016 ; 2017 ) as activation is widely adopted in the literature . Later works introduce new transform designs , such as residual blocks with smaller kernels ( Cheng et al. , 2020 ) , nonlocal ( sigmoid gating ) layers ( Zhou et al. , 2019 ; Chen et al. , 2021 ) , invertible neural networks ( Xie et al. , 2021 ) , and PReLU as an efficient replacement of GDN ( Egilmez et al. , 2021 ) . Vision transformers Although many transform networks are proposed , they are still mainly based on ConvNets . Recently transformers ( Vaswani et al. , 2017 ) have been introduced to the vision domain and have shown performance competitive with ConvNets in many tasks , e.g . object detection ( Carion et al. , 2020 ) , classification ( Dosovitskiy et al. , 2021 ) , image enhancement ( Chen et al. , 2020 ) , and semantic segmentation ( Zheng et al. , 2021 ) . Inspired by their success , in this work we explore how vision transformers work as nonlinear transforms for image and video compression . 3 SWIN-TRANSFORMER BASED TRANSFORM CODING . Among the large number of vision transformer variants , we choose Swin Transformer ( Liu et al. , 2021 ) ( hereafter referred to as SwinT ) to build the nonlinear transforms , mainly because of 1 ) its linear complexity w.r.t . input resolution due to local window attention , and 2 ) its flexibility in handling varying input resolutions at test time , enabled by relative position bias and hierarchical architecture . 3.1 SWINT ENCODER AND DECODER . The original SwinT is proposed as a vision backbone , i.e . an encoder transform with downsampling . As shown in Figure 2 , the SwinT encoder ga contains SwinT blocks interleaved with Patch Merge blocks . The Patch Merge block contains Space-to-Depth ( for downsampling ) , LayerNorm , and Linear layers sequentially . SwinT block performs local self-attention within each non-overlapping window of the feature maps and preserves feature size . Consecutive SwinT blocks at the same feature size shift the window partitioning with respect to the previous block to promote information propagation across nearby windows in the previous block . 3It is reported in ( Wu et al. , 2020 ) ( Table I ) that decoding time of spatial autoregressive models on a 512x768 image range from 2.6s to more than half a minute , depending on the specific designs . Also see Figure 1 . 4The ChARM architecture ( Minnen & Singh , 2020 ) is detailed in Figure 12 of Appendix A.1 . We adopt SwinT encoder as the encoder transform ga in our model , and extend it to SwinT decoder gs by reversing the order of blocks in ga , and replacing the Patch Merge block with a Patch Split block , which contains Linear , LayerNorm , Depth-to-Space ( for upsampling ) layers in sequence . The architectures for hyper transforms ha , hs are similar to ga , gs with different configurations . With these four SwinT transforms , we propose two image compression models , SwinT-Hyperprior and SwinT-ChARM , whose prior and hyper prior models are respectively the same as in ConvHyperprior and Conv-ChARM introduced in Section 2 . The full model architectures are shown in Figure 2 and Figure 13 . 3.2 EXTENSION TO P-FRAME COMPRESSION . To investigate the effectiveness of SwinT transforms for video compression , we study one popular P-frame compression model called Scale-Space Flow ( SSF ) ( Agustsson et al. , 2020 ) . There are three instances of Conv-Hyperprior in SSF , which are respectively for compressing I-frames , scale-space flow and residual . We propose a SwinT variant , referred to as SwinT-SSF , which is obtained by replacing Conv transforms ga , gs in the flow codec and residual codec of SSF with SwinT tranforms . To stabilize training of flow codec in SwinT-SSF , we need to remove all LayerNorm layers and reduce the window size ( e.g . from 8 to 4 ) . The baseline SSF model will be referred as Conv-SSF . Even though we build our solution on top of SSF , we believe this general extension can be applied to other ConvNet-based video compression models ( Rippel et al. , 2021 ; Hu et al. , 2021 ) as well . 4 EXPERIMENTS AND ANALYSIS . 4.1 EXPERIMENT SETUP . Training All image compression models are trained on the CLIC2020 training set . ConvHyperprior and SwinT-Hyperprior are trained with 2M batches . Conv-ChARM and SwinT-ChARM are trained with 3.5M and 3.1M steps . Each batch contains 8 random 256× 256 crops from training images . Training loss L = D + βR is a weighted combination of distortion D and bitrate R , with β being the Lagrangian multiplier steering rate-distortion trade-off . Distortion D is MSE in RGB color space . To cover a wide range of rate and distortion , for each solution , we train 5 models with β ∈ { 0.003 , 0.001 , 0.0003 , 0.0001 , 0.00003 } . The detailed training schedule is in Appendix B.1 . For P-frame compression models , we follow the training setup of SSF . Both Conv-SSF and SwinTSSF are trained on Vimeo-90k Dataset ( Xue et al. , 2019 ) for 1M steps with learning rate 10−4 , batch size of 8 , crop size of 256×256 , followed by 50K steps of training with learning rate 10−5 and crop size 384× 256 . The models are trained with 8 β values 2γ × 10−4 : γ ∈ { 0 , 1 , ... , 7 } . We adopt one critical trick to stablize the training from ( Jaegle et al. , 2021 ; Meister et al. , 2018 ) , i.e . to forward each video sequence twice during one optimization step ( mini-batch ) , once in the original frame order , once in the reversed frame order . Finally we add flow loss5 only between 0 and 200K steps , which we found not critical for stable training but improves the RD . Evaluation We evaluate image compression models on 4 datasets : Kodak ( Kodak , 1999 ) , CLIC2021 testset ( CLIC , 2021 ) , Tecnick testset ( Asuni & Giachetti , 2014 ) , and JPEG-AI testset ( JPEG-AI , 2020 ) . We use BPG and VTM-12.1 to code the images in YUV444 mode , and then calculate PSNR in RGB . For a fair comparison all images are cropped to multiples of 256 to avoid padding for neural codecs . We evaluated P-frame models on UVG ( Mercat et al. , 2020 ) 6 and MCL-JCV ( Wang et al. , 2016 ) , and compare them with the test model implementation of HEVC , referred to as HEVC ( HM ) , and open source library implementation of HEVC , refered to as HEVC ( x265 ) . To align configuration , all video codecs are evaluated in low-delay-P model with a fixed GOP size of 12 . Besides rate-distortion curves , we also evaluate different models using BD-rate ( Tan et al. , 2016 ) , which represents the average bitrate savings for the same reconstruction quality . For image codecs , BD-rate is computed for each image and then averaged across all images ; for video codecs , BD-rate 5We did not observe RD improvement when applying flow loss to Conv-SSF training . 6We use the original 7 UVG sequences that are commonly used in other works ( Agustsson et al. , 2020 ) . is computed for each video and then averaged across all videos . More details on testset preprocessing , and traditional codecs configurations can be found in Appendix B.2 .

Summary: This paper addresses the problem of improving rate-distortion (RD) performance for learned image and video compression without ignoring runtime. Much of the literature on learned compression improves RD performance with more entropy models. Especially for autoregressive models, these models can have excessively slow decode times. This paper uses a less expensive (and generally less powerful) entropy model and explores the use of a transformer in the encoder and decoder transforms instead.  The result is a much faster model that still achieves near-SOTA rate-distortion performance. In particular, the SwinT-ChARM model described in this paper outperforms VVC (via VTM 12.1 a recent version of the reference implementation for H.266) with faster decode times. As far as I know, this is a milestone for learned image compression (other methods have outperformed in terms of RD but only with much slower decoders).