Heavy-tailed denoising score matching

TL;DR

The paper tackles the high-dimensional limitations of Gaussian denoising score matching by introducing heavy-tailed denoising score matching (HTDSM) with generalised normal noise. It shows that DSM objectives remain equivalent when the score is differentiable almost everywhere, analyzes how GN noise concentrates and skews in high dimensions, and proposes a quantile-based noise-scale scheme for annealed Langevin dynamics. Empirically, HTDSM improves score estimation and sampling quality, helps mitigate class-imbalance effects, and yields competitive unconditional generation results on MNIST, Fashion-MNIST, CIFAR-10, and CelebA, with sub-Gaussian diffusion aiding exploration of sparse regions. The work outlines a path toward non-Gaussian diffusion, including a continuous SDE framework and potential Lévy-flight-like sampling for future exploration.

Abstract

Score-based model research in the last few years has produced state of the art generative models by employing Gaussian denoising score-matching (DSM). However, the Gaussian noise assumption has several high-dimensional limitations, motivating a more concrete route toward even higher dimension PDF estimation in future. We outline this limitation, before extending the theory to a broader family of noising distributions -- namely, the generalised normal distribution. To theoretically ground this, we relax a key assumption in (denoising) score matching theory, demonstrating that distributions which are differentiable almost everywhere permit the same objective simplification as Gaussians. For noise vector norm distributions, we demonstrate favourable concentration of measure in the high-dimensional spaces prevalent in deep learning. In the process, we uncover a skewed noise vector norm distribution and develop an iterative noise scaling algorithm to consistently initialise the multiple levels of noise in annealed Langevin dynamics (LD). On the practical side, our use of heavy-tailed DSM leads to improved score estimation, controllable sampling convergence, and more balanced unconditional generative performance for imbalanced datasets.

Figures (21)

• Figure 1: DSM training and LD sampling. In a, $p(\mathbf{x})$ is modelled as an additive mixture of ($k=2$) bivariate Gaussians with 20,000 samples. An MLP is trained to estimate the score from samples noised by $q_{\sigma}(\Tilde{\mathbf{x}}|\mathbf{x})\sim\mathcal{N}(\mathbf{x},\mathbf{I})$. 1,000 sampled paths are evolved in b to demonstrate the decision boundary, its asymmetry (relevant for class imbalance), and the upper bound on approximation accuracy due to the underlying unit noise. Full details in Figure \ref{['fig:dsm example 2']}.
• Figure 2: Multiple noise level DSM training and ALD sampling. The setup and figures are identical to Figure \ref{['fig:dsm example 1']} except that two noise scales, $\sigma_{1}=1.0$ and $\sigma_{2}=0.25$, are used. Full details in Figure \ref{['fig:dsm ald example 2']}.
• Figure 3: The generalised normal distribution for varied $\beta$.
• Figure 4: Laplace DSM with ALD. The setup and figures are identical to Figure \ref{['fig:dsm ald example 1']}, except that Laplace noise is used ($\beta=1$ in the general formulation). a depicts the diamond, rather than circular, noise structure of a diagonal bivariate Laplace distribution. b demonstrates that ALD sampling converges even with Laplace (sub-Gaussian, piece-wise differentiable) diffusion, confirming Theorem \ref{['thm:piecewise differentiable']}. Full details in Figure \ref{['fig:dsm ald laplace 2']}.
• Figure 5: Laplace DSM with ALD. The setup and figures are identical to Figure \ref{['fig:dsm ald 10x example 2']}, except that Laplace noise is used. b demonstrates that Laplace noise compensates for the class imbalance. $29.5\%$ of particles finishing in mode 2 even manages to overcompensate.

Theorems & Definitions (4)

• Theorem 2.1
• proof
• Proposition 2.1
• proof