Filtered Posterior Mean Collections: A Unified Framework for Analytical Models of Diffusion Generalization
Authors: Matthew Niedoba, Berend Zwartsenberg, Frank Wood
Organizations: University of British Columbia · 2Inverted AI · 3Alberta Machine Intelligence Institute
Abstract
The neural-network denoising functions which form the backbone of image diffusion models are remarkably consistent in their generalization behaviour across a wide variety of network architectures and training procedure hyperparameters. A recent line of research has sought to model the outputs of these networks by aggregating posterior weighted averages of training dataset patches. In this work, we consolidate these approaches into a unified model class which we call Filtered Posterior Mean Collections (FPMCs). We define this model class using query precision vectors, response weights, and source distributions, and illustrate that existing methods are recoverable with specific choices of these design axes. Investigating each axis in turn, we find that FPMC performance can be improved with soft relaxations of prior patch-based methods, and through augmentations of source distributions. Applying these findings to an existing FPMC, we demonstrate consistent sample improvement across three natural image datasets.
Recent work has shown that diffusion models trained with the denoising score matching (DSM) objective often violate the Fokker--Planck (FP) equation that governs the evolution of the true data density. Directly penalizing these deviations in the objective function reduces their magnitude but introduces a significant computational overhead. It is also observed that enforcing strict adherence to the FP equation does not necessarily lead to improvements in the quality of the generated samples, as often the best results are obtained with weaker FP regularization. In this paper, we investigate whether simpler penalty terms can provide similar benefits. We empirically analyze several lightweight regularizers, study their effect on FP residuals and generation quality, and show that the benefits of FP regularization are available at substantially lower computational cost. Our code is available at https://github.com/OnnoNiemann/fp_diffusion_analysis.
Onno Niemann, Gonzalo Martínez Muñoz, Alberto Suárez Gonzalez
Pretrained diffusion models represent image distributions through a continuum of progressively smoothed distributions. This multiscale structure organizes generation from global structure to fine detail and supports high-quality, diverse samples. We exploit the same multiscale diffusion prior for linear imaging inverse problems. Rather than using the pretrained model only as a denoiser in an outer iteration, we define a surrogate likelihood whose center is aligned with the clean-image coordinate and whose covariance accounts for residual diffusion uncertainty. This construction defines an explicit surrogate posterior path, from which we derive continuous posterior dynamics. A tunable Langevin component supports target tracking and allows the amount of posterior exploration to be adapted to the application. We prove endpoint consistency and a finite-horizon tracking bound and, in the exact-score setting, first-order weak accuracy. For computation, we derive the Posterior-Dynamics Implicit--Explicit sampler (PD-IMEX), a stable method using one score evaluation per diffusion scale and an implicit data-consistency update. Experiments on deblurring, super-resolution, and inpainting show strong reconstruction quality at 100 score evaluations, coarse-grid stability, and controllable fidelity--diversity behavior.
Steering diffusion models toward conditions unseen during training typically requires either retraining with conditional inputs or per-step gradient computations, both of which incur substantial computational overhead. We present Noise-Aligned RFM Steering (NA-RFM), a general recipe for efficiently steering diffusion models without gradient guidance during inference, enabling fast controllable generation. The method combines two offline-computed signals: noise alignment, a high-noise correction from PCA statistics of the target examples and the full data, and Recursive Feature Machine (RFM) activation steering, which learns a target-discriminative direction from labeled forward-process activations. During sampling, noise alignment provides coarse control at high noise, while the RFM direction is reused over intermediate/late timesteps through lightweight activation edits. Experiments on CIFAR-10, ImageNet, CelebA, and fine-grained bird species show improved target accuracy over gradient-based post-hoc guidance baselines, improved FID on the class-guidance benchmarks, and substantial inference speedups. Code: https://github.com/isotrivial/na-rfm.