One-sentence Summary

The authors propose adapting score matching to settings with partially missing data through two variants: an importance weighting approach with finite sample bounds suited to small-sample, low-dimensional cases, and a variational approach that performs best in high-dimensional settings, demonstrated on graphical model estimation tasks with real and simulated data.

Key Contributions

• Score matching and major score-matching extensions are adapted to estimate full data scores from partially observed samples, using a mask-based formulation for arbitrary missing-coordinate sets and covering missing-completely-at-random and missing-not-at-random scenarios.
• Two missing-data objectives are introduced, an importance-weighting approach and a variational approach, with extensions to truncated, sliced, and denoising score matching and finite-sample bounds for truncated importance-weighted estimation on finite domains.
• The methods are evaluated on simulated and real data for parameter estimation and Gaussian graphical model edge detection, showing importance weighting is strongest in lower-dimensional, small-sample cases and the variational approach is strongest in more complex high-dimensional settings.

Introduction

Score matching is a flexible density-free method for learning the score function of a distribution, supporting applications such as energy based modeling, clustering, and diffusion processes. Prior work on missing data often focuses on neural network diffusion models that require artificial corruption and do not suit explicit parametric score estimation, while related generative and latent variable methods are constrained by model class, convergence uncertainty, or biased estimation. The authors leverage missing score matching, extending score matching to partially observed multidimensional data through two methods: a marginal importance weighting approach with finite sample guarantees and a marginal variational approach for higher dimensional settings. Their framework is compatible with explicit parameterizations and neural networks, and adapts classical, truncated, sliced, and denoising score matching for parameter estimation and graphical model tasks.

Dataset

1. Dataset composition and source

   • The authors use yeast gene expression data introduced by Brem and Kruglyak in 2005.
   • The original dataset contains expression readings for 7086 genes or ORFs measured across 262 yeast segregants.
   • In the analysis, genes are treated as dimensions and segregants are treated as samples.

2. Key subset and filtering rules

   • The full dataset is reduced to 106 genes.
   • A gene is retained only if it is present in at least 95 percent of the samples.
   • This filtering step selects genes with very high availability across the 262 segregants.

3. How the data is used

   • The authors use the filtered gene expression matrix to learn relationships among genes.
   • The data is used in a Gaussian graphical model estimation experiment.
   • Results for this dataset are shown in Figure 7, where the Marg-Var approach is compared with other methods.

4. Processing and metadata details

   • Each sample corresponds to one yeast segregant.
   • Each feature or dimension corresponds to one gene or ORF.
   • The main processing detail is the presence-based subset rule, which keeps only genes observed in at least 95 percent of segregants.
   • No additional cropping strategy is described for this dataset.

Method

The proposed method reframes density estimation under missing observations as a marginal score matching problem. The central idea is to learn an unnormalized model density $q_\theta:\mathcal X\to(0,\infty)$qθ​:X→(0,∞) by training its score

$$\boldsymbol{s}_\theta(\boldsymbol{x})=\nabla_{\boldsymbol{x}}\log q_\theta(\boldsymbol{x}),$$sθ​(x)=∇x​logqθ​(x),

rather than maximizing a marginal likelihood that would require intractable integrals over the unobserved coordinates.

The data are represented by a random vector $X\in\mathcal X\subseteq\mathbb R^d$X∈X⊆Rd. For a subset of observed coordinates $\lambda\subseteq[d]$λ⊆[d], the marginal score associated with $q$q is defined as

$$\boldsymbol{s}_\lambda(\boldsymbol{x}_\lambda) = \nabla_{\boldsymbol{x}_\lambda} \log \int q(\boldsymbol{x})\,\mathrm d\boldsymbol{x}_{-\lambda}.$$sλ​(xλ​)=∇xλ​​log∫q(x)dx−λ​.

This definition ensures that the marginal score of a valid unnormalized density is itself a genuine score function. The corresponding marginal Fisher divergence is

$$F_{\mathrm M}(\theta) = \mathbb E\left[ \left\| \boldsymbol{s}_\Lambda(X_\Lambda) - \boldsymbol{s}_{\Lambda;\theta}(X_\Lambda) \right\|^2 \right],$$FM​(θ)=E[∥sΛ​(XΛ​)−sΛ;θ​(XΛ​)∥2],

where $\Lambda$Λ denotes the random observation pattern induced by missingness. Under regularity and identifiability assumptions, the authors derive an observable surrogate

$$L_{\mathrm M}(\theta) = \mathbb E\left[ 2\nabla_{X_\Lambda}\cdot\boldsymbol{s}_{\Lambda;\theta}(X_\Lambda) + \left\| \boldsymbol{s}_{\Lambda;\theta}(X_\Lambda) \right\|^2 \right] = F_{\mathrm M}(\theta)-C.$$LM​(θ)=E[2∇XΛ​​⋅sΛ;θ​(XΛ​)+∥sΛ;θ​(XΛ​)∥2]=FM​(θ)−C.

