Hierarchical Diffusion Class Pruning

• Hierarchical Diffusion Class Pruning is a family of methods that reduces computational cost in diffusion models by using hierarchical strategies to prune label candidates, training samples, or diffusion branches.
• The approach eliminates redundant evaluations by narrowing the search space via explicit label pruning, two-stage data filtration, or branched reverse diffusion routing.
• Empirical evaluations demonstrate improved runtimes and convergence rates, with significant speed-ups and better FID scores on datasets like ImageNet and remote sensing collections.

Searching arXiv for the specified papers and closely related diffusion hierarchy/pruning work. Hierarchical diffusion class pruning is an umbrella designation for diffusion-model methods that reduce the effective search space through coarse-to-fine structure. In the literature, the most direct instance is the "Hierarchical Diffusion Classifier" (HDC), which prunes label candidates along a hierarchy during zero-shot classification. Closely related work applies a two-stage, hierarchical pruning pipeline to diffusion foundation-model training data in remote sensing, and another line of work organizes class-conditional reverse diffusion through hierarchically branched pathways rather than flat class conditioning (Shanbhag et al., 2024, Wei et al., 29 Dec 2025, Tseng et al., 2022). This suggests that the topic is best understood as a family of hierarchy-aware reduction strategies spanning inference-time class evaluation, training-time subset selection, and branch-specific routing.

1. Scope and conceptual boundaries

The central motivation is that flat diffusion pipelines often scale poorly. In zero-shot diffusion classification, standard classifiers compare an input against all class labels and use multiple diffusion samples for each comparison. In large label spaces, this makes inference expensive. In diffusion foundation-model training, very large remote sensing corpora contain redundancy, noise, class imbalance, and scene homogeneity, so indiscriminate scaling can slow convergence and degrade generation quality. Hierarchical methods intervene by discarding, narrowing, or routing computation before the full diffusion procedure is applied.

Within this landscape, three method families should be distinguished. First, explicit hierarchical class pruning eliminates irrelevant branches of a label tree and evaluates only the surviving leaves. Second, hierarchical data pruning removes low-information samples and then enforces scene-aware coverage within the retained training set. Third, hierarchically branched diffusion models partition reverse diffusion over a class hierarchy; they narrow and route computation, but the source paper does not describe this as pruning. A common misconception is to treat these formulations as interchangeable. They are related in spirit, but they act on different objects: class candidates, training samples, or reverse-diffusion branches (Shanbhag et al., 2024, Wei et al., 29 Dec 2025, Tseng et al., 2022).

2. Diffusion classification and hierarchical pruning of label candidates

HDC begins from the standard diffusion-classifier formulation. Given a conditional diffusion model $p_\theta(\mathbf{x}\mid \mathbf{c}_i)$ and labels $\{\mathbf{c}_i\}_{i=1}^{N_C}$, the posterior is written via Bayes’ rule as

$$p_\theta(\mathbf{c}_i \mid \mathbf{x}) = \frac{p(\mathbf{c}_i)\, p_\theta(\mathbf{x} \mid \mathbf{c}_i)}{\sum_{j=1}^{N_C} p(\mathbf{c}_j)\, p_\theta(\mathbf{x} \mid \mathbf{c}_j)}.$$

Under a uniform prior, this simplifies to

$$p_\theta(\mathbf{c}_i \mid \mathbf{x}) = \frac{p_\theta(\mathbf{x} \mid \mathbf{c}_i)}{\sum_{j=1}^{N_C} p_\theta(\mathbf{x} \mid \mathbf{c}_j)}.$$

The likelihood is approximated by the diffusion model’s noise prediction error using the $\varepsilon$-prediction loss

$$d(\varepsilon, \mathbf{x}, \mathbf{c}) = \|\varepsilon - \varepsilon_\theta(\mathbf{x}, \mathbf{c})\|^2,$$

and the paper also uses the pairwise-difference form

$$p_\theta(\mathbf{c}_i\mid\mathbf{x}) \approx \frac{1}{\sum_{j=1}^{N_C}\exp\left\{\mathbb{E}_{t,\varepsilon}\Delta(\varepsilon,\mathbf{x}_t,\mathbf{c}_i,\mathbf{c}_j)\right\}},$$

with

