Multi-Sample Training for Neural Image Compression
Abstract
This paper considers the problem of lossy neural image compression (NIC).
Current state-of-the-art (sota) methods adopt uniform posterior to approximate quantization noise, and single-sample pathwise estimator to approximate the gradient of evidence lower bound (ELBO).
In this paper, we propose to train NIC with multiple-sample importance weighted autoencoder (IWAE) target, which is tighter than ELBO and converges to log likelihood as sample size increases.
First, we identify that the uniform posterior of NIC has special properties, which affect the variance and bias of pathwise and score function estimators of the IWAE target.
Moreover, we provide insights on a commonly adopted trick in NIC from gradient variance perspective.
Based on those analysis, we further propose multiple-sample
NIC (MS-NIC), an enhanced IWAE target for NIC.
Experimental results demonstrate that it improves sota NIC methods.
Our MS-NIC is plug-and-play, and can be easily extended to other neural compression tasks.
Introduction
Latent variable-based lossy neural image compression (NIC) has witnessed significant success.
The majority of NIC follows the framework proposed by Ballé et al.
[2017]: For encoding, the original image x is transformed into y by the encoder.
Then y is scalar-quantized into integer ¯ y , estimated with an entropy model p ( ¯ y ) and coded.
For decoding, ¯ y is transformed back by the decoder to obtain reconstructed ¯ x .
The optimization target of NIC is R-D cost: R + λD .
R denotes the bitrate of ¯ y , D denotes the distortion between x and ¯ x , and λ denotes the hyper-parameter controlling their trade-off.
During training, the quantization ¯ y =
⌊ y ⌉ is relaxed with ˜ y
= y + ϵ to simulate the quantization noise.
And ϵ is fully factorized uniform noise ϵ ∼ p
( ϵ ) =
(cid:81) U ( − 12 , + 12 ) .
Ballé et al.
[2017] further recognises that such training framework is closely related to variational inference.
Indeed, the above process can be formulated as a graphic model x ← ˜ y .
During encoding, x is transformed into variational parameter y by inference model (encoder), and ˜ y is sampled from variational posterior q ( ˜ y | x ) , which is a unit unifrom distribution centered in y .
The prior likelihood p ( ˜ y ) is computed, and ˜ y is transformed back by the generative model (decoder) to compute the likelihood p ( x | ˜ y ) .
Under such formulation, the prior is connected to the bitrate, the likelihood is connected to the distortion, and the posterior likelihood is connected to the bits-back bitrate (See
Appendix. 2.3), which is 0 in NIC.
Finally, the evidence lower bound (ELBO) is the negative R + λD target (Eq. 1).
Denote the transform function ˜
y
( ϵ ; ϕ ) =
y + ϵ , and sampling ˜
y ∼ q ( ˜ y
| x ) is equivalent to transforming ϵ through ˜ y
( ϵ ; ϕ ) .
Then the gradient of ELBO is estimated via pathwise estimator with single-sample Monte Carlo (Eq. 2).
This is the same as SGVB-1 [Kingma and Welling, 2013].
[Equation]
Ballé et al.
[2018] further extends this framework into a two-level hierarchical structure, with graphic model x ← ˜ y ← ˜ z .
The variational posterior is fully factorized uniform distribution
U ( y − 12 , y + 12 ) U ( z − 12 , z + 12 )
To simulate the quantization noise.
And y , z denote outputs of their inference networks.
[Equation]
The majority of later NIC follows this hierarchical latent framework [Minnen et al., 2018, Cheng et al., 2020].
Some focus on more expressive network architectures [Zhu et al., 2021, Xie et al., 2021], some stress better context models [Minnen and Singh, 2020,
He et al., 2021, Guo et al., 2021a], and some emphasize semi-amortization inference [Yang et al., 2020].
However, there is little research on multiple-sample methods, or other techniques for a tighter ELBO.
On the other hand, IWAE [Burda et al., 2016] has been successful in density estimation.
Specifically,
IWAE considers a multiple-sample lowerbound L k (Eq. 4), which is tighter than its single-sample counterpart.
The benefit of such bound is that the implicit distribution defined by IWAE approaches true posterior as k increases [Cremer et al., 2017].
This suggests that its variational posterior is less likely to collapse to a single mode of true posterior, and the learned representation is richer.
The gradient of L k is computed via pathwise estimator.
Denote the exponential ELBO sample as w i , its reparameterization as w ( ϵ i ; ϕ ) , and its weight ˜ w i =
w i (cid:80) w j .
Then ∇ ϕ L k has the form of importance weighted sum
(Eq. 5).
[Equation]
In this paper, we consider the problem of training NIC with multiple-sample IWAE target (Eq. 4), which allows us to learn a richer latent space.
First, we recognise that NIC’s factorized uniform variational posterior has impacts on variance and bias properties of gradient estimators.
Specifically, we find NIC’s pathwise gradient estimator equivalent to an improved STL estimator [Roeder et al., 2017], which is unbiased even for the IWAE target.
However, NIC’s IWAE-DReG estimator [Tucker et al., 2018] has extra bias, which causes performance decay.
Moreover, we provide insights on a commonly adopted but little explained trick of training NIC from gradient variance perspective.
Based on those analysis and observations, we further propose MS-NIC, a novel improvement of multiplesample IWAE target for NIC.
Experimental results show that it improves sota NIC methods [Ballé et al., 2018, Cheng et al., 2020] and learns richer latent representation.
Our method is plug-and-play, and can be extended into neural video compression.
To wrap up, our contributions are as follows:
• We provide insights on the impact of the uniform variational posterior upon gradient estimators, bits-back coding and a commonly adopted but little discussed trick of NIC training from gradient variance perspective.
• We propose multiple-sample neural image compression (MS-NIC).
It is a novel enhancement of hierarchical IWAE [Burda et al., 2016] for neural image compression.
To the best of our knowledge, we are the first to consider a tighter ELBO for training neural image compression.
• We demonstrate the efficiency of MS-NIC through experimental results on sota NIC methods.
Our method is plug-and-play for neural image compression and can be easily applied to neural video compression.
Gradient Estimation for Neural Image Compression
The common NIC framework (Eq. 1, Eq 3) adopts fully factorized uniform distribution q
( ˜ y , ˜
z | x ) = (cid:81)
U ( y i − 12 , y i + 12 ) (cid:81) U ( z j − 12 , z j + 12 ) to simulate the quantization noise.
Such formulation has the following special properties:
• Property I: q ( ˜ z | ˜ y ) and q ( ˜ y
| x ) ’s support depends on the parameter.
• Property II: log q ( ˜ z | ˜ y )
=
log q
( ˜ y
| x ) = 0 on their support.
The impacts of these two properties are frequently neglected in previous works, which does not influence the results for single-sample pathwise gradient estimators (a.k.a. reparameterization trick in
Kingma and Welling [2013]).
In this section, we discuss the impacts of these two properties upon the variance and biasness of gradient estimators.
Our analysis is based on single level latent (Eq. 1) instead of hierarchical latent (Eq. 3) to simplify notations.
2.1 Impact on Pathwise Gradient Estimators
First, let’s consider the single-sample case.
We can expand the pathwise gradient of ELBO in Eq.into Eq. 6.
As indicated in the equation, ϕ contributes to L in two ways.
The first way is through the reparametrized ˜ y
( ϵ ; ϕ ) (pathwise term), and the other way is through the parameter of log q ( ˜ y | x ) (parameter score term).
Generally, the parameter score term has higher variance than the pathwise term.
The STL [Roeder et al., 2017] reduces the gradient by dropping the score.
It is unbiased since the dropped term’s expectation E q ( ˜ y | x ) [ ∇ ϕ log q ϕ
( ˜ y
| x )] is 0 .
[Equation]
Now let’s consider the STL estimator of multiple-sample IWAE bound (Eq. 4).
As shown in Tucker et al.
[2018], the STL estimation of IWAE bound gradient is biased.
To reveal the reason, consider expanding the gradient Eq.
5 into partial derivatives as we expand Eq.
2 into Eq.
Unlike singlesample case, the dropped parameter score term E p ( ϵ 1: k )
[ (cid:80) ˜
w i ( −∇ ϕ log q ϕ ( ˜ y | x ))] is no longer
due to the importance weight ˜ w i .
This means that STL loses its unbiasness in general IWAE cases.
Regarding NIC, however, the direct pathwise gradient for IWAE bound is automatically an unbiased
STL estimator.
Property II means that variational posterior has constant entropy, which further means that the parameter score gradient is 0 .
So, NIC’s pathwise gradient of IWAE bound is equvailent to an extended, unbiased STL estimator.
2.2 Impact on Score Function Gradient Estimators
In previous section, we show the bless of NIC’s special properties on pathwise gradient estimators.
In this section, we show their curse on score function gradient estimators.
Sepcifically, Property I implies that q ( ˜ z | ˜ y ) and q ( ˜ y | ˜ x ) are not absolute continuous, and hence the score function gradient estimators of those distributions are biased.
For example, consider a univariate random variable x ∼ p θ ( x )
=
U ( θ − 12 , θ + 12 ) .
Our task is to estimate the gradient of a differentiable function f ( x ) .
And consider the θ -independent random variable ϵ ∼ p ( ϵ )
=
U ( − 12 , + 12 ) , the transform x ( ϵ ; θ ) = θ + ϵ .
Under such conditions, the Monte
Carlo estimated pathwise gradient and score function gradient are:
[Equation]
score
[Equation]
Eq.
7 does not equal to Eq. 8, and Eq.8 is wrong.
The score function gradient is only unbiased when the distribution satisfies the absolute continuity condition of [Mohamed et al., 2020].
This reflects that under the formulation of NIC, the equivalence between the score function gradient (a.k.a.
REINFORCE
[Williams, 1992]) and pathwise gradient (a.k.a reparameterization trick in [Kingma and Welling, 2013]) no longer holds.
[Table]
Such equivalence is the cornerstone of many gradient estimators, and IWAE-DReG [Tucker et al., 2018] is one of them.
IWAE-DReG is a popular gradient estimator for IWAE target (Eq. 4) as it resolves the vanish of inference network gradient SNR (signal to noise ratio).
However, the correctness of IWAE-DReG depends on the equivalence between the score function gradient and pathwise gradient, which does not hold for NIC.
Specifically, IWAE-DReG expand the total derivative of IWAE target as Eq. 9 and perform another round of reparameterization on the score term as Eq.
10 to further reduce the gradient variance.
However, Eq.
10 requires the equivalence of pathwise gradient and score function gradient.
[Table]
[Equation]
[Table]
[Equation]
As we show empirically in Tab. 1, blindly adopting IWAE-DReG estimator for multiple-sample NIC brings evident performance decay.
Other than IWAE-DReG, many other graident estimators such as
NVIL [Mnih and Gregor, 2014], VIMCO [Mnih and Rezende, 2016] and GDReG [Bauer and Mnih, 2021] do not apply to NIC.
They either bring some extra bias or are totally wrong.
2.3 Impact on Bits-Back Coding
It is well known that the ELBO L is the minus overall bitrate for bits-back coding in compression [Hinton and Van Camp, 1993, Hinton et al., 1995, Chen et al., 2017], and the entropy of variational posterior is exactly the bits-back rate itself.
For this reason, earlier works [Townsend et al., 2018, Yang et al., 2020] point out
that [Ballé et al., 2018, Minnen et al., 2018] waste bits for not using bits-back coding on z .
However, during training the differential entropy [Cover, 1999] E q ( ˜ z | ˜ y )
[log q ( ˜ z | ˜ y )] is constant.
And this means that this term does not have impact on the optimization procedure.
And due to the deterministic inference, the log q (¯ z | ¯ y ) is 0 , which means that the bitrate saved by bits-back coding is 0 .
In this sense, [Ballé et al., 2018, Minnen et al., 2018] is also optimal in bits-back coding perspective, although no actual bits-back coding is performed.
In fact, there is no space for bits-back coding so long as encoder is deterministic.
Since we can view deterministic encoder as a posterior distribution with mass 1 on a single point.
And then the posterior’s entropy is always 0 .
2.4 The direct-y Trick in Training NIC
In NIC, we feed deterministic parameter y into z inference model instead of noisy samples ˜ y .
This implies that ˜ z is sampled from q ( ˜ z | y ) instead of q ( ˜ z | ˜ y ) .
This trick is initially adopted in Ballé et al.
[2018] and followed by most of the subsequent works.
However, it is little discussed.
In this paper, we refer it to direct-y trick.
Yang et al.
[2020] observes that feeding ˜ y instead of y causes severe performance decay.
We confirm this result in Tab. 2.
Thus, direct-y trick is essential to train hierarchical NIC.
[Table]
[Table]
One explanation is to view q ( ˜ z | y ) as q ( ˜ z | x ) , and q ( ˜ y , ˜
z | ˜ x ) factorized as q ( ˜ y | x ) q ( ˜ z | x ) (See Fig.
(a)-(c)).
A similar trick of feeding mean parameter can be traced back to the Helmholtz machine [Dayan et al., 1995].
However, this provides a rationale why direct-y is fine to be adopted but does not explain why samping ˜ z from q ( ˜ z | ˜ y ) fails.
We provide an alternative explanation from the gradient variance perspective.
Specifically, q ( ˜ z | ˜ y ) has two stochastic arguments that could cause high variance in the gradient of z inference model, and make its convergence difficult.
To verify this, we follow
Rainforth et al.
[2018] to compare the gradient SNR, which is the absolute value of the empirical mean divided by standard deviation.
We trace the gradient SNR during different training stages as model converges (See Sec.
5.1 for detailed setups).
As demonstrated in Tab. 3, the gradient SNR of z inference model of standard 2-level VAE (without direct y ) is indeed significantly lower than Ballé et al.
[2018] (with direct y ) during all 3 stage of training.
This result reveals that the z inference model is more difficult to train without direct-y .
And such difficulty could be the source of the failure of NIC without direct-y trick.
Multiple-sample Neural Image Compression
In this section, we consider the multiple-sample approach based on the 2-level hierarchical framework by Ballé et al.
[2018], which is the de facto NIC architecture adopted by many sota methods.
To simplify notations, log q ( ˜ z | ˜ y ) and log q ( ˜ y | x ) in ELBO are omitted as they are 0 .
First, let’s consider directly applying 2-level IWAE to NIC without direct-y trick (See Fig.
1 (d)).
Regarding a k sample
IWAE, we first compute parameter y of q ( ˜ y | x ) and sample ˜ y 1: k from it.
Then, we compute parameter z 1: k of q ( ˜ z 1: k | ˜ y 1: k ) and samples ˜
z 1: k from it.
Afterward, ˜ y 1: k and ˜ z 1: k are fed into the generative model and compute w 1: k .
Finally, we follow Eq 5 to compute the gradient and update parameters.
In fact, this is the standard 2-level IWAE in the original IWAE paper.
However, the vanilla 2-level IWAE becomes a problem for NIC with direct-y trick.
Concerning a k sample IWAE, we sample ˜ y 1: k from q ( ˜ y | x ) .
Due to the direct-y trick, we feed y instead of ˜ y 1: k into z inference network, and our q ( ˜ z | y ) has only one parameter z other than k parameter z 1: k .
If we follow the 2-level IWAE approach, only one sample ˜ z is obtained, and w 1: k can not be computed.
One method is to limit the multiple-sample part to ˜ y related term only and optimize other parts via single-sample SGVB-1, which produces our MS-NIC-MIX (See Fig 1 (e)).
Another method is to sample another l samples of ˜ z
j from q ( ˜ z | y ) and nest it with MS-NIC-MIX, which generates our
MS-NIC-DMS (See Fig 1 (f)).
3.1 MS-NIC-MIX: Multiple-sample NIC with Mixture
One way to optimize multiple-sample IWAE target of NIC with direct-y trick is to sample ˜
y k times to obtain ˜
y 1: k and ˜
z only 1 time.
Then we perform k sample log mean of p ( x | ˜
y i ) p
( ˜ y i | ˜ z ) to obtain a multiple-sample estimated log p ( x | ˜ z ) , add it with single-sample log p ( ˜ z ) .
This brings a L MIXk with the form of a mixture of 1-level VAE and 1-level IWAE ELBO:
[Equation]
Moreover, L MIXk is a reasonably preferable target over ELBO as it satisfies the following properties (See Appendix.
A.2 for proof):
1. log p ( x ) ≥ L MIXk 2.
L MIXk ≥ L MIXm for k ≥ m
Although L MIX k does not converge to true log p ( x ) as k grows, it is still a lower bound of log p ( x ) and tighter than ELBO (as L MIX= ELBO).
Its gradient can be computed via pathwise estimator.
Denote the per-sample integrand p ( x | ˜ y i ) p
( ˜ y i | ˜ z ) as w MIXi , and its relative weight as ˜ w MIXi , then the gradient ∇ ϕ L MIXk can be estimated as Eq. 13.
[Equation]
Another way to understand MS-NIC-MIX is to view the y inference/generative model as a single level
IWAE, and the z inference/generative model as a large prior of ˜ y which is optimized via SGVB-1.
This perspective is often taken by works in NIC context model [Minnen et al., 2018, He et al., 2021], as the context model of NIC is often limited to ˜ y .
3.2 MS-NIC-DMS: Multiple-sample NIC with Double Multiple Sampling
An intuitive improvement over MS-NIC-MIX is to add another round of multiple-sample over ˜ z .
Specifically, we sample ˜
z l times, nest it with L MIXk to obtain L DMSk,l :
[Equation]
l
[Equation]
And we name it MS-NIC-DMS as it adopts multiple sampling twice.
Moreover, L DMSk,l is a reasonably better target for optimizaion over ELBO and L MIXk , as it satisfies the following properties (See proof in Appendix.
A.2):
1. log p ( x ) ≥ L DMS k,l 2.
L DMSk,l ≥ L DMSm,n for k ≥ m, l ≥ n 3.
L DMSk,l ≥ L MIXk 4.
L DMSk,l → log p ( x ) as k, l → ∞ , under the assumption that log p ( x | ˜ y i ) p ( ˜ y
i | ˜
z j ) q ( ˜ y
i | x )
and
log p ( x | ˜ z j ) p ( ˜ z j ) q ( ˜ z
j | x ) are bounded.
In other words, the target L DMSk,l is a lowerbound of log p ( x ) , converging to log p ( x ) as k, l → ∞ , tighter than L MIXk and tighter than ELBO (as L DMS 1 ,= ELBO).
However, its Monte Carlo estimation is biased due to the nested transformation and expectation.
Empirically, we find that directly adopting biased pathwise estimator works fine.
And its gradient can be estimated by pathwise estimator similar to original IWAE target (See Eq. 5).
[Equation]
Another interpretation of MS-NIC-DMS is to view it as a multiple level IWAE with repeated local samples.
The L DMSk,l
Monte Carlo pathwise estimator has the form of IWAE with k × l samples.
However, there are multiple repeated samples that contain the same ˜
y
i and ˜ z j .
For example, the samples w IWAE 1: of 2 level IWAE with sample size 6 look like Eq.
While the samples w DMS 1:2 , 1:3 of
MS-NIC-DMS with 2 × 3 samples look like Eq.
We can see that in IWAE, we have 6 pairs of independently sampled
˜ y and ˜ z , while in MS-NIC-DMS, we have 2 independent ˜ y and 3 independent ˜ z , they are paired to generate 6 samples in total.
Note that this is only applicable to NIC as ˜ y and
˜ z are conditionally independent given ˜ x
due to direct-y trick.
[Equation]
, p ( x | ˜ y 2 ) p ( ˜ y 2 | ˜ z 2 ) p ( ˜ z 2 ) , p ( x | ˜ y 3 ) p ( ˜ y 3 | ˜ z 3 ) p ( ˜ z 3 ) , p ( x | ˜ y 4 ) p ( ˜ y 4 | ˜ z 4 ) p ( ˜ z 4 ) , p ( x | ˜ y 5 ) p ( ˜ y 5 | ˜
z 5 ) p ( ˜ z 5 ) , p ( x | ˜ y 6 ) p ( ˜ y 6 | ˜
z 6 ) p ( ˜ z 6 ) } (16) w DMS 1:2 , 1:3 =
{ p ( x | ˜ y 1 ) p ( ˜ y 1 | ˜ z 1 ) p ( ˜ z 1 ) , p ( x | ˜ y 1 ) p ( ˜ y 1 | ˜ z 2 ) p ( ˜ z 2 ) , p ( x | ˜ y 1 ) p ( ˜ y 1 | ˜ z 3 ) p ( ˜ z 3 ) , p ( x | ˜ y 2 ) p ( ˜ y 2 | ˜ z 1 ) p ( ˜ z 1 ) , p ( x | ˜ y 2 ) p ( ˜ y 2 | ˜ z 2 ) p ( ˜ z 2 ) , p ( x | ˜ y 2 ) p ( ˜ y 2 | ˜ z 3 ) p ( ˜ z 3 ) } (17)
Related Work 4.1 Lossy Neural Image and Video Compression
Ballé et al.
[2017] and Ballé et al.
[2018] formulate lossy neural image compression as a variational inference problem, by interpreting the additive uniform noise (AUN) relaxed scalar quantization as a factorized uniform variational posterior.
After that, the majority of sota lossy neural image compression methods adopt this formulation [Minnen et al., 2018, Minnen and Singh, 2020, Cheng et al., 2020, Guo et al., 2021a, Gao et al., 2021, He et al., 2022].
And Yang et al.
[2020], Guo et al.
[2021b] also require a AUN trained NIC as base.
Moreover, the majority of neural video compression also adopts this formulation [Lu et al., 2019, 2020, Agustsson et al., 2020, Hu et al., 2021, Li et al., 2021], implying that MS-NIC can be extended to video compression without much pain.
Other approaches to train NIC include random rounding [Toderici et al., 2015, 2017] and straight through estimator (STE)
[Theis et al., 2017].
Another promising approach is the VQ-VAE [Van
Den Oord et al., 2017].
By the submission of this manuscript, one unarchived work [Zhu et al., 2022] has shown the potential of VQ-VAE in practical NIC.
Our MS-NIC does not apply to the approaches mentioned in this paragraph, as the formulation of variational posterior is different.
4.2 Tighter Lower Bound for VAE
IWAE [Burda et al., 2016] stirs up the discussion of adopting tighter lower bound for training VAEs.
However, at the first glance it is not straightforward why it might works.
Cremer et al.
[2018] decomposes the inference suboptimality of VAE into two parts: 1)
The limited expressiveness of interence model.
The gap between ELBO and log likelihood.
However, this gap refers to inference not training.
The original IWAE paper empirically shows that IWAE can learn a richer latent representation.
And Cremer et al.
[2017] shows that the IWAE target converges to ELBO under the expectation of true posterior.
And thus the posterior collapse is avoided.
From the information preference [Chen et al., 2017] perspective, VAE prefers to distribute information in generative distribution than autoencoding information in the latent.
This preference formulates another view of posterior collapse.
And it stems from the gap between ELBO and true log likelihood.
There are various approaches alleviating it, including soft free bits [Theis et al., 2017] and KL annealing [Serban et al., 2017].
In our opinion, IWAE also belongs to those methods, and it is asymptotically optimal.
However, we have not found many works comparing IWAE with those methods.
Moreover, those approaches are rarely adopted in NIC community.
Many follow-ups of IWAE stress gradient variance reduction [Roeder et al., 2017, Tucker et al., 2018, Rainforth et al., 2018], discrete latent [Mnih and Rezende, 2016] and debiasing IWAE target [Nowozin, 2018].
Although the idea of tighter low bound training has been applied to the field of neural joint source channel coding [Choi et al., 2018, Song et al., 2020], to the best of our knowledge, no work in NIC consider it yet.
4.3 Multi-Sample Inference for Neural Image Compression
Theis and Ho [2021] considers the similar topic of importance weighted NIC.
However, it does not consider training of NIC.
Instead, it focuses on achieving IWAE target with an entropy coding technique named softmin , just like BB-ANS
[Townsend et al., 2018] achieving ELBO.
It is alluring to apply softmin to MS-NIC, as it closes the multiple-sample training and inference gap.
However, it requires large number of samples (e.g. 4096 ) to achieve slight improvement for 64 × 64 images.
The potential sample size required for practical NIC is forbidding.
Moreover, we believe the stochastic lossy encoding scheme [Agustsson and Theis, 2020] that Theis and Ho [2021] is not yet ready to be applied (See Appendix.
A.7 for details).
Experimental Results 5.1 Experimental Settings
Following He et al.
[2022], we train all the models on the largest 8000 images of ImageNet [Deng et al., 2009], followed by a downsampling according to Ballé et al.
And we use Kodak [Kodak,
1993] for evaluation.
For the experiments based on Ballé et al.
[2018] (include Tab. 1, Tab. 2), we follows the setting of the original paper except for the selection of λ s, For the selection of λ s, we set λ ∈ { 0 . 0016 , 0 . 0032 , 0 .
08 } as suggested in Cheng et al.
And for the experiments based on Cheng et al.
[2020], we follows the setting of original paper.
More detailed experimental settings can be found in Appendix.
A.
[Table]
And when comparing the R-D performance of models trained on multiple λ s
[Table]
we use Bjontegaard metric (BD-Metric) and Bjontegaard bitrate (BD-BR)
[Bjontegaard, 2001], which is widely applied when comparing codecs.
More detailed experimental results can be found in Appendix.
A.5.
[Table] [Table]
IWAE [Burda et al., 2016] - - - -
MS-NIC-MIX
[Table] [Table]
MS-NIC-DMS
[Table]
[Table] [Table]
5.2 R-D Performance
We evaluate the performance of MS-NIC-MIX and MS-NIC-DMS based on sota NIC methods [Ballé et al., 2018, Cheng et al., 2020].
Empirically, we find that MS-NIC-MIX works best with sample size 8, and MS-NIC-DMS with sample size 16.
The experimental results on sample size selection can be found in Appendix.
A.3.
Without special mention, we set the sample size of MS-NIC-MIX to 8 and
MS-NIC-DMS to 16.
For Ballé et al.
[2018], MS-NIC-MIX saves around 4% of bitrate compared with single-sample baseline (See Tab. 4).
And MS-NIC-DMS saves around 5% of bitrate.
On the other hand, the original
IWAE suffers performance decay as it is not compatible with direct-y trick.
For Cheng et al.
[2020], we find that both MS-NIC-MIX and NS-NIC-DMS suppress baseline in PSNR.
However, it is not as evident as Ballé et al.
Moreover, the MS-SSIM is slightly lower than the baseline.
This is probably due to the auto-regressive context model.
Besides, the original IWAE without direct-y trick suffers from severe performance decay in both cases.
The BD metric of IWAE on Cheng et al.
[2020] can not be computed as its R-D is not monotonous increasing, we refer interested readers to
Appendix.
A.5 for details.
5.3 Latent Space Representation of MS-NIC
To better understand the latent learned by MS-NIC, we evaluate the variance and coefficient of variation (Cov) of per-dimension latent distribution mean parameter y
( i ) ,
z ( i ) , with regard to input distribution p ( x ) .
As we are also interested in the discrete representation, we provide statistics of rounded mean ¯ y ( i ) , ¯ z ( i ) .
These metrics show how much do latents vary when input changes, and a large variation in latents means that there are useful information encoded.
A really small variation indicates that the latent is "dead" in that dimension.
As shown in Tab. 10 of Appendix.
A.6, the latent of multiple-sample approaches has higher variance than those of single-sample approach.
Moreover, the Cov ( y ) of multiple-sample approaches is around 4 − 5 times higher than single-sample approach.
Although the Cov ( z ) of multiple-sample approaches is around 2 times lower, the main contributor of image reconstruction is y , and z only serves to predict y ’s distribution.
Similar trend can be concluded from quantized latents ¯ y , ¯ z .
From the variance and Cov perspective, the latent learned by MS-NIC is richer than single-sample approach.
It is also noteworthy that although the variance and Cov of y , ¯ y of MS-NIC is significantly higher than single-sample approach, the bpp only varies slightly.
[Table]
Limitation & Discussion
A major limitation of our method is that the improvement in R-D performance is marginal, especially when based on Cheng et al.
Moreover, evaluations on more recent sota methods are also helpful to strengthen the claims of this paper.
In general, we think that the performance improvement of our approach is bounded by how severe the posterior collapse is in neural image compression.
We measure the variance in latent dimension according to data in Fig.
A.6.
And from that figure it might be observed that the major divergence of IWAE and VAE happens when the variance is very small.
And for the area where variance is reasonably large, the gain of IWAE is not that large.
This probably indicates that the posterior collapse in neural image compression is only alleviated to a limited extend.
See more discussion in why the result on Cheng et al.
[2020] is negative in Appendix.
A.
Conclusion
In this paper we propose MS-NIC, a multiple-sample importance weighted target for training NIC.
It improves sota NIC methods and learns richer latent representation.
A known limitation is that its R-D performance improvement is limited when applied to models with spatial context models (e.g. Cheng et al.
Despite the somewhat negative result, this paper provides insights to the training of
NIC models from VAE perspective.
Further work could consider improving the performance and extend it into neural video compression.
Acknowledgments and Disclosure of Funding
This work is supported by SenseTime Research.
The content is solely the responsibility of the authors and does not necessarily represent the official views of SenseTime Research.
Checklist
For all authors... (a) Do the main claims made in the abstract and introduction accurately reflect the paper’s contributions and scope?
[Yes] (b) Did you describe the limitations of your work?
[Yes] (c) Did you discuss any potential negative societal impacts of your work?
[Yes] See
Appendix
(d) Have you read the ethics review guidelines and ensured that your paper conforms to them?
[Yes] 2.
If you are including theoretical results... (a) Did you state the full set of assumptions of all theoretical results?
[Yes] (b) Did you include complete proofs of all theoretical results?
[Yes] In Appendix 3.
If you ran experiments... (a) Did you include the code, data, and instructions needed to reproduce the main experimental results (either in the supplemental material or as a URL)?
[No] (b) Did you specify all the training details (e.g., data splits, hyperparameters, how they were chosen)?
[Yes] (c) Did you report error bars (e.g., with respect to the random seed after running experiments multiple times)?
[No] It is really expensive to train a NIC model.
The error bar is rarely reported out of toy size model and dataset.
(d) Did you include the total amount of compute and the type of resources used (e.g., type of GPUs, internal cluster, or cloud provider)?
[Yes] 4.
If you are using existing assets (e.g., code, data, models) or curating/releasing new assets... (a) If your work uses existing assets, did you cite the creators?
[Yes] (b) Did you mention the license of the assets?
[N/A] (c) Did you include any new assets either in the supplemental material or as a URL?
[N/A] (d) Did you discuss whether
and how consent was obtained from people whose data you’re using/curating?
[N/A] (e) Did you discuss whether the data you are using/curating contains personally identifiable information or offensive content?
[N/A] 5.
If you used crowdsourcing or conducted research with human subjects... (a) Did you include the full text of instructions given to participants and screenshots, if applicable?
[N/A] (b) Did you describe any potential participant risks, with links to Institutional Review Board (IRB) approvals, if applicable?
[N/A] (c)
Did you
include the estimated hourly wage paid to participants and the total amount spent on participant compensation?
[N/A]
[-In fact, the real distribution of quantization noise is pretty much a highly concentrated distribution around 
-]