基于蒸馏的模块神经架构搜索的轻量级扩散模型
计算机视觉与模式识别
2023-11-17 v2 机器学习
摘要
扩散模型近来展现出卓越的生成能力，在许多任务中取得了最先进的性能。然而，高计算成本仍是扩散模型的一个棘手问题。为解决这个问题，我们提出利用所提出的扩散蒸馏模块级神经架构搜索(DiffNAS)自动去除扩散模型中的结构冗余。具体而言，给定一个较大的预训练教师模型，我们利用DiffNAS搜索最小架构，其性能可达到与教师模型相当甚至更优。考虑到当前扩散模型基于天然具有模块级结构的UNet，我们在每个模块中独立执行神经架构搜索，大幅缩减了搜索空间。与以往模块级NAS方法不同，DiffNAS包含一种模块级局部搜索策略和带有联合动态损失的重训练策略。具体地，在搜索过程中，我们按模块选择最佳子网络，以避免以往工作所用全局搜索策略带来的不公平性。重训练搜索到的架构时，我们采用动态联合损失来维持超网训练与子网络重训练间的一致性，这也对每个模块提供了信息丰富的目标并缩短了梯度传播路径。我们证明该联合损失能有效提升模型性能。我们也证明了该损失动态调整的必要性。实验表明，我们的方法能实现显著的计算削减，尤其在隐扩散模型上约减少50%的MACs和参数量。
引用
@article{arxiv.2311.04950,
title = {Lightweight Diffusion Models with Distillation-Based Block Neural Architecture Search},
author = {Siao Tang and Xin Wang and Hong Chen and Chaoyu Guan and Yansong Tang and Wenwu zhu},
journal= {arXiv preprint arXiv:2311.04950},
year = {2023}
}