$$\Delta(\varepsilon,\mathbf{x}_t,\mathbf{c}_i,\mathbf{c}_j) = d(\varepsilon,\mathbf{x}_t,\mathbf{c}_i)-d(\varepsilon,\mathbf{x}_t,\mathbf{c}_j).$$

The resulting bottleneck is the need to compare against all $N_C$ classes with multiple diffusion samples, yielding a cost that grows roughly as $O(N_C \cdot M)$ per image (Shanbhag et al., 2024).

HDC addresses this bottleneck by replacing flat exhaustive scoring with coarse-to-fine traversal of a hierarchy

$\{\mathbf{c}_i\}_{i=1}^{N_C}$0

where $\{\mathbf{c}_i\}_{i=1}^{N_C}$1 is the set of nodes, $\{\mathbf{c}_i\}_{i=1}^{N_C}$2 the set of edges, $\{\mathbf{c}_i\}_{i=1}^{N_C}$3 the tree depth, and each node $\{\mathbf{c}_i\}_{i=1}^{N_C}$4 has a label $\{\mathbf{c}_i\}_{i=1}^{N_C}$5. Inference begins at the root and initializes

$\{\mathbf{c}_i\}_{i=1}^{N_C}$6

At each depth, the method evaluates the children of the currently selected nodes, computes

$\{\mathbf{c}_i\}_{i=1}^{N_C}$7

and prunes according to a threshold induced by the pruning ratio $\{\mathbf{c}_i\}_{i=1}^{N_C}$8: $\{\mathbf{c}_i\}_{i=1}^{N_C}$9 When the surviving nodes are leaves, HDC applies the original diffusion classifier only to that reduced set, with final prediction

$$p_\theta(\mathbf{c}_i \mid \mathbf{x}) = \frac{p(\mathbf{c}_i)\, p_\theta(\mathbf{x} \mid \mathbf{c}_i)}{\sum_{j=1}^{N_C} p(\mathbf{c}_j)\, p_\theta(\mathbf{x} \mid \mathbf{c}_j)}.$$0

The hierarchy used in the experiments is the WordNet hierarchy underlying ImageNet-1K, modified by merging ambiguous nodes, collapsing subtrees containing only one leaf, and fixing the final hierarchy depth at $$p_\theta(\mathbf{c}_i \mid \mathbf{x}) = \frac{p(\mathbf{c}_i)\, p_\theta(\mathbf{x} \mid \mathbf{c}_i)}{\sum_{j=1}^{N_C} p(\mathbf{c}_j)\, p_\theta(\mathbf{x} \mid \mathbf{c}_j)}.$$1. Pruning begins from level $$p_\theta(\mathbf{c}_i \mid \mathbf{x}) = \frac{p(\mathbf{c}_i)\, p_\theta(\mathbf{x} \mid \mathbf{c}_i)}{\sum_{j=1}^{N_C} p(\mathbf{c}_j)\, p_\theta(\mathbf{x} \mid \mathbf{c}_j)}.$$2 rather than level $$p_\theta(\mathbf{c}_i \mid \mathbf{x}) = \frac{p(\mathbf{c}_i)\, p_\theta(\mathbf{x} \mid \mathbf{c}_i)}{\sum_{j=1}^{N_C} p(\mathbf{c}_j)\, p_\theta(\mathbf{x} \mid \mathbf{c}_j)}.$$3, because starting from level $$p_\theta(\mathbf{c}_i \mid \mathbf{x}) = \frac{p(\mathbf{c}_i)\, p_\theta(\mathbf{x} \mid \mathbf{c}_i)}{\sum_{j=1}^{N_C} p(\mathbf{c}_j)\, p_\theta(\mathbf{x} \mid \mathbf{c}_j)}.$$4 did not improve error discrimination but did increase runtime. The paper studies two pruning strategies: fixed pruning with $$p_\theta(\mathbf{c}_i \mid \mathbf{x}) = \frac{p(\mathbf{c}_i)\, p_\theta(\mathbf{x} \mid \mathbf{c}_i)}{\sum_{j=1}^{N_C} p(\mathbf{c}_j)\, p_\theta(\mathbf{x} \mid \mathbf{c}_j)}.$$5 at all depths, and dynamic pruning that retains nodes within two standard deviations of the minimum error at a given level (Shanbhag et al., 2024).

3. Empirical behavior of hierarchical diffusion classifiers

