IGSM: Improved Geometric and Sensitivity Matching for Finetuning Pruned Diffusion Models

Abstract: Diffusion models achieve realistic outcomes across a wide range of generative tasks, but their high computational cost remains a major barrier to deployment. Model pruning has emerged as a promising strategy to reduce inference cost and enable lightweight diffusion models. While effective, pruned diffusion models are proned to quality reduction due to limited capacity. A key limitation of current pruning approaches is that pruned models are finetuned using the same objective as the dense model, typically denoising score matching (DSM). Since the dense model is accessible during finetuning, it warrants a more effective approach for knowledge transfer from the dense to the pruned model. Motivated by this aim, we revisit the finetuning stage and propose IGSM (\textbf{I}mproved \textbf{G}eometric and \textbf{S}ensitivity \textbf{M}atching), a general-purpose finetuning framework that introduces a second-order Jacobian projection loss inspired by Finite-Time Lyapunov Exponents (FTLE). IGSM efficiently captures and aligns the geometric and the temporal dynamics of pruned models with their dense teachers using scalable second-order projections. Our approach is architecture-agnostic and applies to both U-Net- and Transformer-based diffusion models. Experiments on CIFAR-10, CelebA, LSUN-Church, and LSUN-Bedroom show that IGSM consistently narrows the performance gap between pruned and dense models, substantially improving sample quality. Code is available on GitHub: https://github.com/FATE4869/IGSM-Official