Where the Score Lives: What Wavelets Reveal About Diffusion Models

The Paradox of Generalization in Diffusion Models

Diffusion models have demonstrated a remarkable ability to generate novel, high-quality samples, yet a significant theoretical question remains: why don’t they simply memorize the training data? In theory, the optimal denoiser should learn the empirical distribution of the data, which would lead it to “snap” a noisy input back to the nearest training sample. This suggests a tendency towards memorization rather than generalization.

Two primary schools of thought attempt to explain this phenomenon:

1. Architectural Inductive Biases: Work by researchers like Kam and Ganguli suggests that the architecture of the score network, typically a U-Net or CNN, introduces crucial inductive biases. Properties like a limited receptive field and translation equivariance lead to a “patch mosaic” form of creativity. The model learns to denoise local patches by matching them to a vast dataset of patches from all training images, regardless of their original spatial location. This can explain some generative behaviors but also failures, such as generating a hand with six fingers when the model’s local view doesn’t account for the global context.

2. Inherent Data Properties: Other research argues that generalization is a property of the image statistics themselves, independent of the architecture. Studies show that for natural images, the sensitivity of an output pixel to input pixels is highly localized at low noise scales and expands at higher noise scales. This behavior is observed across different architectures (U-Nets, Transformers) and even mirrors that of an optimal linear denoiser derived from the data’s covariance matrix. This suggests that the data’s structure, where high-variance principal components correspond to low-frequency features, naturally encourages local denoising.

A Wavelet-Based Framework for Interpretation

To reconcile and investigate the interplay between these two factors, this work proposes an analytically solvable, wavelet-based parameterization of the score function. Wavelets serve as a natural tool for this analysis due to their long history in image processing and their ability to provide a multi-scale basis that captures both location and scale (frequency). Unlike the Fourier transform, which struggles with localized signals, wavelets are ideal for representing image features compactly.

The core idea is to express the score function as a linear combination of wavelet basis functions: Score(xt) = Σ c_i * W_i(xt)

Instead of learning the coefficients c_i with a black-box neural network, they are modeled as a linear function of features φ(xt) extracted from the noisy image: c_i = α_i * φ(xt)

This construction allows for a closed-form analytical solution. By plugging this parameterization into the standard diffusion loss and leveraging the orthonormality of the wavelet basis, the loss function decouples. Using Stein’s identity to eliminate the unknown true score, the optimal parameters α can be estimated directly from the data without iterative training.

Probing Model Behavior with Feature Engineering

The framework’s power lies in its flexibility; by carefully choosing the feature vector φ(xt), one can explicitly model and test different architectural assumptions. Three model classes were investigated:

1. Independent Baseline: Models each wavelet coefficient independently. This serves as a baseline to measure the value of cross-coefficient structures. The results are reasonable but show diminishing returns with increased model complexity (polynomial degree) at high noise levels.

2. Band-Tied Model: Models interactions between different wavelet orientations (horizontal, vertical, diagonal) at a fixed scale and location. This is designed to mimic the channel-mixing inductive bias in CNNs, which helps in identifying edges and corners. While it can improve perceptual quality for details, it tends to introduce more artifacts than it solves at high noise regimes, where such fine details are obscured.

3. Local Coupled Model: Allows wavelet coefficients within a local spatial neighborhood to interact. This directly tests the impact of locality. This model showed the most consistent and significant performance improvements across all noise levels. The gains did not diminish with increased model complexity, reinforcing the idea that local interactions are the most strongly informative and robust inductive bias for denoising.

Key Takeaways and Future Directions

The experiments demonstrate that the wavelet-based models can be competitive with trained denoisers, particularly at low noise scales. This suggests that complex, trained models are likely learning similar local dependency structures. As expected, the performance gap widens at high noise, where the simplified models struggle.

The main conclusion is that this wavelet framework provides a powerful and interpretable “probe” to understand what aspects of a model’s architecture and the data’s statistics contribute most to generalization.

Future work aims to:

• Develop more robust theoretical guarantees for the estimators.
• Apply the framework to higher-resolution color images and larger datasets.
• Incorporate multi-scale dependencies to better mimic U-Net architectures.
• Compare performance against optimal linear denoisers to further disentangle linear and non-linear effects.

• Tags:
• Generalization
• Feature
• Diffusion Models
• Denoising
• Memorization
• Wavelets
• Score-Based Models
• Interpretable AI