T OWARDS R EAL -T IME N EURAL I MAGE C OMPRESSION WITH M ASK D ECAY
Anonymous authors
Paper under double-blind review
A BSTRACT
Neural image compression has surpassed state-of-the-art traditional codecs (H.266/VVC) for rate-distortion (RD) performance, but suffers from large complexity and separate models for different rate-distortion trade-offs.
In this paper, we propose an efﬁcient single-model variable-bit-rate network, which is able to run at 30 FPS with 768x512 input images and still outperforms VVC for the
RD performance.
By further reducing both encoder and decoder complexities, our small model even achieves 30 FPS with 1920x1080 input images.
To bridge the performance gap between our different capacities models, we meticulously design the mask decay, which transforms the large model’s parameters into the small model automatically.
And a novel sparsity regularization loss is proposed to mitigate shortcomings of L p regularization.
Our algorithm signiﬁcantly narrows the performance gap by 50% and 30% for our medium and small models, respectively.
At last, we advocate the scalable encoder for neural image compression.
The encoding complexity is dynamic to meet different latency requirements.
We propose decaying the large encoder multiple times to reduce the residual representation progressively.
Both mask decay and residual representation learning greatly improve the RD performance of our scalable encoder.
Our code will be public.
I NTRODUCTION
The image compression based on deep learning has achieved extraordinary rate-distortion (RD) performance compared to traditional codecs (H.266/VVC) (Bross et al., 2021).
However, two main issues limit its practicability in real-world applications.
One is the large complexity.
Most stateof-the-art (SOTA) neural image codecs rely on complex models, such as the large capacity backbones (Zhu et al., 2022; Zou et al., 2022), the sophisticated probability model (Cheng et al., 2020), and the parallelization-unfriendly auto-regressive model (Minnen et al., 2018).
The large complexity easily results in unsatisﬁed latency in real-world applications.
The second issue is the inefﬁcient rate-control.
Multiple models need to be trained and stored for different rate-distortion trade-offs.
For the inference, it requires loading speciﬁc model according to the target quality.
For practical purposes, a single model is desired to handle variable RD trade-offs.
In this paper, we try to solve both issues to make a single real-time neural image compression model, while its RD performance is still on-par with that of other SOTA models.
Inspired by recent progress, we design an efﬁcient framework that is equipped with Depth-Conv blocks (Liu et al., 2022) and the spatial prior (He et al., 2021; Li et al., 2022a; Qian et al., 2022).
For variable RD trade-offs, we introduce an adjustable quantization step (Chen & Ma, 2020; Li et al., 2022a) for the representations.
All modules within our framework are highly efﬁcient and GPU friendly, different from recent Transformer based models (Zou et al., 2022).
Encoding and decoding (including the arithmetic coding) achieves 30 FPS for the 768 × 512 inputs.
Compared with other SOTA models, ours enjoys comparable RD performance, low latency, and a single model for all RD trade-offs.
To further accelerate our model, we reduce the complexity of both the encoder and decoder.
Three series models are proposed: Large, Medium, and Small (cf.
Appendix Tab. 3).
Note that our small model even achieves 30 FPS for 1920 × 1080 inputs.
Recent work (Zhu et al., 2022) points out that simply reducing the model’s capacity results in signiﬁcant performance loss.
However, this crucial problem has not been solved effectively.
In this paper, we advocate training small neural compression models with a teacher to mitigate this problem.
In particular, mask decay is proposed
[Figure]
[Table]
(cid:9)(cid:10)(cid:11)(cid:12) (cid:9)(cid:13)(cid:14)(cid:1)(cid:15)(cid:16)(cid:1)(cid:9)(cid:13)(cid:14) (cid:12)(cid:14)(cid:17)(cid:18)(cid:11)(cid:18)(cid:19)(cid:14)(cid:20)(cid:14)(cid:13)(cid:21)(cid:9)(cid:22)(cid:14)(cid:11)(cid:12) (cid:23)(cid:24)(cid:25)(cid:13)(cid:26)(cid:1)(cid:27)(cid:16)(cid:17)(cid:14)(cid:11)(cid:17)(cid:11)(cid:25)(cid:9)(cid:11)
[Figure]
(Zhu et al., 2022) for comparison.
to transform the large model’s parameters into the small model automatically.
Speciﬁcally, mask layers are ﬁrst inserted into a pretrained teacher model.
Then we sparsify these masks until all layers become the student’s structure.
Finally, masks are merged into the raw layers meanwhile keeping the output unchanged.
For the sparsity regularization, L p -norm based losses are adopted by most previous works (Li et al., 2016;
Ding et al., 2018; Zhang et al., 2021).
But they hardly work for neural image compression (cf.
Fig. 7).
In this paper, we propose a novel sparsity regularization loss to alleviate the drawbacks of L p so that optimize our masks more effectively.
With the help of our large model, our medium and small models are improved signiﬁcantly by 50% and 30%, respectively (cf.
Fig.
1a).
That demonstrates the effectiveness of our mask decay and reusing a large model’s parameters is helpful for training a small neural image compression model.
In addition, considering the various device capabilities in real-world codec applications, the encoder scalability of supporting different encoding complexities while with only one decoder is also a critical need.
To achieve this, we propose compressing the cumbersome encoder multi-times to progressively bridge the performance gap.
Both the residual representation learning (RRL) and mask decay treat the cumbersome encoder as a reference implicitly, and encourage the diversity of different small encoders.
Therefore ours achieves superior performance than training separate encoders (cf.
Fig.
1b).
And compared to SlimCAE (Yang et al., 2021) which is inspired by slimmable networks (Yu et al., 2019), ours enjoys a simple framework and better RD performance.
Our contributions are as follows.
• We propose an efﬁcient framework for neural image compression.
It enjoys a single model for different RD trade-offs.
Our model is able to run at 30 FPS for the 768 × 512 inputs, while is on-par with other SOTA models for the RD performance.
Our small model even achieves 30 FPS for the 1920 × 1080 inputs.
• We propose mask decay, an effective method to improve the student image compression model with the help of the teacher.
A novel sparsity regularization loss is also introduced, which alleviates the shortcomings of L p regularization.
Thanks to mask decay, our medium and small models are signiﬁcantly improved by 50% and 30%, respectively.
• We enable the encoding scalability for neural image compression.
With residual representation learning and mask decay, our scalable encoder signiﬁcantly narrows the performance gap from the teacher and achieves a superior RD performance than previous SlimCAE.
R ELATED WORKS
Neural image compression is in a scene of prosperity.
Ball ´ e et al.
(2017) proposes replacing the quantizer with an additive i.i.d.
uniform noise, so that the neural image compression model enjoys the end-to-end training.
Then, the hyperprior structure was proposed by Ball´e et al.
After that, numerous methods are proposed to improve the entropy model and the backbone (both encoder and decoder).
Auto-regressive component (Minnen et al., 2018) and Gaussian Mixture
Model (Cheng et al., 2020) are proposed to improve the probability estimation in entropy model.
Informer (Kim et al., 2022) proposes using an attention mechanism to exploit both global and local
[Figure]
information.
Recently, more powerful backbones (e
[Figure]
g., INNs (Xie et al., 2021), Transformer (Qian et al., 2022) and Swin (Zhu et al., 2022)) are introduced to replace ResNet (He et al., 2016) in the encoder and decoder.
However, recent SOTA neural image codec models suffer from high latency and separate models for each RD trade-off.
In this paper, we propose an efﬁcient neural image compression framework that aims at encoding and decoding in real-time.
Especially, our framework can also handle variable rate-distortion trade-offs by the adjustable quantization step (Chen & Ma, 2020;
Cui et al., 2021; Li et al., 2022a).
The RD performance of our models is on-par with other SOTA neural image compression models, and surpasses the SOTA traditional image codec (H.266/VVC).
Model compression and knowledge distillation (KD) are two prevalent techniques for accelerating neural networks.
Model compression, especially ﬁlter level pruning (He et al., 2017), aims at obtaining an efﬁcient network structure from a cumbersome network.
Some works (Li et al., 2016;
Luo & Wu, 2020; He et al., 2019; Liebenwein et al., 2019; Molchanov et al., 2019) calculate the importance of each ﬁlter, and prune the unimportant ﬁlters.
Other methods directly penalize ﬁlters during training.
They treat some parameters as the importance implicitly, such as the scale in batch normalization (Liu et al., 2017), or the inserted redundancy layers (Ding et al., 2021; Wang & Wu, 2022).
L p -norm based sparsity losses are used in these methods (Ding et al., 2018; Wen et al., 2016; Zhang et al., 2021).
For a pruned network, KD (Hinton et al., 2015) is utilized to boost the performance by mimicking the teacher’s knowledge (e.g., the logit (Zhao et al., 2022), or the feature (Wang et al., 2021)).
All these techniques are mainly studied on semantic tasks, but rarely advocated for neural image compression.
This paper ﬁrstly tries to improve a small neural image compression model with the help of a large model.
And we claim that, for neural image compression models, reusing the teacher’s ﬁlters is an effective way to improve the student.
To achieve this, we propose mask decay with a novel sparsity regularization loss.
With all proposed techniques, our student models are greatly improved.
M ETHODOLOGY
We ﬁrst introduce our efﬁcient framework brieﬂy.
Then, three models with different complexities are proposed:
Large, Medium, and Small.
Next, we describe how to improve our small model (student) with the help of our large model (teacher).
Mask decay is proposed to transform the teacher’s parameters into the student automatically.
At last, the scalable encoder is advocated.
With residual representation learning and mask decay, several encoders are trained progressively to bridge the performance gap from the teacher.
3. O UR FRAMEWORK FOR IMAGE COMPRESSION
Our framework is shown in Fig.
Given an input image, the encoder generates its representation y .
Following Chen & Ma (2020);
Li et al.
(2022a),
to achieve variable bit-rates within a single model, we introduce adjustable quantization steps for y .
q globalstep and q chstep are two parameters that quantifying y globally and channel-wisely, respectively.
And different RD trade-offs can be achieved by adjusting the magnitude of q globalstep .
More details of our framework are introduced in Appendix Sec.
A.
Fig.
2 also presents MACs of each module with the 1920 × 1088 input image.
Obviously, most computational costs come from the encoder and the decoder.
Hence, we focus on accelerating these two modules, while keeping other parts’ architectures unchanged.
Fig.
3 shows the detail structure of these two modules.
Different from recent SOTA models (Zhu et al., 2022;
Qian et al., 2022;
[Table]
(cid:1)(cid:2)(cid:3)(cid:4)(cid:5)(cid:17) (cid:7)(cid:8)(cid:9)(cid:10)(cid:5)(cid:17) (cid:1)(cid:2)(cid:3)(cid:4)(cid:5)(cid:24) (cid:7)(cid:8)(cid:9)(cid:10)(cid:5)(cid:24) (cid:7)(cid:1)(cid:8)(cid:9)(cid:5)(cid:10)(cid:11)(cid:12)(cid:13)(cid:14)(cid:12)(cid:4)(cid:3)(cid:15) (cid:16)(cid:1)(cid:17)(cid:18)(cid:19)(cid:13)(cid:20)(cid:4)(cid:2)(cid:21)(cid:13)(cid:14)(cid:12)(cid:4)(cid:3)(cid:15)
(cid:1)(cid:2)(cid:3)(cid:4)(cid:5)(cid:6)(cid:7)(cid:5) (cid:8)(cid:9)(cid:10)(cid:11)(cid:5)(cid:6)(cid:7)(cid:5) (cid:1)(cid:2)(cid:3)(cid:4)(cid:5)(cid:12)(cid:7)(cid:5) (cid:1)(cid:2)(cid:3)(cid:4)(cid:5)(cid:6)(cid:7)(cid:5)
[Figure]
a) presents our pruning method.
The mask layer is inserted between two conv.
layers.
These two cumbersome conv.
layers will be transformed into efﬁcient ones by merging the mask layer into them.
(b) compares the gradient ∂ L sparse ∂x of different sparsity loss functions.
Zou et al., 2022), we adopt residual blocks and depth-wise convolution blocks (Liu et al., 2022) rather than Transformer.
Hence, ours is more GPU-friendly and enjoys lower latency.
C 1 , C 2 , C 3 , and C 4 are pre-deﬁned numbers for different computational budgets.
Three models are proposed with channels [64 , 64 , 128 , 192] (Small), [128 , 128 , 192 , 192] (Medium), and [192 , 192 , 192 , 192] (Large).
Tab. 3 in the appendix summarizes models’ #Params and MACs.
Note that Large is our cumbersome teacher model while Medium and Small are efﬁcient student models.
3. I MPROVE THE STUDENT BY MASK DECAY
Due to the limited capacity, the student performs badly if training from scratch (Zhu et al., 2022).
How to improve the student is a crucial but rarely studied problem for neural image compression.
In this paper, we claim that reusing the teacher’s parameters is helpful.
To achieve it, we propose mask decay to transform parameters in an end-to-end manner.
Compress one layer’s representation by mask decay.
As showed in Fig.
4a, a mask layer is inserted between two layers to reduce the representation’s channels (from N 2 to N s ).
This mask layer contains N 2 parameters to prune corresponding channels.
In the forward computation, the representation multiplies these parameters channel-wisely.
During training, these parameters will become sparse by optimizing our sparsity loss.
After training, indicated by the mask, useless parameters in
Conv #1 and #2 are pruned.
And the mask is eliminated by multiplying itself into Conv #2.
To sparsify parameters m in the mask layer, we design a novel sparsity regularization loss.
Most previous works (Li et al.,
[Figure]
norm ( 12 (cid:107) m (cid:107) 2 ) based functions.
Note that with the gradient descend, m is updated by sign ( m ) and m , respectively.
Miserably, both of them fail for neural image compression models (cf.
Fig. 7).
We argue that, for a parameter with a positive value, the gradient derived by the L1-norm is a constant without considering its own magnitude.
And for the L2-norm, when the parameter is approaching zero, the gradient is too small to sparsify the mask.
In addition, both of them have no stationary points when parameters are positive.
To alleviate these problems, We design a novel sparsity regularization loss from the gradient view.
As showed in Fig.
4b, the gradient is deﬁned by
[Equation]
L sparse ( x
[Equation]
where x ≥ 0 .
For x > 1 , the parameter with a large magnitude suffers a large gradient to decay, which likes the L2-norm.
And for 0 ≤ x < 1 , a smaller parameter enjoys a larger gradient to approach zero fast.
In addition, there is a stationary point at x = 1 , which means ∂ L sparse ( x ) ∂x | x =1 = 0 .
When it comes to L1 and L2, there always is a positive gradient for x > 0 .
To make the model converge to stable, it needs a negative gradient from the task loss to balance this positive gradient.
Thanks to our sparsity loss, only a small negative gradient is needed when x is close to 1 .
L sparse can be calculated by integrating Equation 1, which is
[Equation]
For training, inspired by AdamW (Loshchilov & Hutter, 2018), we decouple L sparse from the task loss L RD as weight decay.
Therefore, the task loss is not affected by our sparsity loss explicitly.
So we name our algorithm as mask decay .
For the mask m , the updating formula is
[Equation]
where γ is the learning rate and η is the rate for mask decay.
Mask decay for our encoder and decoder.
As illustrated in Fig.
3, to transform the teacher into the student, we insert mask layers in some positions if necessary.
Our principle is to insert mask layers as few as possible.
In one residual block, Mask #1 and #2 are inserted to prune channels inside and outside the residual connection, respectively.
In the depth-conv block, Mask #1 determines channels of the depth convolution layer, while Mask #3 prunes representations between Conv #3 and #4.
Note that the output channels of these two blocks are shared, so Mask #2 and #4 in the depth-conv block share weights with Mask #2 in the residual block.
And we place Mask #2 and #4 behind the residual adding to guarantee the identity gradient back-propagation.
Training with mask decay.
The training begins with a pretrained teacher model.
Then, we insert mask layers and initialize them with 1 .
The whole network is optimized by minimizing
L RD =
R + λD , where λ steers the trade-off between the rate R and the distortion D .
Our masks are decayed before each iteration, likely weight decay (Loshchilov & Hutter, 2018).
During training, we stop mask decay for one mask if it is sparse enough.
When all masks meet the sparsity requirements, the whole sparse network can transform into the student by merging masks, meanwhile keeping the output unchanged.
At last, the student will be ﬁnetuned for a few epochs.
3. T HE SCALABLE ENCODER
At last, we advocate the scalable encoder in neural image codecs.
A trivial solution for the scalable encoder is directly training separate small encoders with a uniform decoder.
For a speciﬁc encoding
[Table]
complexity, the corresponding number of encoders are chosen to encode together, and then the best bitstreams will be sent to the decoder.
However, these separate encoders suffer from the homogenization problem.
We propose residual representation learning (RRL) to mitigate this problem.
As illustrated in Fig.
[Table]
we compress the cumbersome encoder several times to learn the residual representations progressively.
In each step, only parameters in the compressing encoder are optimized.
That makes sure previous encoders generate decodable bitstreams.
For a speciﬁc encoding complexity, the ﬁrst k encoders are applied.
Our RRL encourages the encoders’ diversity.
Both RRL and mask decay treat the cumbersome teacher as a reference, which makes the training more effective.
SlimCAE
(Yang et al., 2021) uses slimmable layers (Yu & Huang, 2019) to obtain an image compression model with dynamic complexities.
Our framework is more simple to implement and achieves superior RD performance.
CBANet (Guo et al., 2021) proposes a multi-branch structure for dynamic complexity.
It can be considered as our baseline.
More details are in Appendix C.
E XPERIMENTS
Training & Testing.
The training dataset is the training part of Vimeo-90K (Xue et al., 2019) septuplet dataset.
During training, the raw image is randomly cropped into 256 × 256 patches.
For a fair comparison, all models are trained with 200 epochs.
For our method, it costs 60 epochs for the teacher to transform into the student with mask decay, then the student is ﬁnetuned by 140 epochs.
Testing datasets contain Kodak (Franzen, 1999), HEVC test sequences (Sharman & Suehring, 2017),
Tecnick (Asuni & Giachetti, 2014).
BD-Rate (Bjontegaard, 2001) for peak signal-to-noise ratio (PSNR) versus bits-per-pixel (BPP) is our main metric to compare different models.
More details about experimental settings are introduced in Appendix Sec.
D.
4. C OMPARISON WITH STATE - OF - THE - ART
Fig.
6 compares our models with other SOTA models on Tecnick and Kodak.
Our large model outperforms the SOTA traditional image codec (“VTM 17.2”) and is on-par with other SOTA neural image codec models.
Note that our model handles different RD trade-offs by only one model, while other SOTA methods train separate models for each trade-off.
Tab. 1 compares latency of different models.
For Entroformer (Qian et al., 2022) and STF (Zou et al., 2022), we measure the latency by running their ofﬁcial code.
Note that the latency contains arithmetic coding , but ignores the disk
I/O time.
SwinT (Zhu et al., 2022) is also compared in Appendix Tab. 6.
Our models surpass these
SOTA models with a signiﬁcant margin.
Our large model runs at 30 FPS on A100 for the 768 ×inputs, while our small model even achieves 30 FPS for the 1080P inputs.
For more previous but high-impact neural image codecs (e.g., Cheng
(Cheng et al., 2020) and Xie2021 (Xie et al., 2021)), we make comparisons in Appendix F.
4. M ASK DECAY
Previous work (Zhu et al., 2022) ﬁnds small models suffer from poor RD performance.
This problem is alleviated by our mask decay.
We accelerate the encoder and the decoder respectively to study more cases.
Note that the large model is the teacher and the decay rate is 4 e − 5 .
Fig.
1a and
Tab. 2 summarize experimental results.
In Tab. 2, we also calculate the relative improvement as a
[Table]
[Figure]
[Table]
percentage.
For example, if BD-Rate of the baseline student, our student, and the baseline teacher are 1 .
1% , − 0 .
4% , and − 1 .
1% , respectively.
The relative improvement is 1 .
1 − ( − 0 . 4) 1 . 1 − ( − 1 . 1)
= 68% .
This metric indicates how much the performance gap between the student and the teacher our method narrows.
RD curves are shown in Appendix Fig.
From experimental results, we conclude that:
• Students are improved by mask decay with a signiﬁcant margin.
For medium and small models, relative improvements are about 50% and 30% , respectively
This demonstrates the effectiveness of our method
and the student beneﬁts from
reusing
the
teacher’s
parameters.
Compared with SwinT-Hyperprior in Fig.
1a,
ours achieve a superior trade-off between the RD performance and the complexity.
In addition,
our
models enjoy small sizes.
• The encoder is more redundant than the decoder.
Note that our encoder and decoder are symmetric.
We surprisingly ﬁnd compressing the encoder enjoys less RD performance loss, compared with the decoder.
With the large decoder, compressing the encoder from Large to Small loses 2 .
7% BD-Rate on Kodak.
On the other hand, only compressing the decoder loses 5 . 1% BD-Rate.
[Table]
(cid:2)(cid:1)(cid:1) (cid:12) (cid:13) (cid:5) (cid:14) (cid:15) (cid:16) (cid:6) (cid:17) (cid:9) (cid:10) (cid:11) (cid:1)(cid:2) (cid:5)(cid:12)(cid:16)(cid:6)(cid:17)(cid:12)(cid:16)(cid:6)(cid:2)(cid:18)(cid:7)(cid:12)(cid:16)(cid:5) (cid:1)(cid:2)(cid:3)(cid:2) (cid:1)(cid:4)(cid:3)(cid:2) (cid:5)(cid:6)(cid:3)(cid:7)
(cid:8)(cid:5)(cid:3)(cid:7)
[Figure]
Appendix Fig.
20 visualizes our models’ reconstructions.
Reconstructions of different capacities models are consistent.
Our mask decay does not result in extra obvious artifacts.
And our encoding and decoding time are dramatically shorter than VTM’s.
4. T HE SCALABLE ENCODER
Fig.
1b presents results for the scalable encoder.
First, training separate encoders suffers from poor performance, which is due to the homogenization problem.
Then, with our framework, training these encoders one-by-one to bridge the residual representations greatly improves the performance.
Our
RRL encourages these encoders’ diversity.
Finally, the mask decay treats the cumbersome encoder as a reference and makes training more efﬁcient.
More details can be found in Appendix Sec.
C.
Fig.
6 shows RD curves for our scalable encoder (“Scalable”).
Ours outperforms SlimCAE
(Yang et al., 2021) signiﬁcantly by 15 . 5% BD-Rate.
Note that SlimCAE is a powerful model equipped with hyperprior (Ball ´ e et al., 2018), conditional convolutions (Choi et al., 2019), and an autoregressive context model (Minnen et al., 2018).
Overall, our scalable encoder is on-par other SOTA models.
4. A BLATION STUDIES
Sparsity losses and decay rates.
Fig.
7 summarizes experimental results for mask decay with different sparsity regularization losses and decay rates.
For sparsity, 0% and 100% denote the teacher’s and the student’s structures, respectively.
We adjust the decay rate η for each sparsity loss to get reasonable sparsity decay curves.
BD-Rate is calculated over VTM on Kodak.
For our sparsity loss, all decay rates achieve better RD performance compared with the baseline.
A suitable decay rate ( 4 e − 5 ) achieves the best BD-Rate.
A large decay rate ( 8 e − 5 ) results in fast decay which destroys the model’s performance in early epochs.
This model suffers from the insufﬁcient decay process, and the performance is easily saturated.
Small decay rates ( 1 e − 5 and 2 e − 5 ) waste too many epochs for mask decay and the model lacks ﬁnetuning.
For L1, the model suffers from poor performance with both large and small decay rates.
With a suitable rate ( 1 e − 5 ), it is still worse than ours with a signiﬁcant margin.
For L2, we ﬁnd it hardly works because the gradient vanishes when the mask is close to zero.
Discussions for important channels.
In the ﬁeld of model pruning, many works (Li et al., 2016;
He et al., 2019; Liebenwein et al., 2019; Molchanov et al., 2019) study how to identify unimportant channels for pruning.
However, recent works (Liu et al., 2018; Li et al., 2022b) indicate randomly pruning channels results in comparable or even superior performance, which demonstrates there may be no importance difference among all channels.
Previous works study mainly semantic tasks.
We want to extend this conclusion for low-level
tasks, i.e., image compression
First, we notice that our mask decay selects similar channels with different decay rates in Figure 7a.
We use C i to represent the set of chosen channels, where i ∈ F = { 1 e − 5 , 2 e − 5 , 4 e − 5 , 8 e − 5 } means the decay rate η .
We check masks in the ﬁrst two blocks, and #channels of the teacher and the student are 192 and 64 , respectively.
Hence, about 33% channels need to be kept after pruning.
However, we ﬁnd ∩ i ∈ F C i and ∪ i ∈ F C
i contain 24% and 44% channels, respectively, which means
[Table]
[Figure]
all of them only consider 84 channels from 192 channels.
Are these 44% channels important?
Our answer is negative.
To verify our opinion, we propose mask decay with avoiding speciﬁc channels.
We divide all channels into two groups: freedom F and avoidance A .
To avoid channels in A , we set a large decay rate for them.
Therefore, the loss of our Avoidable Mask Decay (AMD) is
[Equation]
(4) where 1 condition is 1 if the condition is true, 0 otherwise.
Defaultly, we set η a = 10 η f .
Hence, avoidable channels enjoy a large decay rate to approach zero quickly.
We conduct two experiments that avoid selecting channels in ∪ i ∈ F C i for the ﬁrst one or two blocks, respectively.
Figure 8a illustrates our AMD indeed avoid channels.
Speciﬁcally, we calculate C AkB ∩ C F and C AkB ∪ C F , where k = 1 , 2 denotes avoiding channels in the ﬁrst k blocks and C F = ∪ i ∈ F C i .
At most 1% channels are chosen by both A1D and F in the ﬁrst block, and A2B avoids channels of F in both two blocks.
Figure 8b shows the BD-Rate for each epoch.
From it, we ﬁnd AMD results in comparable or a little bit superior performance compared to the raw mask decay.
Overall, if some channels are important, we can avoid these channels to achieve comparable performance.
That demonstrates there are no important channels.
For future works, this conclusion guides us to pay more attention to the optimization process rather than the importance of each channel.
C ONCLUSIONS AND FUTURE WORKS
We proposed an efﬁcient framework for image compression.
With 768 × 512 inputs, our model is able to run at 30 FPS, while the RD performance is on-par with SOTA neural codecs and superior to the SOTA traditional codec (H.266/VVC).
With an adjustable quantization step, our model handles variable RD trade-off within only one model.
These three signs of progress make a new milestone in the ﬁeld of image codec.
By reducing the complexities of both encoder and decoder, a small model is proposed, which achieves 30 FPS for 1080P inputs.
We proposed mask decay for training this small model with the help of the cumbersome model.
And a novel sparsity regularization loss is designed to alleviate the drawbacks of L1 and L2 losses.
Thanks to our algorithm, our medium and small models are improved signiﬁcantly by 50% and 30%, respectively.
At last, the scalable encoder is advocated.
With only one decoder, there are several encoders for different complexities.
With residual representation learning and mask decay, the scalable encoder narrows the performance gap between the teacher.
And ours outperforms previous SlimCAE by a signiﬁcant margin.
We believe this paper will inspire many future works.
First, we point out the encoder is more redundant than the decoder.
Hence, improving the decoder is promising.
Second, this paper mainly focuses on the pre-deﬁned student structure.
How to determine the complexity of each layer is still an open problem for image compression models.
Third, it is also encouraged to ﬁnd a more powerful pruning algorithm.
But according to our AMD, we advise paying more attention to the optimization process rather than the importance criterion for each ﬁlter.
B M ASK DECAY
The pipeline of our model acceleration method.
Figure 12 shows how to accelerate a cumbersome network by our mask decay.
First, this cumbersome network is trained and treated as the pretrained teacher.
Then, mask layers are inserted in the teacher and are optimized to become sparse.
After training, masks are merged into neighbor layers, so that the cumbersome teacher transforms into an efﬁcient student network.
Compress one layer’s representation by mask decay.
In Fig.
4a
[Table]
our mask contains N 2 parameters which are initialized by 1 .
After training, N s parameters remain to be positive while N 2 − N s parameters become zeros.
Then, we conduct two steps to “merge” this mask.
First, for N 2 − N s zero channels, corresponding ﬁlters in Conv #1 and #2 are pruned safely.
Second, for N s positive parameters, they are multiplied into Conv #2.
After these two steps, the mask can be removed.
Claim 1 guarantees that these two steps will not change the outputs.
Hence, our sparse mask is able to be eliminated safely.
Claim 1.
Our mask can be merged into neighbor convolution layers while keeping the outputs unchanged.
Proof.
As showed in Fig.
4a, assume W 1 and b 1 are parameters in pruning Conv #1, and W 2 and b 2 are parameters in pruning Conv #2.
The LeakyReLU applies the element-wise function:
[Equation]
Then given the input X ∈ R N 1 × H × W , the forward computation of the pruning model (the left part in Fig. 4a) is
[Equation]
where (cid:126) denotes the convolution operator, and (cid:12) c means the channel-wise product.
Note that N s elements of m are positive, while other N 2 − N s entries are zeros. In the ﬁrst step, these N 2 − N s zero channels are pruned safely, and the forward computation becomes [Equation] Then
[Equation]
in the second step
[Equation]
N s positive entries in m (cid:48) will be multiplied into W (cid:48) 2 .
We have
[Equation]
The last equation holds for the linearity property of the convolution operator.
Notice that Eq.denotes the computation process in the pruned model (the right part in Fig. 4a).
Hence, the pruned model also outputs Y 4 with the input X .
That means the outputs are unchanged. By merging the mask m , parameters within Conv #1 and # 2 become ( W (cid:48) 1 , b (cid:48) 1 ) and ( W (cid:48) 2 (cid:12) c m (cid:48) , b 2 ), respectively.
Mask decay for our encoder and decoder.
How to prune the channels outside the residual connection is a crucial problem.
Fig.
13 shows two manners.
Note that Mask #1, #2, and #3 share weights.
If we insert Mask #2 and #3 after the residual adding (as in Fig. 13b), the identity output will be degraded by these masks, because most masks are in [0 , 1] .
Inserting masks as in Fig.
13a mitigates this problem and guarantees the identity gradient back-propagation.
C T HE SCALABLE ENCODER
In this part, we introduce more details for different methods to achieve a scalable encoder.
End-to-end.
CBANet (Guo et al., 2021) proposed a multi-branch structure for a scalable network.
The computation complexity is dynamic by using different numbers of branches in testing.
This method can be considered as the baseline in our paper.
That is directly training our framework end-to-end (cf. Fig. 14).
During training, all encoders generate bitstreams for decoding, and the corresponding loss is calculated.
However, this baseline requires a large GPU memory for training and [Table] [Figure] [Figure]
[Figure]
(cid:1) (cid:1)(cid:2)(cid:3) (cid:1) (cid:1)(cid:2)(cid:3)(cid:4)(cid:7) (cid:1)(cid:2)(cid:3)(cid:4)(cid:5)(cid:6)(cid:7)(cid:4) (cid:1)(cid:2)(cid:3)(cid:4)(cid:8) (cid:1)(cid:2)(cid:3)(cid:4)(cid:5)(cid:6)(cid:7)(cid:4) (cid:1)(cid:2)(cid:3) (cid:1) (cid:1)(cid:2)(cid:3) (cid:1) (cid:1)(cid:2)(cid:3)(cid:4)(cid:5) (cid:1)(cid:2)(cid:3) (cid:1) (cid:1)(cid:2)(cid:3)(cid:4)(cid:5) (cid:1)(cid:2)(cid:3)(cid:4)(cid:6) (cid:1)(cid:2)(cid:3) (cid:1) (cid:1)(cid:2)(cid:3)(cid:4)(cid:5) (cid:1)(cid:2)(cid:3)(cid:4)(cid:6) (cid:1)(cid:2)(cid:3)(cid:4)(cid:7) (cid:1)(cid:2)(cid:3)(cid:4)(cid:5) (cid:1)(cid:2)(cid:3)(cid:4)(cid:6) (cid:1)(cid:2)(cid:3)(cid:4)(cid:7) (cid:1)(cid:2)(cid:3)(cid:4)(cid:8) [Figure] [Table]
performs badly demonstrated by our experiments.
End-to-end training requires optimizing all parameters of the whole network.
And it does not utilize the knowledge in the pretrained cumbersome model.
Hence, the training is inefﬁcient.
Separate.
Another baseline method is directly training separate small encoders with a uniform decoder.
As shown in Fig.
15, four small encoders are trained with the frozen decoder and entropy model.
For testing, different numbers of encoders can be chosen for the speciﬁc encoding complexity.
In the right part of Fig. 15, all four encoders are chosen.
The ensemble module will receive all bitstreams from encoders and choose the best bitstream to send.
Note that the raw image is available in encoding.
The ensemble module is able to reconstruct each bitstream and compute the RD score according to the raw image.
However, these separate encoders suffer from the homogenization problem.
With the uniform decoder, these encoders generate similar bitstreams.
In Fig.
1b, we ﬁnd that adding more encoders improves the RD performance marginally.
One-by-one (RRL).
We propose residual representation learning (RRL) to mitigate the homogenization problem.
As illustrated in Fig.
16, we train small encoders progressively.
In each step, a new small encoder is added and trained from scratch.
Note that other parts are frozen.
The new encoder is encouraged to learn the residual representations which effectively alleviates the homogenization problem and enlarges the diversity of these small encoders.
Fig.
1b shows “one-by-one” outperforms “separate encoders” by a signiﬁcant margin.
Ours (RRL + mask decay).
At last, mask decay further improves the RD performance.
As illustrated in Fig.
5, mask decay directly utilizes knowledge of the cumbersome model.
In each step, the small encoder is obtained by pruning from the cumbersome teacher.
Here the teacher means the large encoder in Tab. 3, whereas the student means the small encoder.
In Fig.
5, the encoder with dash boundaries denotes the teacher encoder equipped with masks.
During training, these masks will become sparse so that the network is able to transform into the student structure (i.e., the student encoder).
Both RRL and mask decay treat the cumbersome teacher as a reference, which makes the training more effective.
Fig.
1b shows “ours” (RRL + mask decay) achieves superior RD performance and signiﬁcantly narrows the performance gap from the teacher.
Tab. 4 summarizes results for our scalable encoders with the large decoder.
We also conduct experiments for the small decoder, and Tab. 5 summarizes results. [Table]
D D ETAILS FOR EXPERIMENTS
Training.
The training dataset is the training part of Vimeo-90K (Xue et al., 2019) septuplet dataset.
During training, the raw image is randomly cropped into 256 × 256 patches.
For a fair comparison, all models are trained with 200 epochs.
For our method, it cost 60 epochs for the teacher to transform into the student with mask decay, then the student is ﬁnetuned by 140 epochs.
AdamW is used as the optimizer with batch size 16.
The initial learning rate is 2 e − 4 , and decays by 0.5 at 50, 90, 130, 170 epochs.
The default decay rate η for our mask decay is 4 e − 5 .
All models are trained on a computer with 8 V100 GPUs.
Our large and small model cost about 33 and 16 hours for training, respectively.
Testing.
BD-Rate (Bjontegaard, 2001) for peak signal-to-noise ratio (PSNR) is our main metric to compare different models.
We use VTM 17.2 as the anchor, while negative numbers indicate bitrate saving and positive numbers indicate bitrate increase.
Testing datasets contain Kodak (Franzen, 1999), HEVC test sequences (Sharman & Suehring, 2017), Tecnick (Asuni & Giachetti, 2014).
Kodak contains 24 images with the 768 × 512 resolution.
And Tecnick consists of 100 images with the 1200 × 1200 resolution.
HEVC is a video dataset that consists of 4 classes: B (1080P), C (480P),
D (240P), E (720P).
For testing image codecs on HEVC, only the ﬁrst frame of each video sequence is used.
We report BD-Rate on each HEVC class and average all for convenient comparison.
For one testing image, we pad zeros on the boundary to make the resolution as multiples of 64.
We test the latency on two computers.
One is equipped with one 2080Ti GPU and two Intel(R)
Xeon(R)
E5-2630 v3 CPUs.
Another is equipped with one A100 GPU and one AMD Epyc 7V13 CPU.
Both encoding and decoding times are gathered.
Note that the latency contains arithmetic coding , but ignores the disk I/
O time.
E I MPLEMENT DETAILS FOR OTHER IMAGE CODECS
VTM-17.2 1 is the reference software for the state-of-the-art traditional codec (H.266/VVC), and it is released on Jul 17, 2022.
We use CompressAI-1.1.0 2 to gather evaluation results.
We install these two softwares according to the ofﬁcial instructions.
The default conﬁguration of VTM is used.
All results are gathered by running the following command: python −m compressai.utils.bench vtm [path to image folder] −c [path to VTM folder]/cfg/ encoder intra vtm.cfg −b [path to VTM folder]/bin −j 8 −q 24 26 28 30 32 34
Entroformer 3 is proposed in Qian et al.
Two series models are released: parallel, nonparallel.
We evaluate the RD performance of the parallel model on Tecnick and Kodak, respectively.
But when we measure its latency, we ﬁnd it requires a large GPU memory for generating bitstreams.
[Table]
on numbers in for
SwinT-ChARM is proposed in Zhu et al.
Because the authors did not release their code, we cite numbers in their paper and the webpage 5 .
The latency is reported in their paper, and Tab. compares ours with theirs.
Our models achieve lower latency.
ELIC is proposed in He et al.
Uneven channel-conditional adaptive coding is proposed in this work.
He et al.
(2022) designed two models ELIC and ELIC-sm, where ELIC-sm is adapted from ELIC with a slim architecture.
Because the authors did not release their code, we make comparisons with numbers in their paper.
Cheng2020 is proposed in Cheng et al.
Attention modules and Gaussian Mixture Model are proposed to improve the RD performance.
Due to its high impact, it is always considered as the baseline method for neural image codecs.
However, the autoregressive module is adopted in their raw model.
And inference on GPU is not recommended for the autoregressive models.
When testing on CPU, it takes more than one second to process one image.
He et al.
(2021) proposed the checkerboard module to mitigate this problem.
We compare ours with this model in Tab. 7.
Xie2021 is proposed in Xie et al.
The invertible neural networks are proposed and used as the stronger backbone.
But the same as Cheng2020, Xie2021 is also equipped with the autoregressive module.
It suffers from high latency due to this module.
We compare ours with this model in Tab. 7.
Minnen2018 and Minnen2020 are proposed in Minnen et al.
and Minnen & Singh (2020), respectively.
The autoregressive context model was ﬁrst introduced in Minnen et al.
(2018) to boost the RD performance of neural image codecs.
Minnen & Singh (2020) further improved it by considering channel-wise redundancy.
Both two works have a signiﬁcant impact in the neural image compression and are always considered as the baseline method for comparisons.
We also compare ours with them in Tab. 7.
Ball´e2018 is proposed in Ball ´ e et al.
The hyperprior structure was introduced, which is a milestone in the ﬁeld of image codec.
Most latter models adopt this structure, including ours.
Without a context module, it achieves a high inference speed.
But without using recent proposed techniques, it suffers from a poor RD performance compared with recent SOTA models.
F M ORE EXPERIMENTAL RESULTS
Comparison to more methods.
He et al.
(2022) proposed two models for efﬁcient learned image compression, namely ELIC and ELIC-sm, respectively.
However, they did not release their codes and models and it is difﬁcult for us to make a fair comparison.
We follow their testing manner to compare numbers in their paper.
First, in their supplementary material, they mentioned the deterministic inference mode was not enabled when testing the model speeds.
We test latencies in deterministic/non-deterministic modes, respectively.
We ﬁnd that non-determinism reduces the latency by average 16% but easily results in catastrophic failures for decoding (Ball´e et al., 2019).
Second, different from testing end-to-end latency with arithmetic coding in our main paper, here we test each module’s latency as that in He et al.
Hence, without considering arithmetic coding, [Table] [Table] the latency is lower than that in our main paper.
Third, unfortunately, we do not have TITAN Xp
GPUs
[Table]
and we test our models on 2080Ti 7 .
Tab. 7 summarizes results.
Compared with previous SOTA methods (Ball´e2018, Minnen2018,
Cheng2020, Minnen2020, Xie2021), our large model (“EncL-DecL”) achieves better RD performance and inference latency.
And our model is also comparable to ELIC-sm.
Note that ours enjoys a single model to handle all RD trade-offs, while ELIC trained different models for each trade-off.
Speed (megapixel/s) comparison.
In Tab. 8, we also compute models’ speeds as megapixel/s according to Tab. 1.
Our models outperform previous models for this metric, too.
Comparison for our different models.
Fig.
17 presents RD curves on Kodak.
Our mask decay “MD” improves our baseline models signiﬁcantly.
In these ﬁgures, “L”, “M”, and “S” denote Large,
Medium, and Small, respectively.
And “EncS-DecL” means the encoder is small and the decoder is large.
[Table] [Figure] [Table] Tab. 9 presents these model’s results on CLIC2021 8 and Tecnick.
Mask decay (Ours) outperforms the baselines by signiﬁcant margins.
Sparsity losses and decay rates.
For easier to read, we replicate Fig.
7 for larger sizes.
Figs.and 19 are these replications.
Disscusions for important channels.
Tab.
10 presents the speciﬁc numbers in Fig.
8a.
There are four masks in the ﬁrst two blocks.
We analyze the ﬁrst mask (ResB.1.mask1) for example.
Mask decay is required to choose 64 channels from the teacher’s 192 channels.
With four different decay rates, only 83 channels are considered and they choose 44 channels together.
It seems like these channels are important for superior performance.
We treat these 83 channels as C F and let our algorithm avoid choosing these channels.
A1B (avoid speciﬁc channels in the ﬁrst block) only selects 2 channels in C F , while A2B selects 3.
First, it demonstrates our AMD indeed avoids speciﬁc channels.
Next, we compare the RD performance of these methods.
Surprisingly, AMD results in comparable performance.
That indicates that there are no important channels for superior performance.
[Table]
(cid:7)(cid:2)(cid:3)(cid:4) (cid:1)(cid:2)(cid:3)(cid:2) (cid:4)(cid:3)(cid:5) (cid:6)(cid:3)(cid:7) (cid:8)(cid:3)(cid:6) (cid:6)(cid:3)(cid:2) [Figure] [Table] [Figure]
(cid:2)(cid:1)(cid:1)
[Figure]
(cid:15) (cid:16) (cid:6)
[Figure]
(cid:9) (cid:10)
[Figure]
(cid:1)(cid:2) (cid:5)(cid:12)(cid:16)(cid:6)(cid:17)(cid:12)(cid:16)(cid:6)(cid:2)(cid:18)(cid:7)(cid:12)(cid:16)(cid:5) (cid:1)(cid:2)(cid:3)(cid:2) (cid:1)(cid:4)(cid:3)(cid:2)
[Figure]
These are replications of Fig. 19a and 19b.
[Table]
G V ISUALIZATIONS
Fig.
20 visualizes our models’ reconstructions.
Comparing our small model and large model, the reconstructions are almost the same.
That demonstrates our mask decay does not result in extra artifacts.
Compared with VTM, ours enjoy fast encoding and decoding.