Pluggable Pruning with Contiguous Layer Distillation for Diffusion Transformers

AI summary

Overview

Research area: Model compression for diffusion transformers in computer vision, specifically structured pruning and knowledge distillation of Multi-Modal Diffusion Transformers (MMDiTs) used for text-to-image generation.

Technical level: Advanced. The paper assumes familiarity with transformer block structure, knowledge distillation, Centered Kernel Alignment (CKA) representation analysis, and least-squares linear probing.

Scope: The paper introduces PPCL, a structured pruning framework that locates contiguous redundant layer intervals in MMDiT models via linear probing and CKA trend analysis, then applies non-sequential inter-layer distillation plus width-wise pruning to produce deeply compressed, pluggable text-to-image models.

What This Paper Is About

State-of-the-art diffusion transformers such as SD3.5, FLUX.1, and Qwen-Image deliver strong image fidelity and text-image alignment but carry 8–20 billion parameters, which makes them expensive to train and deploy. Existing structured pruning methods for these models are limited: they generalize poorly across MMDiT variants, offer little flexibility in which layers can be removed, and rely on a weak understanding of how much interdependence exists between layers. The paper's goal is to compress MMDiT models aggressively while keeping generation quality close to the original, and to make the pruning reversible at inference time without retraining.

Key Contributions

1. Empirical characterization of layer redundancy in MMDiT. The authors run controlled removal experiments on Qwen-Image (a 60-layer MMDiT with 20B parameters), removing 1, 2, or 3 layers in contiguous and non-contiguous configurations, and find that redundancy is depth-wise continuous — contiguous removal consistently outperforms non-contiguous removal.

2. A linear-probe-based redundancy detection strategy. Each teacher layer is approximated by a lightweight residual linear probe initialized by a closed-form least-squares solution and trained with an alignment loss. Layer substitutability is then measured through CKA between teacher layer outputs and surrogate-model outputs, with the first-order difference of that CKA used to find where a redundant interval ends.

3. A plug-and-play depth-wise pruning framework with non-sequential distillation. Instead of the conventional sequential pipeline where early-layer compression errors compound, student layers receive the teacher's immediately preceding depth outputs, so each pruned interval is optimized independently. This permits arbitrary activation or bypassing of specific layers at inference without retraining.

4. Width-wise pruning for stream-level and FFN redundancy, and validation across multiple MMDiT models. Redundant text-stream components are replaced by two lightweight linear projectors, and redundant FFNs in both streams are replaced by compact linear projectors, combining with depth-wise pruning to reach higher compression ratios.

Main Findings

• Contiguous redundancy beats random removal. On Qwen-Image evaluated with LongText-Bench, removing one to three layers has minimal effect on generation quality, and contiguous layer removal consistently outperforms non-contiguous removal, indicating that redundancy is continuous across depth rather than scattered.

• MMDiT has two distinct breadth-wise redundancy sources. CKA heatmaps show the text stream has high cross-layer similarity and low inter-layer variability (large redundancy), while the image stream exhibits a smooth diagonal similarity decay reflecting sequential feature evolution with minimal redundancy. FFNs in both streams are described as significantly over-parameterized.

• Strong compression on Qwen-Image with small degradation. The full PPCL pipeline (fine-tuned, 10B parameters) achieves an average performance drop of 3.29% on Qwen-Image, with inference latency dropping from 2625 ms to 1462 ms and GPU memory falling from 100% to 66.9% of the base model's usage — a reduction from 20B to 10B parameters, a 50% cut.

• Plug-and-play layer swapping works without retraining. Starting from the trained 10B model, the authors construct 12B and 14B variants by replacing selected student layers with the corresponding teacher layers. The 14B variant shows only a 0.42% average performance drop (87.9 DPG, 0.847 GenEval, 0.929 LongText EN) and the 12B variant a 3.03% drop — comparable to the fully fine-tuned 10B model's 3.29%.

