Generalization in diffusion models

The density model is as good as the denoiser, so studying the density means studying the denoiser. Substitute $s_\theta=(f_\theta(y)-y)/\sigma^2$ into Miyasawa's identity. By the Pythagorean decomposition around $\mathbb E[x\mid y]$, the score error becomes a denoising gap, $\mathbb E\|\nabla\log p_\sigma-s_\theta\|^2=\sigma^{-4}\big(\mathrm{MSE}(f_\theta,\sigma^2)-\mathrm{MSE}(f^\star,\sigma^2)\big)$, and the likelihood bound for score-based models, for the variance-exploding SDE, gives $$D_{\rm KL}(p\,\|\,p_\theta)\le\int_0^\infty\big(\mathrm{MSE}(f_\theta,\sigma^2)-\mathrm{MSE}(f^\star,\sigma^2)\big)\,\sigma^{-3}\,d\sigma .$$ So learning the density is optimal denoising at every noise level, with small noise weighted heavily. The second-order identity says what optimal denoising looks like: $\mathrm{MSE}(f^\star,\sigma^2)=\mathbb E\,\mathrm{tr}\,\mathrm{Cov}[x\mid y]=\sigma^2\,\mathbb E\,\mathrm{tr}\,\nabla f^\star(y)$, so a small error means a Jacobian with few non-negligible eigenvalues. SURE states the same trade-off for any denoiser: fit the noisy image while keeping the Jacobian's trace, an effective rank, small. The bound is never evaluated for real images, where $f^\star$ is unknown.

Denoisers trained on more images fit their training set worse and agree with each other more. UNets are trained on $N\in\{1,10,\dots,10^5\}$ CelebA faces. At $N=1$ and 10 training PSNR rises with unit slope while test PSNR is low and flat, which is what a denoiser that returns stored images produces; at $N=10^5$ the curves coincide, training PSNR having fallen as test PSNR rose. Two networks trained on disjoint sets and sampled from the same noise agree with this: for $N\le100$ each returns an image from its own training set, and at $N=10^5$ nearly the same image as the other network, unlike the nearest training images. Two things follow. The networks are far from the exact empirical minimizer, whose training error vanishes at small noise, so architecture and optimization set their behavior; the objective alone would memorize. And agreement across disjoint data shows vanishing model variance directly, which train-test parity alone does not.

A bias-free network is exactly adaptive shrinkage, and its basis can be read off. With ReLUs and no additive constants the map is positively homogeneous and piecewise linear, so $f(y)=\nabla f(y)\,y$. Diagonalizing the Jacobian, $$f(y)=\sum_k\lambda_k(y)\,\langle y,e_k(y)\rangle\,e_k(y),$$ a shrinkage whose factors $\lambda_k(y)$ and basis $e_k(y)$ both depend on the noisy image. For a UNet trained on $10^5$ faces, the clean image's coefficients $\langle x,e_k(y)\rangle$ decay quickly, so the image lies mostly in the span of the directions the denoiser keeps, and the factors fall with $k$ from slightly above 1 toward 0. The eigenvectors are oscillations that follow the contours of the face and fill its smooth regions, with frequency rising as $\lambda_k$ falls.

Approximation theory rates a basis by the decay of the image's coefficients in it. Fix an orthonormal basis and let an oracle choose $\lambda_k$ from the clean image. Its risk $\sum_k(1-\lambda_k)^2c_k^2+\lambda_k^2\sigma^2$, with $c_k=\langle x,e_k\rangle$, is minimized at $\lambda_k=c_k^2/(c_k^2+\sigma^2)$, the Bayes linear factor for a coefficient of prior variance $c_k^2$, and the minimal risk is within a factor 2 of an $M$-term approximation error, $$\sum_k\frac{\sigma^2c_k^2}{c_k^2+\sigma^2}\ \asymp\ \sum_k\min(c_k^2,\sigma^2)=M\sigma^2+\|x-x_M\|^2,$$ where $M$ counts coefficients above the noise. If the sorted coefficients obey $c_k^2\asymp k^{-(\alpha+1)}$, then $M\asymp\sigma^{-2/(\alpha+1)}$, both terms are of order $\sigma^{2\alpha/(\alpha+1)}$, and the PSNR curve has asymptotic slope $\alpha/(\alpha+1)$; soft thresholding of the noisy coefficients reaches this rate up to a log factor (Donoho and Johnstone). Adapting the basis to the image can only steepen the slope, but the basis must be chosen from $y$. Choosing the best basis for $y$ from a dictionary of polynomially many vectors and thresholding attains the oracle rate, while unlimited adaptivity such as $e_1(y)=y/\|y\|$ has too much variance. The question for networks is which dictionary their bias implements.

Where a GAHB is optimal, the networks nearly reach the optimal rate. Geometric $C^\alpha$ images are $\alpha$-regular on regions separated by $\alpha$-regular contours. Their optimal sparse representation is a bandlet basis, harmonics in the smooth regions and along the contours, whose sorted coefficients decay as $k^{-(\alpha+1)}$, so by Step 4 the optimal PSNR slope is $\alpha/(\alpha+1)$. UNets trained on $10^5$ such images come close to it (the slopes in Main findings), and their top eigenvectors are harmonics on the two regions and along the boundary, a bandlet-like basis with more flexible geometry. Model bias is therefore small on this class in the sense of the rate, though the slope fixes only the exponent and says nothing about the constant.

Where a GAHB is not optimal, the networks still use one, which identifies the bias. Disks with random center, radius and foreground and background intensities form a five-dimensional curved manifold. For noise small against the curvature the optimal denoiser projects onto the five-dimensional tangent space, with error about $5\sigma^2$ and PSNR slope 1. A UNet trained on $10^5$ disks gets the first five eigenvectors right (they match the tangent basis) but keeps further GAHB eigenvectors with non-negligible $\lambda_k$, oscillating along the rim and across the background, whose number grows as $\sigma$ decreases, so the slope stays below 1. Faces scrambled by one fixed pixel permutation have exactly the optimal error of ordinary faces (unscramble, denoise, rescramble), yet the network trained on them has a much lower PSNR slope and noisier bases. So the preference is specifically for geometry and locality in pixel space; low-dimensional structure in the data is not enough.