按时间步动态跳过冗余计算，让扩散模型生成更快更省资源。
Timestep-Aware Block Masking for Efficient Diffusion Model Inference
- 为每个去噪步骤学习专属跳过策略，智能决定哪些模块该执行
- 在不损失图像质量前提下，推理速度提升显著，支持多种主流模型
- 无需重训练，适配性强，适合部署时追求高效生成的团队
扩散概率模型（DPMs）在图像生成中表现卓越，但因迭代去噪过程导致推理延迟高。针对去噪轨迹中特征动态变化的特点，本文提出一种基于时间步的块掩码优化框架，对预训练DPM的计算图进行逐时间步优化。通过学习时间步特定的掩码，方法在每个推理阶段动态决定哪些模块可跳过或复用特征，实现计算资源的精准分配。与依赖全链反向传播的全局优化方法不同，本方法独立优化各时间步掩码，避免高昂内存开销。为指导优化过程，引入时间步感知的损失缩放机制，在敏感去噪阶段优先保障特征保真度，并采用知识引导的掩码修正策略，去除冗余时空依赖。该方法具备架构无关性，在DDPM、LDM、DiT和PixArt等多种模型上均实现显著效率提升。实验表明，将去噪过程视为一系列优化后的计算路径，可在保持生成质量的同时显著提升采样速度。代码将公开。
原文摘要 · Abstract (English)
Diffusion Probabilistic Models (DPMs) have achieved great success in image generation but suffer from high inference latency due to their iterative denoising nature. Motivated by the evolving feature dynamics across the denoising trajectory, we propose a novel framework to optimize the computational graph of pre-trained DPMs on a per-timestep basis. By learning timestep-specific masks, our method dynamically determines which blocks to execute or bypass through feature reuse at each inference stage. Unlike global optimization methods that incur prohibitive memory costs via full-chain backpropagation, our method optimizes masks for each timestep independently, ensuring a memory-efficient training process. To guide this process, we introduce a timestep-aware loss scaling mechanism that prioritizes feature fidelity during sensitive denoising phases, complemented by a knowledge-guided mask rectification strategy to prune redundant spatial-temporal dependencies. Our approach is architecture-agnostic and demonstrates significant efficiency gains across a broad spectrum of models, including DDPM, LDM, DiT, and PixArt. Experimental results show that by treating the denoising process as a sequence of optimized computational paths, our method achieves a superior balance between sampling speed and generative quality. Our code will be released.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品；中文卡片由大模型生成，请以原文为准。