ARIA: Adaptive Region-Based Importance Allocation for Conditional Diffusion Distillation

2026-06-22

AI Summary

This work addresses the limitations of conditional diffusion model distillation, which suffers from low training efficiency and poor generalization under out-of-distribution settings and large conditional spaces. The authors propose an adaptive region importance allocation mechanism that introduces, for the first time, a region-based online discrepancy estimation to dynamically track the divergence between teacher and student models across different regions of the conditional space. Coupled with an adaptive sampling strategy, this approach optimizes the allocation of training resources while preserving the original distillation objective. The method significantly enhances modeling capability for unseen and low-frequency conditional regions. Extensive experiments demonstrate consistent superiority over existing approaches across diverse architectures and settings, with particularly notable gains in out-of-distribution scenarios, supported by theoretical guarantees.

Abstract

Distilling conditional diffusion models aims to transfer the behavior of a large teacher to a smaller student while preserving alignment across conditioning inputs. Unlike recognition tasks, knowledge distillation in conditional diffusion often struggles to transfer knowledge beyond the training distribution, since the predicted noise strongly depends on the conditioning signal. As a result, effective distillation requires exploring a large conditioning space. In practical settings, this creates a major bottleneck. Paired image-condition data may be limited, and generating synthetic images for every available condition is often computationally infeasible, while the pool of conditions, such as text prompts, can be extremely large. Recent work addresses this issue by switching conditions during training, exposing the student to a broader conditioning space without changing the distillation objective. Yet this raises a complementary question: once a large conditioning corpus is available, how should the training effort be allocated? In this work, we introduce ARIA, a framework that adaptively allocates training effort across coarse regions of the conditioning space. By maintaining online estimates of teacher-student discrepancy at the region level, ARIA focuses updates where misalignment persists while preserving the original distillation objective. Empirically, ARIA improves over RC across most architectures and settings, with the clearest gains observed in unseen and underrepresented regimes. We also provide a theoretical analysis showing that the proposed tracking mechanism follows the evolving discrepancy during training under bounded variance and drift assumptions.

Problem

• conditional diffusion
• knowledge distillation
• conditioning space
• training allocation
• teacher-student discrepancy

Innovation

• conditional diffusion distillation
• adaptive importance allocation
• region-based training
• teacher-student discrepancy
• knowledge distillation