Enhancing Consistency-Based Image Generation via Adversarialy-Trained Classification and Energy-Based Discrimination

Shelly Golan,Roy Ganz,Michael Elad
DOI: https://doi.org/10.48550/arXiv.2405.16260
2024-05-25
Abstract:The recently introduced Consistency models pose an efficient alternative to diffusion algorithms, enabling rapid and good quality image synthesis. These methods overcome the slowness of diffusion models by directly mapping noise to data, while maintaining a (relatively) simpler training. Consistency models enable a fast one- or few-step generation, but they typically fall somewhat short in sample quality when compared to their diffusion origins. In this work we propose a novel and highly effective technique for post-processing Consistency-based generated images, enhancing their perceptual quality. Our approach utilizes a joint classifier-discriminator model, in which both portions are trained adversarially. While the classifier aims to grade an image based on its assignment to a designated class, the discriminator portion of the very same network leverages the softmax values to assess the proximity of the input image to the targeted data manifold, thereby serving as an Energy-based Model. By employing example-specific projected gradient iterations under the guidance of this joint machine, we refine synthesized images and achieve an improved FID scores on the ImageNet 64x64 dataset for both Consistency-Training and Consistency-Distillation techniques.
Computer Vision and Pattern Recognition,Machine Learning
What problem does this paper attempt to address?
The problem that this paper attempts to solve is to improve the perceptual quality of images generated based on Consistency Models while maintaining the speed of image generation. Specifically, although Consistency Models can quickly generate high - quality images, their sample quality is usually slightly inferior to that of diffusion models. To this end, the author proposes a new post - processing technique to enhance the perceptual quality of images generated by Consistency Models through a classifier - discriminator joint model with adversarial training. ### Background of the Paper - **Diffusion Models**: Diffusion models are efficient generative models that generate high - quality images from random noise by gradually denoising. However, these models require a large amount of computational resources when generating samples, resulting in slow inference speed and limiting real - time applications. - **Consistency Models**: Consistency models aim to accelerate the generation process of diffusion models and achieve fast one - step or few - step generation by directly mapping from noise to data. Despite the fast speed, the quality of generated images is limited by the complex distillation process. ### Methods in the Paper - **Joint Classifier - Discriminator Model**: The author proposes a joint classifier - discriminator model, in which the classifier is used to evaluate the probability that an image belongs to a specified category, and the discriminator part uses softmax values to evaluate the closeness of the input image to the target data manifold, thus acting as an energy model. - **Adversarial Training**: Through adversarial training, the classifier and discriminator are jointly optimized, making the generated images more realistic. Specifically, the Projected Gradient Descent (PGD) algorithm is used to iteratively modify the generated images to maximize the conditional probability of the target category and minimize the gap between the generated images and the real data distribution. - **Early - Stopping Mechanism**: To prevent over - optimization, an early - stopping mechanism based on the model's probability output is introduced to ensure that the generated images are not misidentified as being more "real" than the real images. ### Experimental Results - **Performance Improvement**: The experimental results show that this method significantly improves the FID score on the ImageNet 64x64 dataset, increasing by 27.48% and 20.96% under Consistency - Training and Consistency - Distillation techniques respectively. - **Comparison with Existing Methods**: Compared with methods using only robust classifiers (such as BIGROC), this method shows stronger advantages in terms of perceptual quality and the realism of generated images. - **Generality**: This method is not only applicable to Consistency Models, but can also be applied to other generative models, showing good generalization ability. ### Summary This paper proposes a new post - processing technique. Through the joint classifier - discriminator model with adversarial training, it significantly improves the perceptual quality of images generated based on Consistency Models while maintaining the generation speed. The experimental results on multiple generative models verify its effectiveness and generality.