• PPCL outperforms TinyFusion and HierarchicalPrune at matched parameter counts. On Qwen-Image at 14B, TinyFusion records an 8.75% average drop and HierarchicalPrune 6.49%, versus 0.42% for PPCL. On FLUX.1-dev at 8B, PPCL records a 4.03% average drop versus 13.80% (TinyFusion) and 13.38% (HierarchicalPrune).

• Chroma1-HD has the lowest average drop but is much slower. Chroma1-HD, at 8.9B, is the best in the FLUX.1-dev group on average performance drop (1.02%), but its inference latency of 1714 ms is more than double the base model's 715 ms.

• Dense2MoE is far off the base model. Reported directly from its original paper, Dense2MoE runs at 312 ms with DPG 73.6 and an average performance drop of 21.52%.

• PPCL can compress an already-pruned model. Applying PPCL to FLUX.1 Lite removes 1.5B parameters (from 8B to 6.5B) with an average performance drop of 0.07%, at 428 ms latency versus FLUX.1 Lite's 572 ms.

• Ablations isolate the value of each component. In Table 2, the baseline (CKA importance selection plus sequential distillation, 25 layers pruned) averages 0.706; adding linear-probing-based interval selection raises it to 0.761; replacing sequential with non-sequential distillation raises it to 0.848 (described in the text as nearly 9 points higher); adding text-stream linear replacement reaches 0.860; adding FFN replacement reaches 0.850; and full-parameter fine-tuning reaches 0.870 versus the original model's 0.894, which the text describes as 2.61% below the original. The ablation text cites a baseline average of 0.731 with an 18.2% decline, while Table 2 lists 0.706 for that row — the two are inconsistent in the paper as presented.

• The CKA first-order-difference rule beats threshold selection. Variants using a simple CKA similarity threshold (LP-a, average 0.718) or a widened interval upper bound (LP-b, average 0.726) both underperform the first-order-difference interval detection (LP, average 0.761).

• The authors state two explicit limitations. The inflection-point detection via first-order difference of CKA similarity lacks rigorous theoretical foundations and is largely a successful engineering heuristic; and INT4 quantization produces unsatisfactory results, presumably because pruning reduces network redundancy and narrows the quantization fault-tolerant space.

Methodology in Plain English

The approach has two pruning stages plus a short fine-tuning pass.

Stage 1.1 — Probing each layer's substitutability. For every MMDiT block in the teacher, the authors build a small residual linear probe and initialize it with the closed-form least-squares solution that best maps the block's input to its output. Because MMDiT has residual connections, each probe is given a residual structure, and each probe is trained only on that layer's own input so its modeling is independent of other layers. A stack of linear maps is itself linear, which is what allows consecutive layers to be modeled as substitutable.

Stage 1.2 — Finding contiguous redundant intervals. On a calibration

Authors’ abstract

Diffusion Transformers (DiTs) have shown exceptional performance in image generation, yet their large parameter counts incur high computational costs, impeding deployment in resource-constrained settings. To address this, we propose Pluggable Pruning with Contiguous Layer Distillation (PPCL), a flexible structured pruning framework specifically designed for DiT architectures. First, we identify redundant layer intervals through a linear probing mechanism combined with the first-order differential trend analysis of similarity metrics. Subsequently, we propose a plug-and-play teacher-student alternating distillation scheme tailored to integrate depth-wise and width-wise pruning within a single training phase. This distillation framework enables flexible knowledge transfer across diverse pruning ratios, eliminating the need for per-configuration retraining. Extensive experiments on multiple Multi-Modal Diffusion Transformer architecture models demonstrate that PPCL achieves a 50\% reduction in parameter count compared to the full model, with less than 3\% degradation in key objective metrics. Notably, our method maintains high-quality image generation capabilities while achieving higher compression ratios, rendering it well-suited for resource-constrained environments. The open-source code, checkpoints for PPCL can be found at the following link: https://github.com/OPPO-Mente-Lab/Qwen-Image-Pruning.