A Monte Carlo estimate is then obtained from incomplete samples $\{X_{\Lambda_i}^{(i)}\}_{i=1}^n${XΛi​(i)​}i=1n​:

$$\hat\theta = \arg\min_\theta \frac{1}{n} \sum_{i=1}^n \left[ \nabla_{X_{\Lambda_i}^{(i)}}\cdot \boldsymbol{s}_{\Lambda_i,\theta}(X_{\Lambda_i}^{(i)}) + \left\| \boldsymbol{s}_{\Lambda_i;\theta}(X_{\Lambda_i}^{(i)}) \right\|^2 \right].$$θ^=argθmin​n1​i=1∑n​[∇XΛi​(i)​​⋅sΛi​,θ​(XΛi​(i)​)+​sΛi​;θ​(XΛi​(i)​)​2].

This formulation generalizes classical score matching to settings where the density of interest must be marginalized over unobserved variables. It also extends naturally to compactly supported distributions through a truncated marginal objective. For compact $\mathcal X$X with boundary $\partial\mathcal X$∂X, one introduces a coordinatewise weighting function $g_\Lambda$gΛ​ that vanishes near the boundary. The marginal truncated Fisher divergence uses weights to suppress the influence of boundary regions:

$$F_{\mathrm{TM}}(\theta) = \mathbb E\left[ \left\| g_\Lambda(X_\Lambda)^{1/2} \circ \left( \boldsymbol{s}_\Lambda(X_\Lambda) - \boldsymbol{s}_{\Lambda;\theta}(X_\Lambda) \right) \right\|^2 \right].$$FTM​(θ)=E[​gΛ​(XΛ​)1/2∘(sΛ​(XΛ​)−sΛ;θ​(XΛ​))​2].

The corresponding training objective contains both score curvature terms and boundary correction terms:

$$L_{\mathrm{TM}}(\theta) = \mathbb E\left[ \sum_{j\in\Lambda} g_\Lambda(X_\Lambda)_j \left( s_{\Lambda;\theta}(X_\Lambda)_j^2 + 2\partial_j s_{\Lambda;\theta}(X_\Lambda)_j \right) + \sum_{j\in\Lambda} 2\partial_j g_\Lambda(X_\Lambda)_j s_{\Lambda;\theta}(X_\Lambda)_j \right].$$LTM​(θ)=E​j∈Λ∑​gΛ​(XΛ​)j​(sΛ;θ​(XΛ​)j2​+2∂j​sΛ;θ​(XΛ​)j​)+j∈Λ∑​2∂j​gΛ​(XΛ​)j​sΛ;θ​(XΛ​)j​​.

This adaptation is important because standard score matching typically requires the density to vanish at infinity, while truncated variants are needed when the support is bounded or compact.

The main computational challenge is that the marginal score

$$\boldsymbol{s}_{\lambda;\theta}(\boldsymbol{x}_\lambda) = \nabla_{\boldsymbol{x}_\lambda} \log \int q_\theta(\boldsymbol{x}_\lambda,\boldsymbol{x}_{-\lambda})\,\mathrm d\boldsymbol{x}_{-\lambda}$$sλ;θ​(xλ​)=∇xλ​​log∫qθ​(xλ​,x−λ​)dx−λ​

is usually intractable. The method addresses this in two complementary ways: importance weighting and a gradient-first variational estimator.

The first approach estimates the marginal integral using an auxiliary proposal density $p'$p′ over the unobserved coordinates:

