Score-based generative models commonly employ linear forward diffusions with Gaussian reference distributions, limiting their ability to directly incorporate structural information into the noising dynamics. We introduce latent nonlinear denoising score matching (LNDSM), which combines nonlinear, structure-informed latent diffusions with the variational autoencoder (VAE) training framework. In our construction, the forward diffusion is designed to admit a prescribed reference distribution, such as a Gaussian mixture model, that reflects qualitative features of the latent data distribution. We derive a computable approximation to a cross-entropy upper bound in the VAE objective using Euler-Maruyama discretization and the resulting Gaussian conditional transitions. To stabilize training, we remove two zero-mean control variates whose variances become large for small time steps. The resulting objective enables joint training of the encoder, decoder, and latent score network. Numerical experiments on benchmark image datasets demonstrate the effectiveness of LNDSM in learning structured distributions, outperforming conventional latent score-based generative modeling in sample quality and mode balance across several settings, while achieving improved training efficiency. These results highlight the potential of nonlinear, structure-informed latent diffusions for generative modeling.
Figures & tables
Networks
# parameters (MNIST)
# parameters (CIFAR-10)
NVAE
207,648
299,518
NCSNpp
895,076
2,177,608
NVAE + NCSNpp
1,102,724
2,477,126
U-Net in NDSM-SGM
1,115,297
/
Table 1: Number of trainable parameters
Figure 1: The average log-likelihood plateaus at a terminal time of T=0.5 , indicating that the latent-space diffusion process has brought the latent variable distribution sufficiently close to the target GMM.
Figure 2: (a): For the MNIST dataset, LNDSM-SGM only needs a few training epochs to convergence, whereas the benchmark LSGM requires 500 epochs. The total training time of the proposed LNDSM-SGM is lower. (b): For the CIFAR-10 dataset, the total training time is capped at 48 hours. LNDSM-SGM is trained for 100 training epochs to convergence. The benchmark LSGM is trained for 436 epochs to convergence.
Figure 3: LSGM and LNDSM-SGM trained using the full MNIST dataset. (a) and (d): Snapshot and fraction of different digits of the training samples. (b) and (e): Snapshot and fraction of different digits of the 10,000 samples generated by the LSGM. The KL divergence from the fraction of the training samples to the fraction of the generated samples: 0.008. (c) and (f): Snapshot and fraction of different digits of the 10,000 samples generated by the LNDSM-SGM. The KL divergence of the fraction of the training samples from the fraction of the generated samples: 0.00114.
Table 2: FID and IS results on MNIST and CIFAR-10
Figure 4: LSGM and LNDSM-SGM trained using the CIFAR-10 dataset. (a) and (d): Snapshot and fraction of different classes of the training samples. (b) and (e): Snapshot and fraction of different classes of the 10,000 samples generated by the LSGM. The KL divergence from the fraction of the training samples to the fraction of the generated samples: 0.294. (c) and (f): Snapshot and fraction of different classes of the 10,000 samples generated by the LNDSM-SGM. The KL divergence from the fraction of the training samples to the fraction of the generated samples: 0.227.
Appendix figures & tables3 assets
Supplementary material from the paper’s appendix.
Appendix
Figure 5: LSGM and LNDSM-SGM trained using the low data MNIST dataset. (a) and (d): Snapshot and fraction of different digits of the training samples. (b) and (e): Snapshot and fraction of different digits of the 10,000 samples generated by the LSGM. The KL divergence from the fraction of the training samples to the fraction of the generated samples: 0.00883. (c) and (f): Snapshot and fraction of different digits of the 10,000 samples generated by the LNDSM-SGM. The KL divergence from the fraction of the training samples to the fraction of the generated samples: 0.00439.
model
FID ( ↓ )
IS ( ↑ )
KL ( ↓ )
LSGM
27.1
9.12
0.00883
LNDSM-SGM
15.0
9.57
0.00439
Appendix
Table 3: Low data MNIST ( N=3000 )
Figure 6: LSGM and LNDSM-SGM trained using the full Approx.-C2-MNIST dataset. (a) and (d): Snapshot and fraction of different digits of the training samples. (b) and (e): Snapshot and fraction of different digits of the 10,000 samples generated by the LSGM. (c) and (f): Snapshot and fraction of different digits of the 10,000 samples generated by the LNDSM-SGM.
Score-based generative modeling learns data-distribution gradients to generate samples through reverse diffusion, flows, and related stochastic dynamics.
RIKEN AIP & The Institute of Statistical Mathematics · University of Sydney · Agency for Science, Technology and Research & The Institute of Statistical Mathematics +4
State Key Laboratory of General Artificial Intelligence, School of Intelligence Science and Technology, Peking University · University of Chinese Academy of Sciences · Tencent +1