Revisiting Diffusion Autoencoder Training for Image Reconstruction Quality

TL;DR

This work revisits diffusion autoencoder training to improve image reconstruction quality. It argues that spending excessive steps in high-noise regions wastes capacity that could refine details, and thus introduces a two-phase training regime, a shifted cosine noise schedule, and a switch from $\varepsilon$-prediction to $v$-prediction. The proposed method first forces the encoder–decoder to autoencode by using $t=1$ in Phase 1, then refines details in Phase 2 with a schedule biased toward low-noise regimes; across CIFAR-10, CelebA, LSUN Bedroom, and ImageNet, this approach yields sharper reconstructions while preserving latent-space interpolability and manipulability. The results demonstrate improved reconstruction metrics (PSNR, SSIM, LPIPS) and competitive or improved FID, along with retained latent-code information content and smooth latent traversals. Overall, the two-phase method offers a practical path to higher-fidelity reconstructions without sacrificing the latent-space properties that enable editing and manipulation.

Abstract

Diffusion autoencoders (DAEs) are typically formulated as a noise prediction model and trained with a linear-$β$ noise schedule that spends much of its sampling steps at high noise levels. Because high noise levels are associated with recovering large-scale image structures and low noise levels with recovering details, this configuration can result in low-quality and blurry images. However, it should be possible to improve details while spending fewer steps recovering structures because the latent code should already contain structural information. Based on this insight, we propose a new DAE training method that improves the quality of reconstructed images. We divide training into two phases. In the first phase, the DAE is trained as a vanilla autoencoder by always setting the noise level to the highest, forcing the encoder and decoder to populate the latent code with structural information. In the second phase, we incorporate a noise schedule that spends more time in the low-noise region, allowing the DAE to learn how to perfect the details. Our method results in images that have accurate high-level structures and low-level details while still preserving useful properties of the latent codes.

Figures (11)

• Figure 1: Some noise schedules in the shifted cosine family and the linear-$\beta$ noise schedule.
• Figure 2: Some (A) CIFAR-10 images and their reconstructions by (B) a vanilla autoencoder trained to minimize the L2 loss, (C) stochastic autoencoding with a DAE with 10 DDIM sampling steps, and (D) stochastic autoencoding with the same DAE at 50 steps and (E) 1000 steps. We show (F) the absolute pixel differences between (D) and (E) to highlight the subtle differences between them. We also show the FID (computed by the torch-fidelity (tf) library Obukhov:2020 and the clean-fid (cf) library Parmar:2022 in that order) and averaged PSNR, SSIM, and LPIPS values computed from the reconstructions against the dataset. We can see that PNSR and SSIM overly award blurry images (Column B and C). LPIPS does not fall for the same trap but scores the sharpest (as measured by FID) reconstructions at 1000 steps worse than those at 50 steps. The FID, on the other hand, does not measure alignment between the reconstructed and the original, so it scores the most structurally accurate images in Column B the worst.
• Figure 3: Some (A) original CIFAR-10 images and their reconstruction by the (B) $\varepsilon$-predicting, (C) $x$-predicting, and (D) $v$-predicting DAEs in Table \ref{['table:prediction-type-ablation']}. The $\varepsilon$-predicting DAE simply fails at being an accurate autoencoder.
• Figure 4: Reconstructions of a frog image from CIFAR-10 by some of the DAEs from Table \ref{['table:one-phase-noise-schedule-ablation']}. Notice that the images become sharper and noisier but less accurate as one moves to the right.
• Figure 5: Image reconstruction on CelebA and ImageNet datasets.