Speculative draft trees speed up diffusion model generation
Accelerating Diffusion Sampling via Speculative Draft Trees
Machine LearningArtificial Intelligence
Summary
Diffusion models are a way to create things like images or sounds by gradually changing random noise into a final result, but this process can be slow because it needs many careful steps. The authors found a new method that speeds up this process by looking ahead at several possible next steps at once, organizing these guesses in tree-like structures rather than a simple chain. This approach lets the model check and correct its guesses more efficiently, using ideas from a coding method called relative entropy coding. Tests showed this method can make generation about 8% faster while still producing exactly correct results.
What this means in practice
- •For ai developers: Integrate draft tree speculative sampling to reduce compute time for generating images or other outputs from diffusion models.
- •For machine learning engineers: Enhance diffusion model pipelines with improved coupling methods to speed up sampling without compromising output quality.
Authors
Marcello Bullo, Yanxiao Liu, Öykü Sıla Güner, Arpan Mukherjee, Deniz Gündüz
Abstract
Speculative sampling accelerates diffusion model generation by drafting inexpensive candidate states and correcting them under a coupling that preserves the target distribution exactly, reducing the number of expensive target evaluations. Existing diffusion samplers, notably those based on reflection maximal coupling, are topologically constrained: their lookahead drafts form a chain graph, a single linear sequence, which inherently limits the acceptance rate per target evaluation. We connect speculative sampling in diffusion models to relative entropy coding (REC). This perspective shows the lookahead need not be linear and motivates our central contribution, draft trees, which enrich the candidates considered per round and lower the target function evaluations. We further adopt greedy rejection sampling, an REC algorithm, as the draft-target coupling, improving acceptance while guaranteeing exact target samples. Experiments across diverse target and draft models demonstrate up to 8.3% acceleration over the reflection coupling baseline in practical settings.