Year: 2021

Abstract:
Neural data compression based on nonlinear transform coding has made great progress over the last few years, mainly due to improvements in prior models, quantization methods and nonlinear transforms. A general trend in many recent works pushing the limit of rate-distortion performance is to use ever more expensive prior models that can lead to prohibitively slow decoding. Instead, we focus on more expressive transforms that result in a better rate-distortioncomputation trade-off. Specifically, we show that nonlinear transforms built on Swin-transformers can achieve better compression efficiency than transforms built on convolutional neural networks (ConvNets), while requiring fewer parameters and shorter decoding time. Paired with a compute-efficient Channel-wise AutoRegressive Model prior, our SwinT-ChARM model outperforms VTM-12.1 by 3.68% in BD-rate on Kodak with comparable decoding speed. In P-frame video compression setting, we are able to outperform the popular ConvNet-based scalespace-flow model by 12.35% in BD-rate on UVG. We provide model scaling studies to verify the computational efficiency of the proposed solutions and conduct several analyses to reveal the source of coding gain of transformers over ConvNets, including better spatial decorrelation, flexible effective receptive field, and more localized response of latent pixels during progressive decoding.

heading: 1 INTRODUCTION
text: 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.

heading: 2 BACKGROUND & RELATED WORK
text: 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.
cannot 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.

heading: 3 SWIN-TRANSFORMER BASED TRANSFORM CODING
text: 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.

heading: 3.1 SWINT ENCODER AND DECODER
text: 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.

heading: 3.2 EXTENSION TO P-FRAME COMPRESSION
text: 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.

heading: 4 EXPERIMENTS AND ANALYSIS

heading: 4.1 EXPERIMENT SETUP
text: 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.

heading: 4.2 RESULTS
text: RD and BD-rate for image codecs The RD curves for all compared image codecs evaluated on Kodak are shown in Figure 3a, and the relative rate reduction of each codec compared to VTM-12.1 at a range of PSNR levels is shown in Figure 3b 7.
As can be seen from Figure 3, SwinT transform consistently outperforms its convolutional counterpart; the RD-performance of SwinT-Hyperprior is on-par with Conv-ChARM, despite the simpler prior; SwinT-ChARM outperforms VTM-12.1 across a wide PSNR range. In the Appendix (Figure 28 and Figure 30), we further incorporate the results from existing literature known to us for a complete comparison. Particularly, our Conv-Hyperprior is much better than the results reported in (Minnen et al., 2018) (no context), and Conv-ChARM is on par with (Minnen & Singh, 2020).
In Table 1, we summarize the BD-rate of image codecs across all four dataset with VTM-12.1 as anchor. On average SwinT-ChARM is able to achieve 3.8% rate reduction compared to VTM12.1. The relative gain from Conv-Hyperprior to SwinT-Hyperprior is on-average 12% and that from Conv-ChARM to SwinT-ChARM is on-average 9%. Further gain over VTM-12.1 can be obtained by test-time latent optimization (Campos et al., 2019) or full model instance adaptation (van Rozendaal et al., 2021), which are out of the scope of this work.
7The relative rate-saving curves in Figure 3b is generated by first interpolating the discrete RD points (averaged across the testset) with a cubic spline, and then compare bitrate of different models at fixed PSNR.
8RD plot for the other three datasets can be found in Appendix (Figure 14, Figure 15, and Figure 16)
RD and BD-rate for video codecs For P-frame compression, we evaluated SwinT-SSF on UVG and MCL-JCV, with RD comparison shown in Figure 4. Again, SwinT transform leads to consistently better RD. Table 2 summarizes BD-rate with our reproduced Conv-SSF model as anchor. We can see that SwinT-SSF achieves an average of 11% rate saving over Conv-SSF. Additionally, we show that if SwinT transform is only applied to residual-autoencoder (labeled as SwinT-SSF-Res), it can only get about 4.6% gain, which indicates
that both flow and residual compression benefit from SwinT as encoder and decoder transforms. Note that SwinT-SSF still lags behind HM, suggesting lots of room for improvement in neural video compression. For per-video breakdown of BD-rate, see Figure 18 and Figure 17 in the Appendix.
Decoding complexity We evaluate the decoding complexity of 4 image codecs on 100 images of size 768 × 512 and show the metrics in Table 3, including decoding time, GMACs and GPU peak memory during decoding and total model parameters. The models run with PyTorch 1.9.0 on a workstation with one RTX 2080 Ti GPU. From the table, the inference time of SwinT decoder is less than that of Conv decoder. The entropy decoding time of ChARM prior is about twice than the factorized prior. The total decoding time of SwinT-based models is less than Conv-based models. In ablation study A5, we show a smaller SwinT-Hyperprior with 20.6M parameters has almost the same RD as the SwinT-Hyperprior profiled here. For details on encoding complexity, profiling setup, scaling to image resolution, please refer to Table 4 and Section C.3 in the Appendix.
Scaling behavior To see how the BD-rate varies with model size, we scale SwinTHyperprior and Conv-Hyperprior to be twice or half of the size of the base models (i.e. medium size)9. The result is shown in Figure 5. For both types of models, as we reduce the base model size, there is a sharp drop in performance, while doubling model size only leads to marginal gain. Noticeably, SwinT-Hyperpriorsmall is on-par with Conv-Hyperprior-medium even with half of the parameters, and SwinT transforms in general incur fewer MACs per parameter.
In Figure 1, we further consolidate the decoding latency and scaling behavior study into a single plot and show that SwinT-ChARM runs at comparable speed as VTM-12.1 while achiev-
ing better performance,10 as opposed to state-of-the-art neural codecs with spatial autoregressive prior that decodes orders of magnitude slower.

