Timestep-Aware Block Masking for Efficient Diffusion Model Inference

• URL: http://arxiv.org/abs/2603.19939v1
• Date: Fri, 20 Mar 2026 13:35:03 GMT
• Title: Timestep-Aware Block Masking for Efficient Diffusion Model Inference
• Abstract summary: Diffusion Probabilistic Models (DPMs) have achieved great success in image generation but suffer from high inference latency.<n>We propose a novel framework to optimize the computational graph of pre-trained DPMs on a per-timestep basis.
• Score: 48.764748095994996
• License: http://arxiv.org/licenses/nonexclusive-distrib/1.0/
• 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.

Related papers

• Causal Autoregressive Diffusion Language Model [70.7353007255797]
  CARD reformulates the diffusion process within a strictly causal attention mask, enabling dense, per-token supervision in a single forward pass.<n>Our results demonstrate that CARD achieves ARM-level data efficiency while unlocking the latency benefits of parallel generation.
  arXiv  Detail & Related papers  (2026-01-29T17:38:29Z)
• Mixture of Distributions Matters: Dynamic Sparse Attention for Efficient Video Diffusion Transformers [13.366686736005699]
  We present MOD-DiT, a sampling-free dynamic attention framework.<n>It accurately models evolving attention patterns through a two-stage process.<n>It overcomes the computational limitations of traditional sparse attention approaches.
  arXiv  Detail & Related papers  (2026-01-14T16:25:39Z)
• Real-Time Iteration Scheme for Diffusion Policy [23.124189676943757]
  We introduce a novel approach inspired by the Real-Time Iteration (RTI) Scheme to accelerate inference.<n>We propose a scaling-based method to effectively handle discrete actions, such as grasping, in robotic manipulation.<n>The proposed scheme significantly reduces runtime computational costs without the need for distillation or policy redesign.
  arXiv  Detail & Related papers  (2025-08-07T13:49:00Z)
• One-Step Diffusion Model for Image Motion-Deblurring [85.76149042561507]
  We propose a one-step diffusion model for deblurring (OSDD), a novel framework that reduces the denoising process to a single step.<n>To tackle fidelity loss in diffusion models, we introduce an enhanced variational autoencoder (eVAE), which improves structural restoration.<n>Our method achieves strong performance on both full and no-reference metrics.
  arXiv  Detail & Related papers  (2025-03-09T09:39:57Z)
• PRISM: Self-Pruning Intrinsic Selection Method for Training-Free Multimodal Data Selection [68.8373788348678]
  Visual instruction tuning adapts pre-trained Multimodal Large Language Models to follow human instructions.<n>PRISM is the first training-free framework for efficient visual instruction selection.<n>It reduces the end-to-end time for data selection and model tuning to just 30% of conventional pipelines.
  arXiv  Detail & Related papers  (2025-02-17T18:43:41Z)
• PnP-Flow: Plug-and-Play Image Restoration with Flow Matching [9.589644878034553]
  We introduce Plug-and-Play Flow Matching, an algorithm for solving inverse imaging problems.<n>We evaluate its performance on denoising, superresolution, and inpainting tasks.
  arXiv  Detail & Related papers  (2024-10-03T12:13:56Z)
• TMPQ-DM: Joint Timestep Reduction and Quantization Precision Selection for Efficient Diffusion Models [40.5153344875351]
  We introduce TMPQ-DM, which jointly optimize timestep reduction and quantization to achieve a superior performance-efficiency trade-off. For timestep reduction, we devise a non-uniform grouping scheme tailored to the non-uniform nature of the denoising process. In terms of quantization, we adopt a fine-grained layer-wise approach to allocate varying bit-widths to different layers based on their respective contributions to the final generative performance.
  arXiv  Detail & Related papers  (2024-04-15T07:51:40Z)
• DITTO: Diffusion Inference-Time T-Optimization for Music Generation [49.90109850026932]
  Diffusion Inference-Time T-Optimization (DITTO) is a frame-work for controlling pre-trained text-to-music diffusion models at inference-time. We demonstrate a surprisingly wide-range of applications for music generation including inpainting, outpainting, and looping as well as intensity, melody, and musical structure control.
  arXiv  Detail & Related papers  (2024-01-22T18:10:10Z)
• Towards More Accurate Diffusion Model Acceleration with A Timestep Tuner [112.99126045081046]
  A diffusion model, which is formulated to produce an image using thousands of denoising steps, usually suffers from a slow inference speed.<n>We propose a textbftimestep tuner that helps find a more accurate integral direction for a particular interval at the minimum cost.<n>Experiments show that our plug-in design can be trained efficiently and boost the inference performance of various state-of-the-art acceleration methods.
  arXiv  Detail & Related papers  (2023-10-14T02:19:07Z)
• AdaDiff: Accelerating Diffusion Models through Step-Wise Adaptive Computation [32.74923906921339]
  Diffusion models achieve great success in generating diverse and high-fidelity images, yet their widespread application is hampered by their inherently slow generation speed. We propose AdaDiff, an adaptive framework that dynamically allocates computation resources in each sampling step to improve the generation efficiency of diffusion models.
  arXiv  Detail & Related papers  (2023-09-29T09:10:04Z)
• Learning to Efficiently Sample from Diffusion Probabilistic Models [49.58748345998702]
  Denoising Diffusion Probabilistic Models (DDPMs) can yield high-fidelity samples and competitive log-likelihoods across a range of domains. We introduce an exact dynamic programming algorithm that finds the optimal discrete time schedules for any pre-trained DDPM.
  arXiv  Detail & Related papers  (2021-06-07T17:15:07Z)

This list is automatically generated from the titles and abstracts of the papers in this site.

This site does not guarantee the quality of this site (including all information) and is not responsible for any consequences.