The experimental setting for HDC is ImageNet-1K with Stable Diffusion models. The baseline diffusion classifier achieves Top-1 64.90% with runtime 1600 s. Under Strategy 1, HDC reports Top-1 65.16% class-wise and 64.90% overall in one table presentation, runtime 980 s, and speed-up 38.75%. Under Strategy 2, HDC reports Top-1 63.33%, runtime 650 s, and speed-up 59.38%. The paper also reports Baseline Top-3 84.30% and Top-5 89.70%; Strategy 1 Top-3 81.80% and Top-5 86.30%; and Strategy 2 Top-3 82.30% and Top-5 86.30% (Shanbhag et al., 2024).

Configuration	Top-1	Time
Baseline diffusion classifier	64.90%	1600 s
HDC Strategy 1	65.16% class-wise, 64.90% overall	980 s
HDC Strategy 2	63.33%	650 s

The reported backbone comparison indicates that SD 2.0 gives the best overall tradeoff. Under Strategy 1, SD 2.0 reaches 65.16% class-wise accuracy at 980 s. Under Strategy 2, the fastest runtime is achieved with SD 1.4 at 710 s, but accuracy drops to 54.77%. Prompt wording also matters: the default prompt, “A photo of a <class label>”, performs best across strategies, with 64.90% under Strategy 1 and 63.20% under Strategy 2. These findings support the interpretation that hierarchical pruning is not merely a runtime heuristic; it also interacts with backbone choice and prompt specification in the zero-shot diffusion-classification regime (Shanbhag et al., 2024).

4. Two-stage hierarchical pruning of diffusion training data

A second, distinct use of hierarchy appears in RS-Prune for remote sensing diffusion foundation models. The motivating claim is that these models rely on large amounts of globally representative data, but such corpora often contain redundancy, noise, and class imbalance, and existing remote sensing diffusion foundation models typically aggregate multiple classification datasets or apply simplistic deduplication. RS-Prune therefore proposes a training-free, two-stage data pruning approach that selects a high-quality subset under high pruning ratios and jointly considers local information content with global scene-level diversity and representativeness (Wei et al., 29 Dec 2025).

The first stage computes grayscale Shannon entropy for each image,

$$p_\theta(\mathbf{c}_i \mid \mathbf{x}) = \frac{p(\mathbf{c}_i)\, p_\theta(\mathbf{x} \mid \mathbf{c}_i)}{\sum_{j=1}^{N_C} p(\mathbf{c}_j)\, p_\theta(\mathbf{x} \mid \mathbf{c}_j)}.$$6

where $$p_\theta(\mathbf{c}_i \mid \mathbf{x}) = \frac{p(\mathbf{c}_i)\, p_\theta(\mathbf{x} \mid \mathbf{c}_i)}{\sum_{j=1}^{N_C} p(\mathbf{c}_j)\, p_\theta(\mathbf{x} \mid \mathbf{c}_j)}.$$7 is the input image, $$p_\theta(\mathbf{c}_i \mid \mathbf{x}) = \frac{p(\mathbf{c}_i)\, p_\theta(\mathbf{x} \mid \mathbf{c}_i)}{\sum_{j=1}^{N_C} p(\mathbf{c}_j)\, p_\theta(\mathbf{x} \mid \mathbf{c}_j)}.$$8 the number of possible intensity levels, and $$p_\theta(\mathbf{c}_i \mid \mathbf{x}) = \frac{p(\mathbf{c}_i)\, p_\theta(\mathbf{x} \mid \mathbf{c}_i)}{\sum_{j=1}^{N_C} p(\mathbf{c}_j)\, p_\theta(\mathbf{x} \mid \mathbf{c}_j)}.$$9 the empirical probability of intensity level $$p_\theta(\mathbf{c}_i \mid \mathbf{x}) = \frac{p_\theta(\mathbf{x} \mid \mathbf{c}_i)}{\sum_{j=1}^{N_C} p_\theta(\mathbf{x} \mid \mathbf{c}_j)}.$$0. Images with $$p_\theta(\mathbf{c}_i \mid \mathbf{x}) = \frac{p_\theta(\mathbf{x} \mid \mathbf{c}_i)}{\sum_{j=1}^{N_C} p_\theta(\mathbf{x} \mid \mathbf{c}_j)}.$$1 are discarded. In the paper’s interpretation, low-entropy images are often cloud-covered, open ocean, large deserts, very homogeneous scenes, or invalid/noisy acquisitions. This stage is explicitly described as a fast and inexpensive filter that reduces the computational burden for later clustering and improves the quality of the remaining candidate set (Wei et al., 29 Dec 2025).

