Source-linked AI summary
GAN-Diff : Coupling Pretrained WGAN-GP Features with Conditional Diffusion U-Nets
Saif Ahmed, Ashadulla Hil Galib, S. M. Riaz Rahman Antu, Ahmed Faizul Haque Dhrubo, Souvik Pramanik, Mohammad Abdul Qayum, Mohsin Sajjad, Mohammad Ashrafuzzaman Khan
TL;DR
The paper addresses whether a pretrained WGAN-GP can provide an effective feature-level prior for conditional diffusion image restoration. It couples frozen generator features to a diffusion U-Net through cross-attention and reports improvements on CelebA Gaussian denoising and 2× super-resolution over input baselines.
Problem
The paper examines whether a pretrained WGAN-GP can provide an adequate feature-level prior for conditional diffusion-based image restoration.
Method
A frozen WGAN-GP generator supplies intermediate spatial features through cross-attention to a conditional diffusion U-Net during DDIM restoration.
Results
The framework consistently improves degraded and low-resolution CelebA images on Gaussian denoising and 2× super-resolution tasks.
Takeaways & Limitations
A frozen GAN feature prior can guide diffusion models toward stable and effective image restoration within the evaluated CelebA tasks.
Takeaways & Limitations
Evaluation is limited to five tracked faces, lacks final-pipeline perceptual metrics, and uses one fixed Gaussian-noise level for denoising.
Abstract
from arXiv · showhide
Generative adversarial networks (GANs) can provide efficient image generation, while diffusion models offer high-quality image restoration but require iterative sampling. This paper presents a hybrid GAN-guided diffusion framework that uses a pretrained Wasserstein GAN with gradient penalty (WGAN-GP) as a feature prior for conditional diffusion-based image restoration. Intermediate features from the frozen WGAN-GP generator are incorporated into a diffusion U-Net through cross-attention and remain fixed during the DDIM sampling process. The framework is evaluated on two restoration tasks, Gaussian denoising and 2Xsuper-resolution, using CelebA face images. During development, several sources of instability were identified and addressed, including adversarial learning-rate imbalance, inappropriate diffusion initialization, excessive corruption, and insufficient parameter averaging. The resulting framework consistently improves the quality of both degraded and low-resolution images. In particular, it improves denoising performance by 4.40 dB in PSNR and super-resolution performance by 3.70 dB over their respective input baselines. These results demonstrate the potential of a frozen GAN feature prior to guide diffusion models toward stable and effective image restoration.
I. INTRODUCTION
The paper combines a frozen pretrained WGAN-GP feature prior with conditional diffusion restoration, using cross-attention to guide Gaussian denoising and 2× super-resolution. It also addresses training and inference instabilities identified during iterative development.
- I. INTRODUCTION: Development compared DCGAN, WGAN-GP, and StyleGAN-lite, selecting WGAN-GP for stability and quality reasons.
- I. INTRODUCTION: The hybrid model uses intermediate spatial features from a frozen WGAN-GP generator as a feature-level prior for a conditional diffusion U-Net.The features are integrated through cross-attention while preserving the diffusion model’s restoration procedure.
- I. INTRODUCTION: The study diagnoses instability related to generator–critic learning rates, DDIM initialization, corruption strength, and exponential moving-average smoothing.These issues affect both training and inference in the hybrid pipeline.
- I. INTRODUCTION: The framework targets Gaussian image denoising and 2× super-resolution on CelebA face images.
- I. INTRODUCTION: The GAN feature map is calculated once per sampling chain and kept fixed throughout DDIM sampling.This design avoids changing the conditioning signal across reverse-diffusion steps.
- II. RELATED WORK: The approach differs from prior GAN–diffusion hybrids by conditioning the diffusion U-Net on an intermediate GAN feature map rather than replacing reverse diffusion with GAN output.
III. DATASET
The experiments use a 50,000-image CelebA subset prepared at 64 × 64 resolution with normalized channels. Five fixed faces support consistent checkpoint-wise comparisons.
- III. DATASET: The dataset uses a 50,000-image subset of CelebA celebrity face images.
- III. DATASET: Images are center-cropped and resized to 64 × 64 using bicubic interpolation.
- III. DATASET: Pixel values are normalized to the [−1, 1] range per channel.
- III. DATASET: Five fixed benchmark samples enable direct before-and-after comparisons across training checkpoints.The tracked indices are 0, 2501, 8000, 15000, and 31337.
IV. METHODOLOGY
The methodology uses a sequential pipeline that extracts structural priors from a pretrained generative model and supplies them to conditional diffusion restoration paths.
- IV. METHODOLOGY: The framework sequentially extracts robust structural priors from a pretrained generative model and uses them in conditional diffusion restoration paths.
- IV. METHODOLOGY: The WGAN-GP generator maps z ∼ N(0, I128) to 64 × 64 RGB images through five transposed-convolution, BatchNorm, and ReLU blocks.Spatial resolution increases from 4 × 4 to 64 × 64, followed by a tanh output.
- IV. METHODOLOGY: The critic uses strided convolutions with spectral normalization, instance normalization, and LeakyReLU activations.
- IV. METHODOLOGY: The generator is trained with a WGAN-GP critic objective while an EMA copy is maintained for lower inference-time parameter variance.The generator updates every n_critic = 5 critic updates.
2) Phase 2 : GAN-Guided Conditional U-Net and Diffusion Optimization:
The diffusion stage conditions U-Net restoration on a fixed WGAN-GP feature map through cross-attention and performs deterministic DDIM inference from task-specific initialization.
- 2) Phase 2 : GAN-Guided Conditional U-Net and Diffusion Optimization:: The diffusion U-Nets use FiLM-style timestep conditioning, bottleneck self-attention, and cross-attention to a fixed WGAN-GP feature map.
- 2) Phase 2 : GAN-Guided Conditional U-Net and Diffusion Optimization:: The feature map g is computed once for each sampling chain and reused during the reverse process.
- 2) Phase 2 : GAN-Guided Conditional U-Net and Diffusion Optimization:: Optimization uses a cosine noise schedule with T = 1000.
- 2) Phase 2 : GAN-Guided Conditional U-Net and Diffusion Optimization:: Training samples a timestep t and Gaussian noise ε, then constructs x_t through the forward diffusion process.
- 2) Phase 2 : GAN-Guided Conditional U-Net and Diffusion Optimization:: Cross-attention forms queries from x_t and keys and values from g.
- 2) Phase 2 : GAN-Guided Conditional U-Net and Diffusion Optimization:: DDIM inference initializes x_tstart from corrupted-input injection for denoising or noise for super-resolution, then iteratively predicts x̂_0 and updates deterministically.
B. Experimental Setup
The restoration pipeline was evaluated using stable configurations and optimized hyperparameters, with training settings documented in Table II.
- The implementation was evaluated across stable configurations and optimized hyperparameters to address previously identified pipeline instabilities.
- Table II lists the training hyperparameters used for the restoration pipeline.
C. System Architecture
The system combines a pretrained generator prior with conditional diffusion restoration and evaluates outputs using PSNR and SSIM. GAN-derived features are used as a fixed conditioning signal during sampling.
- System Architecture: A pretrained generator acts as a reusable prior in the hybrid restoration architecture.
- System Architecture: The system architecture is presented in Figure 4.
- Evaluation Metrics: PSNR measures reconstruction quality through the ratio between maximum image power and corrupting noise power.
- Evaluation Metrics: MAXI denotes the maximum possible pixel value, while MSE is the mean squared error between clean ground truth and evaluated output.
- Evaluation Metrics: SSIM evaluates image-quality changes by combining luminance, contrast, and structural terms.
- Evaluation Metrics: The metrics compare five tracked faces’ clean ground truth with degraded inputs and final diffusion-restored outputs.
A. Training Dynamics
Training curves indicate stable adversarial and diffusion optimization, while restoration metrics improve across the tracked CelebA faces for both denoising and super-resolution.
- Training Dynamics: The critic loss increased from −23.99 at epoch 1 to −4.71 at epoch 50, while generator loss decreased from 36.65 to 18.44.
- Training Dynamics: Denoising U-Net loss decreased from 0.079 to 0.037 in 30 epochs, while super-resolution U-Net loss decreased from 0.060 to 0.014 in 25 epochs.
- Training Dynamics: Generator validation showed no mode collapse across ten WGAN-GP samples, with sample standard deviation increasing from 0.430 at epoch 10 to 0.476 at epoch 50.
- Restoration Quality: PSNR and SSIM were averaged across five tracked faces before and after diffusion restoration for both tasks.
- Restoration Quality: Every evaluated face exhibited a PSNR gain, with no degenerate or negative improvement case observed.
- Restoration Quality: Super-resolution had higher PSNR and SSIM values before and after restoration because 2× bicubic downsampling is milder than Gaussian noise with σ = 0.15.
C. Qualitative Observations
Visual examples show that GAN-guided diffusion removes visible denoising noise while retaining identity features and synthesizes realistic super-resolution details. The discussion attributes consistent restoration to fixed GAN conditioning during DDIM sampling and stabilized adversarial learning rates.
- The denoising outputs eliminate visible speckle noise while retaining eyes, hair contours, and mouth-region identity features.
- The super-resolution outputs synthesize realistic high-frequency details, including fine hair strands.
- The GAN conditioning signal g is computed once before the reverse loop and held constant throughout all 50 DDIM steps.Resampling g at each step would make the trajectory inconsistent by providing contradictory semantic instructions.
- The corrected DDIM initialization and fixed conditioning support the GAN-guided diffusion restoration shown for Gaussian denoising and 2× super-resolution.The figures compare clean references, degraded inputs, and restored outputs for five tracked faces.
- Returning to the symmetric 2 × 10−4/2 × 10−4 critic-generator learning-rate ratio resolves the instability without architectural modification.The paper attributes the instability to an overly fast critic whose gradients can become meaningless or saturated.
VII. LIMITATIONS
The evaluation is restricted to five tracked CelebA faces and one Gaussian-noise level, while omitting final LPIPS/FID reporting and cross-attention ablations. The authors identify larger held-out evaluations, perceptual metrics, broader corruption testing, and conditioning ablations as needed follow-up work.
- Evaluation is limited to five tracked faces, preventing assessment of restoration quality across a larger population.The authors suggest testing hundreds or thousands of faces in a larger test set.
- The final pipeline reports no LPIPS or FID scores, despite their closer alignment with human perception of face realism than PSNR/SSIM.
- Denoising robustness is tested at only one fixed Gaussian-noise level, σ = 0.15.
- No ablation study isolates the individual effect of GAN cross-attention conditioning.