让小模型更懂复杂条件,自动分配训练资源提升效果
ARIA: Adaptive Region-Based Importance Allocation for Conditional Diffusion Distillation

- 根据条件空间区域动态分配训练精力,聚焦差异大的部分
- 在未见和低频条件下性能提升明显,优于现有方法
- 适合大规模条件场景下的高效模型蒸馏,如文本生成
条件扩散模型的蒸馏旨在将大模型(教师)的行为迁移到小模型(学生),同时保持不同条件输入的一致性。与识别任务不同,条件扩散蒸馏常因预测噪声强烈依赖条件信号,导致知识迁移难以超出训练分布。因此,有效蒸馏需探索广阔的条件空间,但在实际中面临数据瓶颈:成对图像-条件数据有限,为每个条件生成合成图像计算成本过高,而条件集合(如文本提示)可能极为庞大。近期工作通过训练中切换条件,使学生接触更广条件空间,但未解决训练资源如何分配的问题。本文提出ARIA框架,自适应地在条件空间粗粒度区域间分配训练资源。通过在线跟踪各区域的师生差异,ARIA集中优化持续存在偏差的区域,同时保持原有蒸馏目标。实验证明,ARIA在多数架构和设置下优于基线方法,尤其在未见和低频条件下表现突出。理论分析表明,在方差和漂移有界条件下,该追踪机制能有效跟随训练中差异演化。
原文摘要 · Abstract (English)
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.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。