Distillation of Discrete Diffusion through Dimensional Correlations

TL;DR

This paper introduces scalable mixture models for discrete diffusion that capture dimensional correlations, enabling distillation of slow, iterative models into faster, few-step versions for image and language generation.

Contribution

It proposes novel mixture models and loss functions for distilling discrete diffusion models, effectively reducing sampling steps while preserving quality.

Findings

Effective distillation of pretrained models in image and language domains

Significant reduction in sampling steps with maintained performance

Scalable approach for modeling dependencies in high-dimensional discrete data

Abstract

Diffusion models have demonstrated exceptional performances in various fields of generative modeling, but suffer from slow sampling speed due to their iterative nature. While this issue is being addressed in continuous domains, discrete diffusion models face unique challenges, particularly in capturing dependencies between elements (e.g., pixel relationships in image, sequential dependencies in language) mainly due to the computational cost of processing high-dimensional joint distributions. In this paper, (i) we propose "mixture" models for discrete diffusion that are capable of treating dimensional correlations while remaining scalable, and (ii) we provide a set of loss functions for distilling the iterations of existing models. Two primary theoretical insights underpin our approach: First, conventional models with element-wise independence can well approximate the data distribution, but essentially require {\it many sampling steps}. Second, our loss functions enable the mixture models to distill such many-step conventional models into just a few steps by learning the dimensional correlations. Our experimental results show the effectiveness of the proposed method in distilling pretrained discrete diffusion models across image and language domains. The code used in the paper is available at https://github.com/sony/di4c .