Latent Target Score Matching, with an application to Simulation-Based Inference

TL;DR

This work addresses high-variance diffusion-based score matching when latent variables are present by introducing Latent Target Score Matching (LTSM), which leverages the latent-target identity $\nabla_{\theta_t} \log p_t(\theta_t) = \frac{1}{\alpha(t)} \mathbb{E}_{\theta_0,z|\theta_t}[\nabla_{\theta_0} \log p(\theta_0,z)]$ to provide low-variance supervision of the marginal score. It extends Target Score Matching to latent-variable settings and couples LTSM with a diffusion-time dependent DSM mixture, enabling robust performance across noise scales. The authors derive a variance-minimizing optimal mixture weight $w_t^*$ and also train a learned schedule to approximate it, forming a MIX objective $y_{MIX} = w_t y_{DSM} + (1-w_t) y_{LTSM}$. Across simulation-based inference tasks with gray-box simulators, MIX consistently improves score accuracy and posterior quality, particularly under limited simulator budgets, demonstrating improved sample efficiency for diffusion-based SBI.

Abstract

Denoising score matching (DSM) for training diffusion models may suffer from high variance at low noise levels. Target Score Matching (TSM) mitigates this when clean data scores are available, providing a low-variance objective. In many applications clean scores are inaccessible due to the presence of latent variables, leaving only joint signals exposed. We propose Latent Target Score Matching (LTSM), an extension of TSM to leverage joint scores for low-variance supervision of the marginal score. While LTSM is effective at low noise levels, a mixture with DSM ensures robustness across noise scales. Across simulation-based inference tasks, LTSM consistently improves variance, score accuracy, and sample quality.

Figures (5)

• Figure 1: [Left] Regression target variance vs. diffusion time $t$ for DSM and LTSM across three tasks (\ref{['sec:exps']}). DSM rises as $t\to0$, while LTSM stays low at small $t$, increasing for larger $t$. The time-dependent mixture retains the best of both methods. [Right] Score estimation error, the DSM+LTSM mixture yields the best results.
• Figure 2: MMD (lower is better) vs. simulator budget (i.e. size of the dataset used for training) for each task (rows). Left column averages over five observations $x$; remaining columns show results for different potential values for the observations $x$. The Mixture improves over DSM, with the largest gap happening for the lower number of simulator calls.
• Figure 3: Toy Galton board depiction, extracted from Brehmer et al. brehmer2020mining. The blue and green dark dots on the top figure represent potential paths of a ball dropped from the top. The positioning of the "nails" is defined by the parameters $\theta$, and the final position where the ball lands is denoted by $x$. For different parameters $\theta$ the distribution of positions for the ball's landing changes.
• Figure 4: Variance-optimal DSM weight $w_t^\ast$ vs. time $t$. Computed via \ref{['prop:optimal-w']}. Here $w_t^\ast$ is the coefficient on the DSM target in the mixture: $w_t^\ast{=}0$ means pure LTSM, $w_t^\ast{=}1$ means pure DSM. As expected, $w_t^\ast$ is low near $t\!\approx\!0$ and increases toward $1$ as $t\!\to\!1$.
• Figure 5: Learned DSM weight $w_t$ vs. time $t$. $w_t$ is the DSM coefficient in the mixture (0 = LTSM, 1 = DSM), trained jointly with the score network. It is low near $t\!\approx\!0$ and rises toward $1$ as $t\!\to\!1$, closely tracking $w_t^\ast$.

Theorems & Definitions (6)

• Proposition 3.1: Latent Target Score Identity
• Proposition 3.2: Optimal mixture weight
• proof
• proof
• Proposition B.1: Latent Target Score Matching (LTSM)
• proof