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 exact estimator, neglecting the inherent statistical uncertainty of the denoising process. In this work, we propose Forward-Process Aligned Diffusion prediction (\textbf{DiFA}), a training-free framework that reframes inference-time data prediction refinement as a sequential state estimation problem. Rather than reusing past outputs solely for numerical integration, DiFA treats iterative data predictions along the reverse trajectory as correlated observations to build a forward-aligned temporal consensus. Inspired by Kalman filtering, this consensus aggregates historical predictions according to structural consistency and noise-level compatibility. To counteract the over-smoothing tendency of temporal consensus, we introduce a deviation guidance mechanism to adaptively preserve 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.

• The paper introduces a training-free, solver-agnostic framework that aligns recent clean-signal predictions through temporal consensus and anchor-relative refinement before an unchanged diffusion integrator.
• The method improves few-step generation without additional network evaluations, reducing CIFAR-10 FID from 8.40 to 4.15 at 8 NFE and ImageNet-256 Euler FID from 52.64 to 27.61 at 5 NFE.
• The results show consistent gains across diffusion and flow-matching models, while highlighting open challenges in modeling correlated denoiser errors and adapting refinement settings to noise level and inference budget.

Overview and Motivation

DiFA is a training-free inference framework that reframes the refinement of clean-signal predictions in diffusion sampling as a sequential state-estimation problem (2607.17972). The authors' central diagnosis is that prevailing inference frameworks treat generation as numerical integration of the probability-flow ODE, implicitly assuming the learned denoiser is an exact estimator. Because the network is trained with an MSE objective, its clean-data predictions $$\hat{\boldsymbol{x}}_0^{(t)} = D_\theta(\boldsymbol{x}_t, t)$$ constitute a smoothed, biased approximation of the posterior mean, and standard samplers—deterministic or stochastic—discard the temporal redundancy embedded in the sequence of these predictions. In high-curvature regions of the generative manifold, linearized solver steps accumulate per-step estimation bias and injected-noise variance into irreversible trajectory drift, particularly under few-step inference.

The proposed remedy is not a new solver but a plug-and-play wrapper at the interface between the denoiser and the downstream integrator: historical clean predictions are fused into a forward-aligned temporal consensus, which serves as a reference anchor against which the current prediction is refined. DiFA introduces no additional network function evaluations (NFEs); its overhead is $O(Kd)$ per step for a window of size $K$ and prediction dimension $d$.

The Forward-Process Alignment Principle

The theoretical foundation rests on the signal-normalized representation of the forward process. Since $$\bar{\boldsymbol{x}}_t = \boldsymbol{x}_t/\alpha_t = \boldsymbol{x}_0 + \operatorname{SNR}(t)^{-1/2}\boldsymbol{\epsilon}$$, all forward states share a common clean-data anchor $\boldsymbol{x}_0$, with $\operatorname{SNR}(t)$ exactly characterizing the observation precision at noise level $t$. DiFA transfers this geometry to reverse inference: the sequence of clean predictions along the trajectory is treated as temporally correlated observations of a locally stable trajectory-implied anchor.

Under an idealized independent-view observation model in which each prediction is an unbiased, mutually independent Gaussian observation of a static anchor with covariance $R_i = c/\operatorname{SNR}(t_i)$, the Best Linear Unbiased Estimation (BLUE) analysis yields the precision-weighted fusion rule

$$\hat{\boldsymbol{x}}_{0,\mathrm{ideal}}^{\star} = \frac{\sum_i \operatorname{SNR}(t_i)\, \boldsymbol{y}_i}{\sum_i \operatorname{SNR}(t_i)},$$

with fused covariance $O(Kd)$0, strictly smaller than any individual observation's covariance for $O(Kd)$1 (Proposition 3.1). The authors further prove (Theorem 3.2) that this estimator admits an exact recursive realization via static-state Kalman updates, connecting the consensus construction to classical filtering.

Two caveats are stated plainly by the authors: the independent-view model is an analytical proxy, not an exact model of learned denoiser errors, which are in practice biased and temporally correlated; and the static-anchor assumption holds only locally. The practical algorithm therefore uses the BLUE/Kalman analysis as motivation for a structural design rather than as a guarantee.

Algorithmic Instantiation

The practical DiFA procedure operates in three stages per reverse step:

