Latent Nonlinear Denoising Score Matching (LNDSM)

• The paper introduces LNDSM, which integrates nonlinear latent dynamics with control variates to stabilize training and enhance generative performance.
• It employs a VAE framework augmented with a Gaussian mixture prior to capture multimodal latent representations for mode-balanced sampling.
• Empirical results on MNIST variants demonstrate state-of-the-art improvements in FID, IS, and generative diversity while reducing gradient variance.

Latent Nonlinear Denoising Score Matching (LNDSM) is a training paradigm for score-based generative models that combines nonlinear stochastic dynamics in latent space with variational autoencoder (VAE) architectures. Its core innovation is to embed structured priors—such as multimodalities and approximate symmetries—into the forward diffusion mechanism, yielding sample-efficient, mode-balanced generation on tasks where latent structure departs from simple isotropic Gaussianity. LNDSM achieves numerically stable optimization by identifying and subtracting high-variance, zero-mean control variates in the discretized denoising score matching (DSM) loss, and enables end-to-end training of encoder, decoder, and score networks. The approach has demonstrated improvements over baseline conditional score-based models and structure-agnostic latent SGMs, with state-of-the-art results for sample quality, mode coverage, and generative diversity on tasks such as MNIST and its symmetry-perturbed variants (Birrell et al., 24 May 2024, Shen et al., 7 Dec 2025).

1. Foundations in Score-Based Generative Modeling

Score-based generative models (SGMs) define a diffusion process in data or latent space $\mathbb{R}^d$, typically specified by a stochastic differential equation (SDE):

$dx_t = f(x_t, t)\,dt + g(t)\,dw_t,$

where $f$ is a drift term, $g(t)$ a scalar noise schedule, and $w_t$ a standard Wiener process. The generative task involves running the reverse-time SDE, which requires estimation of the score $\nabla_x \log p_t(x)$ at each time $t$. Score matching trains a neural network $s_\theta(x,t)$ to approximate this score via objectives derived from denoising score matching (DSM).

Latent SGMs (LSGM) operate in the latent space of a VAE, decoupling inference from high-dimensional data and allowing tractable modeling of more complex distributions. Classical methods rely on linear drifts (Ornstein-Uhlenbeck process), enforcing isotropic Gaussian priors.

2. LNDSM: Nonlinear Forward Dynamics in Latent Space

LNDSM generalizes latent diffusion to nonlinear drifts by replacing the linear OU process with

$dz_t = f(z_t, t)\,dt + g(t)\,dW_t,$

where $f(z, t) = -\nabla_z V(z)$ and $V(z)$ is a potential encoding latent structure, typically the negative log-density of a fitted $K$-component Gaussian mixture:

$$V(z) = -\log \eta_*(z),\quad \eta_*(z) = \sum_{i=1}^K w_i \mathcal{N}(z; \mu_i, \Sigma_i).$$

This choice yields overdamped Langevin dynamics with multi-well stationary density reflecting multimodal clusters or approximate symmetries in the latent representations.

3. Variational Objective and Discretization

The VAE loss for a datapoint $x$ combines reconstruction and KL divergence:

$$L(\phi,\theta,\psi;x) = \mathbb{E}_{z_0\sim q_\phi(z_0|x)}[-\log p_\psi(x|z_0)] + KL[q_\phi(z_0|x)\,\|\,p_\theta(z_0)].$$

LNDSM reinterprets the cross-entropy term via an SDE-induced Markov chain and Euler–Maruyama discretization. For time partition $0 = t_0 < \ldots < t_{n_f} = T$, latent transitions are simulated:

$$z_n = z_{n-1} + f(z_{n-1}, t_{n-1}) \Delta t_{n-1} + g(t_{n-1}) \sqrt{\Delta t_{n-1}} U_{n-1},$$

$U_{n-1} \sim \mathcal{N}(0, I)$, with $n_f$ steps per batch.

Expectation over a uniformly sampled index $N$ substitutes time averaging, reducing computational complexity and variance.

4. Control Variates and Stable Training

Direct computation of score terms in the DSM loss introduces variance-exploding components $(U_{N-1}/\sigma_{N-1})$, which scale as $1/\sqrt{\Delta t}$ and destabilize stochastic gradients as discretization becomes fine. LNDSM applies two analytic interventions:

• Subtraction of zero-mean control variate terms: Expansions of $s_\theta(z_N)$ and $f(z_N)$ about $\mu_{N-1}$ reveal high-variance, zero-mean terms, which are removed to stabilize training.
• Neural control variate parameterization: Additional auxiliary networks $\epsilon_\varphi(t)$ further minimize the variance of gradient estimates via standard control-variate fitting.