heading: 4.3 ANALYSIS
text: Latent correlation One of the motivating principles of transform coding is that simple coding can be made more effective in the transform domain than in the original signal space (Goyal, 2001; Ballé et al., 2021). A desirable transform would decorrelate the source signal so that simple scalar quantization and factorized entropy model can be applied without constraining coding performance. In most mature neural compression solutions, uniform scalar quantization is adopted together with a learned factorized or conditionally factorized Gaussian prior distribution. It is critical, then, to effectively factorize and Gaussianize the source distribution so that coding overhead can be minimized.
Specifically, in hyperprior based models (Ballé et al., 2018), ȳ , (y − µ)/σ is modeled as a standard spherical normal vector. The effectiveness of the analysis transform ga can then be evaluated by measuring how much correlation there is among different elements in ȳ. We are particularly interested in measuring the correlation between nearby spatial positions, which are heavily correlated in the source domain for natural images. In Figure 6, we visualize the normalized spatial correlation of ȳ averaged over all latent channels, and compare Conv-Hyperprior with SwinT-Hyperprior at β = 0.001. It can be observed that while both lead to small cross-correlations, Swin-Transformer does a much better job with uniformly smaller correlation values, and the observation is consistent
9detailed model configurations are provided in Appendix A.3. 10Note that it is difficult to fairly compare the decoding time of VTM and neural codecs since they run on different hardware. For more discussion please refer to Appendix C.3. 11The value with index (i, j) corresponds to the normalized cross-correlation of latents at spatial location (w, h) and (w + i, h+ j), averaged across all latent elements of all images on Kodak.
with other β values, which are provided in Figure 20 in the Appendix. This suggests that transformer based transforms incur less redundancy across different spatial latent locations compared with convolutional ones, leading to an overall better rate-distortion trade-off. The larger spatial correlation (and thus redundancy) in Conv-Hyperprior also explains why a compute-heavy spatial auto-regressive model is often needed to improve RD with convolutional based transforms (Minnen et al., 2018; Lee et al., 2019; Ma et al., 2021; Guo et al., 2021; Wu et al., 2020). Figure 6 also reveals that most of the correlation of a latent comes from the four elements surrounding it. This suggests that a checkerboard-based conditional prior model (He et al., 2021) may yield further coding gain.
Effective receptive field Intra prediction in HEVC or AV1 only rely on left and top boarders of the current coding block (Sullivan et al., 2012; Chen et al., 2018), except for intra block copy for screen content (Xu et al., 2016). We would like to see how large the effective receptive field (ERF) (Luo et al., 2017) of SwinT encoders compared to Conv encoders. The theoretical receptive field of the encoders (ga, ha ◦ ga) in SwinT-based codecs is much larger than that of Conv-based codecs. However comparing Figure 7a with 7e and Figure 7b with 7f, the ERF of SwinT encoders after training is even smaller than Conv encoders. When we examine the ERF of the released Swin transformers for classification, detection and segmentation tasks, they are all spanning the whole input image. This contrast suggests that (natural) image compression with rate-distortion objective is a local task, even with transformer-based nonlinear transforms. We further look into P-frame compression models, particularly the ERF of two types of transforms in flow codec and residual codec, as shown in Figure 7d & 7h, and Figure 7c & 7g. Clearly for flow codec, SwinT transform has much larger ERF than the convolution counterpart. For residual codec, the ERF of SwinT transforms is similar to image (I-frame) compression case. This shows of flexibility of SwinT encoders to attend to longer or shorter range depending on the tasks. To get a better picture of the behavior of attention layers in SwinT transforms, we also show the attention distance in each layer in Figure 22.
Progressive decoding The ERF in the previous section shows the behavior of the encoder transforms, here we further investigate the decoder transforms through the lens of progressive decoding (Rippel et al., 2014; Minnen & Singh, 2020; Lu et al., 2021). Initialized with the prior mean, the input to the decoder is progressively updated with the dequantized latent ŷ in terms of coding units, leading to gradually improved reconstruction quality. For the latent with shape (C,H,W ), we consider three types of coding units, i.e. per channel (1, H,W ), per pixel (C, 1, 1), per element (1, 1, 1). The coding units are ordered by the sum of prior std of all elements within each unit. The RD curves of progressive decoding for SwinT-Hyperprior and Conv-Hyperprior are shown in Figure 8a, which closely follow each other when ordered by channel or element, but significantly apart when ordered by pixel (spatial dim). Particularly, we show an extreme case when the half pixels in the latent (masked by checkerboard pattern) are updated with dequantized values, corresponding to the two scatter points in Figure 8a. One visual example (CLIC2021 test) is shown in Figure 8b under
this setup, where we can clearly see SwinT decoder achieves better reconstruction quality than the Conv decoder, mainly in terms of more localized response to a single latent pixel. This is potentially useful for region-of-interest decompression. More visual examples are shown in Figure 26.
4.4 ABLATION STUDY
Relative position bias There are two sources of positional information in SwinT transforms, namely the Space-to-Depth modules and the additive relative position bias (RPB). Even when the RPB is removed, SwinT-Hyperprior still outperforms Conv-Hyperprior across all bitrates, which indicates image compression may not require accurate relative position.
Shifted window The original motivation of shifted window design is to promotes the inter-layer feature propagation across nonoverlapping windows. Image compression performance drops slightly when there is no shifted window at all. This further suggests image compression requires local information.
The details of ablations A3-A5 in Figure 9 can be found in Section E.

