Nonlinear denoising score matching for enhanced learning of structured distributions
This work addresses challenges in generative modeling for datasets with multimodality or symmetries, offering a method that is more data-efficient and computationally cheaper, though it is incremental as it builds on existing score-based approaches.
The paper tackles the problem of training score-based generative models on structured distributions by introducing nonlinear noising dynamics, which improves learning with less data and reduces computational cost, as demonstrated on low-dimensional examples, high-dimensional images, and latent space representations.
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.