Sample what you cant compress
Vighnesh Birodkar Gabriel Barcik James Lyon Sergey Ioffe David Minnen Joshua V. Dillon



[
    [
     
=====



§ ABSTRACT 
For learned image representations, basic autoencoders often produce blurry results. Reconstruction quality can be improved by incorporating additional penalties such as adversarial (GAN) and perceptual losses. Arguably, these approaches lack a principled interpretation. Concurrently, in generative settings diffusion has demonstrated a remarkable ability to create crisp, high quality results and has solid theoretical underpinnings (from variational inference to direct study as the Fisher Divergence). Our work combines autoencoder representation learning with diffusion and is, to our knowledge, the first to demonstrate the efficacy of jointly learning a continuous encoder and decoder under a diffusion-based loss. 
We demonstrate that this approach yields better reconstruction quality as compared to GAN-based
autoencoders while being easier to tune. 
We also show that the resulting representation is easier to model
with a latent diffusion model as compared to the representation obtained from a state-of-the-art GAN-based loss.
Since our decoder is stochastic, it can generate details not encoded in the otherwise deterministic latent representation; we therefore name our approach “Sample what you can't compress”, or  for short.





[1]Primary contact.
[2]Work done while at Google.



§ INTRODUCTION





Image autoencoders ultimately necessitate a pixel-level loss to measure and minimize distortion. A common choice is to use mean squared error (MSE). This is a problem for image and video models because MSE favors low frequencies over high frequencies
<cit.>. Although generalized robust loss functions have been developed <cit.>, they are insufficient on their own for avoiding blurry reconstructions. A popular fix is to augment a pixel-level loss with additional penalties. Typically, MSE is still used because it is easy to optimize due to its linear gradient.

For example, <cit.> use a combination of MSE, perceptual
loss, and adversarial loss. <cit.> noted that an adversarial loss
helps them get high-quality images with realistic textures.
Unfortunately GANs remain challenging to train; which was most
recently noted by <cit.>, when they couldn't naively scale up 
their architecture.
The diversity of their outputs is also limited, because modern GAN based decoders
are deterministic, and thus lack the capacity to sample multiple different possibilities.


0.47
    
    < g r a p h i c s >

     figureReconstruction distortion (lower is better) as a function of compression for  and GAN based auto-encoders.
     

0.47
    
    < g r a p h i c s >

     figureClass conditional generation  quality (lower is better) as a function of compression for  and GAN based auto-encoders.
     




As an alternative, this paper describes a technique for using a diffusion loss to learn an autoencoder. The diffusion loss is sensible because it is a proper scoring rule with favorable theoretical properties such as being formally connected to KL divergence.
It has proven itself capable of generating crisp results with high perceptual quality as reflected by human evaluation studies <cit.>.

To demonstrate its simplicity, we take a popular encoder architecture
<cit.> and marry it with a U-net decoder, a popular denoising
architecture for diffusion <cit.>.
With some additional details outlined in <ref>,
we show that this approach achieves lower distortion at all
compression levels as measured by the CMMD metric <cit.>.
Because our decoder is able to sample details at test-time that are not encoded in the latents, we call our approach “Sample what you can't compress” or  for short.


This work will show that,

    
  * SWYCC achieves lower reconstruction distortion at all tested compression levels vs SOTA GAN-based autoencoders (<ref>).
    
  * SWYCC representations enable qualitatively better latent diffusion generation results at alltested compression levels vs SOTA GAN-based autoencoders (<ref>).
    
  * Splitting the decoder into two parts improves training dynamics (<ref> and <ref>).






§ METHOD


Eliding its various parametrizations for brevity, the standard diffusion loss is characterized by the Monte Carlo approximation of the following loss,

    ℓ(x) _ε∼(0,I_h· w· 3), t∼[0,1][w_t(⃗x) - D(α_t (⃗x)+σ_tε, t)_2^2].

Herein x ∈^h× w× 3 denotes a natural image and D is a neural network fit using gradient descent and which serves to denoise the corrupted input x_tα_t x + σ_t ε at a given noise-level t.
Let the corruption process be the cosine schedule <cit.>, σ_t^21-α_t^2 and α_t^2(-2logtan(at+b)) where aarctan(e^-1/2(-20)) - b and barctan(e^-1/2 (20)).



We extend this definition to the task of autoencoding by simply allowing the denoising function to take an additional argument, (E(x)), itself having access to the uncorrupted input x but only through the bottlenecking function E. The result is,

    ℓ_AE(x) _ε∼(0,I_h· w· 3), t∼[0,1][w_tx - (α_t x+σ_tε, t, (E(x))) _2^2].


As the notation suggests E is an encoder which, notably, is learned 
jointly with “diffusion decoder”  and secondary decoder :→^h× w× 3. The specification of  is largely a convenience but also merits secondary advantages. By mapping z=E(x) back into x-space, we can simply concatenate the corrupted input x_t and its “pseudo reconstruction,” (z). Additionally, we find that directly penalizing (z), as described below, speeds up training.




 §.§ Architectural Details

Encoder: We use a fully convolutional encoder in all of our
experiments, whose specifics we borrow from MaskGIT<cit.>.
The encoder consists of multiple ResNet<cit.> blocks stacked on top
of each other, with GeLU<cit.> for its non-linearities and GroupNorm<cit.> for training stability.
The ResNet blocks are interspersed with strided convolutions
with stride 2 which achieves a 2 × downsampling
by itself.
To get the 8×8 patch size, we use 4 ResNet blocks with 3 downsampling
blocks.
The encoder architecture is common for all of our experiments,
and we only change the number of channels at the output layer
to achieve the desired compression ratio.

Decoder: For the decoder in the GAN baseline and  we use an architecture that is 
the reverse of the encoder. For up-sampling,
we use the depth-to-space operation.
Just like the encoder, we have 
4 ResNet blocks interspersed with 
3 depth-to-space operations. 
For  we use a U-Net as defined
by <cit.>.
The U-Net has 4 ResNet blocks for downsampling
and corresponding 4 ResNet blocks for upsampling with residual connections
between blocks of the same resolution.
After 4 downsampling stages, when 
resolution is 16×16, we
use a self-attention block to give
the network additional capacity.




 §.§ Speeding-Up Training





We find that additional direct penalization of (E(x) leads to faster training and improved CMMD and FID. This was achieved by minimizing a composite loss
containing terms with favorable Hessian (<ref>) and perceptual characteristics (<ref>),

    ℓ_Totalℓ_AE + λ_pℓ_Perceptual
    + λ_mℓ_MSE

where,

    ℓ_MSE x - (E(x))_2^2

and,

    ℓ_Perceptualf_Frozen(x) - f_Frozen( (E(x)))_2^2.

The function f_Frozen is an unlearnable standard ResNet, itself trained on 
ImageNet and used for both the baseline and . We found the best hyper-parameter setting for <ref> is λ_m = 1 and λ_p = 0.1
Setting λ_p > 0 was particularly important to be competitive at reconstruction with GAN based methods in Table <ref>. 
The visual impact of perceptual loss is shown in Figure <ref>. We maintain
that the main advantage of using auxiliary losses is to speed-up training,
because given enough data and training time, our
diffusion decoder should converge to the true distribution.

For generating reconstructions (recall that the  decoder is stochastic) we used classifier-free guidance during inference <cit.>; for the unconditional model the U-Net was trained with (E(x)) dropped-out, i.e., randomly zeroed out on 10% of training instances.







§ EXPERIMENTS



In this section we explore how the GAN-based loss compares to our approach.
Without loss of generality, we define the relative compression
ratio of 1 to be a network that maps 8 × 8 RGB patches to an 8 dimensional
latent vector. Effectively, this means for our encoder 
E if and x∈^256 × 256 × 3,
then E(x) ∈^32 × 32 × C
where C=8.
In general, to achieve a relative compression ratio of 
K we set C = 8/K. The effect of increasing
the compression ratio is plotted in Figure <ref>.
Observe that distortion degrades much more rapidly for the
GAN based auto-encoder as measured by CMMD <cit.> which
Imagen-3 <cit.> showed better correlates with
human perception.

Not only is our approach better at all compression levels; the gap between the GAN based
autoencoder and  widens as we increase the relative compression ratio.
Using the much simpler diffusion formulation under Equation <ref>,
we are able to reconstruct crisp looking images with detailed textures
(See Figure <ref>). Our method has the added benefit that
we do not need to tune any GAN related hyper-parameters, and can scale up effectively 
using the large body of diffusion literature (<cit.>, <cit.>).








 §.§ Impact of 



Observing Equation <ref> and Figure <ref>, we note
that the output of  is an intermediate tensor that is not strictly
required for the diffusion loss or for generating the output. We
show using Table <ref> that this piece is crucial for
achieving performance comparable to the GAN based autoencoder.
The perceptual loss term in particular has a large impact in
reducing distortion. Visual examples are shown in Figure <ref>.



 §.§ Analysis of sampling in 


In Figure <ref> we show the qualitative difference number of
sampling steps makes to reconstructions. We can see that even with just
2 steps the high level structure of the image is present. In Figure <ref>
we study the impact of number of sampling steps by using the CMMD metric.
Figure <ref> shows what sampling in  actually ends up changing.
We can see that only regions with high-frequency components and detailed textures
are changed between samples, while regions containing similar colors over large
areas are left untouched.


Figure <ref> studies the effect of classifier-free guidance <cit.>
as used in . We ablate the guidance with a model trained at a relative compression
factor of 4 (See Section <ref> for definition) and find
that a guidance scale of 0.5 works the best. This is not to be
confused with the guidance scale of the latent diffusion model that may
be trained on top of our autoencoder, which is a completely separate parameter
to be tuned independently. 




 §.§ Modeling latents for diffusion


We use a DiT model <cit.> to model the latent space of our models
for the task of class-conditional image generation.
In Figure <ref> we compare our latents with those of a GAN based autoencoder.
Crucially, note that despite using an identical encoder architecture,
our approach can improve FID <cit.> by over 12%.
We hypothesize that our latent space is easier to model
because our encoder does not need to model high-frequency texture
information because the decoder can independently sample those details.



 §.§ Exploring better perceptual losses

In Table <ref> we showed the large impact perceptual loss
has on reconstruction quality. This begs the question; are there
better auxiliary losses we can use?
We compare the perceptual loss as described by in VQGAN <cit.> and replace it with the DISTS <cit.>.
DISTS loss differs from perceptual loss in 2 important ways. a)
It uses SSIM <cit.> instead of mean squared error and b) It
uses features are multiple levels instead of using only
the activation's from the last layer.
The results are shown in Figure <ref>.
At lower relative compression ratios, DISTS loss helps GANs
and . But at higher relative compression ratios,
GAN based autoencoders perform even worse than the perceptual
loss based baseline. We think this is an avenue for future exploration. 
We perform all other experiments with perceptual loss <cit.>,
since it is more prevalent in literature and it helps
GAN based auto-encoders at higher relative compression ratios.








 §.§ Architecture and hyper-parameters

Autoencoder training hyper-parameters
We train all of our models on the ImageNet dataset resized at 256 × 256
resolution. During training, we resize the image such that the shorter side measures
256 pixels and take a random crop in that image of size 256 × 256. For measuring reference statistics on the validation set, we take the largest
possible center square crop. All of our models are trained at a batch size of 256 for
10^6 steps which roughly equals 200 epochs. 

GAN baseline:
We use the popular convolutional encoder-decoder architecture popularized
by MaskGIT <cit.> in our GAN-based baselines. This autoencoder design is used by many image models, including FSQ <cit.> and GIVT <cit.>, and was extended to the video domain by MAGVIT-v2 <cit.>, VideoPoet <cit.>, and WALT <cit.>. We take advantage of decoder improvements developed by MAGVIT-v2 (see Section 3.2 in <cit.>) to improve reconstruction quality.
Note that while the video models enhance the base autoencoder architecture with 3D convolution to integrate information across time, the discriminator and perceptual loss are still applied on a per-frame basis and thus are essentially unchanged in our model.



: In our experiments we keep the architecture of  identical
to the decoder used in the GAN baseline. For  we use the U-Net architecture
as parameterized by <cit.>. We borrow the U-Net 256 architecture
and make the following modifications:

We use the Adam optimizer to learn our parameters. The learning rate is warmed up
for 10^4 steps from 0 to a maximum value of 10^-4 and cosine decayed
to 0. We use gradient clipping with global norm set to 1.

Latent Diffusion:
We train for 2 × 10^6 steps with a batch size of 256. 
We use a classifier-free guidance scale of 1 and dropout the class embedding
10% of the time during training.
All of our experiments are done with a DiT-XL architecture.




§ RELATED WORK

Autoencoders for 2-stage generation: For discrete representation learning, <cit.> showed the usefulness of the 2-stage modeling approach. In this broad category, the first stage fits an autoencoder to the training data with the goal of learning a compressed representation useful for reconstructing images. This is followed by a second stage where the encoder is frozen and a generative model is trained to predict the latent representation based on a conditioning signal. This approach regained popularity when <cit.> showed that it can be used for zero-shot 
text generation, and is now the dominant approach for image and video generation <cit.>.

Adversarial losses: <cit.> extended the autoencoder from <cit.> with
two important new losses, the perceptual loss and the adversarial loss,
taking inspiration from the works of <cit.>
and <cit.>. The perceptual loss is usually defined as the L2 loss between a latent representation of the original and reconstructed image.
The latent representation, for example, can be extract from
the final layer activations of a ResNet optimized to classify ImageNet images.
The adversarial loss is a patch-based discriminator that uses a discriminator
network to predict at a patch-level whether it is real or fake.
This encourages the decoder to produce realistic looking textures.

Latent Diffusion:  <cit.> popularized text-to-image
generation using latent diffusion models. They kept the autoencoder
from <cit.> intact and simply removed the quantization layer.
This accelerated diffusion model research in the community owing
to the fact that the latent space was much smaller than the
pixel space, which allows fast training and inference compared
to diffusion models like Imagen that sample pixels directly <cit.>.

Compression and diffusion: <cit.> 
showed that diffusion models can be used for compression. Crucially,
compared to our approach they use a frozen autoencoder, and do not
train their autoencoder end-to-end. They also use an objective based on modified flow matching. In contrast, we did not modify the loss or the sampling algorithm.

In a similar context, <cit.> developed an end-to-end optimized compression model using a diffusion decoder. They show improved perceptual quality compared to earlier GAN-based compression methods at the expense of higher distortion (pixel-level reconstruction accuracy). Different from our approach, they use a discrete latent space, which is required for state-of-the-art compression rates achieved via entropy coding. This limits the reconstruction quality but is required for a compression model that ultimately seeks to minimize a rate-distortion objective, not just a reconstruction and sampling quality objective.




§ CONCLUSION

We have described a general autoencoder framework that uses a diffusion based
decoder. Compared to decoders that use GANs, our system is much more easier
to tune and has the same theoretical underpinnings as diffusion models.
We showed our method produces sigificantly less distortions as compared
to GAN based autoencoders in Figure <ref> and are better
behaved as latent spaces for diffusion in Figure <ref>.
In Section <ref> and <ref> we studied
the hyper-parameter settings on the 2 major components of our decoder,
 and .

Possible extensions: The autoencoder technique we describe is fairly
general and can be extended to any other continuous modality like audio, video or 
3D-point clouds. In addition, all improvements to diffusion algorithms like those
by <cit.> can be carried over.

Limitations: The main limitation of our method is the increase in inference
cost during decoding. This can be partly mitigated by using fewer steps like in Figure <ref>. In addition, techniques used to improve diffusion
sampling time like Progressive distillation<cit.> and Instaflow<cit.>
are also prudent.






iclr2025_conference



§ APPENDIX



 §.§ Visual examples