heading: 5 CONCLUSION
text: In this work we propose Swin transformer based transforms for image and video compression. In the image compression setting, SwinT transform consistently outperforms its convolutional counterpart. Particularly, the proposed SwinT-ChARM model outperforms VTM-12.1 at comparable decoding speed, which, to the best of our knowledge, is the first in learning-based methods. We also show the effectiveness of SwinT transforms when extended to the P-frame compression setting. Compared with convolution transforms, SwinT transforms can spatially decorrelate the latent better, have more flexible receptive field to adapt to tasks that requires either short-range (image) and long-range (motion) information, and better progressive decoding of latent pixels. While pushing the neural image compression to a new level in terms of rate-distortion-computation trade-off, we believe it is only the starting point for developing more efficient transformer-based image and video codecs.

heading: Rethinking_Semantic_Segmentation_From_a_Sequence-to-Sequence_
text: Perspective_With_Transformers_CVPR_2021_paper.html.
Lei Zhou, Zhenhong Sun, Xiangji Wu, and Junmin Wu. End-to-end optimized image compression with attention mechanism. In CVPR workshops, pp. 0, 2019.

heading: A MODELS

heading: A.1 CONVOLUTION BASELINES
text: Conv-Hyperprior and Conv-ChARM The architecture of Conv-Hyperprior and ConvChARM are shown in Figure 10 and Figure 11. For both architecture, our base model (i.e. medium size) has the following hyperparameters: (C1, C2, C3, C4, C5, C6, C7) = (320, 320, 320, 320, 192, 192, 192).

