Enhancing Variational Autoencoders with Smooth Robust Latent Encoding

TL;DR

The paper addresses the robustness of VAEs used in diffusion-based generation, proposing SRL-VAE, a post-training latent-space adversarial framework that smooths representations while preserving original latent structure via an originality regularization term. It formulates a min–max objective in the latent space, leveraging PGD-based perturbations and a combination of reconstruction and perceptual losses, plus an originality loss that anchors the encoder to pre-trained latent statistics. Empirical results show SRL-VAE improves image reconstruction fidelity and diffusion-generation quality, while significantly enhancing robustness against Nightshade poisoning and perturbation-based defenses in image editing tasks. The approach requires minimal additional computation, integrates with existing diffusion pipelines, and highlights a new pathway for robust generative modeling by shaping a stable, high-fidelity latent space. Overall, SRL-VAE demonstrates that adversarial training can simultaneously boost fidelity and robustness in VAEs, offering practical guidance for robust diffusion-based content creation.

Abstract

Variational Autoencoders (VAEs) have played a key role in scaling up diffusion-based generative models, as in Stable Diffusion, yet questions regarding their robustness remain largely underexplored. Although adversarial training has been an established technique for enhancing robustness in predictive models, it has been overlooked for generative models due to concerns about potential fidelity degradation by the nature of trade-offs between performance and robustness. In this work, we challenge this presumption, introducing Smooth Robust Latent VAE (SRL-VAE), a novel adversarial training framework that boosts both generation quality and robustness. In contrast to conventional adversarial training, which focuses on robustness only, our approach smooths the latent space via adversarial perturbations, promoting more generalizable representations while regularizing with originality representation to sustain original fidelity. Applied as a post-training step on pre-trained VAEs, SRL-VAE improves image robustness and fidelity with minimal computational overhead. Experiments show that SRL-VAE improves both generation quality, in image reconstruction and text-guided image editing, and robustness, against Nightshade attacks and image editing attacks. These results establish a new paradigm, showing that adversarial training, once thought to be detrimental to generative models, can instead enhance both fidelity and robustness.

Figures (16)

• Figure 1: Concept figure of SRL-VAE. Compared to SD-VAE, SRL-VAE maintains similar representations for clean examples while achieving robust representation against perturbed examples.
• Figure 2: Training objective of Smooth Robust Latent Variational Autoencoders (ours). A novel adversarial training approach in the latent space of VAE with originality regularization.
• Figure 3: Concept of smooth latent space. A smooth latent space ensures that perturbed examples are mapped closely to their original counterparts, enabling the VAE to extract robust features.
• Figure 5: Visual examples of reconstruction under various perturbations. SD-VAE struggles to reconstruct images with added perturbed noise, whereas ours robustly handles both clean and various perturbed images.
• Figure 6: Comparison of image-to-image editing results under defensive perturbations. SRL-VAE produces valid and high-quality edited outputs based on the given prompts, while the baseline VAE fails to preserve the original semantics.