T RANSFORMER - BASED T RANSFORM C ODING
A BSTRACT
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.
Speciﬁcally, we show that nonlinear transforms built on
Swin-transformers can achieve better compression efﬁciency than transforms built on convolutional neural networks (ConvNets), while requiring fewer parameters and shorter decoding time.
Paired with a compute-efﬁcient 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-ﬂow model by 12 .
35% in BD-rate on UVG.
We provide model scaling studies to verify the computational efﬁciency of the proposed solutions and conduct several analyses to reveal the source of coding gain of transformers over ConvNets, including better spatial decorrelation, ﬂexible effective receptive ﬁeld, and more localized response of latent pixels during progressive decoding.
I NTRODUCTION
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.
All three components can be enhanced by deep neural networks: autoencoder networks are adopted as ﬂexible nonlinear transforms, deep generative models are used as powerful learnable entropy models, and various differentiable quantization schemes are proposed to aid end-to-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 (Ball´e 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-the-art 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.
However, the rate-distortion improvement from context based prior modeling often comes with a hefty price tag 1 in terms of decoding complexity.
Noteably, all existing works that claimed on-par or better performance than
VTM (Cheng et al., 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?
[Figure]
Our Swin transformer based image compression models land in a favorable region of rate-distortion-computation trade-off that has never been achieved before.
See Section 4.2 for more results and evaluation setup
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 ﬁrst in the neural compression literature.
As main contributions, we 1) extend SwinTransformer (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 scalespace-ﬂow, 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.
B ACKGROUND & R ELATED W ORK
Conv-Hyperprior
The seminal hyperprior architecture (Ball ´ e et al., 2018; Minnen et al., 2018) is a two-level hierarchical variational autoencoder, consisting of a pair of encoder/decoder g a , g s , and a pair of hyper-encoder/hyper-decoder h a , h s .
Given an input image x , a pair of latent y =
g a ( x ) and hyper-latent z
= h a ( 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 ( µ , σ ) =
h s (ˆ z ) .
The quantized version of the latent ˆ y =
Q
( y − µ )
+ µ is then entropy coded and passed through decoder g s to derive reconstructed image ˆ x = g s (ˆ y ) .
The tranforms g a , g s , h a , h s 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 y 1 , . .
, y S ), and the Gaussian prior p ( y s | ˆ z , ˆ y <s ) is made autoregressive across groups where the mean/scale of y s 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 cannot be parallelized along spatial dimensions, leading to impractical 3 decoding 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 (Ball´e 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 efﬁcient 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), classiﬁcation (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.
S WIN -T RANSFORMER BASED T RANSFORM C ODING
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 ﬂexibility in handling varying input resolutions at test time, enabled by relative position bias and hierarchical architecture.
I n p u t
I m a g e
S w i n T r a n s f o r m e r
B l o c k
P a t c h
M e r g e
P a t c h
M e r g e
S w i n T r a n s f o r m e r
B l o c k
P a t c h
M e r g e
S w i n T r a n s f o r m e r
B l o c k
P a t c h
M e r g e
S w i n T r a n s f o r m e r
B l o c k
Q
S w i n T r a n s f o r m e r
B l o c k
P a t c h
M e r g e
P a t c h
M e r g e
S w i n T r a n s f o r m e r
B l o c k
Q
R e c o n s t r u c t
i o n
S w i n T r a n s f o r m e r
B l o c k
P a t c h
S p li t
P a t c h
S p li t
S w i n T r a n s f o r m e r
B l o c k
P a t c h
S p li t
S w i n T r a n s f o r m e r
B l o c k
P a t c h
S p li t
S w i n T r a n s f o r m e r
B l o c k
P a t c h
S p li t
S w i n T r a n s f o r m e r
B l o c k
P a t c h
S p li t
S w i n T r a n s f o r m e r
B l o c k
C h a nn e l - w
i s e
A u t o R e g r e ss i v e
M o d e l
AE
AD
Bits Bits
F a c t o r i z e d
M o d e l
AE
AD _ +
A.1).
3. S WIN T 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 g a 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.
We adopt SwinT encoder as the encoder transform g a in our model, and extend it to SwinT decoder g s by reversing the order of blocks in g a , 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 h a , h s are similar to g a , g s with different conﬁgurations.
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. E XTENSION TO P- FRAME C OMPRESSION
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 ﬂow and residual.
We propose a SwinT variant, referred to as SwinT-SSF , which is obtained by replacing Conv transforms g a , g s in the ﬂow codec and residual codec of SSF with SwinT tranforms.
To stabilize training of ﬂow 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.
E XPERIMENTS AND A NALYSIS 4. E XPERIMENT S ETUP
Training All image compression models are trained on the CLIC2020 training set.
ConvHyperprior and SwinT-Hyperprior are trained with 2 M 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 ﬂow loss 5 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 conﬁguration, all video codecs are evaluated in low-delay-P model with a ﬁxed 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 is computed for each video and then averaged across all videos.
More details on testset preprocessing, and traditional codecs conﬁgurations can be found in Appendix B.2. [Figure]
4. R ESULTS
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.
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).
[Table]
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.
[Figure]
[Table]
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 ﬂow and residual compression beneﬁt 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.
[Table]
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 proﬁled here.
For details on encoding complexity, proﬁling setup, scaling to image resolution, please refer to Table 4 and Section D.3 in the Appendix.
[Figure]
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 achieving better performance, 10 as opposed to state-of-the-art neural codecs with spatial autoregressive prior that decodes orders of magnitude slower.
4. A NALYSIS
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;
Ball ´ e 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
[Figure]
Speciﬁcally, in hyperprior based models (Ball´e et al., 2018), ¯ y (cid:44) ( y − µ ) / σ is modeled as a standard spherical normal vector.
The effectiveness of the analysis transform g a 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 .
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 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.
[Figure]
Effective receptive ﬁeld 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 ﬁeld (ERF)
(Luo et al., 2017) of SwinT encoders compared to Conv encoders.
The theoretical receptive ﬁeld of the encoders ( g a , h a ◦ g a ) 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 classiﬁcation, 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 ﬂow codec and residual codec, as shown in Figure 7d & 7h, and Figure 7c & 7g. Clearly for ﬂow 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 ﬂexibility 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 signiﬁcantly 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.
[Figure]
4. A BLATION STUDY
[Figure]
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 F of the appendix.
C ONCLUSION
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 ﬁrst 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 ﬂexible receptive ﬁeld 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 efﬁcient transformer-based image and video codecs.
A CKNOWLEDGMENTS
We would like to thank Amir Said for developing entropy coding and great advice on data compression in general.
We would also appreciate the helpful discussions from Reza Pourreza and Hoang
Le, and draft reviews from Auke Wiggers and Johann Brehmer.
Appendix Table of Contents
A Models D More Results E Analysis
[Table]
[Table]
[Table]
[Table]
A M ODELS A. C ONVOLUTION BASELINES
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:
[Equation]
SwinT-Hyperprior, SwinT-ChARM For both SwinT-Hyperprior and SwinT-ChARM, we use the same conﬁgurations: ( w g , w h ) =
(8 , 4) , ( C 1 , C 2 , C 3 , C 4 , C 5 , C 6 ) = (128 , 192 , 256 , 320 , 192 , 192) , ( d 1 , d 2 , d 3 , d 4 , d 5 , d 6 ) = (2 , 2 , 6 , 2 , 5 , 1) where C , d , and w are deﬁned in Figure 13 and Figure 2. The head dim is 32 for all attention layers in SwinT-based models.
SwinT-SSF For SwinT transforms used in SSF variant, the ﬁrst Patch Merge block is with downsample rate of 4 and two other Patch 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 ﬂow codec and 8 for residual codec.
[Figure]
[Figure]
d e =
C o n v 5 x
S t r i d e =
Q
R e c o n s t r u c t
i o n
C o n v T 5 x
S t r i d e =
C o n v T 5 x
S t r i d e =
C o n v T 5 x
S t r i d e =
C o n v T 5 x
S t r i d e =
C o n v T 5 x
S t r i d e =
C o n v T 5 x
S t r i d e = 2 C h a nn e l - w i s e
A u t o R e g r e ss i v e
M o d e l ( C h
A R
M )
AE
AD
Bits Bits
F a c t o r i z e d
M o d e l
AE
AD _ +
C o n v 3 x 3 s t r i d e
R e L U
C o n v 3 x 3 s t r i d e
G D N G D N
I G D N I G D N
I G D N
R e L U R e L U R e L U
SwinT-SSF-Res This is a variant where only residual autoencoder uses SwinT transforms. Same architecture as the residual autoencoder in SwinT-SSF.
A. M ODEL CONFIGURATIONS FOR MODEL SIZE SCALING STUDY A.3. S WIN T-H YPERPRIOR
Set of model hyperparameters that are common to all experiments: ( d 1 , d 2 , d 3 , d 4 , d 5 , d 6 ) =
(2 , 2 , 6 , 2 , 5 , 1) ( w g , w h )
= (8 , 4)
SwinT-Hyperprior (small) ( C 1 , C 2 , C 3 , C 4 , C 5 , C 6 ) =
SwinT-Hyperprior [Table] A.3. C ONV -H YPERPRIOR
[Table]
[Figure]
I n p u t
I m a g e
S w i n T r a n s f o r m e r
B l o c k
P a t c h
M e r g e
P a t c h
M e r g e
S w i n T r a n s f o r m e r B l o c k P a t c h
M e r g e
S w i n T r a n s f o r m e r
B l o c k
P a t c h
M e r g e
S w i n T r a n s f o r m e r
B l o c k
Q
S w i n T r a n s f o r m e r B l o c k P a t c h M e r g e
P a t c h
M e r g e
S w i n T r a n s f o r m e r
B l o c k
Q
R e c o n s t r u c t
i o n
S w i n T r a n s f o r m e r
B l o c k
P a t c h
S p li t
P a t c h
S p li t
S w i n T r a n s f o r m e r
B l o c k
P a t c h
S p li t
S w i n T r a n s f o r m e r
B l o c k
P a t c h
S p li t
S w i n T r a n s f o r m e r B l o c k P a t c h
S p li t
S w i n T r a n s f o r m e r
B l o c k
P a t c h
S p li t
S w i n T r a n s f o r m e r
B l o c k
AE
AD
Bits Bits
F a c t o r i z e d
M o d e l
AE
AD _ +
Conv-Hyperprior (large) ( C 1 , C 2 , C 3 , C 4 , C 5 , C 6 , C 7 ) =
B T RAINING AND E VALUATION B. T RAINING
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 2 M batches.
Each batch contains 8 patches of size 256 ×randomly cropped from the training images.
Learning rate starts at 10 − 4 and is reduced to 10 − 5 at 1 .
8 M step.
For Conv-ChARM, we ﬁrst 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 ﬁnetune 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 −and crop size 12 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 beinstead of 8.
Finally we add ﬂow 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.
B. T RADITIONAL CODEC EVALUATION
In this section, we provide evaluation script used to generate results for traditional codecs.
B.2. I MAGE 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.
Speciﬁcally, the following command is issued to gather VTM-12.1 image compression evaluation results: python -m compressai.utils.bench vtm [ image folder]
BPG
BPG software is obtained https bellard.org/bpg the following commands are used for encoding and decoding.
bpgenc -e x265 -q [0 to 51] -f-o [encoded heic file]
[original png file] bpgdec -o
[decoded png file] [encoded heic file]
B.2. V IDEO 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= 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]
HEVC (HM)
[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-q [12, 17, 22, 27, 32, 37, 42]
C BD RATE COMPUTATION
def Bjontegaard_Delta_Rate( #rateandpsnrinascending order rate_ref, psnr_ref, # reference rate_new, psnr_new, # new result ): min_psnr = max(psnr_ref[0], psnr_new[0], 30) max_psnr = min(psnr_ref[-1], psnr_new[-1], 44) log_rate_ref = log(rate_ref) log_rate_new = log(rate_new) spline_ref = scipy.interpolate.
CubicSpline( psnr_ref, log_rate_ref, bc_type=’not-a-knot’, extrapolate=
True, ) spline_new = scipy.interpolate.
CubicSpline( psnr_new, log_rate_new, bc_type=’not-a-knot’, extrapolate=
True, ) delta_log_rate
= ( spline_new.integrate(min_psnr, max_psnr) - spline_ref.integrate(min_psnr, max_psnr) ) delta_rate
= exp(delta_log_rate / (max_psnr - min_psnr)) return 100 * (delta_rate - 1)
C.
BD RATE FOR IMAGE CODEC # Evaluate BD-rate on an image dataset bd_rates = list() for image in image_dataset: # evaluate rate and psnr on reference and new codec # for this image with different qualities rate_ref, psnr_ref = ReferenceCodec(image, qp=[...]) rate_new, psnr_new = NewImageCodec(image, beta=[...]) bd_rates.append(
Bjontegaard_Delta_Rate( rate_ref,psnr_ref, rate_new,psnr_new, ) )
# BD is computed per image and then averaged bd_rate = bd_rates.mean
C.
BD RATE FOR VIDEO CODEC # Evaluate BD-rate on a video dataset bd_rates = list() for video in video_dataset: # evaluate rate and psnr on reference and new codec # for this video with different qualities rate_ref, psnr_ref = ReferenceCodec(video, qp
[Table]
[Table]
new = NewVideoCodec(video, beta
[Table]
[Table]
[Table]
D M ORE R ESULTS D. I MAGE COMPRESSION
Additional rate-distortion results on CLIC2021, Tecnick, and JPEG-AI are provided in Figure 14,
Figure [Figure] and Figure 16.
For a complete comparison of results from existing literatures [Figure] 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 [Figure] we plot the percentage rate saving with BPG444 and VTM-12.1 as reference, respectively.
[Figure]
[Figure]
[Figure]
D. V IDEO C OMPRESSION
In Figure 17 and Figure
[Figure]
we provide performance comparison of Conv-SSF, SwinT-SSF, HEVC (x265) [Figure] and AVC (x264) with per-video breakdown. [Figure]
[Figure]
D. C ODING COMPLEXITY
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 [Figure] 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 deterministic 13 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,
GPU peak memory, and model parameter count are proﬁled using get model profile function of deepspeed proﬁler 14 .
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.
[Figure]
For completeness, we also report the proﬁling for CPU coding time in Table 5 and Table 6. The evaluation setup is the same as the GPU proﬁling case, except models are run on the CPU instead (same host machine with Intel(R)
Xeon(R)
W-2123 CPU @ 3.60GHz). Table 7 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 by models trained with β = 0 .
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 [Table] [Table] [Table]
[Table]
E A NALYSIS E. S PATIAL CORRELATION OF LATENT
We visualize the spatial correlation map for Conv-Hyperprior and SwinT-Hyperprior at different β in Figure 20.
E. E FFECTIVE R ECEPTIVE F IELD
See Figure 21 for the effective receptive ﬁeld for the composed encoding tranforms h a ◦ g a
and
Figure 22 for the mean attention distance visualization of each head within each transformer layer
. E. R ATE DISTRIBUTION ACROSS LATENT CHANNELS
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 ﬁnd 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 β ).
E. C ENTERED KERNEL ALIGNMENT
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 E.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).
Identiﬁcation 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 identiﬁed SwinT channel, we ﬁnd 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 identiﬁed top 8 channels in Figure 25.
The channels are indeed highly similar, up to a sign ﬂip, even through the two model architectures are quite different.
This empirical result is relevant to the literature on the identiﬁability of generative models (Khemakhem et al., 2020).
E. P ROGRESSIVE DECODING
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 signiﬁcantly different channel (the ﬁrst 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 ﬁrst channel, the progressive reconstruction with the rest leadingchannels only leads to PSNR around 16dB, instead of 26dB shown in the ﬁgure.
F M ORE ABLATION STUDIES
Local self-attention To see if local self-attention is the most important component in transformers, we replace it by depthwise separable convolution block 15 (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 ﬁnd the last block in each of the last two encoder stages has about half of its attention heads degenerate to attending to ﬁxed 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 (Ball ´ e 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 on this type of transforms did not report results at higher bitrates
[Figure]
[Figure]
[Figure]
-Hyperprior(Beta
[Figure]
[Figure]
[Figure] [Figure]
[Figure]
[Figure]
Conv-Hyperprior(Beta=0.001) Kodak: PSNR=33.86dB, BPP
[Figure]
[Figure]
[Figure]
SwinT-Hyperprior(Beta
[Figure]
.001)Kodak: PSNR=34.49dB, BPP=0.
[Figure] [Figure]
[Figure] [Figure]
[Figure]
Conv-Hyperprior(Beta=0.003)
Kodak: PSNR=31.09dB, BPP
[Figure]
[Figure] [Figure]
[Figure] [Figure]
[Figure] [Figure]
-Hyperprior(Beta=0.003)Kodak: PSNR=31.66dB, BPP
[Figure] [Figure]
[Figure]
[Figure] [Figure]
[Figure] [Figure]
[Figure]
r c
[Figure]
nn
[Figure]
l o v e r K o
[Figure]
a
[Figure]
[Figure]
n T l a t
[Figure]
n t c h a n n
[Figure]
l i n d e x ( s o r t e d b
y m e a n B PP )
[Figure] [Figure]
[Figure] [Figure]
[Figure]
i n T l a t e n t c h a nn e l i n d e x (
[Figure]
o r t e d b y m e a n B PP )
[Figure]
(f
[Figure]
S w i n T l a t e n t c h a nn e l i n d e x ( s o r t e d b y m e a n B PP ) 0.2 0.
[Figure] [Table]
[Figure] [Figure]
[Figure] [Figure]
[Figure]
c o d e c
c
c c
c c
[Table]
m e a n s c a l e bpp=0.0123 c o d e b i t m a p bpp=0.
bpp=0.
bpp=0. bpp=0.
bpp=0. bpp=0.
bpp=0.
psnr=16.92, bpp=0.0123 r e c o n w / c
o d e [ i ] psnr=20.93, bpp=0.
psnr=24.31, bpp=0.
psnr= [Table] 27, bpp=
[Table]
psnr=26.06, bpp=
[Table]
psnr [Table] 26.23, bpp=0.
[Table]
psnr=26.33, bpp=0.
psnr=26.51, bpp=0.
bpp=0.
r e c o n w / c
o d e [ : i + 1 ] bpp=0.
bpp=0. bpp=0.
bpp=0. bpp=0.
bpp=0. bpp=0.
[Table] [Table]
[Table] [Table]
[Table]
[Table]
[Table]
[Table]
[Table]
c o d e c
c
c c
c c
c
c m e a n s c a l e bpp=0.0001 c o d e b i t m a p bpp=0.
bpp=0.
bpp=0. bpp=0.
bpp=0. bpp=0.
bpp=0.
psnr=13.00, bpp=0.0001 r e c o n w / c
o d e [ i ] psnr=17.38, bpp=0.
psnr=20.89, bpp=0.
psnr=23.84, bpp=0.
psnr=25.00, bpp=0.
psnr=26.03, bpp=0.
psnr=26.24, bpp=0.
psnr=26.30, bpp=0.
bpp=0.0013 r e c o n w / c
o d e [ : i + 1 ] bpp=0.
bpp=0. bpp=0.
bpp=0. bpp=
[Table]
bpp=0.
[Table]
bpp=0.
[Table]
[Table] [Table]
[Table] [Table]
[Table] [Table]
[Figure] [Figure]
[Figure]
R a t e s a v
i n g ( % ) o v e r
B P G 444 ( l a r g e
r
i s b e tt e r )
[Figure] [Figure]
R a t e s a v
i n g ( % ) o v e r
V T M12 . 1 ( l a r g e r i s b e tt e r )
[Figure]