Towards a Mechanistic Explanation of Diffusion Model Generalization

General Comment

We'd like to thank the reviewer for their thorough review. We have summarized and responded to the key points of your review below in the limited available space. If we have missed any of your points, we welcome further discussion.

Methods & Evaluation

Network errors for $t>30$

The reviewer is right. U-ViT & NCSN++ MSE is higher for $t> 30$. As U-VIT predictis $\epsilon$, we compute $x$ as $x = z - t \epsilon_\theta(z, t)$. This scales $\epsilon_\theta$ errors by $t$ (see L162). NCSN++ errors fluctuate near $t=40$, but are ~4x smaller than peak errors at t=3. We believe this is a training artifact. Further, Fig. 9 implies high-$t$ errors have little impact on generalization. Both U-ViT& NCSN++ have SSCD similarities > 0.7 with every model, indicating near-copy samples.

"Averaging effect"

Beyond MSE, Fig. 1, 2, & 13-20 provide evidence that network outputs are qualitatively similar for the same input. To quantify this output similarity, we have measured output cosine similarity at $t=3.2$ for shared forward process inputs [Link].

Lack of other networks baselines past Fig. 2

This assertion is incorrect. Past Fig 2, Fig 9 and 13-20 all compare various models. However, for transparency and completeness, we will include the following modified versions of Fig 3, 4, and 6 to the appendix:

• Fig 3 All models have a similar trend between gradient concentration and $t$.
• Fig. 4 Patch posterior MSE is similar for all models, except high-$t$ U-ViT.
• Fig. 6. We compare our methods against each network baseline. For UViT & DiT, we only have CIFAR checkpoints.

We are happy to provide additional figures upon request.

On the impact of attention

We do not believe attention layers are the cause of our observations. To test this, we trained an attention-free DDPM++ model on CIFAR. We find the MSE of PSPC is mostly unaffected[Link]. Qualitatively, the attention-free model's outputs are similar to other models [Link]. Similar to our work, Kamb & Ganguli find their patch-based denoiser performs better against attention free models.

On the inclusion of U-Net and BFCNN

This paper's aim was to investigate generalization in modern image diffusion models, which overwhelmingly utilize attention layers. In addition to the lack of attention, the networks of Kadkhodaie et al. (2023) have unique properties such as no bias terms, and no $t$ conditioning. Due to these architectural quirks, we believe that analyzing commonly used U-Net and DiT architectures is of greater value to the community than BFCNN.

Other Comments

Q1

Niedoba et al. (Fig. 3) finds network denoisers are biased for $t \in [0.3, 10]$, which matches our Fig 2 findings. For clarity, we suggest the change:

Across all architectures, we observe similar behaviour to Niedoba et al. (2024), Figure 3 - network denoisers exhibit low MSE for both small and large values of $t$, but substantial error for $t \in [0.3, 10]$.

Q2

We use "local inductive bias" to describe a preference for denoising functions where outputs at a spatial position (x,y) are more influenced by nearby input pixels than distant ones. Goyal & Bengio (2022) define inductive bias as a tendency for learning algorithms to favour solutions with certain properties. Kamb & Ganguli (2024) provide a more formal definition of the bias, calling it "locality." However, since Sec 3.1 mainly motivates our hypothesis on generalization via local denoising, we felt a formal definition was unnecessary. We propose the following changes:

• Replacing “local bias” with “local inductive bias” throughout
• Rephrasing the intro of Sec. 3.1 (L142) as: “One potential inductive bias (Goyal & Bengio 2022) of network denoisers is local inductive bias, where denoiser outputs are more sensitive to spatially local perturbations of $z$ than distant ones”

Q3

The networks analyzed have high performance but are not SoTA. Here are their FID scores

Q4

Pixel (x,y) is the output pixel at spatial location x,y where $x \in \\{1, ..., w\\}, y \in \\{1, ..., h\\}$ for image height $h$ & width $w$. We'll update this notation for clarity

Q5

We'll fix this error.

Q6

The subscript x,y,c denotes indexing at position (x,y) and channel $c \in \{1, 2, 3\}$. We'll reword this to improve clarity.

Q7

For each $t$ value, we draw 10k samples from the forward process $z \sim p_t(z | x^{(i)}) p_D(x)$, with $p_t(z | x^{(i)})$ defined on L75.

General Response

We'd like to thank the reviewer for taking the time to read and review or work. Please find our responses to your review below. If there are any items which you believe have not been addressed, we welcome this feedback.

Other Comments Or Suggestions:

Are the labels in Figure 8 correct? Which row corresponds to DDPM++ ?

The labels in Figure 8 are correct, the middle row of each subplot corresponds to DDPM++ (An EDM architecture) evaluated on the $\mathbf{z}$ value in the top row. The left column of subplots show sampling trajectories using the DDPM++ denoiser, while the right column shows sampling trajectories using the PSPC-Flex denoiser.

