Hierarchically branched diffusion models leverage dataset structure for class-conditional generation

TMLR Paper2419 Authors

24 Mar 2024 (modified: 05 Jun 2024)Under review for TMLREveryoneRevisionsBibTeXCC BY-SA 4.0
Abstract: Diffusion models have attained state-of-the-art performance in generating realistic objects, including when conditioning generation on class labels. Current class-conditional diffusion models, however, implicitly model the diffusion process on all classes in a flat fashion, ignoring any known relationships between classes. Class-labeled datasets, including those common in scientific domains, are rife with internal structure. To take advantage of this structure, we propose hierarchically branched diffusion models as a novel framework for class-conditional generation. Branched diffusion models explicitly leverage the inherent relationships between distinct classes in the dataset to learn the underlying diffusion process in a hierarchical manner. We highlight several advantages of branched diffusion models over the current state-of-the-art methods for class-conditional diffusion. Firstly, they can be easily extended to novel classes in a continual-learning setting at scale. Secondly, they enable more sophisticated forms of conditional generation, such as analogy-based conditional generation (i.e. transmutation). Finally, they offer a novel interpretability into the class-conditional generation process. We extensively evaluate branched diffusion models on several benchmark and large real-world scientific datasets, spanning different data modalities (images, tabular data, and graphs). We particularly highlight the advantages of branched diffusion models on a single-cell RNA-seq dataset, where our branched model leverages the intrinsic hierarchical structure between human cell types.
Submission Length: Regular submission (no more than 12 pages of main content)
Assigned Action Editor: ~Ekin_Dogus_Cubuk1
Submission Number: 2419
Loading