heading: A.2 SWIN-TRANSFORMER BASED COMPRESSION MODELS
text: SwinT-Hyperprior, SwinT-ChARM For both SwinT-Hyperprior and SwinT-ChARM, we use the same configurations: (wg, wh) = (8, 4), (C1, C2, C3, C4, C5, C6) = (128, 192, 256, 320, 192, 192), (d1, d2, d3, d4, d5, d6) = (2, 2, 6, 2, 5, 1) where C, d, and w are defined in Figure 13 and Figure 2.
SwinT-SSF For SwinT transforms used in SSF variant, the first Patch Merge block is with downsample rate of 4 and two other Path Merge blocks with downsampling rate of 2. Thus the downsampling rate for the encoder is still 16, the same as the image compression models. There are only 3 transformer stages with depths 2, 4, 2. The embedding dim is 96. The number of latent and hyper latent channels are all 192. The window size is 4 for flow codec and 8 for residual codec.
SwinT-SSF-Res This is a variant where only residual autoencoder uses SwinT transforms. Same architecture as the residual autoencoder in SwinT-SSF.

heading: A.3 MODEL CONFIGURATIONS FOR MODEL SIZE SCALING STUDY

heading: A.3.1 SWINT-HYPERPRIOR
text: Set of model hyperparameters that are common to all experiments: (d1, d2, d3, d4, d5, d6) = (2, 2, 6, 2, 5, 1) (wg, wh) = (8, 4)
SwinT-Hyperprior (small) (C1, C2, C3, C4, C5, C6) = (96, 128, 160, 192, 96, 128)
SwinT-Hyperprior (medium) (C1, C2, C3, C4, C5, C6) = (128, 192, 256, 320, 192, 192)
SwinT-Hyperprior (large) (C1, C2, C3, C4, C5, C6) = (160, 256, 352, 448, 192, 256)

heading: A.3.2 CONV-HYPERPRIOR
text: Conv-Hyperprior (small) (C1, C2, C3, C4, C5, C6, C7) = (192, 192, 192, 192, 128, 128, 128)
Conv-Hyperprior (medium) (C1, C2, C3, C4, C5, C6, C7) = (320, 320, 320, 320, 192, 192, 192)
Conv-Hyperprior (large) (C1, C2, C3, C4, C5, C6, C7) = (448, 448, 448, 448, 256, 256, 256)

heading: B TRAINING AND EVALUATION

heading: B.1 TRAINING
text: All image compression models are trained on CLIC2020 training set, which contains both professional and mobile training sets, in total 1,633 high resolution natural images. Conv-Hyperprior and SwinT-Hyperprior are trained with 2M batches. Each batch contains 8 patches of size 256 × 256 randomly cropped from the training images. Learning rate starts at 10−4 and is reduced to 10−5 at 1.8M step.
For Conv-ChARM, we first train a model Conv-ChARM at β = 0.0001 from scratch for 2M steps, with it as the starting point, we continue to train other beta values Conv-ChARM-β, β ∈ B for 1.5M steps. For SwinT-ChARM-β, we load the transform weights from the checkpoint at 2M step of the pretrained SwinT-Hyperprior-β, then finetune the transforms together with the random initialized ChARM prior for 1.1M steps. Learning rate starts at 10−4 and is reduced to 10−5 for the last 100K steps.
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 β ∈ B = {0.003, 0.001, 0.0003, 0.0001, 0.00003}. Usually we need to train longer for the model with larger bitrates (i.e. smaller β) to converge. Particularly for the results presented in this paper, We train SwinT-Hyperprior-0.00003 for 2.5M steps instead of 2M steps for the other 4 lower bitrates.
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 size12 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. When this trick is used, we set the batch size to be 4 instead of 8. Finally we add flow loss only between 0 and 200K steps, which we found not critical for stable training but helps improve the RD.
For all model training, Adam optimizer is used without weighted decay. Training for 2M steps takes about 10 days and 14 days respectively for Conv-Hyperprior and SwinT-Hyperprior on a single Nvidia V100 GPU. Total training time is about 7.5 days on a single Nvidia V100 GPU.
For all models, we use mixed quantization during training (Minnen & Singh, 2020), i.e. adding uniform noise to the continuous latent before passing to the prior model, subtracting prior mean from the continuous latent followed by rounding before passing to the decoder transform.
12We did not use the crop size 384 × 384 during the second stage as in the original paper because the resolution of Vimeo dataset is 448 × 256. We found in our case increasing crop size from 256 × 256 to 384× 256 in the second stage does not improve RD.

