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] Fig.
A.1 shows the generative models of Fig. 1 (d)-(f). Note that we have only one unique sample of x inside the n plate, but it is repeated k times for IWAE, MS-NIC-MIX and k × l times for
MS-NIC-DMS.
Similarly, k samples of ˜ y is repeated l times and l samples of ˜ z is repeated k times for MS-NIC-DMS. A.2 Proof on the Properties of MS-NIC-MIX and MS-NIC-DMS In this section, we add the q ( ˜ z | ˜ y ) , q ( ˜ y | x ) back to equations for clarity of the proof.
This makes the notations slightly different from Eq.
12 and Eq.
Note that divide by q ( ˜ z | ˜ y ) , q ( ˜ y
| x ) does not effect the value of equation, and add log q ( ˜ z | ˜ y ) , log q ( ˜ y | x ) does not effect the value of equation.
For MS-NIC-MIX to be a reasonably better approach to apply over [Ballé et al., 2018], we show that
L MIXk satisfies following properties:
1. log p ( x ) ≥ L MIXk 2. L MIX k ≥ L MIXm for k ≥ m We can show 1. log p ( x ) ≥ L MIXk by applying Jensen’s inequality twice:
[Equation]
We can show 2.
L MIXk ≥ L MIXm for k ≥ m by borrowing the Theorem 1 from IWAE paper:
[Equation]
Applying Eq.
19 to the internal part of L MIXk , when k ≥ m , we have:
[Equation]
For MS-NIC-DMS to be a reasonably better approach to apply over [Ballé et al., 2018] and MS-NICMIX, we show that L DMSk,l statisfies following properties:
1. log p ( x ) ≥ L DMSk,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.
Similar to MS-NIC-DMS, we can show 1.
log p ( x ) ≥ L DMSk,l by applying Jensen’s inequality twice:
L DMSk,l
[Equation]
Also similar to MS-NIC-MIX
[Equation]
we can borrow conclusion from IWAE (Eq. 19) and apply it twice to show 2.
L DMSk,l ≥ L DMSm,n for k ≥ m, l ≥ n :
[Equation]
x )] ≥ E q ϕ ( ˜ z 1: n | x ) [log 1 n n (cid:80) j exp (
[Equation]
ϕ (
[Equation]
: m
[Equation]
x ) [log 1 m m (cid:80) i p ( x | ˜ y
[Equation]
p (
[Equation]
y
[Equation]
| ˜
[Equation]
j )
[Equation]
[Equation]
y i | x ) | ˜ z j ]) p ( ˜ z j ) /q
( ˜ z j | x )]
=
L DMSm,n (22)
With 2.
L DMSk,l ≥ L DMSm,n for k ≥ m, l ≥ n holds, we can show 3.
L DMSk,l ≥ L MIXk immediately as
L DMSk,l ≥ L DMSk,= L MIXk .
To show 4.
L DMSk,l → log p ( x ) as k, l → ∞ , we first define intermediate variables W k , ˜ M k,l , M k,l :
[Equation]
Under the assumption that log p ( x | ˜ y i ) p ( ˜ y i | ˜
z j ) /q
( ˜ y i | x ) is bounded, from the strong law of large number, we have W k a.s. −−→ p ( x | ˜ z
j ) (Eq. 24).
Then we have E [log W k | ˜
z j ] → log p ( x | ˜ z j ) .
[Equation]
Moreover, as E [log W k | ˜ z j ] → log p ( x | ˜ z j ) , we have M k,l → ˜ M k,l .
This means that ∀ ϵ > 0 , ∃ k, l, s.t. | M k,l − ˜ M k,l | < ϵ .
And thus we have | E
[ M k,l ] − E [ ˜ M k,l ] | ≤ E [ | M k,l − ˜ M k,l | ] <
ϵ → 0 .
Then we have | E [ M k,l ] − p ( x ) | ≤ | E
[ M k,l ] − E [ ˜ M k,l ] | + | E [ ˜ M k,l ] − p ( x ) | → 0 , and thus
E [ M k,l ] → p ( x ) .
Finally we have E [log M k,l ] =
L DMSk,l → log p ( x ) . A.3 Effects of Sample Size When comparing the R-D performance of models trained with a single λ , we use R-D cost as our metric. The R-D cost is simply computed as bpp + λ MSE, where bpp is a short of bits-per-pixel, and
MSE is a short of mean square error.
The lower the R-D cost is, the better the R-D performance is.
Another way to interpret R-D cost is to view it as the ELBO with constant offset.
Then the λ MSE is connected to the log likelihood of a Gaussian distribution whose mean is the output of decoder and sigma is determined by λ .
Note that R-D cost is only comparable when λ is the same.
Tab. 7 shows the effect of sample size to MS-NIC.
Moreover, we compare the naïve increase of batch size versus multiple importance weighted samples.
As shown by the table, increasing the batch size × 3 − 16 only slightly affects the R-D cost (from 1 .
017 to 1 . 013 ).
However, the MS-NIC-MIX can achieve R-D cost of 0 .
9988 with sample size 8, and MS-NIC-DMS can achieve 0 .
9954 with sample size 16.
This means that MS-NIC is effective over the baseline and vanilla batch size increases.
It is also noteworthy that we have not observed inference model training failure as sample size increase.
While MS-NIC also suffers from gradient SNR vanishing problem, a sample size of 16 is probably not large enough to make it evident.
Limited by computational power, we can not raise sample size by several magnitudes as [Rainforth et al., 2018] does with small model.
A.4 Detailed Experimental Settings
All the experiments are conducted on a computer with Intel(R)
Xeon(R) CPU E5-2620 v4 @ 2.10GHz and 8 × Nvidia(R)
TitanXp.
All the training scripts are implemented with Pytorch 1.7 and CUDA 9.0. For experiments with single-sample, we adopt Adam optimizer with β 1 = 0 . 90 , β 2 = 0 . 95 , lr = 1 e − 4 .
For experiments with multiple-sample/big batch, we scale lr linearly with sample size.
All the models are trained for 2000 epochs with the settings in Sec. 5.1.
For first 200 epochs, we adopt cosine annealing [Loshchilov and Hutter, 2016] to schedule learning rate.
It takes around 1 − 2 days to train models based on Ballé et al.
[2018], and 3 − 5 days to train models on Cheng et al.
Note that our multiple-sample approaches’ training time does not scale linearly with sample size, as we perform sampling on posterior, and the variational encoder only computes parameter of posterior parameters once.
Further, we provide the pytorch style sudo code for implementation guidance of
MS-NIC-MIX and MS-NIC-DMS.
import torch from torch .
nn import f u n c t i o n a l as F def IWAELoss( minus_elbo ) : ’ ’ ’
args−−−− minus_elbo : tensor , [b , k ] , which i s R + \ lambda D
[Table] [Table]
A.5 Detailed Experimental Results
In this section we present more detailed experimental results in Tab. 8 and Tab. [Table] Note that without direct-y trick [Table]
the IWAE for Cheng et al.
[2020] totally fails and we can not produce a valid BD metric from it. [Table] IWAE [Burda et al., 2016] 0. MS-NIC-MIX 0.
MS-NIC-DMS 0.
[Table]
A.6 Distribution of Latent Variance
We show the histogram of latent variance in log space in Fig. A.6. From the histogram we can observe that for latent y , the variance distribution of two MS-NIC approaches is similar and single-sample [Table]
approach is quite different.
MS-NIC has more latent dimensions that have high variance (the right mode), and less with low variance (the left mode).
Moreover, the low variance mode of MS-NIC has less variance than single-sample approach, which indicates that MS-NIC does a better job in separating active and inactive latent dimensions.
Similarly, the low variance mode of ˜ z in MS-NIC approaches is lower than single-sample approach.
[Table] [Figure]
f r e q u e n c y
[Figure]
f r e q u e n c y [Balle et al, 2018] MS-NIC-MIX
MS-NIC-DMS
A.7 Tighter ELBO for Inference Time A.7.1 Inference time ELBO and Softmin Coding [Theis and Ho, 2021]
The inference time tighter ELBO is another under-explored issue. In fact, the training time tighter ELBO and inference time ELBO is independent. We can train a model with tighter ELBO, infer with single sample ELBO.
Or we can also conduct multiple sample infer on a model trained with single sample.
The general idea is:
• The training time tighter ELBO benefits the performance in terms of avoiding posterior collapse. As we state and empirically verify in Sec. We adopt deterministic rounding during inference time, and there is no direct connection between the training time tighter ELBO and inference time R-D cost. However, we indeed end up with a richer latent space (Sec. 5.3), which means more active latent dimensions and less bitrate waste.
• The inference time tighter ELBO sounds really alluring for compression community.
However, there remains two pending issue to be resolved prior to the application of the inference time tighter ELBO: 1) How this inference time multiple-sample ELBO is related to R-D cost remains under-explored.
In other words, whether the entropy coding itself can achieve the R-D cost defined by multiple-sample ELBO is a question.
The inference time multiplesample ELBO only makes sense with stochastic encoder (you can not importance weight the same deterministic ELBO), whose impact on lossy compression remains dubious.
For the first pending issue, the softmin coding [Theis and Ho, 2021] is proposed to achieve multiple sample ELBO based on Universal Quantization (UQ)
[Agustsson and Theis, 2020]. However, it is not a general method and is tied to UQ. Moreover, as we stated in Sec 4.3, its computational cost is forbiddingly high and its improvement is marginal. But those are not the real problem of softmin coding.
Instead, the real problem is the second pending issue: stochastic lossy encoder.
The softmin coding relies on UQ, and UQ relies on stochastic lossy encoder.
And the stochastic lossy encoder is exactly the second issue that we want to discuss.
A.7.2 Stochastic Lossy Encoder and Universal Quantization
It is known to lossless compression community that stochastic lossy encoder benefits compression performance [Ryder et al., 2022] with the aid of bits-back coding [Townsend et al., 2018].
While the bits-back coding is not applicable to lossy compression.
For lossy compression, currently we know that the stochastic encoder degrades R-D performance especially when distortion is measured in MSE [Theis and Agustsson, 2021].
In the original UQ paper, the performance decay of vanilla
UQ over deterministic rounding is obvious ( ≈ 1 db).
When we writing this paper, we also find the performance decay of UQ is quite high.
As shown in Tab. 11, the R-D cost of UQ is significantly higher than deterministic rounding.
This negative result makes softmin coding less promising than it seems as it only obtains a marginal gain over UQ.
In our humble opinion, this performance decay
[Table] of UQ
is partially brought by stochastic encoder itself.
For lossless compression, the deterministic encoder and stochastic encoder are just two types of bit allocation preference:
• The deterministic encoder allocate less bitrate to log p ( ˜ y ) , more to log p ( x | ˜ y ) and 0 to log q ( ˜ y | x ) .
• Given the same quantization step-size, from the differential entropy’s perspective, the stochastic encoder allocate more bitrate to log p ( ˜ y ) , less to log p ( x | ˜ y ) and minus bitrate to log q ( ˜ y | x )
Therefore, for lossless compression, it is reasonable that the bitrate increase to log p ( ˜ y ) and log p ( x | ˜ y ) can be offset by bits-back coding bitrate log q ( ˜ y | x ) .
While for lossy compression, there is no obvious way to bits-back log q ( ˜ y | x ) (as we can not reconstruct q ( ˜ y | x ) without x ).
If the bitrate increase in log p ( ˜ y ) and log p ( x | ˜ y ) , the R-D cost just increases for lossy compression.
Prior to other entropy coding bitrate that is able to achieve R-D cost equals to minus ELBO with E q [log q ] ̸ = 0 becomes mature (such as relative entropy coding [Flamich et al., 2020]), we have no way to implement a stochastic lossy encoder with reasonable R-D performance.
By now, we have no good way to achieve tighter ELBO during inference time.
{+Moreover, whether the quantization error is uniform distribution remains a question. And we think that is another reason why UQ does not work well. In fact, the real distribution of quantization noise is pretty much a highly concentrated distribution around 
+}
(See Fig. A.7.3).
And it is quite far away from uniform distribution, which violates the assumption of Ballé et al.
We also find that this concentrated distribution is caused by that most of latent dimension is quite close to 0 .
The evidence is, if we remove the latent dimension y i ∈ [ − 0 .
5] , then the quantization noise looks similar to a uniform distribution (See Fig. A.7.3).
So, if we apply direct rounding, they are kept as 0 and the latent is sparse.
However, adding uniform noise to it loses this sparsity, which result in bitrate increase.
And from Tab. 11, we can wee that the UQ reduce MSE while increase the bitrate.
From total R-D cost perspective, the deterministic rounding outperforms UQ.
As a matter of fact, the assumption of UQ that resolving training-testing distribution mismatch improves R-D performance does not hold well.
To wrap up, we find that there is some pending issues to be resolved prior to the practical solution of tighter ELBO for inference time.
A.8 The Effect on Training Time
The MS-NIC-MIX and MS-NIC-DMS should be more time efficient than simply increase batchsize.
For MS-NIC-MIX with k samples, the ˜ y encoder q ( ˜ y | x ) is inferred with only 1 sample, and the ˜ z encoder, decoder and entropy model q ( ˜ z | x ) , p ( ˜ y | ˜ z ) , p ( ˜ z ) is inferred with only 1 sample.
And only the ˜ y decoder p ( x | ˜ y ) is inferred k times.
This sample efficiency makes the training time grows slowly with k .
In our experiment, the MS-NIC-MIX with 8 samples only increases the training time by × 1 . 5 , the MS-NIC-MIX with 16 samples only increases the training time by × 3 .
The
MS-NIC-DMS is slightly slower, as the z entropy model and decoder p ( ˜ z ) , p ( ˜ y | ˜ z ) also requires k times inference.
However, it is still much more efficient than batchsize × k as all the encoders q ( ˜ y | x ) , q ( ˜ z | x ) requires only 1 inference.
In fact, sampling from posterior is much cheaper than inferring the posterior parameters.
Similar spirit has also been adopted in improving the efficiency of sampling from Gumbel-Softmax relaxed posterior
[Paulus et al., 2020].
The trade-off between batchsize and sample number is a more subtle issue.
As stated in Rainforth et al.
[2018], the gradient SNR of encoder (inference model) scales with Θ( M/K ) , and the gradient
SNR of decoder (generative model) scales with Θ( M/K ) , where M is the batchsize and K is the sample size.
Another assumption required prior to further discussion is that the suboptimality of VAE mainly comes from inference model [Cremer et al., 2018], which means that the encoder is harder to train than the decoder.
This means that an infinitely large K ruins the convergence of encoder, and solemnly increasing sample number frustrates training.
In practice the overall performance is determined by both inference suboptimality and ELBO-likelihood gap.
In a word, we believe there is no general answer for all problem.
But a reasonable balance between sample size and batchsize is the golden rule to maximize performance (as T ( M ) and T ( K ) grow linearly with batchsize/sample size).
And the obvious case is that neither setting batchsize to M = 1 and give all resources to K , nor setting sample size K = 1 and give all resources to M is optimal.
A.9 More Limitation and Discussion
The cause of negative results on MS-SSIM of Cheng et al.
[2020] is more complicated.
One possible explanation is that the gradient property of Cheng et al.
[2020] is not as good as Ballé et al.
As a reference, the training of [Burda et al., 2016] totally fails on Cheng et al.
[2020] and produces garbage R-D results (See Tab. 9).
This bad gradient property might account for the bad results of
MS-SSIM on Cheng et al.
[2020], as the gradient of IWAE and MS-NIC is certainly trickier than the gradient of single sample approaches.
As evidence, when we are studying the stability of the network in Cheng et al.
[2020], we find that without limitation of entropy model (imagine setting λ to ∞ ) and quantization, Cheng et al.
[2020] produces PSNR of 43 .
27 db, while Ballé et al.
[2018] produces PSNR of 48 .
54 db.
This means that Cheng et al. [2020] is not as good as Ballé et al. [2018] as an auto-encoder.
Moreover, when we finetune these pre-trained model into a lossy compression model, Cheng et al.
[2020] produces nan results while Ballé et al.
[2018] converges.
This result indicates that the backbone of Cheng et al.
[2020]’s gradient is probably more difficult to deal with than Ballé et al.
A.10 Broader Impact
Improving the R-D performance of NIC methods is valuable itself.
It is beneficial to reducing the carbon emission by reducing the resources required to transfer and store images.
And NIC has potential of saving network channel bandwidth and disk storage over traditional codecs.
Moreover, for traditional codecs, usually dedicated hardware accelerators are required for efficient decoding.
This codec-hardware bondage hinders the wide adaptation of new codecs.
Despite the sub-optimal
R-D performance of old codecs such as JPEG, H264, they are still prevalent due to broad hardware support.
While modern codecs such as H266 [Bross et al., 2021] can not be widely adopted due to limited hardware decoder deployment.
However, for NIC, the general purpose neural processors are able to fit to all codecs.
Thus the neural decoders have better hardware flexibility, and the cost to update neural decoder only involves software, which encourages the adoption of newer methods with better R-D performance.