Uni-DAD: Unified Distillation and Adaptation of Diffusion Models for Few-step Few-shot Image Generation

1. Problem Statement

Diffusion Models (DMs) have achieved state-of-the-art (SoTA) performance in image generation but suffer from two critical limitations when applied to new domains:

1. High Computational Cost: Standard DMs require iterative denoising over hundreds of time-steps (e.g., $T \approx 1000$), making inference slow and unsuitable for real-time applications.
2. Data Scarcity in Adaptation: Adapting DMs to novel domains (Few-Shot Image Generation - FSIG) or specific subjects (Subject-Driven Personalization - SDP) using only a handful of images ($N \le 10$) often leads to overfitting, loss of diversity, or "mode collapse."

Current Limitations: Existing solutions typically rely on two-stage pipelines:

• Distill-then-Adapt: Distills a fast student first, then adapts it. This often leads to "over-smoothed" outputs because the student saturates its adaptation capacity, and fine-tuning under the original diffusion loss negates distillation benefits.
• Adapt-then-Distill: Adapts the teacher first, then distills. This risks overfitting to the few-shot target data during the adaptation phase and loses transferable source information during distillation.

Both approaches are complex, prone to quality degradation, and fail to jointly optimize for speed, diversity, and target realism.

2. Methodology: Uni-DAD

The authors propose Uni-DAD, a single-stage pipeline that unifies diffusion model distillation and adaptation. It compresses a high-quality, slow teacher DM into a fast, few-step student while simultaneously adapting it to a few-shot target domain.

Core Architecture

The method couples two complementary training signals to optimize the student generator ($G$):

A. Dual-Domain Distribution Matching Distillation (DMD)

• Concept: Instead of matching the student only to the source teacher ($\epsilon_{src}$), Uni-DAD matches the student to a weighted combination of the source teacher and an optional online target teacher ($\epsilon_{trg}$).
• Mechanism:
  • $\epsilon_{src}$: A frozen source-domain teacher trained on large-scale data (e.g., FFHQ). It preserves diverse, transferable knowledge (e.g., pose, background, general structure).
  • $\epsilon_{trg}$: An online teacher initialized from $\epsilon_{src}$ and fine-tuned on the few-shot target set ($Y$). It guides the student toward the specific structural distribution of the target domain.
  • Objective: The gradient update is a weighted sum:
    $$\nabla_\theta L_{DMD}^{trg+src} = (1-a)\nabla_\theta L_{DMD}^{src} + a\nabla_\theta L_{DMD}^{trg}$$
    where $a \in [0, 1]$ balances the influence of the target domain. This guides the student to a "common area" between the two distributions, preventing overfitting to the small target set while ensuring structural adaptation.

B. Multi-Head Generative Adversarial Network (GAN) Loss

• Concept: To enforce sharp realism and stabilize training in the few-shot regime, a multi-head GAN is employed.
• Mechanism:
  • Discriminator ($D$): Reuses the encoder and middle blocks of a "fake teacher" ($\epsilon_{fk}$) which tracks the student's evolving distribution. Linear heads are attached to every block of the feature extractor.
  • Multi-Scale Discrimination: This allows the discriminator to distinguish real target images from student generations at multiple feature scales (local and global).
  • Benefit: This multi-scale approach significantly mitigates overfitting and mode collapse, which are common when training GANs on very small datasets ($N \le 10$).

C. Training Loop
The training alternates between three updates per iteration:

1. Student ($G$): Optimized via the Dual-Domain DMD loss and the GAN generator loss.
2. Fake Teacher & Discriminator ($\epsilon_{fk}, D$): $\epsilon_{fk}$ is updated to match the student's output (MSE), and $D$ is updated to distinguish real target images from student generations.
3. Target Teacher ($\epsilon_{trg}$): (Optional) Updated via MSE on the target set to better model the target distribution.

3. Key Contributions

1. First Single-Stage Pipeline: Uni-DAD is the first method to jointly distill and adapt diffusion models in a single training loop, eliminating the need for sequential two-stage pipelines.
2. Dual-Domain DMD: Introduces a novel objective that leverages both source and target teachers to preserve source diversity while adapting to structurally distant target domains.
3. Multi-Head GAN for Few-Shot: Proposes a multi-scale discriminator architecture specifically designed to stabilize training and reduce overfitting in extreme few-shot regimes ($N \le 10$).
4. Checkpoint Agnosticism: The framework is flexible; it can initialize the student with a pre-distilled model or use a pre-adapted model as the target teacher without retraining.

4. Experimental Results

The method was evaluated on two comprehensive benchmarks:

• FSIG (Few-Shot Image Generation): Using guided DDPM on FFHQ, adapting to 10-shot sets of Babies, Sunglasses, MetFaces, and Cats.
• SDP (Subject-Driven Personalization): Using Stable Diffusion v1.5 on the DreamBooth benchmark (30 subjects, 4-6 images each).

Key Findings:

• Quality vs. Speed: Uni-DAD achieves better or comparable quality (FID, CLIP-I, DINO scores) to SoTA adaptation methods while using substantially fewer sampling steps (1 to 4 steps vs. 25–1000 steps).
• Outperforming Two-Stage Pipelines: Uni-DAD consistently surpasses Distill-then-Adapt and Adapt-then-Distill baselines in both image quality and diversity.
  • Distill-then-Adapt baselines produced over-smoothed, low-detail images.
  • Adapt-then-Distill baselines suffered from overfitting and loss of source diversity.
• Diversity Preservation: Unlike many distilled models that suffer from diversity collapse, Uni-DAD maintains high intra-domain diversity (measured by Intra-LPIPS) even with 1-step sampling.
• Computational Efficiency:
  • Inference: Reduces sampling time from 35–63 minutes (for 5k images) to **4.2 minutes**.
  • Training: Requires less training compute (2.2–2.8 GPU hours) compared to two-stage pipelines (3+ GPU hours).

5. Significance and Impact

Uni-DAD represents a paradigm shift in deploying diffusion models for real-world, resource-constrained applications.

• Real-Time Personalization: By enabling high-quality, diverse generation in 1–4 steps, it makes personalized image generation feasible for real-time applications (e.g., mobile apps, interactive design tools).
• Robustness to Domain Shift: The dual-domain approach allows models to adapt to structurally distant domains (e.g., faces to cats) without losing the generative capabilities learned from the massive source dataset.
• Simplicity: By unifying distillation and adaptation, it simplifies the engineering pipeline, removing the complexity of managing separate stages and reducing the risk of error propagation between stages.

In conclusion, Uni-DAD demonstrates that distillation and adaptation are not mutually exclusive stages but can be optimized simultaneously to achieve fast, high-fidelity, and diverse image generation under data-scarce conditions.