heading: B.2 TRADITIONAL CODEC EVALUATION
text: In this section, we provide evaluation script used to generate results for traditional codecs.
B.2.1 IMAGE CODECS
VTM-12.1: VTM-12.1 software is built from https://vcgit.hhi.fraunhofer.de/ jvet/VVCSoftware_VTM/-/tags/VTM-12.1 and we use the script from CompressAI (https://github.com/InterDigitalInc/CompressAI/tree/efc69ea24) for dataset evaluation. Specifically, the following command is issued to gather VTM-12.1 image compression evaluation results:
python -m compressai.utils.bench vtm [path to image folder] -c [path to VVCSoftware_VTM folder]/cfg/encoder_intra_vtm.cfg -b [path to VVCSoftware_VTM folder]/bin -q 16, 18, 20, 22, 24, 26, 28, 30, 32, 34, 36, 38, 40
BPG: BPG software is obtained from https://bellard.org/bpg/ and the following commands are used for encoding and decoding.
bpgenc -e x265 -q [0 to 51] -f 444 -o [encoded heic file] [original png file] bpgdec -o [decoded png file] [encoded heic file]
B.2.2 VIDEO CODECS
HEVC (x265)
ffmpeg -y -pix_fmt yuv420p -s [resolution] -r [frame-rate] -i [input yuv420 raw video] -c:v libx265 -preset medium -crf [9, 12, 15, 18, 21, 24, 27, 30] -tune zerolatency -x265-params "keyint=12:min-keyint=12:verbose=1" [output mkv file path]
AVC (x264)
ffmpeg -y -pix_fmt yuv420p -s [resolution] -r [frame-rate] -i [input yuv420 raw video] -c:v libx264 -preset medium -crf [9, 12, 15, 18, 21, 24, 27, 30] -tune zerolatency -x264-params "keyint=12:min-keyint=12:verbose=1" [output mkv file path]

heading: HEVC (HM)
text: [Path to HM folder]/bin/TAppEncoderStatic -c [Path to HM folder]/cfg/encoder_lowdelay_P_main.cfg -i [input yuv raw video] --InputBitDepth=8 -wdt [width] -hgt [height] -fr [frame-rate] -f [number of frames] -o [output yuv video] -b [encoded bitstream bin file] -ip 12 -q [12, 17, 22, 27, 32, 37, 42]

heading: C MORE RESULTS
text: C.1 IMAGE COMPRESSION
Additional rate-distortion results on CLIC2021, Tecnick, and JPEG-AI are provided in Figure 14, Figure 15, and Figure 16.
For a complete comparison of results from existing literatures, we provide a summary RD plot of all neural image codec solutions known to us in Figure 28 on Kodak. In Figure 29 and Figure 30, we plot the percentage rate saving with BPG444 and VTM-12.1 as reference, respectively.
C.2 VIDEO COMPRESSION
In Figure 17 and Figure 18, we provide performance comparison of Conv-SSF, SwinT-SSF, HEVC (x265), and AVC (x264) with per-video breakdown.