1. Causal temporal consensus. A buffer of the most recent $O(Kd)$2 historical clean predictions (default $O(Kd)$3) is maintained. Historical candidates undergo channel-wise affine mean–variance alignment to the current prediction, then are weighted by a softmax over a compatibility logit combining pooled cosine similarity (structural compatibility, sharpness $O(Kd)$4) and a logSNR-proximity penalty (noise-level compatibility). The current prediction serves as a query but is excluded from the anchor, preventing trivial self-reinforcement and preserving a non-degenerate anchor-relative deviation.
2. Anchor-relative deviation guidance. The deviation $O(Kd)$5 is projected orthogonally to the current prediction to suppress magnitude-parallel components, decomposed into low- and high-frequency parts via local pooling, and recombined with a sigmoidal logSNR gate $O(Kd)$6 that admits high-frequency detail only at low noise levels.
3. Additive refinement. The refined prediction $O(Kd)$7 (with $O(Kd)$8, default $O(Kd)$9, i.e., $K$0) is passed to an unmodified downstream solver (DDIM, DPM-Solver++, UniPC, Heun).

Because the integration rule itself is untouched, DiFA is solver-agnostic: any solver expressible in clean-prediction parameterization can consume the refined estimate. The deviation-guidance stage exists specifically to counteract the over-smoothing that naive temporal consensus would induce, which the authors identify as the mechanism by which averaging-based methods typically improve FID at the cost of blurriness.

Empirical Results

The evaluation covers CIFAR-10 (32×32), ImageNet-64, LSUN Bedroom via latent diffusion, and ImageNet 256×256 via SiT-XL/2, using FID, IS, sFID, precision/recall, and FD-DINOv2.

On CIFAR-10, the gains are largest in the few-step regime: at NFE 8, DiFA reduces DPM-Solver++ FID from 8.40 to 4.15 (over 50% relative improvement); at NFE 12 it achieves FID 2.18 versus the baseline's 3.70, also surpassing the trained entropy-aware method EVODiff (2.25 at the same budget); at NFE 20 it reaches 1.96 versus 2.33, breaking the FID 2.0 barrier without training. On ImageNet-64, DiFA reaches FID 1.63–1.64 at 25 NFE with DPM-Solver++ and UniPC, compared to 1.83 and 1.73 for the unmodified solvers—competitive with EDM2-S (1.58 at 63 NFE) using far fewer steps. The most dramatic single result concerns Heun on ImageNet: FID drops from 230.05 to 110.20 at NFE 5, indicating that prediction alignment can stabilize severely drifting high-order solvers at very low budgets.

The flow-matching validation is notable for its breadth. On SiT-XL/2 with ImageNet 256×256 and CFG 1.5, the FM-adapted DiFA reduces Euler 5-NFE FID from 52.64 to 27.61 (−47.5%) and Heun2 5-NFE FID from 14.64 to 7.99 (−45.4%), with IS gains up to +65.2%. At 20 NFE the relative margins shrink but remain positive (e.g., Euler 3.33 → 2.41). A refinement-scale sweep from 1.25 to 1.75 shows consistent improvement across all scales, indicating the gains are not an artifact of a single hyperparameter.

Ablations attribute the bulk of the improvement to historical consensus: with the window collapsed ($K$1), DiFA exactly reproduces the baseline (EDM FID 5.071 at 10 NFE), while the full configuration achieves 2.797. Window-size sensitivity exhibits a U-shaped trend (W=2: 3.219; W=3: 2.797; W=4: 3.791), reflecting a trade-off between exploiting redundancy and incorporating stale predictions. Hyperparameter sweeps over SNR gating thresholds, magnitude-alignment factor, deviation scale, and logSNR-compatibility coefficient show stable improvements across all tested values, though the optimal configuration varies across NFE budgets.

Limitations and Open Questions

The authors acknowledge that the framework relies on clean-prediction parameterizations, fixed (non-adaptive) hyperparameters, and an idealized static-anchor assumption that does not hold exactly for MSE-trained denoisers whose errors are biased and temporally correlated. The theoretical variance-reduction guarantee applies only under the independent-view model; no guarantee is provided for the practical weighting scheme, which substitutes heuristic structural and logSNR compatibility for the ideal inverse-SNR precisions. A mild precision–recall trade-off is observed in the Heun2 20-NFE setting on ImageNet 256×256, and the optimal deviation scale and gating configuration shift across NFE budgets, leaving open the design of adaptive, noise- and budget-aware parameter selection. The paper also leaves unaddressed the extension of the BLUE analysis to explicitly temporally correlated error models, and the application to text-to-image, video generation, inverse problems, and distillation pipelines is deferred as future work.

Conclusion

DiFA contributes a repositioning of inference-time improvement: from refining the numerical integration of a trusted vector field to refining the clean-signal estimate consumed by an unchanged integrator, using the forward process's common-anchor geometry and inverse-SNR reliability structure as organizing principles. The empirical evidence—consistent FID, IS, and FD-DINOv2 improvements across solvers, noise schedules, pixel and latent spaces, and both diffusion and flow-matching paradigms, at zero additional NFE cost—supports the claim that temporal prediction redundancy in pretrained generative models is substantially underexploited by standard samplers. The gap between the idealized estimation-theoretic analysis and the heuristic practical instantiation remains the principal theoretical open point.