Towards a Mechanistic Explanation of Diffusion Model Generalization

TL;DR

<3-5 sentence high-level summary> The paper investigates why diffusion models generalize well beyond their training data by contrasting pretrained network denoisers with the empirically optimal denoiser, uncovering a persistent local inductive bias across architectures. It posits that neural denoisers operate via localized denoising that, when aggregated across patches, approximates the global optimal denoiser for much of the forward process. The authors formalize this intuition with Patch Set Posterior Composite (PSPC) denoisers, including PSPC-Square and PSPC-Flex, which compute patch posterior means over spatial crops and combine them to match network outputs. PSPC exhibits strong alignment with network denoisers in forward and reverse diffusion, yielding samples that resemble neural-network outputs while remaining training-free, with implications for attribution, efficiency, and non-neural diffusion strategies.

Abstract

We propose a simple, training-free mechanism which explains the generalization behaviour of diffusion models. By comparing pre-trained diffusion models to their theoretically optimal empirical counterparts, we identify a shared local inductive bias across a variety of network architectures. From this observation, we hypothesize that network denoisers generalize through localized denoising operations, as these operations approximate the training objective well over much of the training distribution. To validate our hypothesis, we introduce novel denoising algorithms which aggregate local empirical denoisers to replicate network behaviour. Comparing these algorithms to network denoisers across forward and reverse diffusion processes, our approach exhibits consistent visual similarity to neural network outputs, with lower mean squared error than previously proposed methods.

Figures (25)

• Figure 1: Denoiser outputs given shared reverse process noisy inputs from CIFAR-10. Column 1: Optimal empirical denoiser, i.e. what a "perfect" neural network denoiser would output if appropriately parameterized and trained to achieve minimal loss on the diffusion denoising loss (\ref{['eq:score_matching']}). Columns 2-4: Outputs from various denoising neural networks. For $t < 3.3$ all networks deviate from the optimal denoiser in similar ways. Column 5: Our learning-free patch set posterior composite denoiser produces qualitatively similar outputs to the neural network denoisers, suggesting that neural networks may generalize in part via patch denoising and composition.
• Figure 2: Left: Mean squared error between empirical and network denoisers for three architectures on CIFAR-10. Right: Comparison of network and empirical denoiser for a shared $\mathbf{z} \sim p_t(\mathbf{z} | \mathbf{x})$ at three $t$ values. Network estimators have low error for small and large $t$, but large errors around $t=3$. At this point, each network varies substantially from the empirical denoiser in the same way.
• Figure 3: Left: Comparison of average square patch sizes required to capture 50%, 75%, and 95% of the total gradient sensitivity heatmap. As $t$ increases, larger patch sizes are required to capture a fixed percentage of the total gradient sensitivity heatmap. Right: Gradient sensitivity heatmaps for DDPM++ denoiser on CIFAR-10 for output pixel (15,15) across varying $t$.
• Figure 4: Left: Comparison of patch posterior means with varying patch sizes to corresponding patches of the optimal denoiser over forward process samples. For $t<1$, relatively small square patch posterior means exactly match the optimal denoiser. As $t$ increases, larger patch sizes are required to exactly estimate patches of the optimal denoiser. Right: Comparison of patch posterior means with varying patch sizes to DDPM++ denoiser patches on $\mathbf{z}$ drawn from the reverse process. Patch posterior means estimate the network denoiser patches better than the optimal denoiser for all $t<3$. As $t$ decreases, so does the patch size which best estimates network denoiser patches.
• Figure 5: Our PSPC denoiser. First, $\mathbf{z}$ is decomposed into patches using a set of cropping matrices. For each patch, we compute the patch posterior mean via \ref{['eq:patch_posterior']}. Resulting means are then combined into one image and normalized by the by the number of patches that overlap each pixel. Although square patches are visualized, PSPC can be used with any set of cropping matrices.