Learnable Sparsity for Vision Generative Models
Authors: Yang Zhang
Er Jin
Wenzhong Liang
Yanfei Dong
Ashkan Khakzar
Philip Torr
Johannes Stegmaier
Kenji Kawaguchi
Venue: ICLR 2026 Poster (2026)
Keywords: efficiency
diffusion model
pruning
flow matching
Primary area: applications to computer vision, audio, language, and other modalities
TL;DR: We propose a memory-efficient end-to-end pruning framework that removes 20% of SDXL and FLUX parameters in 10 A100 hours.

Abstract: Generative models have achieved impressive advancements in various vision tasks. However, these gains often rely on increasing model size, which raises computational complexity and memory demands. The increased computational demand poses challenges for deployment, elevates inference costs, and impacts the environment. While some studies have explored pruning techniques to improve the memory efficiency of diffusion models, most existing methods require extensive retraining to maintain model performance. Retraining a large model is extremely costly and resource-intensive, which limits the practicality of pruning methods. In this work, we achieve low-cost pruning by proposing a general pruning framework for vision generative models that learns a differentiable mask to sparsify the model. To learn a mask that minimally deteriorates the model, we design a novel end-to-end pruning objective that spans the entire generation process over all steps. Since end-to-end pruning is memory-intensive, we further design a time step gradient checkpointing technique for the end-to-end pruning, a technique that significantly reduces memory usage during optimization, enabling end-to-end pruning within a limited memory budget. Results on the state-of-the-art U-Net diffusion models Stable Diffusion XL (SDXL) and DiT flow models (FLUX) show that our method efficiently prunes 20% of parameters in just 10 A100 GPU hours, outperforming previous pruning approaches.