heading: C.3 CODING COMPLEXITY
text: We evaluate the decoding complexity of all neural image codecs in terms of the time for network inference and entropy coding, peak GPU memory, model size, etc. We select 100 high resolution images from CLIC testset and center crop them to three resolutions (768× 512, 1280× 768, 1792×
1024) to see how those metrics scale with image size. Batch size is one for all model inference. We run the experiment on a local workstation with one RTX 2080 Ti GPU, with PyTorch 1.9.0 and Cuda toolkit 11.1. For bit-exact entropy coding, we need to use deterministic13 convolution. The neural networks run on the single GPU and entropy coding runs on CPU with 8 threads. We follow the standard protocols to measure the inference time and peak memory of neural nets, such as GPU warm up and synchronization. File open/close is excluded from coding time measurement. MACs,
13https://pytorch.org/docs/1.9.0/generated/torch.use_deterministic_ algorithms.html?highlight=deterministic#torch.use_deterministic_ algorithms
GPU peak memory, and model parameter count are profiled using get model profile function of deepspeed profiler14.
We show more details on the coding complexity of neural image codecs in Figure 19, particularly the linear scaling to image resolution of both SwinT-based and Conv-based models. The break-down of encoding complexity is shown in Table 4.
Table 5 reports encoding and decoding time of VTM-12.1 under different quantization parameters (QPs), evaluated on an Intel Core i9-9940 CPU @ 3.30GHz, averaged over 24 Kodak images. As can be seen from the table, decoding time of VTM-12.1 is a function of reconstruction quality, where longer decoding time is observed for higher quality reconstruction. In Figure 1, the reported VTM12.1 decoding speed corresponds to a QP value of 28, where the bpp value is similar to that obtained
14https://www.deepspeed.ai/tutorials/flops-profiler/ #usage-outside-the-deepspeed-runtime
by models trained with β = 0.001. It is worth pointing out that VTM-12.1 encoding process is much slower, ranging anywhere from 1 to 5 minutes per image, whereas neural codec runs much faster.

heading: 16 2.5441 44.09 0.284 300.47

heading: 28 0.7844 36.76 0.249 118.85

heading: 40 0.1557 29.51 0.149 71.10

heading: D ANALYSIS

heading: D.1 SPATIAL CORRELATION OF LATENT
text: We visualize the spatial correlation map for Conv-Hyperprior and SwinT-Hyperprior at different β in Figure 20.

heading: D.2 EFFECTIVE RECEPTIVE FIELD
text: See Figure 21 for the effective receptive field for the composed encoding tranforms ha ◦ ga and Figure 22 for the mean attention distance visualization of each head within each transformer layer.

heading: D.3 RATE DISTRIBUTION ACROSS LATENT CHANNELS
text: It is generally believed that ConvNets learn to extract various features and store them in each channel of the activations. Here we look into the features in the latent channels which are to be quantized and entropy coded to bitstreams. Particularly we order the total bitrate of each channel averaged over Kodak dataset (24 768× 256 images). The result is shown in Figure 23. We find an interesting phenomenon across models under different bitrates: there is a cutoff point of the bitrate-vs-channel curve where the bitrate suddenly drops to zero, which manifest the rate constraint in the loss function. As expected, the cutoff index decreases for the model trained for smaller bitrate (larger β).

heading: D.4 CENTERED KERNEL ALIGNMENT
text: To investigate the difference or similarity between latent features of Conv-based and SwinT-based models, we resort to a commonly used tool in representation learning called centered kernel alignment (CKA) (Kornblith et al., 2019). We evaluate CKA between each of the Conv latent channel and SwinT latent channel (both models are trained under the same β) over the 24 Kodak images. There are 320 channels for both Conv latent and SwinT latent, resulting a 320 × 320 CKA matrix. The result is shown in Figure 24. The latent channel is ordered by the averaged bitrate of each channel over Kodak images (same as in Section D.3). The CKA matrix has clear block structure, where high similarity region corresponds to the latent channels before the bitrate cutoff in the rate distribution curve (Figure 23).
Identification of SwinT and Conv latent channels with CKA Within the block of high similarity (from the CKA matrix), we identify the ‘less’ similar SwinT latent channels with lowest CKA values between this SwinT channel and all other Conv latent channels. For each of the identified SwinT channel, we find the Conv latent channel with the largest CKA value between the two. This way, we are able to identify latent channels of two different models with high similarity. We show the identified top 8 channels in Figure 25. The channels are indeed highly similar, up to a sign flip, even through the two model architectures are quite different. This empirical result is relevant to the literature on the identifiability of generative models (Khemakhem et al., 2020).

