Just Leaf It: Accelerating Diffusion Classifiers with Hierarchical Class Pruning

TL;DR

This paper introduces a Hierarchical Diffusion Classifier that leverages label hierarchies to prune irrelevant categories, significantly reducing inference time while maintaining or improving accuracy in large-scale image classification tasks.

Contribution

The paper presents a novel hierarchical pruning method for diffusion classifiers, enabling faster inference without sacrificing accuracy in large-scale applications.

Findings

Up to 60% faster inference times.

Maintains or improves classification accuracy.

Effective use of hierarchical label structures.

Abstract

Diffusion models, celebrated for their generative capabilities, have recently demonstrated surprising effectiveness in image classification tasks by using Bayes' theorem. Yet, current diffusion classifiers must evaluate every label candidate for each input, creating high computational costs that impede their use in large-scale applications. To address this limitation, we propose a Hierarchical Diffusion Classifier (HDC) that exploits hierarchical label structures or well-defined parent-child relationships in the dataset. By pruning irrelevant high-level categories and refining predictions only within relevant subcategories (leaf nodes and sub-trees), HDC reduces the total number of class evaluations. As a result, HDC can speed up inference by as much as 60% while preserving and sometimes even improving classification accuracy. In summary, our work provides a tunable control mechanism between speed and precision, making diffusion-based classification more feasible for large-scale applications.