Timestep-Aware Block Masking for Efficient Diffusion Model Inference

TL;DR

This paper introduces a timestep-aware block masking framework that dynamically optimizes the computational graph of diffusion models during inference, significantly reducing latency while maintaining high image quality.

Contribution

It proposes a novel, memory-efficient method to learn timestep-specific masks for diffusion models, improving inference speed without sacrificing performance.

Findings

Achieves faster sampling with minimal quality loss.

Demonstrates effectiveness across multiple diffusion architectures.

Introduces a timestep-aware loss and mask rectification strategy.

Abstract

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.