DiFA: Inference-Time Forward-Process Alignment for Diffusion Models
Abstract
The prevailing inference framework for diffusion models formulates generation fundamentally as a problem of numerical integration. This perspective casts the model as an accurate estimator, neglecting the inherent statistical uncertainty of the denoising process. In this work, we propose Forward-Process Aligned Diffusion prediction (DiFA), a training-free framework that reframes inference-time clean-signal prediction refinement within diffusion sampling as a sequential state estimation problem. Instead of discarding historical predictions, DiFA treats the inference trajectory as a sequence of correlated observations with varying uncertainties. Inspired by Kalman estimation, DiFA builds a logSNR-aware temporal consensus to align historical clean predictions. Crucially, to counteract the over-smoothing typically associated with temporal consensus, we introduce a deviation guidance mechanism that adaptively preserves residual details. Empirically, DiFA yields significant improvements on CIFAR-10 and ImageNet across the evaluated metrics, including FID, IS, and FD-DINOv2, demonstrating that aligning inference with the forward statistical structure substantially improves generative fidelity. Code is available at https://github.com/ShiguiLi/DiFA.
Lay Summary
Artificial Intelligence (AI) models that generate high-quality images and videos, like Midjourney or Sora, work by gradually "cleaning up" random noise until a clear picture emerges. However, this step-by-step cleaning process is computationally expensive and slow. Imagine trying to drive to a destination using a GPS that only recalculates your route based on your immediate surroundings at every single intersection—if you take one wrong turn, the errors pile up, leading you off course. To speed up the process, researchers often use "shortcuts" (fast solvers) to skip steps, but taking too many shortcuts usually results in blurry or distorted images. In our paper, we introduce a method called DiFA. Instead of blindly trusting the immediate "GPS instruction" at each step, DiFA acts like a smart co-pilot. It looks at the history of recent predictions to find a reliable consensus and gently corrects the steering direction without requiring the car to slow down or need a new engine (no extra training cost). Our experiments show that DiFA allows these AI models to generate significantly better and sharper images while taking far fewer steps, making advanced AI generation faster and more accessible.