heading: D.5 PROGRESSIVE DECODING
text: More visual examples of reconstructions of checkerboard (spatially) masked latent are provided in Figure 26.
Channel-wise progressive decoding Here we visualize the behavior of channel-wise progressive decoding of Conv and SwinT models. For both models, we order the latent channels with a heuristic importance metric: the reduction of distortion over the increase of rate if one channel is included for decoding. We start with the order of bitrate per channel, we pass the leading channels of bitrate order (zero out all rest channels) to the decoder to obtain reconstruction and calculate distortion. We plot the top 8 channels following this importance order. For each channel, we show 6 maps from top to bottom: latent values, mean prediction from the hyper decoder, standard deviation prediction from the hyper decoder, the bitmap, the reconstruction with only current channel, the reconstruction with up to current channel (all leading channels). The result is shown in Figure 27. For Conv models, usually the top 3 important channels are responsible for lightness, and two color components (blueyellow, red-green), similar to the LAB colorspace. The rest of the latent channels are responsible for adding details like texture and edges. For the Swint models, at low bitrate, there is one significantly different channel (the first column in Figure 27b), which is in a coarse scale (with smooth blocks) and responsible for the reconstruction of a nearly constant image with value close to 120 (the mean value of natural image dataset). This latent channel costs extremely small bitrate but reaches PSNR of 13dB. We tried remove this first channel, the progressive reconstruction with the rest leading 7 channels only leads to PSNR around 16dB, instead of 26dB shown in the figure.

heading: E MORE ABLATION STUDIES
text: Local self-attention To see if local self-attention is the most important component in transformers, we replace it by depthwise separable convolution block15 (Han et al., 2021; El-Nouby et al., 2021), which performs similar spatial feature aggregation as self-attention, while keeping all other components the same as in the SwinT. We found this change only leads to minor degradation in RD. This suggests other components in transformers such as MLPs and skip connections may also play a big role, other than just self-attention, for the leading performance in our work and many other tasks (Dong et al., 2021).
Small depths Upon investigating the mean attention distance as shown in Figure 22, we find the last block in each of the last two encoder stages has about half of its attention heads degenerate to attending to fixed nearly pixels. This suggests redundant transformer blocks at that stage, so we remove those two blocks, i.e. from depths [2, 2, 6, 2] to [2, 2, 5, 1]. The resulting SwinT-Hyperprior has even less parameters (20.6M) than Conv-Hyperprior (21.4M) while with almost no RD loss compared to the larger model. We expect more hyperparameter search will identify models with better RD-complexity trade-off than we currently show in this work.
Deeper Conv encoder Deeper models are usually more expressive (Raghu et al., 2017) and the state-of-the-art Conv-based compression models typically use much deeper layers than the encoder in the original Hyperprior model (Ballé et al., 2018). As a sanity check on whether deeper convolutional transforms can outperform SwinT-based encoder transforms with 12 blocks, we take an existing design (Chen et al., 2021) with residual blocks and attention (sigmoid gating) layers, which has over 50 conv layers in either encoder or decoder, and more parameters than conv baseline. It indeed improves the RD in lower bitrate, but still worse than SwinT-Hyperprior, and gets much worse in higher bitrate. This is probably the reason that compression models based on this type of transforms did not report results at higher bitrates.
15Conv1×1-LayerNorm-ReLU-DSConv3×3-LayerNorm-ReLU-Conv1×1-LayerNorm-ReLU
Kodak

Review summaries:
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.

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). 

This paper introduces the swin-transformer for the learned image and video coding. Instead of developing more advanced and expensive entropy coding approaches, the authors focus on the transform network(i.e., networks of encoder/decoder) and improve the compression performance significantly.  The experimental results are impressive. Although introducing SwinT is not very new, the authors provide a neat and practical solution, which should be encouraged. The authors also have a good balance between the coding gain and complexity.