Lightweight Stable Diffusion via StableKOT: Knowledge Distillation Meets Optimal Transport

Abstract

Stable Diffusion models deliver photorealistic image synthesis but face computational bottlenecks that hinder deployment in resource-constrained environments. However, most existing distillation methods for diffusion models rely on point-wise feature matching, which often fails to preserve global semantic structure, or require additional and costly teacher finetuning. To address these limitations, we propose StableKOT — a novel knowledge distillation framework that leverages Optimal Transport (OT) theory. Still, direct OT alignment of high-dimensional feature maps is computationally prohibitive and sensitive to spatial noise. To overcome this issues, we transform teacher-student knowledge transfer into a distribution matching problem, applying OT to max-pooled features augmented with positional embeddings across U-Net layers. This captures geometric relationships in latent space while reducing computational overhead. Empirically, our method reduces parameters by 32.6% and accelerates inference 1.5×, while maintaining generative fidelity.