Questions For Authors:

1. Please clarify the gap between denoising ability and generalization ability. Why the proposed patch denoiser well-matches the diffusion models in Figure 13-20 but diverges significantly in Figure 21? Such gap is not adequately addressed in the current paper. This raises concerns on how much the locally denoising mechanism contributes to the generalization ability.

The reviewer is correct that the samples produced by PSPC are significantly different to those produced by network denoisers. Although we tried our best to match the output of network denoisers as closely as possible, our method is an imperfect approximation to their behaviour. There is a significant difference between our fully empirical method and a deep neural network trained with gradient descent over millions of $\mathbf{z}$ samples. We believe the remarkable similarity of our method our denoiser presents compelling evidence that network denoisers may employ similar generalization mechanisms.

Figure 1 & 6 illustrate that significant differences remain between PSPC and network denoisers, especially for intermediate $t$. When sampling with PSPC-Flex, these differences in denoiser outputs compound to result in PSPC-Flex PF-ODE trajectories which slowly drift from the PF-ODE trajectories of network denoisers. We discuss this briefly in section 6.2, line 371.

This process of error accumulation is visualized in Figure 8 which illustrates the denoiser outputs of PSPC-Flex and DDPM++ on shared $\mathbf{z}$ inputs drawn from DDPM++ (left) and PSPC-Flex (Right) PF-ODE trajectories. Comparing the network and PSPC denoisers, we can see that for all cases and $t$, the outputs are highly similar. However, differences in two, which are most pronounced in the middle of the trajectory, result in substantially worse samples in the right column of figure 8 (PSPC-Flex) than the left column (DDPM++). Despite the artifacts present in PSPC samples, we highlight that the structure of the final samples in the right column of Figure 8 are similar to the final network samples in the right column.

Notably, even as sample quality degrades (ie right column, t < 0.5) both denoisers produce similar denoiser outputs. These degraded $\mathbf{z}$ samples are obviously outside the training distribution $p_t(\mathbf{z})$ and are clear examples of neural network generalization. The fact that PSPC-Flex outputs are similar to network denoisers in this case provides further evidence that a local denoising mechanism may be partially responsible for the generalization of diffusion denoisers.

2. Please compare the proposed denoiser with more (classical patch based) denoisers.

In the methods & evaluation criteria section of your review, you specifically mention Closed-Form Diffusion Models (Scarvellis et al. 2023) as a missing baseline method. We have therefore implemented this method and added it as a baseline to Figure 6 here. We used hyperparameters $\sigma=0.1, M=2$, which they report in Appendix section C2. for their CelebA results. However, we did not find that CFDM outputs were significantly different to the optimal denoiser with these settings.

If there are other baselines you believe are necessary, please cite them explicitly and we will try our best to implement them by the end of the discussion period.

3. Why PCSC shares the highest similarity with the Gaussian denoiser rather than the diffusion model (Figure 9)?

We are unsure as to why PSPC-Flex samples are more similar to the Gaussian denoiser. Recently, several methods have been proposed to explain diffusion generalization, including our work, the optimal Gaussian denoiser, and the geometrically adaptive harmonic bases of Kadkhodaie et. al (2023). Understanding the similarities between these methodologies is an interesting direction for future research.

General Comment

We'd like to thank the reviewer for taking the time to read and review our paper. Furthermore, we appreciate the additional references that the reviewer has brought to our attention. Below, we have responded to several items of your review. If there are items which you feel have not been sufficiently addressed, we would welcome additional discussion on these matters.

Experimental Designs or Analyses:

There are, however, several missing details that I think are important to mention and discuss:

We have clarified both of the details below. Although we believe that our definitions of $C_s$ on line 199, $C_{G(t, \lambda)}$ on line 316, as well as equations 9,10 and 11 address these details, we will update the text in the camera ready revision to improve clarity on these items

• What was the dataset of patches used to construct each empirical patch denoiser (both for PSPC-patch and for PSPC-square)? For example, when constructing the denoised patch at location (i,j), does the computation involve all the overlapping patches extracted from all the training images, or only the patches at location (i,j) extracted from all the training images (namely, only one patch per training image)?

For both PSPC-Flex and PSPC-Square, we construct an empirical patch set for each patch posterior mean. In the case of PSPC-Flex, this corresponds to one patch set for each pixel in the input image (ie $32^2$ patch sets for CIFAR-10 and $64^2$ patch sets for AFHQ & FFHQ). For PSPC-Square, we create a dataset for each possible square patch location with no padding and a stride of one. The total number of patch sets created for PSPC-Square varies depending on patch size. In general, we create $(h - s + 1)^2$ patch sets, where $h$ is the height of the image in pixels and $s$ is the height of the square patch in pixels. For both PSPC-Flex and PSPC-Square, each patch set is created by applying the same cropping matrix to each image in the training set. This results in patch sets with one spatially localized patch per training set image.