$$\int q_\theta(\boldsymbol{x}_\lambda,\boldsymbol{x}_{-\lambda})\,\mathrm d\boldsymbol{x}_{-\lambda} = \mathbb E_{X'_{-\lambda}\sim p'} \left[ \frac{q_\theta(\boldsymbol{x}_\lambda,X'_{-\lambda})}{p'(X'_{-\lambda})} \right].$$∫qθ​(xλ​,x−λ​)dx−λ​=EX−λ′​∼p′​[p′(X−λ′​)qθ​(xλ​,X−λ′​)​].

Given $r$r independent completions $X'^{(1)}_{-\lambda},\dots,X'^{(r)}_{-\lambda}\sim p'$X−λ′(1)​,…,X−λ′(r)​∼p′, the marginal score is estimated by

$$\hat{\boldsymbol{s}}_{\lambda,r;\theta}(\boldsymbol{x}_\lambda) = \nabla_{\boldsymbol{x}_\lambda} \log \left[ \frac{1}{r} \sum_{k=1}^r \frac{q_\theta(\boldsymbol{x}_\lambda,X'^{(k)}_{-\lambda})} {p'(X'^{(k)}_{-\lambda})} \right].$$s^λ,r;θ​(xλ​)=∇xλ​​log[r1​k=1∑r​p′(X−λ′(k)​)qθ​(xλ​,X−λ′(k)​)​].

These estimated marginal scores are substituted directly into the Monte Carlo marginal score matching objective. This yields a practical estimator for both ordinary marginal score matching and truncated marginal score matching. For the truncated case, finite sample error bounds can be derived with respect to the marginal objective. The bound scales roughly with

$$\sqrt{ \frac{p\log(dnrA/\delta)}{\min\{r,n\}} },$$min{r,n}plog(dnrA/δ)​​,

showing that increasing the number of importance weighting samples $r$r improves the score error. However, matching the statistical rate controlled by $n$n requires $r$r of order $n$n, leading to a quadratic computational cost in the data size. In practice, small values of $r$r are often sufficient for stable training, but the importance weighting estimator can become biased or unstable in high dimensions because the integral estimate sits inside nonlinear operations.

The second approach resolves this by deriving the parameter gradient before approximating expectations. Rather than forming an explicit marginal score estimator and then differentiating it, the authors first write the gradient of the marginal score matching objective in terms of expectations and covariances over conditional completions. A key identity states that

$$\boldsymbol{s}_{\theta;\lambda}(\boldsymbol{x}_\lambda) = \mathbb E'\left[ \boldsymbol{s}_\theta(\boldsymbol{x}_\lambda,X'_{-\lambda})_\lambda \right],$$sθ;λ​(xλ​)=E′[sθ​(xλ​,X−λ′​)λ​],

where $\mathbb E'$E′ denotes expectation with respect to the model conditional completion distribution

$$X'_{-\lambda}\mid X_\lambda=\boldsymbol{x}_\lambda \sim p_\theta(\cdot\mid \boldsymbol{x}_\lambda).$$X−λ′​∣Xλ​=xλ​∼pθ​(⋅∣xλ​).

More generally, for a function $h_\theta$hθ​, gradients of expectations can be decomposed as

$$\nabla \mathbb E'\left[ h_\theta(\boldsymbol{x}_\lambda,X'_{-\lambda}) \right] = \mathbb E'\left[ \nabla h_\theta(\boldsymbol{x}_\lambda,X'_{-\lambda}) \right] + \operatorname{Cov}' \left( \nabla_\theta\log q_\theta(\boldsymbol{x}_\lambda,X'_{-\lambda}), h_\theta(\boldsymbol{x}_\lambda,X'_{-\lambda}) \right).$$∇E′[hθ​(xλ​,X−λ′​)]=E′[∇hθ​(xλ​,X−λ′​)]+Cov′(∇θ​logqθ​(xλ​,X−λ′​),hθ​(xλ​,X−λ′​)).

This transformation is central to the method because the difficult expectations over $p_\theta(\cdot\mid X_\Lambda)$pθ​(⋅∣XΛ​) appear outside the differentiation operator. Once the gradient has been rewritten in this form, the conditional distributions and covariances can be approximated using any generative approximation scheme, rather than being tied to explicit importance weights inside a logarithm.

The resulting gradient-first objective for marginal score matching can be written as

$$\nabla_\theta L_{\mathrm M}(\theta) = \mathbb E_{X_\Lambda\sim p} \left[ 2 \sum_{j\in\Lambda} \left( \Psi_\Lambda \left( s_\theta(.)_j^2+\partial_j s_\theta(.)_j \right) - \mathbb E'\left[ s_\theta(X_\Lambda,X'_{-\Lambda})_j \right] \Psi_\Lambda \left( s_\theta(.)_j \right) \right) \right],$$∇θ​LM​(θ)=EXΛ​∼p​​2j∈Λ∑​(ΨΛ​(sθ​(.)j2​+∂j​sθ​(.)j​)−E′[sθ​(XΛ​,X−Λ′​)j​]ΨΛ​(sθ​(.)j​))​,

where the covariance-augmented expectation operator is