Reviews:
[{"review_id": "bfw0mVzQyM", "reviewer": "Reviewer_kteZ", "summary": "The paper proposes EcoDiff, a learnable sparsity framework for large vision generative models (diffusion and flow-based, e.g., SDXL, SD2, FLUX-dev, FLUX-schnell). The key idea is to learn structural neuron masks (for attention and FFN layers) end-to-end over the full denoising or generation trajectory, rather than relying on heuristic or per-step pruning criteria. Conceptually, EcoDiff can be viewed as an adaptation of differentiable mask learning approaches such as MaskLLM, extended to the diffusion domain. Like those methods, it learns continuous mask parameters using a hard-discrete (L₀-style) relaxation, but it departs in several crucial ways: it optimizes masks over the entire diffusion trajectory using a latent reconstruction objective, incorporates time-step gradient checkpointing to efficiently backpropagate through long Markov chains, and targets spatial–temporal architectures (U-Net or DiT) instead of token-based transformers. This design makes it practical for large-scale vision generative models.\n\nEmpirically, EcoDiff prunes about 20% of parameters in SDXL and FLUX-family models using only ~10 A100 GPU hours and ~100 calibration text prompts, while maintaining image quality close to the original models and outperforming prior pruning baselines such as DiffPruning, BK-SDM, and FLUX-Lite under comparable or lower compute budgets. Overall, the paper presents a general, compute-efficient pruning pipeline for state-of-the-art vision generative models, bridging ideas from learnable sparsity in language models to the more complex setting of diffusion and flow-based architectures.", "questions": "On SSIM:\n\nSection 6.2 reports that SSIM between original and pruned models is consistently below 0.65 and interprets this as EcoDiff producing high-quality images without strictly mimicking the teacher. However, the paper does not provide a detailed analysis of how the outputs change. \nGiven that SSIM is relatively low while FID and CLIP remain strong, could the authors examine more systematically whether certain prompt categories or visual structures are more vulnerable to pruning (for example, text rendering, small details, or crowded scenes)? Additionally, has EcoDiff been observed to alter the style or bias of generated images - for instance, shifting color palettes, composition tendencies, or object frequencies, as sparsity increases? This could partially explain SSIM decrease. \n\nOn missing baselines and comparative scope:\n\nRecent works such as LD-Pruner (Castells et al., 2024) and Efficient Pruning of Text-to-Image Models (Ramesh & Zhao, 2024) have explored structural pruning for latent and text-to-image diffusion models, demonstrating strong trade-offs between sparsity, image quality, and compute efficiency. The paper briefly mentions LD-Pruner in the discussion, but does not address Efficient Pruning of Text-to-Image Models, which reports competitive results using simple baselines such as magnitude and WANDA pruning. Could the authors clarify why these methods were not included in the experimental comparison or discussion, and whether their approaches could be evaluated or extended within the proposed EcoDiff framework?\n\nMoreover, Efficient Pruning of Text-to-Image Models identifies a pruning configuration that not only preserves but in some cases improves generative performance after pruning (at 33.5% sparsity). In contrast, EcoDiff achieves only minimal but consistent degradation relative to the original models. It would be valuable for the authors to discuss or empirically compare these differences, ideally by reproducing the Efficient Pruning method on the same models and datasets used in this paper to contextualize EcoDiff’s trade-offs.\n\nOn the scope of pruning:\n\nThe Efficient Pruning of Text-to-Image Models paper distinguishes between pruning the text encoder (e.g., CLIP/T5) and the image generator (U-Net), whereas EcoDiff focuses solely on pruning the visual backbone (U-Net or DiT) and leaves the text encoder untouched. Could the authors clarify whether EcoDiff could, in principle, be extended to prune the text encoder as well, and what motivated the decision to exclude this component from the current framework?\n\nOn inference-time speedups in real deployments:\n\nCould the authors provide more detailed measurements of inference latency, throughput (images/sec), and memory usage for pruned vs original vs distilled+pruned models, on a standard GPU at typical resolutions ? This would solidify the practical impact of the method.\n\nReferences: \n\nCastells, M., Avraham, T., Ghosh, E., & Hoffmann, J. (2024). LD-Pruner: Efficient Pruning of Latent Diffusion Models using Task-Agnostic Insights. CVPR Workshops. arXiv:2404.11936\n\nRamesh, S. N., & Zhao, Z. (2024). Efficient Pruning of Text-to-Image Models: Insights from Pruning Stable Diffusion. arXiv:2409.02496\n\nFang, G., et al. (2023). MaskLLM: Learnable Semi-Structured Sparsity for Large Language Models. arXiv preprint. arXiv:2305.18191", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 4, "contribution": 2, "strengths": "General framework across architectures and paradigms\n\nEcoDiff is demonstrated on both U-Net diffusion models (SD2, SDXL) and DiT-based flow models (FLUX-dev, FLUX-schnell), with a unified mask-learning approach. This cross-architecture applicability is a strong plus.\n\nMemory-efficient optimization via time-step checkpointing\n\nThe proposed time-step gradient checkpointing significantly reduces memory usage for backprop through the full trajectory, making mask learning feasible for large models on a single 80GB GPU. The complexity analysis and empirical VRAM/runtime plots support this.\nStrong empirical evaluation and ablations\nThe paper includes experiments on multiple large models (SD2, SDXL, FLUX-dev, FLUX-schnell) and standard text-to-image benchmarks (COCO, Flickr30k), using FID, CLIP scores, and SSIM.\n\nCompute efficiency and practicality\n\nAchieving ~20% sparsity with 10 A100 GPU hours and only ~100 calibration prompts is a compelling practical story, especially compared to much more expensive baselines. This positions EcoDiff as something people might actually adopt.\n\nClear visualizations and qualitative analysis:\n\nThe paper presents numerous visual comparisons and ablation figures that effectively illustrate the visual impact of pruning across sparsity levels, model architectures, and prompt types, making the technical claims intuitive and easy to follow.", "weaknesses": "Missing baselines \n\nThe paper omits comparison with recent pruning methods such as LD-Pruner (Castells et al., 2024) and Efficient Pruning of Text-to-Image Models (Ramesh & Zhao, 2024). LD-Pruner proposes its own task-agnostic structured pruning strategy for latent diffusion models, while Efficient Pruning reports strong results using simple baselines like magnitude and WANDA pruning. Including both would better contextualize EcoDiff’s performance relative to recent structured and baseline pruning approaches.\n\nPerformance comparison gap \n\nEfficient Pruning of Text-to-Image Models achieves performance gains at around 33% sparsity, while EcoDiff shows slight but consistent degradation. The paper would benefit from discussing or empirically comparing these contrasting outcomes.\n\nCompute comparisons could be more rigorously normalized\n\nWhile the 10 A100 GPU hours vs 1120 H200 GPU hours comparison is striking, the paper doesn’t fully normalize across:\n\nHardware generations (A100 vs H200/H100);\n\nFLOPs per update\n\nThere are also small inconsistencies between main text (A100) and appendices (H100 mentioned for some configs), which have to be cleared out to get the full picture.\n\nInference-time measurements are underdeveloped\n\nBeyond parameter counts, there is limited wall-clock inference evaluation (latency, throughput, GFLOPS) for pruned vs original vs distilled+pruned models. For a practical efficiency story, this would help a lot.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "sxGohKEQAY", "reviewer": "Reviewer_7mkK", "summary": "This paper introduces EcoDiff, an end-to-end structural pruning framework for vision generative models. The key idea is to learn a differentiable neuron mask that is applied across all denoising steps, enabling efficient pruning without extensive retraining. The authors propose a time-step gradient checkpointing technique to reduce memory usage, making it feasible to prune large models like SDXL and FLUX on a single GPU. The method is evaluated on multiple SOTA models and demonstrates competitive performance at 20% sparsity with only fewer GPU hours and samples. The framework also supports lightweight post-pruning adaptation via LoRA or full fine-tuning.", "questions": "(1) The method applies the same mask across all denoising steps. Is there an empirical or theoretical support for this choice? Have authors experimented with time-varying masks? If so, how does it affect performance and memory? Could a mask scheduling mechanism (e.g., step-wise or block-wise masks) further improve performance or sparsity tolerance?\n\n(2) $L_0$ -approximation encourages sparsity uniformly. However, not all layers or blocks are equally redundant. At high sparsity targets, does this uniform pressure lead to a catastrophic under-allocation of parameters to critical layers? Do authors observe a highly uneven distribution of remaining parameters across layers in aggressively pruned models, and if so, does this misalignment with layer sensitivity explain the performance drop?\n\n(3) The performance degradation at high sparsity is presented as a given. Is this primarily due to the irreversible removal of critical structural components or is it a functional optimization issue where the remaining sub-network has the capacity, but the post-pruning retraining fails to recover? Have authors conducted analysis to distinguish between these two causes?\n\n(4) LoRA and full fine-tuning is applied for recovery. The significant gap at 50% sparsity, even with full fine-tuning, suggests a fundamental limit. Can authors quantify what is being lost? For instance, authors could identify which high-level features become incompressible beyond a certain sparsity threshold?\n\n(5) Pruned models may be more vulnerable to out-of-distribution prompts or adversarial attacks due to reduced capacity. Have authors tested the pruned models on OOD datasets or adversarially crafted prompts? Does the mask learning process inadvertently amplify bias or reduce robustness?\n\n(6) This paper overlooks several relevant studies in efficient diffusion model training. For example, it does not cite prior works [1-2] from a data perspective. Additionally, the paper omits some studies on diffusion training acceleration [3-6].\n\nReference:\n\n[1] Z. Qin, K. Wang, Z. Zheng, J. Gu, X. Peng, Z. Xu, D. Zhou, L. Shang, B. Sun, X. Xie, and Y. You, “Infobatch: Lossless training speed up by unbiased dynamic data pruning,” In ICLR, 2024.\n\n[2] Y. Li, Y. Zhang, S. Liu, and X. Lin, “Pruning then reweighting: Towards data-efficient training of diffusion models,” In ICASSP, 2025.\n\n[3] H. Zheng, W. Nie, A. Vahdat, and A. Anandkumar, “Fast training of diffusion models with masked transformers,” In TMLR, 2024.\n\n[4] Z. Ding, M. Zhang, J. Wu, and Z. Tu, “Patched denoising diffusion models for high-resolution image synthesis,” In ICLR, 2024.\n\n[5] Z. Wu, P. Zhou, K. Kawaguchi, and H. Zhang, “Fast diffusion model,” arXiv preprint arXiv:2306.06991, 2023.\n\n[6] T. Hang, S. Gu, C. Li, J. Bao, D. Chen, H. Hu, X. Geng, and B. Guo, “Efficient diffusion training via min-snr weighting strategy,” In ICCV, 2023.", "limitations": "", "rating": 4, "confidence": 4, "soundness": 2, "presentation": 2, "contribution": 2, "strengths": "(1) The paper is technically sound, with various empirical results across multiple models (SDXL, FLUX) and metrics (FID, CLIP, SSIM). The writing is clear and well-structured.\n\n(2) The work addresses a critical problem—efficient deployment of large generative models—and offers a practical solution with low computational overhead. The method seems model-agnostic and compatible with existing acceleration techniques.", "weaknesses": "(1) The approximation of the $L_0$ regularization to  $L_1$ (Appendix A) is heuristic and lacks theoretical guarantees. More rigorous analysis would strengthen the method.\n\n(2) The method is only validated on image generation tasks. Its applicability to video generation remains unverified. Moreover, the assumption that all denoising steps share the same mask may not hold for models with highly temporal dynamics.\n\n(3) The use of synthetic data from GCC3M for retraining, rather than original training data, may limit the validity of the post-pruning recovery claims.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "CVsuYkjMPY", "reviewer": "Reviewer_2dT4", "summary": "This paper explores pruning of visual generative models, including LDM and Flux. Existing pruning strategies are coarse-grained and fail to balance performance and sparsity. Therefore, in this paper, the authors propose low-cost pruning by proposing a general pruning framework for vision generative models that learns a differentiable mask to sparsify the model. Experiments demonstrate the effectiveness of the proposed method.", "questions": "Please refer to Weaknesses.", "limitations": "", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1. Pruning generative models facilitates their deployment and application.\n\n2. The paper is presented intuitively and clearly.\n\n3. The method is simple and effective.", "weaknesses": "1. The mask isn't actually composed of only 0s and 1s; its value can range from [0,1]. How should this be handled? This would likely result in a performance penalty.\n\n2. While this paper is very effective in engineering, its overall contribution is incremental. Differentiable masks and gradient checkpointing are readily available techniques.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "TXTN91YUdO", "reviewer": "Reviewer_h5Na", "summary": "This paper introduces **EcoDiff**, an efficient and model-agnostic framework for pruning large vision generative models. Its core innovation is an **end-to-end differentiable masking** scheme that learns to remove redundant parameters by optimizing for final output quality.", "questions": "See the Weaknesses section.", "limitations": "", "rating": 4, "confidence": 3, "soundness": 2, "presentation": 2, "contribution": 2, "strengths": "1. **High Efficiency:** The method achieves up to 20% sparsity on billion-parameter models (SDXL, FLUX) using only **100 samples** and **10 A100 GPU hours**, which is highly impressive given the scale.\n2. **End-to-End Pruning Objective:** The proposed holistic loss jointly optimizes all denoising steps, outperforming traditional per-step pruning methods.\n3. **Practicality:** The design is model-agnostic and potentially applicable to both diffusion and flow-based generative models.", "weaknesses": "### **Theoretical Concerns**\n\nThere is a fundamental error in Equation (12). Contrary to what is presented, the correct form of the L₀ regularization term, as established in *Learning Sparse Neural Networks through L₀ Regularization* (Louizos et al., 2018), should be:\n$$\nL_0(\\lambda)=\\Sigma \\text{Sigmoid}\\left(\\lambda_i-\\alpha\\log\\frac{-\\gamma}{\\zeta}\\right),\n$$\nthis error propagates to the derivation in Appendix A, which is consequently invalid. The intended relationship should instead be:\n$$\nL_0(\\lambda)=\\Sigma \\text{Sigmoid}\\left(\\lambda_i-\\alpha\\log\\frac{-\\gamma}{\\zeta}\\right)\\approx \\|e^{\\lambda}\\|_{L^1}.\n$$\nMoreover, the claim that ||e^lambda|| can be bounded by ||\\lambda|| is mathematically unsound. There is no generic upper bound of ||e^lambda|| in terms of ||\\lambda||, since the former grows exponentially while the latter only grows linearly.\n\nFurthermore, the motivation for using this inaccurate approximation is not justified. It is both clearer and more rigorous to simply use the exact expression: \n$$\nL_0(\\lambda)=\\Sigma \\text{Sigmoid}\\left(\\lambda_i-\\alpha\\log\\frac{-\\gamma}{\\zeta}\\right),\n$$\nwhich is well-founded in the literature and avoids unnecessary ambiguity.\n\n### Comment on experiments\n\n**Experimental Fairness**\n\nThe comparison with FLUX-Lite is problematic in terms of computational fairness: FLUX-Lite used 1,120 H200 GPU hours, while EcoDiff required only 10 A100 hours. Although this highlights efficiency, it remains unclear whether FLUX-Lite could achieve similar performance under a comparable computational budget.\n\n**Ablation studies on loss parameters are missing.**\n\nThe paper lacks ablation studies investigating the impact of key parameters in the proposed loss function (Eq. 13), such as the regularization coefficient `β` and the hard concrete parameters (e.g., δ, γ, ζ). The influence of these choices on the trade-off between sparsity and performance is not analyzed.\n\n### Reading difficulty\n\nThe specific implementation of the learnable mask—such as the hard discrete sampling procedure and the selection of parameters (e.g., δ, γ, ζ)—is only briefly described in the main text. Readers must refer to the appendix to fully understand these mechanisms.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (Poster)