HierarchicalPrune: Position-Aware Compression for Large-Scale Diffusion Models

TL;DR

Large diffusion models with 8–11B parameters demand excessive memory and compute for deployment on constrained hardware. The paper proposes HierarchicalPrune, a hierarchy-aware compression framework that combines Hierarchical Position Pruning, Positional Weight Preservation, and Sensitivity-Guided Distillation, plus INT4 weight quantisation, to prune inter-block and intra-block structure while preserving semantic content and texture. It demonstrates 77.5–80.4% memory footprint reduction and 27.9–38.0% latency improvements on SD3.5 Large Turbo and FLUX.1 Schnell, with minimal quality loss (GenEval down to ~2.6% and HPSv2 down to ~7%), corroborated by a human user study. The work shows that leveraging two-fold hierarchy and targeted distillation enables deployable, high-quality SOTA diffusion models on resource-limited devices.

Abstract

State-of-the-art text-to-image diffusion models (DMs) achieve remarkable quality, yet their massive parameter scale (8-11B) poses significant challenges for inferences on resource-constrained devices. In this paper, we present HierarchicalPrune, a novel compression framework grounded in a key observation: DM blocks exhibit distinct functional hierarchies, where early blocks establish semantic structures while later blocks handle texture refinements. HierarchicalPrune synergistically combines three techniques: (1) Hierarchical Position Pruning, which identifies and removes less essential later blocks based on position hierarchy; (2) Positional Weight Preservation, which systematically protects early model portions that are essential for semantic structural integrity; and (3) Sensitivity-Guided Distillation, which adjusts knowledge-transfer intensity based on our discovery of block-wise sensitivity variations. As a result, our framework brings billion-scale diffusion models into a range more suitable for on-device inference, while preserving the quality of the output images. Specifically, combined with INT4 weight quantisation, HierarchicalPrune achieves 77.5-80.4% memory footprint reduction (e.g., from 15.8 GB to 3.2 GB) and 27.9-38.0% latency reduction, measured on server and consumer grade GPUs, with the minimum drop of 2.6% in GenEval score and 7% in HPSv2 score compared to the original model. Finally, our comprehensive user study with 85 participants demonstrates that HierarchicalPrune maintains perceptual quality comparable to the original model while significantly outperforming prior works.

Figures (13)

• Figure 1: HierarchicalPrune achieves 79.5% memory reduction (left) while maintaining image quality. User study with 85 participants (right) demonstrates minimal quality drop (4.8-5.3%) with 95% confidence intervals, contrary to the excessive degradation (11.1-52.2%) of prior methods.
• Figure 2: High-resolution image samples generated by compressed model using HierarchicalPrune, showcasing its superior visual quality across various visual styles, precisely following text prompts, and preserving the ability to draw typography.
• Figure 3: HierarchicalPrune's compression framework leverages MMDiT's two-fold hierarchy (inter-block: early blocks establish semantics, later blocks refine; intra-block: varying subcomponent importance). It comprises (1) Hierarchical Position Pruning (HPP), maintaining early blocks while pruning later ones, (2) Positional Weight Preservation (PWP), freezing critical early blocks during distillation, and (3) Sensitivity-Guided Distillation (SGDistill), applying inverse weights—minimal updates to sensitive blocks and subcomponents. The resulting framework enables effective compression while preserving model capabilities.
• Figure 4: Fine-grained contribution analysis of SD3.5 Large Turbo on the HPSv2 dataset by removing either an entire MMDiT block (a), following prior depth pruning approaches lee2024koalakim2024bksdmfang2024tinyfusion, or an intra-block subcomponent (b, c and see Fig. \ref{['fig:contribution_analysis_full']} in Appendix \ref{['app:additional_analysis']} for full set of analysis). We report the performance drop compared to the original model. The discrepancy in performance drop patterns reveals the different patterns of importance of each subcomponent.
• Figure 5: Impact of removing MMDiT blocks at different positions. Compared to original outputs (f), removing earlier layers leads to high impact on image structure (a-c), whereas removing later blocks affects mainly fine details (d, e).