$$\Psi_\Lambda(h) = \mathbb E'\left[ \nabla_\theta h(X_\Lambda,X'_{-\Lambda}) \right] + \operatorname{Cov}' \left( \nabla_\theta\log q_\theta(X_\Lambda,X'_{-\Lambda}), h(X_\Lambda,X'_{-\Lambda}) \right).$$ΨΛ​(h)=E′[∇θ​h(XΛ​,X−Λ′​)]+Cov′(∇θ​logqθ​(XΛ​,X−Λ′​),h(XΛ​,X−Λ′​)).

A parallel expression is obtained for truncated marginal score matching by inserting the boundary weights $g_\Lambda$gΛ​ and their derivatives:

$$\nabla_\theta L_{\mathrm{TM}}(\theta) = \mathbb E \left[ 2 \sum_{j\in\Lambda} \left\{ g_\Lambda(X_\Lambda)_j \left[ \Psi_\Lambda \left( s_\theta(.)_j^2+\partial_j s_\theta(.)_j \right) - \mathbb E'\left[ s_\theta(X_\Lambda,X'_{-\Lambda})_j \right] \Psi_\Lambda \left( s_\theta(.)_j \right) \right] + \frac{1}{2} \partial_j g_\Lambda(X_\Lambda)_j \Psi_\Lambda \left( s_\theta(.)_j \right) \right\} \right].$$∇θ​LTM​(θ)=E​2j∈Λ∑​{gΛ​(XΛ​)j​[ΨΛ​(sθ​(.)j2​+∂j​sθ​(.)j​)−E′[sθ​(XΛ​,X−Λ′​)j​]ΨΛ​(sθ​(.)j​)]+21​∂j​gΛ​(XΛ​)j​ΨΛ​(sθ​(.)j​)}​.

The variational module approximates the model conditional distribution $p_\theta(x_{-\lambda}\mid x_\lambda)$pθ​(x−λ​∣xλ​) with a tractable generative conditional

$$p'_\phi(x_{-\lambda}\mid x_\lambda).$$pϕ′​(x−λ​∣xλ​).

Because the target conditional depends on the current model parameters $\theta$θ, the variational parameters $\phi$ϕ must be updated as $\theta$θ changes. The variational model is trained to minimize either a KL-style discrepancy

$$J_{\mathrm{KL}}(\phi,\theta) = \mathbb E \left[ \log \left( \frac{ p'_\phi(X'_{-\Lambda}\mid X_\Lambda) }{ q_\theta(X_\Lambda,X'_{-\Lambda}) } \right) \right],$$JKL​(ϕ,θ)=E[log(qθ​(XΛ​,X−Λ′​)pϕ′​(X−Λ′​∣XΛ​)​)],

or a score-fitting Fisher-style objective

$$J_{\mathrm{F}}(\phi,\theta) = \mathbb E \left[ \left\| \nabla_{X'_{-\Lambda}} \log \left( \frac{ p'_\phi(X'_{-\Lambda}\mid X_\Lambda) }{ q_\theta(X_\Lambda,X'_{-\Lambda}) } \right) \right\|^2 \right],$$JF​(ϕ,θ)=E[​∇X−Λ′​​log(qθ​(XΛ​,X−Λ′​)pϕ′​(X−Λ′​∣XΛ​)​)​2],

where the expectation is taken over incomplete observations $X_\Lambda$XΛ​ drawn from the corrupted data distribution and completions drawn from the current variational conditional:

$$(X_\Lambda,X'_{-\Lambda}) \sim p'_\phi(X'_{-\Lambda}\mid X_\Lambda)\,p(X_\Lambda).$$(XΛ​,X−Λ′​)∼pϕ′​(X−Λ′​∣XΛ​)p(XΛ​).

The score-form objective has the advantage that it avoids requiring the normalizing constant of the variational approximation, because gradients of log densities are invariant to multiplicative normalizing factors.

The full training process is therefore alternating and nested. First, for a batch of incomplete observations, the method samples several completions of the missing coordinates from $p'_\phi(\cdot\mid X_\Lambda)$pϕ′​(⋅∣XΛ​). These completions are used to estimate the expectations and covariances appearing in the gradient-first objective. The gradient of the marginal score matching loss is then evaluated and used to update the score model parameters $\theta$θ. Since changing $\theta$θ changes the conditional completions that $p'_\phi$pϕ′​ is trying to match, the variational parameters $\phi$ϕ are also updated, typically by taking several inner gradient steps after each outer $\theta$θ step. This keeps the variational conditional accurate enough to approximate the increasingly intractable model conditional as the score model changes.

The same gradient-first variational structure can be adapted to other score matching variants, including sliced and denoising score matching, and to truncated marginal objectives. The essential architectural pattern is the coupling of three components: an unnormalized score model $q_\theta$qθ​, an auxiliary or variational completion model $p'_\phi$pϕ′​, and a Monte Carlo gradient estimator that avoids differentiating through intractable marginal integrals.

Experiment

The experiments evaluate the proposed importance weighting and variational marginal score matching methods against EM and a naive zeroed score approach on truncated Gaussian, non-Gaussian, and Gaussian graphical model tasks with simulated and real missing data. Across these settings, the proposed methods are generally reliable, with the variational approach showing the clearest advantage in higher-dimensional and strongly structured problems, while the naive zeroed strategy struggles as interdependence increases. These results support the conclusion that principled marginalization enables more accurate score estimation and dependency recovery under data corruption, although simpler approximations remain competitive in lower-structure or truncated Gaussian cases.