U NIVERSAL R ATE -D ISTORTION -P ERCEPTION R EPRESENTATIONS FOR L OSSY C OMPRESSION
A BSTRACT
In the context of lossy compression, Blau & Michaeli (2019) adopt a mathematical notion of perceptual quality and deﬁne the rate-distortion-perception function, generalizing the classical rate-distortion tradeoff.
We consider the notion of (approximately) universal representations in which one may ﬁx an encoder and vary the decoder to (approximately) achieve any point along the perception-distortion tradeoff.
We show that the penalty for ﬁxing the encoder is zero in the Gaussian case, and give bounds in the case of arbitrary distributions, under MSE distortion and W 22 ( · , · ) perception losses.
In principle, a small penalty refutes the need to design an end-to-end system for each particular objective.
We provide experimental results on MNIST and SVHN to suggest that there exist practical constructions that suffer only a small penalty, i.e. machine learning models learn representation maps which are approximately universal within their operational capacities.
I NTRODUCTION
A lossy compression system consists of an encoder producing a low-rate representation of a source and a decoder reconstructing an approximation.
Unlike lossless compression, the decoder in a lossy compression system has ﬂexibility in how it would like to reconstruct the source.
It is conventional to optimize for the reconstruction minimizing some distortion measure between the original and the reconstruction such as mean squared error, PSNR or SSIM/MS-SSIM (Wang et al., 2003; 2004).
Accordingly, lossy compression algorithms are analyzed through rate-distortion theory, wherein the objective is to minimize the amount of distortion for a given rate.
However, low distortion is not necessarily synonymous with high perceptual quality; indeed, deep learning based image compression has inspired works in which authors have noted that increased perceptual quality may come at the cost of increased distortion (Blau & Michaeli, 2018;
Agustsson et al., 2019).
This culminated in the work of (Blau & Michaeli, 2019) who propose the rate-distortion-perception theoretical framework.
The main idea was to introduce a third perception axis which more closely mimics what humans would deem to be visually pleasing.
Unlike distortion, judgement of perceptual quality is inherently no-reference; indeed, the mathematical proxy for perceptual quality is deﬁned by a constraint between the source and reconstruction distributions.
Leveraging generative adversarial networks (Goodfellow et al., 2014) in the training procedure has made such a task possible for complex data-driven settings, with efﬁcacy even at very low rates (Tschannen et al., 2018).
Naturally, this induces a tradeoff between optimizing for perceptual quality and optimizing for distortion.
But in designing a lossy compression system, one may wonder where exactly this tradeoff lies: is the objective tightly coupled with optimizing the representations generated by the encoder, or can most of this tradeoff be achieved by simply changing the decoding scheme?
Our contributions are as follows.
We prove that there is no loss in ﬁxing the representation map across the perception-distortion tradeoff for the Gaussian distribution, then provide a bound on the penalty incurred for arbitrary source distributions.
We perform experiments on MNIST and SVHN to demonstrate that learned representation maps can be approximately universal within their operational
[Figure]
[Figure]
P e r ce p t i o n L o ss ( W 22 )
R = 0 :
R = 0 :
R = 0 : 25 R = 0 :
R = 2 :
Rate (bits) (
[Figure]
Perception-distortion universality of Gaussian sources implies the existence of a representation rate-distortion-perception tradeoffs.
This has important practical implications as ﬂexibility is a highly desirable property in designing representations.
It may not be plausible to design an entire end-to-end system for each potential objective in advance; our results show that representations from universal encoders are reusable across a family of lossy compression objectives without much loss.
U NIVERSAL R EPRESENTATIONS
The process of quantizing a source X ∼ p
X leads to imperfect reconstruction as measured by some distortion function ∆ : R → R .
Minimizing an appropriate choice of ∆( · , · ) should cause the output to look more like the input as perceived by humans.
This methodology however does not explicitly capture how the distribution of the reconstructions compares to the distribution of the source, which Blau & Michaeli (2019) argue is a better proxy for what humans would perceive to be natural (disregarding the original source image).
Accordingly, the conventional rate-distortion function is augmented with some ”distance” between probability measures d ( · , · ) as follows 1 .
Deﬁnition 1.
The information rate-distortion-perception function for a source X ∼ p
X is deﬁned as
[Equation]
≤ P. Popular choices for d ( · , · ) include Wasserstein distances due to their empirical success, which we will assume henceforth unless otherwise stated.
We ﬁrst give a solution to the Gaussian case to provide an example of a distribution with the perception-distortion universality property.
Theorem 1.
For X ∼ N ( µ X , σ 2 X ) , the rate-distortion-perception function under squared error distortion and squared W 2 distance is achieved by some ˆ X jointly Gaussian with X and is given by
[Equation]
Due to joint Gaussianity of ( X, ˆ X ) , this expression admits the equivalent form
[Equation]
which describes the rate-distortion-perception function in terms of the minimum mean squared error estimator for X given ˆ X .
In principle, different ( D, P ) pairs are associated with different encoding schemes for X .
One may use (1) to show that different ( D, P ) pairs associated with the same R can be achieved by linearly scaling the representation
ˆ X produced by a single encoder.
This representation x f z gggˆxˆxˆxh
W 1 ( p X ,p ˆ X 1 ) h 2 ···
W 1 ( p X ,p ˆ X 2 ) h 3 ···
W 1 ( p X ,p ˆ X 3 ) is universal in the sense that it can achieve optimality across all perception and distortion tradeoffs for a given rate; to move between the points, it is sufﬁcient to modify only the decoding scheme.
It turns out that general distributions satisfy an approximate version of universality.
Let X ∼ p
X be a source and ˆ X (1) be an optimal representation in the conventional rate-distortion framework (i.e.,
I ( X ; ˆ X (1) ) =
R ( D 1 , P 1 ) , E [( X − ˆ X (1) ) 2 ] ≤ D 1 for some D 1 ∈ (0 , σ 2 X ) and P 1 = ∞ ).
Theorem 2.
(Approximate universality for general sources)
Under MSE and W 22 ( · , · ) , ˆ X (1) can meet distortion constraint D 2 = 2 D 1 and any perception constraint P 2 via suitable decoding.
This result implies that the optimal representation in the conventional rate-distortion sense can be transformed to meet any perception constraint possibly with a distortion penalty of no more than 2 D 1 − D 2 .
It is worth noting that this special case gives rise to Theorem 2 ( R ( D, ∞ ) ≥ R (2 D, 0) ) in Blau & Michaeli (2019).
Namely, the numerical connection between R ( D, ∞ ) and R (2 D, 0) is a manifestation of the existence of approximately universal representations.
E XPERIMENTAL R ESULTS
The rate-distortion-perception tradeoff was observed in the context of GAN-enhanced image compression (Blau & Michaeli, 2019), while training an entire model for each desired setting over distortion and perception.
In practice, it is undesirable to develop an entire system from scratch for each objective.
Our theoretical results have already established a bound on the loss of optimality for using universal representations.
We provide experimental evidence to show that the loss is small in practice.
Concretely, in an end-to-end model the encoder and decoder are trained jointly for an objective so as to utilize the ﬂexibility in learning a representation map whereas in a universal model, the representation map is ﬁxed so only the decoder can be trained.
The encoder for a universal model may be re-used across many models and be designed separately from the decoders; we refer to it as a universal encoder.
The architecture used in the experiments is a stochastic autoencoder augmented with a Wasserstein GAN and follows closely the design of (Blau & Michaeli, 2019).
The task of generating both the end-to-end and universal models is broken into two stages.
The initial stage trains an end-to-end model, consisting of an encoder f producing latent representation z
= f ( x ) and decoder g producing reconstruction ˆ x = g ( z ) .
Here f incorporates both the deterministic quantized encoding step and the addition of stochastic noise u to support high perceptual quality.
In practice, the quantized versions of the representations are used for compression then u can be added before decoding.
For the distortion metric we use MSE and for the perception metric we use
[Figure]
e r s t e i n - 1 P
e r c e p t i o n
L
[Figure]
Wasserstein-1 distance.
The is a weighted sum the given by
[Equation]
for some initial λ to be chosen.
W 1 is estimated using the Kantorovich-Rubinstein dual form
[Equation]
Here, F is the set of all bounded 1-Lipschitz functions.
In practice, we use a critic network as h and the Lipschitz condition is approximated with a gradient penalty (Gulrajani et al., 2017).
Optimization alternates between minimizing f, g with h ﬁxed and maximizing h with f, g ﬁxed.
In essence, g is optimized for both low distortion and high perception, so it acts as both a decoder and a generator.
Afterwards, the encoders of end-to-end models can be used to construct universal models.
To do the latter, we freeze the parameters of f and introduce a new decoder g 1 and critic h 1 trained to minimize
[Equation]
where λ 1 is a different tradeoff parameter and p ˆ X 1 is the distribution induced by pushing X through g 1 ◦ f .
The rest of the training procedure follows that of the ﬁrst stage.
This is repeated over many different parameters { λ i } to generate a tradeoff curve.
Figure 3 shows rate-distortion-perception curves at multiple rates on MNIST and SVHN.
Note that the rate for each individual curve is ﬁxed through using the same quantizer across all models.
The rates are chosen to be low so that the tension between distortion and perception is most visible.
Points outlined in black are losses for end-to-end models, while the other points correspond to the universal models sharing an encoder.
The universal models are able to achieve a tradeoff which is very close to the end-to-end models (with outputs that are visually comparable).
For any ﬁxed rate, decreasing the perception loss P reduces the output blur, but the reconstruction is less faithful to the original input.
As noted by (Blau & Michaeli, 2019), this is especially evident at very low rates in which the compression system appears to act as a generative model.
However, this does not necessarily come at the expense of the representation map learned by the encoder, i.e. the embeddings generated by an encoder trained for small P can also be used to produce a low-distortion reconstruction by training a new decoder.
Conversely, training a decoder to produce reconstructions with high perceptual quality on top of an encoder trained only for distortion loss is also possible as the decoder is sufﬁciently expressive to act purely as a generative model.
D ISCUSSION
The use of deep generative models in data compression has highlighted the tradeoff between optimizing for low distortion and high perceptual quality.
With a suitable encoding scheme, we have shown that modifying the decoder is sufﬁcient to achieve this tradeoff in practice.
Future directions include one-shot theoretical analysis, and employing the scheme to high-resolution images and videos.
A T HEORETICAL R ESULTS
Deﬁnition 1 .
Let ∆( · , · ) be a non-negative distortion measure satisfying ∆( x, y )
=
if and only if x = y , and d ( · , · ) be a (non-negative) divergence between probability measures, likewise with d ( p, q ) =
if and only if p = q .
The information rate-distortion-perception function for a source
X ∼ p X is deﬁned as
[Equation]
We let R ( D ) =
R ( D, ∞ ) denote the conventional rate-distortion function.
We assume the use of MSE distortion loss and squared Wasserstein-2 perception loss in the theoretical results unless otherwise stated.
Recall that
[Equation]
where the inﬁmum is over all joint distributions of ( X, ˆ X ) with marginals p X and p ˆ X .
When p X and p ˆ X are both Gaussian, W 2 2 ( p X , p ˆ X ) =
( µ X − µ ˆ X )
+ ( σ X − σ ˆ X ) 2 .
A. P ROOF OF T HEOREM
Theorem 1 .
For X ∼ N ( µ X , σ 2 X ) [Equation] the rate-distortion-perception function under squared error distortion and squared W 2 distance is achieved by some ˆ X jointly Gaussian with X and is given by
[Equation]
We ﬁrst restate a useful result from estimation theory.
Let ˆ
X be a random variable with E [ ˆ X ] = µ ˆ X ,
Var( ˆ X ) =
σ 2ˆ X and Cov( X, ˆ X )
= θ .
Let ˆ
X G be a random variable jointly Gaussian with X with the same ﬁrst and second order statistics as ˆ X .
Lemma 1.
Given µ ˆ X , σ 2ˆ X , and θ , we have that
[Equation]
Proof of Theorem 1.
We shall show that there is no loss of optimality in assuming that ˆ X is jointly
Gaussian with X .
It is clear that E [( X − ˆ X ) 2 ] =
E [( X − ˆ X G ) 2 ] , as the ﬁrst and second order statistics are all given.
Moreover, since every coupling of p X and p ˆ X induces a Gaussian coupling of p X and p
ˆ X G with the same covariance, it follows that
[Equation]
Finally, we have
[Equation]
where (a) is because the Gaussian distribution maximizes differential entropy for a given variance, (b) follows from Lemma 1 and (c) is because the estimation error is independent of ˆ X G .
Thus, it sufﬁces to solve the problem
[Equation]
Note that
[Equation]
so there is no loss of optimality in assuming µ ˆ X = µ X .
Accordingly,
[Equation]
When √ P < σ X − (cid:112) | σ 2 X − D | , both P and D are active, and consequently we have σ 2ˆ X =
( σ X − √ P ) 2 and θ =
σ 2 X + σ 2ˆ X − D. Noting that the other case is simply the solution to R ( D ) , this concludes the proof.
Remark.
The proof of Theorem 1 can be easily modiﬁed to handle to the case d ( p X , p ˆ X ) =
KL ( p X , p ˆ X ) , where KL ( p X , p ˆ X ) = (cid:82) p ˆ X ( x ) log p ˆ X ( x ) p X ( x ) d x is the KL-divergence between p X and p
ˆ X .
Note that given ( µ ˆ X , σ 2ˆ X ) , KL ( p X , p ˆ X ) is minimized when p ˆ X is a Gaussian distribution.
Given p X , one may see that there is a one-to-one correspondence between
[Equation]
which implies the rate-distortion-perception functions under KL ( p X , · ) and W 22 ( p X , · ) also share a one-to-one correspondence in P .
Remark.
One may see that the Gaussian distribution is universally representable as follows.
Let ˆ Xand ˆ X 2 be the induced outputs associated with ( D 1 , P 1 ) and ( D 2 , P 2 ) , respectively.
Note that we have equality of ﬁrst moments, E [ X ] =
E [ E [ X | ˆ X 1 ]]
=
E [ E [ X | ˆ X 2 ]] .
Since X is jointly Gaussian with ˆ X 1 , it follows from standard facts that (a) E
[ X | ˆ X 1 ] is independent of the estimation error
X − E
[ X | ˆ X 1 ] , and (b) E [ X | ˆ X 1 ] is a linear function of ˆ X 1 (and thus is itself Gaussian).
The same is true for ˆ X 2 .
Furthermore, after simplifying R ( D 1 , P 1 )
=
R ( D 2 , P 2 ) we get that
[Equation]
One may expand (3) and apply (a) to obtain equality of the cross terms, i.e. E [ X E [ X | ˆ X 1 ]] =
E [ X E
[ X | ˆ X 2 ]] .
Plugging this back into (3) shows that E [ E [ X | ˆ X 1 ] 2 ] =
E [ E [ X | ˆ X 2 ] 2 ] .
Since Gaussians are uniquely parameterized by their ﬁrst and second moments, ( X, E
[ X | ˆ X 1 ]) and ( X, E [ X | ˆ X 2 ]) must be identically distributed.
From (b), E [ X | ˆ X 1 ] and E [ X | ˆ X 2 ] are linear in ˆ X 1 and ˆ X 2 , respectively, so ˆ X 1 and ˆ X 2 must themselves be linearly related.
A. P ROOF OF T HEOREM
Recall that X ∼ p
X is an information source and ˆ X (1) is an optimal representation in the conventional rate-distortion framework (i.e., I ( X ; ˆ X (1) )
=
R ( D 1 , P 1 ) and E [( X − ˆ X (1) ) 2 ] ≤ D 1 for some
D 1 ∈ (0 , σ 2 X ) and P 1 = ∞ ).
Theorem 2 .
(Approximate universality for general sources)
Via suitable decoding, representation ˆ X (1) can meet distortion constraint D 2 = 2 D 1 and any perception constraint P 2 .
In fact, we ﬁrst prove the following more general result.
Representation ˆ X (1) can meet distortion constraint D 2 + δ D and perception constraint P 2 for any ( D 2 , P 2 ) , where δ D
= D 1 − D 2 + inf p ˆ X (2) :
d ( p X ,p ˆ X (2) ) ≤ P 2 W 22 ( p ˆ X (1) , p ˆ X (2) ) .
Proof of Theorem 2.
Let ˆ X (2) be an arbitrary representation jointly distributed with ( X, ˆ X (1) )
such that X − ˆ X (1) − ˆ X (2) form a Markov chain and E [( ˆ X (1) − ˆ X (2) ) 2 ]
= W 22 ( p ˆ X (1) , p ˆ X (2) ) .
Note that ˆ X (2) can be obtained from ˆ X (1) via a determinstic/stochastic transformation.
We have
[Equation]
where (a) is because ˆ X (1)
=
E [ X | ˆ X (1) ] almost surely and the MMSE estimation error is orthogonal to any function of ˆ X (1) .
Minimizing W 22 ( p ˆ X (1) , p ˆ X (2) ) subject to the constraint d ( p X , p ˆ X (2) )
≤ P 2 ﬁnishes the proof.
Observe that
[Equation]
because we may choose X (2) to have the same distribution as X .
This yields the original statement.
B E XPERIMENTAL D ETAILS
Training lasted 30 epochs for MNIST and 80 epochs for SVHN, and alternates between training the encoder and decoder with the critic ﬁxed and training the critic with the encoder and decode ﬁxed.
The learning rate was decayed by a factor of 5 after 20 epochs for MNIST, and likewise afterepochs for SVHN.
All models were trained with the Adam optimizer.
The batch size used was 64.
Training a single model takes about 30 minutes and 100 minutes for MNIST and SVHN, respectively.
All training was performed on an RTX 2070 GPU.
The architectures used for the experiments are given as follows.
Here each row represents a group of layers.
Noise is added for stochasticity after the output of the encoder.
d denotes the latent dimension and L the number of quantization levels per dimension, with R =
d log L used to upper bound the true rate (for MNIST, we found that this was within 15% of optimality).
The quantizer performs hard nearest-neighbour quantization on the forward pass and uses a soft relaxation given by Equation (3) in (Mentzer et al., 2018) during the backward pass.
The bin centers for quantization are spaced evenly in [ − 1 , 1] for each dimension.
The type of compression systems are denoted by E for end-to-end and
U for (perception-distortion) universal.
We found the results to be robust against model architectures when equipped with hyperparameters with good training performance.
Moreover, we found that using universal encoders trained with larger λ performed best in practice, though using small λ is nearly as effective also.
B. MNIST
The universality experiments build off of the encoders produced by the end-to-end experiments of the same rate with λ = 0 . 015 .
System
R d
L
E+U 4 .
E+U
E+U E+U
[Table] [Table]
[Table]
SVHN
The experiments are similar to MNIST, with the main difference being in the encoder architecture.
The universality experiments build off of the encoders produced by the end-to-end experiments of the same rate with λ = 0 .
[Table]
[Table] [Table]
[Table]
B. E XPERIMENTAL S AMPLES
As the emphasis on perception loss λ i increases
[Table]
the output blurriness is reduced
[Table]
but the reconstruction is less faithful to the original (in extreme cases even changing the identity of the digit).
The visual quality of both the end-to-end and universal models are on average comparable for each λ
i
[Table]
[Figure] [Figure]