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 proposed to replace the typical cnn-based transform in image and video compression networks by Swin-transform. Experiments show its better performance and lower computational complexity over cnn-basd methods. Some extensive analysis are conducted to explore the differences between convolution and transformers.