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.

翻译：本文提出了一种训练基于分数的生成模型的新方法，该方法利用非线性加噪动力学来改进对结构化分布的学习。通过推广至非线性漂移项，可以在动力学中融入额外的结构信息，从而使训练过程更好地适应数据特性，例如在多模态分布或（近似）对称性场景中。此类结构可通过低成本的数据预处理步骤从数据中提取。非线性动力学的引入为训练带来了新的挑战，我们通过两种方式应对：1）开发了一种新的非线性去噪分数匹配方法；2）引入神经控制变量以降低NDSM训练目标的方差。我们在多个示例中验证了该方法的有效性：a）针对潜在空间聚类问题设计的一系列低维示例；b）高维图像数据，解决了模态不平衡、训练集规模小及近似对称性问题——后者对于基于等变神经网络的方法尤为挑战，因其要求精确对称性；c）高维数据的潜在空间表示，展示了在显著降低计算成本的同时提升性能的效果。本方法通过灵活融入数据集中涌现的结构特征，能够以更少的数据量学习基于分数的生成模型。