KDC-Diff: A Latent-Aware Diffusion Model with Knowledge Retention for Memory-Efficient Image Generation

TL;DR

Diffusion-based text-to-image models are highly computationally intensive, limiting real-world deployment. KDC-Diff addresses this by introducing a lightweight UNet backbone, a dual-layer knowledge distillation scheme, and latent-space replay-based continual learning to preserve quality with lower resources. The approach achieves strong performance while reducing parameters to $482$M and FLOPs to $228.85$ GMac, delivering competitive FID, CLIP, KID, and LPIPS scores on benchmark datasets. This work enables practical diffusion-based generation in low-resource environments and lays groundwork for edge-ready, scalable diffusion systems.

Abstract

The growing adoption of generative AI in real-world applications has exposed a critical bottleneck in the computational demands of diffusion-based text-to-image models. In this work, we propose KDC-Diff, a novel and scalable generative framework designed to significantly reduce computational overhead while maintaining high performance. At its core, KDC-Diff designs a structurally streamlined U-Net with a dual-layered knowledge distillation strategy to transfer semantic and structural representations from a larger teacher model. Moreover, a latent-space replay-based continual learning mechanism is incorporated to ensure stable generative performance across sequential tasks. Evaluated on benchmark datasets, our model demonstrates strong performance across FID, CLIP, KID, and LPIPS metrics while achieving substantial reductions in parameter count, inference time, and FLOPs. KDC-Diff offers a practical, lightweight, and generalizable solution for deploying diffusion models in low-resource environments, making it well-suited for the next generation of intelligent and resource-aware computing systems.

Figures (5)

• Figure 1: Overview of the KDC-Diff framework: (a) illustrates the memory buffer mechanism used to store latent representations of previously learned classes, enabling efficient replay during training on subsequent classes; (b) depicts the overall training architecture, integrating a streamlined U-Net design with CL strategies and knowledge distillation to ensure both efficiency and knowledge retention.
• Figure 2: Comparison between the baseline U-Net and the optimized KDC-Diff U-Net architecture. Modules outlined in red represent components removed from the teacher model to facilitate a more efficient design, while the highlighted blocks indicate layers eliminated from the original U-Net to derive the lightweight student network.
• Figure 3: Illustration of sequential training in a diffusion model using latent-space replay for continual learning. Latent representations from multiple classes are stored in a fixed-size buffer and replayed during training on new tasks. This mitigates catastrophic forgetting by preserving class diversity, enabling efficient memory use and stable generation quality across tasks.
• Figure 4: Visual Comparison of Original and Generated Images for Different Classes in Both Datasets. Our proposed methods produced high-quality, color-enhanced images from scratch under computational constraints.
• Figure 5: Attention map visualization during image generation. The grayscale heatmap highlights spatial regions with higher attention weights, indicating where the model focused most while generating the corresponding image. These heatmaps captured the semantic structure so well from the reference images.