L EARNED TRANSFORM COMPRESSION WITH OPTIMIZED ENTROPY ENCODING .
A BSTRACT
We consider the problem of learned transform compression where we learn both, the transform as well as the probability distribution over the discrete codes.
We utilize a soft relaxation of the quantization operation to allow for back-propagation of gradients and employ vector (rather than scalar) quantization of the latent codes.
Furthermore, we apply similar relaxation in the code probability assignments enabling direct optimization of the code entropy.
To the best of our knowledge, this approach is completely novel.
We conduct a set of proof-of concept experiments conﬁrming the potency of our approaches.
I NTRODUCTION
We consider the problem of compressing data x ∈ X sampled i.i.d.
according to some unknown probability measure (distribution) x ∼ µ x .
We take the standard transform coding (Sayood, 2012) approach where we ﬁrst transform the data x by a learned non-linear function, an encoder E θ : X → Z , into some latent representation z ∈ Z .
We then quantize the transformed data z
= E θ ( x ) using a quantization function Q E : Z → C parametrized by learned embeddings E (codebook/ dictionary) so that the discrete codes composed of indexes of the embedding vectors c = Q E ( z ) can be compressed by a lossless entropy encoding and transmitted.
The received and losslessly decoded integer codes are then used to index the embedding vectors and dequantized back to the latent space
Q E : C → (cid:98)
Z ⊂ Z introducing a distortion due to mapping the codes only to
the discrete subset (cid:98) Z ⊂ Z corresponding to the quantization embeddings.
The dequantized data ˆ z = Q E ( c ) are then decoded by a learned non-linear decoder D φ : (cid:98) Z → X to obtain the reconstructions ˆ x = D φ (ˆ z ) .
Our aim is to learn the transform (encoder/decoder) as well as the quantization so as to minimize the expected distortion E µ x d ( x , ˆ x ) 1 while, at the same time, minimizing the expected number of bits transmitted (the rate) when passing on the discrete codes E µ c l ( c ) , where l is the length of the bit-encoding.
The two competing objectives are controlled via a hyper-parameter λ
[Equation]
The optimal length of encoding a symbol c ∼ µ c is determined by Shannon’s self informationi c ( c ) =
− log p c ( c ) (Cover & Thomas, 2006), where p c is the discrete probability mass 3 of the distribution µ c .
Consequently, the expected optimal description length for the discrete code c can be bounded by its entropy H µ c ( c ) =
− E µ c log p c ( c ) as H µ c ( c ) ≤ E µ c l ( c ) ∗ < H µ c ( c ) + 1 4 .
To minimize the rate we therefore minimize the entropy of the discrete code H µ c ( c ) so that
[Equation]
respect counting measure
Q UANTIZATION
We employ the soft relaxation approach to quantization proposed in Agustsson et al.
(2017) simpliﬁed similarly to Mentzer et al.
However, instead of the scalar version Mentzer et al.
Habibian et al.
we use the vector formulation of the quantization as in van den Oord et al.
(2017) in which the k quantization centers are the learned embedding vectors { e ( j ) } kj =1 , e i ∈ R m , columns of the m × k embedding matrix E =
[ e (1) , . . .
, e ( k ) ] .
The quantizer Q E ﬁrst reshapes the transformed data
5 z ∈ Z ⊆ R d into a m × d/m matrix Z
=
[ z (1) , . . .
, z ( d/m ) ] , then ﬁnds for each column z ( i
) ∈
R m its nearest embedding and replaces it by the embedding vector index to output the d/m dimensional vector of discrete codes c = Q E ( z )
Q E
[Equation]
After transmission, the quantized latent representation ˆ Z
= [ˆ z (1) , . .
, ˆ z ( d/m ) ] is recovered from c by indexing to the shared codebook E and decoded ˆ x = D φ (ˆ z ) , ˆ z = ﬂatten( ˆ Z ) .
In practice, the quantized latent ˆ z can be used directly by the decoder at training in the forward pass without triggering the c indexing operation.
The ﬁnite quantization operation in equation 3 is non-differentiable.
To allow for the ﬂow of gradients back to E θ and Q E we use a differentiable soft relaxation for the backward pass
˜ z ( i ) =
k (cid:88)
[Equation]
where instead of the hard encoding ˆ¯ z ( i ) picking the single nearest embedding vector, the soft ˜ z ( i ) is a linear combination of the embeddings weighted by their (softmaxed) distances.
The distortion loss is thus formulated as d ( x , ˆ x ) =
d ( x , D φ [sg(ˆ z − ˜ z ) + ˜ z ]) , where sg is the stopgradient operator.
The hard/soft strategy is different from the approach of van den Oord et al.
(2017) where they use a form of straight-through gradient estimator and dedicated codebook and commitment terms in the loss to train the embeddings E .
This is also different from Williams et al.
(2020), where they use the relaxed formulation of equation 4 for both forward and backward passes in a fully stochastic quantization scheme aimed at preventing the mode-dropping effect of a deterministic approximate posterior in a hierarchical vector-quantized VAE.
M INIMIZING THE CODE CROSS - ENTROPY
Though the optimal lossless encoding is decided by the self-information of the code − log p c ( c ) , it cannot be used directly since p c is unknown.
Instead we replace the unknown p c by its estimated approximation q c , derive the code length ˆ l ( c ) from ˆ i c
= − log q c ( c ) , and therefore minimize the expected approximate self-information, the cross-entropy H µ c | q c ( c ) = − E µ c log q c ( c ) .
This, however, yields inefﬁciencies as H µ c | q c ( c ) ≥ H µ c ( c ) due to the decomposition
[Equation]
where D KL ≥ 0 is the Kullback-Leibler divergence between p c and q c which can be interpreted as the expected additional bits over the optimal rate H µ c ( c ) caused by using q c instead of the true p c .
In addition to E θ , D φ and Q E we shall now therefore train also a probability estimator P ψ : { c } ni → q c by minimizing the cross-entropy H µ c | q c ( c ) so that the estimated q c is as close as possible to the true p c , the D KL is small, and the above mentioned inefﬁciencies disappear.
As we cannot evaluate the cross-entropy over the unknown µ c , we instead learn P ψ by its empirical estimate over the sample data which is equivalent to minimizing the negative log likelihood (NLL)
[Equation]
Similar strategy has been used for example in Theis et al.
(2017) and Ball´e et al.
both using some form of continuous relaxation of q c as well as in Mentzer et al.
(2018) using an autoregressive
PixelCNN as P ψ to model q c ( c ) =
(cid:81)
i q c
( c i |
c i − 1 , . . .
.c 1 ) .
There is one caveat to the above approach.
In the minimization in equation 6 the sampling distribution µ c is treated as ﬁxed.
Minimizing the cross-entropy in such a regime minimizes the D KL and hence the additional bits due to q c (cid:54) =
p c but not the entropy H µ c ( c ) (see equation 5) which is treated as ﬁxed and therefore not optimized for low rate as explained in section 2.
This may seem natural and even inevitable since the samples c are the result of sampling the data x from the unknown yet ﬁxed distribution µ x .
Yet, the distribution µ c is not ﬁxed .
It is determined by the learned transformation T E ,θ = Q E ◦ E θ as the push-forward measure of µ x
[Equation] [Equation]
where A ∈ C and T − 1 is the inverse image 6 deﬁned as T − 1 ( A
[Equation]
:= { x ∈ X : T ( x ) ∈ A } .
Changing the parameters of the encoder E θ and the embeddings E will change the measure µ c and hence the entropy H µ c ( c ) and the cross-entropy H µ c | q c ( c ) even with the approximation q c ﬁxed.
We therefore propose to optimise the encoder and the embeddings so as to minimize the crossentropy not only through learning better approximation q c but also through changing µ c to achieve overall lower rate.
Since the discrete sampling operation from µ c is non-differentiable, we propose to use a simple continuous soft relaxation similar the one described in section 2.
Instead of using
[Equation]
we use the differentiable soft relaxation
[Equation]
Our ﬁnal objective is the minimization of the empirical loss composed of three terms: the distortion, the soft cross-entropy, and the hard cross-entropy
[Equation]
The distortion may be the squared error d ( x , ˆ x )
= (cid:107) x − ˆ x (cid:107) 2 or other application-dependent metric (e.g. multi-scale structural similarity index for visual perception in images).
Through the distortion we optimize the parameters of the decoder D φ and using the relaxation described in section 2 for the backward pass also the parameters of the encoder E θ and the quantization embeddings Q E .
The hard cross-entropy loss is
[Equation]
where we treat the dimensions of the vector c
= [ c (1) , . . .
, c ( d/m ) ] as independent so that the approximation q c can be used directly as the entropy model for the lossless arithmetic coding (expects the elements of the messages to be sampled i.i.d.
from a single distribution).
The loss simpliﬁes to the ﬁnal form due to the 0/1 probabilities of the deterministic code-assignments in equation 8.
Through the hard cross-entropy loss we learn the parameters of the probability model P ψ outputting the q c distribution.
The soft cross-entropy loss is
[Equation]
which uses the differentiable soft relaxation ˆ p c of equation
This allows for back-propagating the gradients to the encoder E θ and the quantizer Q E .
We use the sg operator here to treat q c as ﬁxed in this part of the loss preventing further updating of the parameters of the probability model P ψ .
E XPERIMENTS
As a proof of concept we conducted a set of experiments on the tiny 32x32 CIFAR-10 images (Krizhevsky, 2009).
We use similar architecture of the encoder E θ and decoder D φ as Mentzer et al.
(2018) (without the spatial importance mapping), with the downsampling and upsampling strideconvolutions with kernels of size 4, and 10 residual blocks with skip connections between every 3.
We ﬁx the annealing parameter σ = 1 and the loss hyper-parameter β = 1 7 .
We use ADAM with default pytorch parameters, one cycle cosine learning rate and train for 15 epochs.
The code is available at: https://bitbucket.org/dmmlgeneva/softvqae/ .
We ﬁrst compare the vector quantization (VQ) approach where the codebook is composed of mlong vectors versus the scalar (SQ) approach where it contains scalars m = 1 as e.g. in Mentzer et al.
By construction the VQ version needs to transmit shorter messages for the same level of downsampling.
For example, with 8-fold downsampling to 4 × 4 latents z with 8 channels the discrete codes c of SQ have d/m = 128 elements.
In VQ, the channel dimension forms the rows of the matrix Z with m = 8 and the c messages to be encoded and transmitted have only d/m =elements.
On the other hand, in the scalar version each of the 8 channels is represented by its own code and therefore allows for more ﬂexibility compared to a single code for the whole vector.
Our preliminary experiments conﬁrm the superiority of the VQ.
In ﬁgure 1 left we plot the ratedistortion for comparable parts of the trade-off space.
The VQ models use 2-folds downsampling resulting in d/m = 16 × 16 = 256 long messages c .
The two curves are for embeddings (latent z channels) with size m = 8 and m = 16 respectively and the points around the curves are the result of increasing the size of the dictionary E as k =
[8 , 16 , 32 , 64 , 128] from left to right.
The SQ models use 8-folds downsampling with latent z channels 8 and 16 resulting in d/m = 4 × 4 × 8 = 128 and d/m = 4 × 4 × 8 = 256 ( m = 1 here).
We observe that VQ clearly achieves better trade-offs being in the bottom-left of the plots.
We next conﬁrm the effectiveness of the soft cross-entropy term s ( c ) in our ﬁnal loss formulation in equation 10.
Increasing the α hyper-parameter should put more importance on the rate minimization (through the entropy) as compared to the distortion.
In the right plot of ﬁgure 1 we compare the points k
= [8 , 32 , 128] of the ‘vq, 8’ curve for values α = [0 . 01 , 0 .
001 , 0] from left to right.
With the highest α = 0 . 01 , the objective trade-off searches for the lowest rate tolerating higher distortion.
The lower the α , the less we push for low rates which allows for smaller distortion.
This behaviour corresponds well to the expected and desirable one where, as formulated in equation 1, we can now directly control the trade-off between the two competing objectives by setting the hyper-parameter α .
In the appendix we provide examples of the learned q c histograms for ‘vq, 8, k=32’ with different values of α showing the concentration of the measure into a few points in the support for high α .
A A PPENDIX
Histograms q c learned by the probability model P ψ for the ‘vq, 8’ models with the number of embedding vectors k =
≈ the red line in the right graph in Figure 1 with increasing α ∈ { 0 , 0 .
For the highest α = 0 .
01 the distribution is concentrated into a few points in the support resulting in the lowest entropy (and therefore the best rate) but highest distortion.