The final stable cross-entropy approximation is

$$\begin{align*} CE(q_\phi\,\|\,p_\theta) &\approx H[q_\phi(z_{n_f}|x)] \ &\qquad + \frac{1}{2} \mathbb{E}_{N,z_N}[g^2(t_N) \|s_\theta(z_N,t_N)\|^2] \ &\qquad + \mathbb{E}_{N,z_{N-1},z_N} \left[ g^2(t_N) \frac{U_{N-1}}{\sigma_{N-1}} \cdot (s_\theta(z_N, t_N) - s_\theta(\mu_{N-1}, t_N)) \right.\ &\qquad \left. - \frac{U_{N-1}}{\sigma_{N-1}} \cdot (f(z_N, t_N) - f(\mu_{N-1}, t_N)) \right]. \end{align*}$$

This loss can be computed in closed form for Gaussian transitions.

5. Integration into VAE Framework and Sampling Procedure

LNDSM is implemented by several preparatory and iterative steps:

1. Preprocess the dataset $\{x_i\}$ via a fixed or learned encoder $q_\phi$, collect latent codes $\{z_i\}$.
2. Fit a GMM to latent codes, freeze parameters to define structured reference $\pi(z)$.
3. Simulate latent SDE forward chains per batch, apply the stable LNDSM objective.
4. Jointly backpropagate and update encoder ($\phi$), decoder ($\psi$), and score net ($\theta$). In cases where entropy of $\pi$ is intractable, encoder learning rate is reduced and $\pi$ parameters are frozen.

The generative sampling procedure consists of:

• Sampling $z_T$ from the structured prior $\pi(z)$.
• Integrating the reverse-time SDE from $t = T$ to $0$ using learned $s_\theta(z, t)$.
• Decoding $x$ via $p_\psi(x|z_0)$.

6. Empirical Results and Comparative Metrics

Experiments on MNIST and approximate symmetry (C₂)-MNIST demonstrate LNDSM's advantages in both abundant and scarce-data regimes (Shen et al., 7 Dec 2025). Key quantitative metrics are tabulated below:

Model	FID ↓	IS ↑
NDSM-SGM	36.1	8.93
LSGM	6.7	9.42
LNDSM	2.2	9.75

On low-data (N=3,000) MNIST: | Model | FID ↓ | IS ↑ | |------------|-------|------| | LSGM | 27.1 | 9.12 | | LNDSM | 15.0 | 9.57 |

On approx-C₂-MNIST: | Model | FID ↓ | IS ↑ | |------------|-------|------| | LSGM | 12.4 | 9.01 | | LNDSM | 5.3 | 9.48 |

LNDSM more accurately recovers target class proportions, with Kullback-Leibler divergence of label frequencies reduced from 0.008 to 0.00114 (full MNIST) and from 0.00883 to 0.00439 (low data). Qualitative assessment finds sharper samples, fewer artifacts, and stronger generative diversity. Synthesis in latent space is faster ($\approx$0.068 s/batch with 100 EM steps, latent 128-dim) compared to pixel-space NDSM ($0.35$ s/batch).

7. Properties, Limitations, and Extensions

Theoretical analysis establishes consistency: as network capacity increases and discretization step decreases, minimizers of the LNDSM objective converge to the true transition score. Control-variate removal reduces gradient variance scaling from $O(1/\Delta t)$ to $O(1)$, enabling stable training without importance sampling in diffusion time.

LNDSM is agnostic to hard equivariant constraints, accommodating approximate symmetries not captured by equivariant architectures. The method leverages inexpensive preprocessing to fit the latent GMM, integrating domain-specific structure with minimal overhead.

A plausible implication is that LNDSM can be generalized beyond the tested VAE-SGM family, provided structured priors in latent space are computable or estimable. However, the approach assumes latent representations remain sufficiently expressive for target data, and that fitted latent GMMs do not inadvertently bias sampling against rare but valid modes.

LNDSM unites the efficiency of latent SGMs with the structural flexibility of nonlinear SGMs. By discretizing the nonlinear latent diffusion, reformulating the VAE cross-entropy in terms of Gaussian transitions, and removing two zero-mean exploding control variates, one obtains a stable, importance-sampling-free training loss. Empirical results confirm gains in FID, IS, mode coverage, and overall training/sampling efficiency versus both LSGM and full-space NDSM-SGM (Shen et al., 7 Dec 2025, Birrell et al., 24 May 2024).