The second stage is scene-aware clustering with stratified sampling. Because remote sensing lacks an ImageNet-like universal benchmark for scene diversity, RS-Prune uses multiple expert-curated scene-classification datasets as a reference bank: NWPU-RESISC45, UC Merced Land-Use, AID, WHU-RS19, and RSD46-WHU. Features are extracted with Git-RSCLIP and normalized as

$$p_\theta(\mathbf{c}_i \mid \mathbf{x}) = \frac{p_\theta(\mathbf{x} \mid \mathbf{c}_i)}{\sum_{j=1}^{N_C} p_\theta(\mathbf{x} \mid \mathbf{c}_j)}.$$2

after which $$p_\theta(\mathbf{c}_i \mid \mathbf{x}) = \frac{p_\theta(\mathbf{x} \mid \mathbf{c}_i)}{\sum_{j=1}^{N_C} p_\theta(\mathbf{x} \mid \mathbf{c}_j)}.$$3-means clustering on the unit hypersphere produces centroids

$$p_\theta(\mathbf{c}_i \mid \mathbf{x}) = \frac{p_\theta(\mathbf{x} \mid \mathbf{c}_i)}{\sum_{j=1}^{N_C} p_\theta(\mathbf{x} \mid \mathbf{c}_j)}.$$4

The paper reports that $$p_\theta(\mathbf{c}_i \mid \mathbf{x}) = \frac{p_\theta(\mathbf{x} \mid \mathbf{c}_i)}{\sum_{j=1}^{N_C} p_\theta(\mathbf{x} \mid \mathbf{c}_j)}.$$5 is a robust choice.

For each unlabeled image $$p_\theta(\mathbf{c}_i \mid \mathbf{x}) = \frac{p_\theta(\mathbf{x} \mid \mathbf{c}_i)}{\sum_{j=1}^{N_C} p_\theta(\mathbf{x} \mid \mathbf{c}_j)}.$$6, cosine similarity to every centroid is computed: $$p_\theta(\mathbf{c}_i \mid \mathbf{x}) = \frac{p_\theta(\mathbf{x} \mid \mathbf{c}_i)}{\sum_{j=1}^{N_C} p_\theta(\mathbf{x} \mid \mathbf{c}_j)}.$$7 and each sample is assigned to the cluster

$$p_\theta(\mathbf{c}_i \mid \mathbf{x}) = \frac{p_\theta(\mathbf{x} \mid \mathbf{c}_i)}{\sum_{j=1}^{N_C} p_\theta(\mathbf{x} \mid \mathbf{c}_j)}.$$8

With total budget $$p_\theta(\mathbf{c}_i \mid \mathbf{x}) = \frac{p_\theta(\mathbf{x} \mid \mathbf{c}_i)}{\sum_{j=1}^{N_C} p_\theta(\mathbf{x} \mid \mathbf{c}_j)}.$$9, the method sets

$\varepsilon$0

and allocates a quota to each cluster. Within each cluster pool $\varepsilon$1, samples are ranked by similarity and selected by

$\varepsilon$2

If $\varepsilon$3, then $\varepsilon$4, and the remaining budget is reallocated from the global remainder. The intended balance is between cluster-level uniformity and sample representativeness: diversity is maintained by quota allocation, while representativeness is favored by centroid-prioritized selection (Wei et al., 29 Dec 2025).

5. High-ratio pruning, convergence, and downstream performance

RS-Prune is evaluated on three large-scale training datasets: Git-10M with 10.5M satellite images, RS5M with about 1.04M curated images after preprocessing, and USA-1m with 8.77M four-channel multispectral images at 1 m resolution. The pretrained diffusion backbone is evaluated on Super-resolution (SR) and Semantic image synthesis (SIS) using FID and LPIPS. The reported training setup uses a DiT-XL/2 backbone, 256×256 input size, batch size 256, AdamW, learning rate $\varepsilon$5, training on 4 NVIDIA H100 GPUs, 40K–100K diffusion steps for pretraining, and 5,000 fine-tuning steps for downstream tasks (Wei et al., 29 Dec 2025).

