Nonlinear denoising score matching for enhanced learning of structured distributions

TL;DR

This paper introduces a nonlinear denoising score matching approach for training score-based generative models that better captures data structure, improving learning efficiency and performance on complex, high-dimensional datasets.

Contribution

It proposes a novel nonlinear noising dynamics framework with neural control variates to enhance score-based generative modeling, especially for structured and multimodal data.

Findings

Effective on low-dimensional clustering tasks

Improves image generation with mode balance and symmetry handling

Reduces computational cost in high-dimensional latent space

Abstract

We present a novel method for training score-based generative models which uses nonlinear noising dynamics to improve learning of structured distributions. Generalizing to a nonlinear drift allows for additional structure to be incorporated into the dynamics, thus making the training better adapted to the data, e.g., in the case of multimodality or (approximate) symmetries. Such structure can be obtained from the data by an inexpensive preprocessing step. The nonlinear dynamics introduces new challenges into training which we address in two ways: 1) we develop a new nonlinear denoising score matching (NDSM) method, 2) we introduce neural control variates in order to reduce the variance of the NDSM training objective. We demonstrate the effectiveness of this method on several examples: a) a collection of low-dimensional examples, motivated by clustering in latent space, b) high-dimensional images, addressing issues with mode imbalance, small training sets, and approximate symmetries, the latter being a challenge for methods based on equivariant neural networks, which require exact symmetries, c) latent space representation of high-dimensional data, demonstrating improved performance with greatly reduced computational cost. Our method learns score-based generative models with less data by flexibly incorporating structure arising in the dataset.