• Were denoised patches constructed only for the patches that are fully contained within the 64x64 image, or did patches at the boundaries also contain zero-padded regions?

We did not use padding in our computations.

Essential References Not Discussed:

The paper draws connections to classical patch-based image restoration methods. But there were also quite a few patch-based image generation methods. Most of them in the context of learning from a single image. Starting from the classical texture-generation paper:

• J. De Bonet, "Multiresolution sampling procedure for analysis and synthesis of texture images", SIGGRAPH'97.

To more recent papers, like:

• N. Granot, et al. "Drop the GAN: In defense of patches nearest neighbors as single image generative models", CVPR'22.

These are also related to methods that learn patch statistics using GANs (e.g. SinGAN) or using diffusion models (e.g. SinFusion, SinDiffusion, SinDDM).

We thank the reviewer for drawing our attention to these highly relevant references. We will include the references mentioned in the camera ready version of the paper.

Other Strengths And Weaknesses:

A weakness is that the paper does not explain, and does not provide supporting experiments, how local processing induces coherent global structure. If when denoising a patch, the denoiser doesn't know from where in the image that patch was extracted, how does the denoiser know whether this patch is more likely to contain sky (if it was extracted from the upper part of the image) or grass (if it was extracted from the lower part of the image)?

As mentioned in the experimental design, we use spatially consistent patches for each patch posterior mean. This means that for example, when denoising patches in the upper portion of the image, the patch dataset contains patches which are more likely to contain sky than grass. We found this detail to be especially important for localized datasets such as FFHQ where facial features (noses, eyes, mouths etc.) are located in specific spatial regions of the image.

Another key finding of our work is that the ideal patch size and therefore the degree of locality is anti-correlated with the level of noise. When generating a sample with PSPC-flex or PSPC-square, we initially denoise using large patches before moving to smaller patches. The larger patches therefore condition the posterior means of the later, smaller patch denoising operations. We believe this large to small patch hierarchy is important for ensuring coherent global structure

Other Comments Or Suggestions:

Typo: It seems that in Eq. (8), the input to D_\\theta should be $\\mathbf{z}$ rather than $\\mathbf{x}$

Thank you for this observation. You are correct, the input to D_\\theta is $\\mathbf{z}$. We will correct this mistake in the camera ready.

General Comment

We'd like to thank the reviewer for their thoughtful consideration of our work. Below, we have responded to specific components of your review. If you have additional questions, we would welcome further discussion.

Relation To Broader Scientific Literature:

This paper makes a major contribution to the study of diffusion model generalization/memorization, which is itself a fairly major subfield of diffusion models. It also contributes meaningfully to literature on how generative models generalize and produce 'creative' output. Their main claim is kind of similar to that of Kamb and Ganguli, which they mention, but the details of their approach are a bit different, and it is likely that the two works developed somewhat independently.

Although our method shares many similarities to Kamb & Ganguli, our work is entirely concurrent and independent. There are some differences between our methodologies and experimental results which we will highlight here, but plan to elaborate on in our camera ready revision.

The primary difference between our works is that Kamb & Ganguli utilize shared patch distributions whereas we utilize patch distributions which are localized to specific image locations. We found that this localization was essential for FFHQ & AFHQ where the subject is centered within the image. In addition, Kamb & Ganguli utilize square patches exclusively, while PSPC-Flex uses localized, adaptive patch shapes and sizes. Experimentally, while Kamb & Ganguli restrict their analysis to only convolutional networks, we find similar generalization patterns exist even in convolution free architectures such as DiT and UViT.

Other Strengths And Weaknesses:

As a very minor point, it could be helpful to include some additional discussion of what these findings may mean. This inductive bias ('locality', to use Kamb and Ganguli's terminology) seems to be true for at least U-net- and vision-transformer-based diffusion models trained on images. What about other kinds of architectures or data sets?

Please see our response to this point in the Questions for Authors section

Other Comments Or Suggestions:

line ~158: "plot four such heatmaps in Figure 3" needs a period

line ~359: "can be found in Appendix D" needs a period

We'd like to thank the reviewer for identifying these items. We will fix both in the camera ready revision.

Questions For Authors:

What kind of architectures does one expect to exhibit this inductive bias? If one trained a fully-connected MLP-like architecture to denoise images, would it exhibit this bias?

This is an interesting open question in our opinion. Although more rigorous study is required, we would guess that the inductive biases are probably due to interplay between the diffusion process and the correlations in the data. For example, in datasets where the features are not locally correlated, (ie shuffling the data dimensions), we would not expect this local inductive bias to be useful.

Our observations in this paper seem to suggest that network architectures do not significantly affect the generalization behaviour. We would therefore expect MLPs to have similar properties to DiT & U-Net architectures when trained on the same data. This hypothesis is somewhat supported by prior art. For example, Li et al. (2024) show that the optimal linear denoiser is the Gaussian denoiser which produces outputs which are quite similar to our method.