The paper studies high pruning ratios of 70% pruning, 85% pruning, and 90% pruning. It states that even after pruning 85\% of the training data, the method significantly improves convergence and generation quality, and that diffusion foundation models trained with the method consistently achieve state-of-the-art performance across downstream tasks. On Git-10M, retaining only a subset can outperform the full dataset. At 70% pruning, RS-Prune gives the best overall results among compared methods. For Git-10M, full-dataset generation FID is 45.30, whereas RS-Prune reaches 28.46. For USA-1m, full-dataset generation FID is 242.52, whereas RS-Prune reaches 175.93. At 85% pruning, stage II improves stage I alone by about 10% on both Git-10M and RS5M in the reported analysis. At 90% pruning, performance degrades slightly compared with 85% but remains better than training on the full dataset (Wei et al., 29 Dec 2025).

Setting	Reported outcome
Git-10M, full dataset	FID 45.30
Git-10M, RS-Prune at 70% pruning	FID 28.46
USA-1m, full dataset	FID 242.52
USA-1m, RS-Prune at 70% pruning	FID 175.93

The entropy-only ablation on Git-10M shows that retaining the top 30% by entropy gives the best FID in that ablation, outperforming the full dataset substantially. The runtime comparison further indicates why the two-stage design is described as efficient: clustering the full unlabeled dataset takes 4630.3 s, entropy-pruned unlabeled clustering takes 308.4 s, and reference-guided clustering with 5 datasets takes 115.1 s. These numbers support a specific interpretation of hierarchical pruning in generative training: for diffusion foundation models, the right subset is not necessarily the largest one, but the one that is informative, diverse, and scene-representative (Wei et al., 29 Dec 2025).

6. Related hierarchical routing and adjacent pruning paradigms

Hierarchically branched diffusion models provide an important adjacent formulation because they exploit class structure during generation without performing explicit class elimination. The forward diffusion process is unchanged, but the reverse process is partitioned by a class hierarchy derived from branch points between classes. A branch is defined as

$\varepsilon$6

where $\varepsilon$7 is a diffusion-time interval and $\varepsilon$8 is a subset of classes. For $\varepsilon$9 classes, the paper states that there are $$d(\varepsilon, \mathbf{x}, \mathbf{c}) = \|\varepsilon - \varepsilon_\theta(\mathbf{x}, \mathbf{c})\|^2,$$0 branches. Reverse diffusion then uses a multi-task neural network with one output head per branch: $$d(\varepsilon, \mathbf{x}, \mathbf{c}) = \|\varepsilon - \varepsilon_\theta(\mathbf{x}, \mathbf{c})\|^2,$$1 This design decomposes the class set into branch-specific subsets, narrows the active class structure at each timestep, and routes generation through branch-specific heads. The source paper emphasizes extension to novel classes in a continual-learning setting, analogy-based conditional generation through transmutation, and interpretability of branch points; it also reports branched sampling faster than linear sampling when generating one batch per class, with MNIST: 78.73 s linear vs 37.30 s branched, Letters: 110.42 s linear vs 67.54 s branched, and single-cell RNA-seq: 275.81 s linear vs 132.37 s branched (Tseng et al., 2022).

An adjacent but non-diffusion-specific reference point is module-aware pruning for hierarchical vision transformers. DIMAP is not a diffusion method, but it is relevant insofar as it addresses a recurring issue in hierarchical systems: weights or computations that are locally important can be pruned incorrectly when compared at a global level. The method defines modules such as QKV-M, PRJ-M, and MLP-M, introduces a data-independent weight metric based solely on weights, and reports strong ImageNet-1K results on Swin models, including a Top-5 accuracy drop of only 0.07% when removing 52.5% FLOPs and 52.7% parameters of Swin-B (He et al., 2024). A plausible implication is that hierarchical pruning in diffusion systems may likewise benefit from respecting local structure rather than applying global thresholds indiscriminately.

Taken together, these related paradigms clarify the boundaries of hierarchical diffusion class pruning. Explicit pruning is exemplified by HDC, which removes label candidates during diffusion classification. RS-Prune extends the hierarchy principle to training-set construction for diffusion foundation models. Hierarchically branched diffusion narrows and routes reverse diffusion without hard class removal. DIMAP shows that hierarchy-aware pruning is also a broader design principle in neural architectures beyond diffusion.