Score-Based Diffusion Modeling for Nonparametric Empirical Bayes in Heteroscedastic Gaussian Mixtures

Abstract

We propose a generalized score-based diffusion framework for learning multivariate Gaussian mixture models with homoscedastic or heteroscedastic noise. Our goal is to nonparametrically estimate the latent location distribution and denoise the observations. Departing from the conventional maximum likelihood approach, we reinterpret each observation as a temporal slice of a family of stochastic diffusion processes. This modeling choice enables a principled characterization of the additive noise structure and supports a multi-step denoising procedure grounded in reverse-time dynamics. We introduce a score-based objective that explicitly models the latent distribution and accommodates observation-specific noise covariances. Theoretically, we establish that the score estimation error with $n$ independent observations achieves a near-parametric error rate of $\raisebox{1ex}{$\text{polylog}(n)$}\!\left/ \!\raisebox{-1ex}{$n$}\right.$, improving upon existing results in the diffusion literature. Empirically, our method outperforms the nonparametric maximum likelihood estimator in both density estimation and denoising fidelity, especially in high-dimensional settings. These findings suggest a promising direction for integrating nonparametric empirical Bayes with diffusion-based generative modeling for latent structure recovery.

1. Introduction

In many large-scale scientific domains, such as the analysis of high-dimensional astronomical data, measurements are often corrupted by heteroscedastic noise, where error variances vary across dimensions and differ across observations [Kelly, 2012, Akritas and Bershady, 1996, Anderson et al., 2018]. While this noise complicates downstream inference, it is frequently well-characterized through the data collection process. A central goal in such settings is to recover the true latent signals and their underlying distribution from noisy observations by leveraging structural information inherent in the data. This problem can be naturally modeled using heteroscedastic Gaussian location mixtures, described by:

$$X_i={\theta }_i^{*}+Z_i,\text{with}{\theta }_i^{*}\stackrel{\text{i.i.d}}{~}{G}^{*}\text{and}{Z}_i\stackrel{\text{ind}}{~}𝒩\left(0,{\sum }_i\right),\text{for}i=1,\dots ,n,$$ (1)

where ${\left\{X_i\right\}}_{i=1}^n\subseteq {\mathbb{R}}^d$ are the observed data with distinct error distributions, ${\left\{{\theta }_i^{*}\right\}}_{i=1}^n$ are the unobserved latent signals drawn from a common unknown prior measure $G^{*}$, and ${\left\{Z_i\right\}}_{i=1}^n$ are independent Gaussian noise with known positive-definite covariance matrices ${\left\{{\sum }_i\right\}}_{i=1}^n$.

Model (1) falls within the empirical Bayes (EB) framework [Robbins, 1951, 1983], where the goal is to estimate the prior distribution $G^{*}$ and recover the latent vectors ${\left\{{\theta }_i^{*}\right\}}_{i=1}^n$. One way to achieve the goal is through the G-modeling approach [Efron, 2014], which is directly modeling the prior on the parameter scale and then plugged into the posterior calculation. A principled and widely used approach to estimate $G^{*}$, without imposing any parametric assumption on it, is the nonparametric maximum likelihood estimator (NPMLE) [Robbins, 1950, Kiefer and Wolfowitz, 1956, Jiang and Zhang, 2009, Gu and Koenker, 2016, Polyanskiy and Wu, 2020]. The NPMLE seeks a distribution $\widehat{G}\in 𝒫\left({\mathbb{R}}^d\right)$ that maximizes the marginal likelihood of the observed data ${\left\{X_i\right\}}_{i=1}^n$:

$${\widehat{G}}_{\text{MLE}}\in \underset{G\in 𝒫\left({\mathbb{R}}^d\right)}{\text{arg}\text{max}}\frac{1}{n}\sum _{i=1}^n\text{log}{f}_{G,{\sum }_i}\left(X_i\right),$$ (2)

where $𝒫\left({\mathbb{R}}^d\right)$ is the family of all probability measures on ${\mathbb{R}}^d$ and $f_{G,{\sum }_i}(x)≔\int {\phi }_{{\sum }_i}(x-\theta )dG(\theta )$ is the marginal density of $X_i$ under $G$ with ${\phi }_{{\sum }_i}(y)≔{\left(\text{det}\left(2\pi {\sum }_i\right)\right)}^{-1/2}\text{exp}\left(-\frac{1}{2}{y}^{\top }{\sum }_i^{-1}y\right)$. Once $\widehat{G}$ is obtained, each latent signal ${\theta }_i$ can be estimated using the empirical Bayes posterior mean:

$${\widehat{\theta }}_i≔{\mathbb{E}}_{\widehat{G}}\left[\theta \mid X_i,{\sum }_i\right]\stackrel{(*)}{=}{X}_i+{\sum }_i{\nabla }_x\text{log}{f}_{\widehat{G},{\sum }_i}\left(X_i\right),$$ (3)

where $(*)$ is due to the Tweedie’s formula [Efron, 2011], linking posterior means to ${\nabla }_x\text{log}{f}_{\widehat{G},{\sum }_i}\left(X_i\right)$, the score function of the marginal density. This nonparametric approach flexibly adapts to latent structures without strong modeling assumptions on $G^{*}$. Theoretical guarantees for the NPMLE in this model are well established in the homoscedastic setting ${\sum }_i=\sum$ or the univariate setting $d=1$: it achieves a near-parametric sample complexity rate $O\left(n^{-1}(\text{log}n{)}^{O(d)}\right)$, i.e., parametric up to logarithmic multiplicative factors, for both density estimation and Bayes regret risk [Jiang and Zhang, 2009, Saha and Guntuboyina, 2020]. Recently, the results have further been extended to multivariate and heteroscedastic setting in Soloff et al. [2025].

Despite these guarantees, the NPMLE and its associated Tweedie-based denoising procedure suffer from inherent statistical and computational limitations that render their practical effectiveness in high-dimensional settings: (a) Statistically, Tweedie’s posterior mean estimator often has an over-shrinkage bias. As an example, consider a joint Gaussian model for $\left(X,{\theta }^{*}\right):\left[\begin{array}{c}X\\ {\theta }^{*}\end{array}\right]~𝒩\left(\left[\begin{array}{l}0\\ 0\end{array}\right],\left[\begin{array}{cc}2I_d& I_d\\ I_d& I_d\end{array}\right]\right)$. It follows that marginally ${\theta }^{*}~𝒩\left(0,I_d\right)$, and $X\mid {\theta }^{*}~𝒩\left({\theta }^{*},I_d\right)$. However, the posterior mean of ${\theta }^{*}$ is $\mathbb{E}\left[{\theta }^{*}\mid X\right]=X/2~𝒩\left(0,I_d/2\right)$, thus underestimating the variance of the latent signal. (b) Computationally, although problem (2) is a convex optimization in $G$, it is infinite-dimensional and difficult to solve in practice. Since it is known that there exists a solution ${\widehat{G}}_{\text{MLE}}$ with support on at most $n$ atoms, one common way to address the issue is to restrict $G$ to be a discrete measure of the form $\sum _{i=1}^m{w}_i{\delta }_{{\mu }_i}$, with $\sum _{i=1}^m{w}_i=1$ for some $m\le n$. Optimizing over ${\left\{\left(w_i,{\mu }_i\right)\right\}}_{i=1}^m\subseteq {\mathbb{R}}_{+}\times {\mathbb{R}}^d$ leads to a finite dimensional but nonconvex optimization, often tackled via the expectation-maximization (EM) algorithm [Laird, 1978]. However, EM suffers from slow convergence and may converge to local maxima due to the nonconvexity. To mitigate these issues, prior work has suggested fixing the atoms ${\left\{{\mu }_i\right\}}_{i=1}^m$ using either grid points over a compact region [Koenker and Mizera, 2014, Feng and Dicker, 2018] or the observed data itself [Lashkari and Golland, 2007]. Yet the grid-based approach scales poorly with dimension, and the data-driven strategy is suboptimal under heteroscedastic noise [Soloff et al., 2025]. Figure 1 further illustrates the over-shrinkage and computational challenges that arise from these limitations.

Figure 1:

Left: 5, 000 two-dimensional observations $X_i$ (in blue) sampled from $X_i\mid {\theta }_i^{*}~𝒩\left({\theta }_i^{*},(0.3{)}^2{I}_2\right)$ with ${\theta }^{*}$ uniformly distributed over a unit circle (in black). The posterior means (in red), obtained via the NPMLE and (3), are over-shrunk relative to the true unit circle. Right (adapted from Soloff et al. [2025, Figure 7]): a real-world 19-dimensional astronomy data (in black). The estimated posterior means (in blue) from NPMLE poorly reveal the underlying chemical structure (see Section 5 for the denoising results from our methods).

1.1. Our Methods and Contributions

To address the statistical and computational limitations of the classical NPMLE and Tweedie-based posterior denoising, we propose a novel perspective that reinterprets the data generation mechanism in model (1) through the lens of stochastic diffusion processes. Our method is built on two core ideas: replacing the NPMLE with score matching to estimate the prior within the G-modeling approach, and applying multi-step posterior denoising to recover latent signals. Below, we present an overview of our approach in the simplified homoscedastic setting; the full treatment of the heteroscedastic case is developed in detail in Section 2.

Modeling the data measurement process.

Consider the homoscedastic version of model (1), where ${\sum }_i=\sum$ for all $i=1,\cdots ,n$. In this special case, the observed data ${\left\{X_i\right\}}_{i=1}^n$ are i.i.d. samples from the convolution of the latent distribution $G^{*}$ with the Gaussian kernel $𝒩(0,\sum )$:

$$X_i~G^{*}*𝒩(0,\sum ),\text{for}i=1,\dots ,n.$$ (4)

Rather than treating each observed data $X_i$ as the result of a static additive noise model, we propose to interpret ${\left\{X_i\right\}}_{i=1}^n$ as snapshots from a forward stochastic diffusion process. Specifically, we model the measurement process as a zero-drift stochastic differential equation (SDE) that continuously perturbs samples from the latent prior $G^{*}$:

$$d{X}^{(t)}=g(t){\sum }^{1/2}d{B}_t,\text{with}{X}^{(0)}\stackrel{d}{=}{\left\{{\theta }_i^{*}\right\}}_{i=1}^n~G^{*},$$ (5)

where $B_t\in {\mathbb{R}}^d$ denotes a standard Brownian motion and $g:{\mathbb{R}}_{+}\to {\mathbb{R}}_{+}$ is the diffusion coefficient controlling the noise scaling. The function $g(t)$ can be adapted to reflect assumptions on the measurement process; for simplicity, one natural choice is $g(t)=\sqrt{\frac{2}{e^2-1}}{e}^t\Rightarrow \alpha (t)≔\left({\int }_0^{t}g(s{)}^{2}ds\right)=\frac{e^{2t}-1}{e^2-1}$. Under this model, the marginal distribution of $X^{(t)}$ evolves according to:

$$X^{(t)}~G^{*}*𝒩\left(0,{\sum }^{(t)}\right),\text{where}{\sum }^{(t)}≔\alpha (t)\sum =\frac{e^{2t}-1}{e^2-1}\sum .$$

In particular, at time $t=1$, we have ${\sum }^{(1)}=\sum$. Therefore, we can indeed view each $X_i$ as sampled from the diffusion process (5) at $t=1$, i.e., ${\left\{X_i\right\}}_{i=1}^n~X^{(1)}$.

Denoising the observations.

The forward-time process in (5) is associated with a reverse-time diffusion process, which can be approximated by a deterministic flow that shares the same marginal distributions over time [Anderson, 1982, Maoutsa et al., 2020, Song et al., 2021a]. This flow is governed by the following ordinary differential equation (ODE):

$$d{\bar{X}}^{(t)}=-\frac{1}{2}g(t{)}^2\sum {\nabla }_x\text{log}{f}_{G^{*},{\sum }^{(t)}}\left({\bar{X}}^{(t)}\right)dt.$$ (6)

Given an estimator $\widehat{G}$ of the prior $G^{*}$, this reverse time ODE naturally defines a generalized multi-step Tweedie-type empirical Bayes denoising procedure. Specifically, replacing $G^{*}$ with $\widehat{G}$ and discretizing the ODE backward in time from $t=1$ to $t=0$ with step size $\Delta t$ yields the iterate:

$${\widehat{\theta }}_i^{(t-\Delta t)}={\widehat{\theta }}_i^{(t)}+\frac{\Delta t}{2}g(t{)}^2\sum {\nabla }_x\text{log}{f}_{\widehat{G},{\sum }^{(t)}}\left({\widehat{\theta }}_i^{(t)}\right)\text{with}\text{initial}\text{condition}{\widehat{\theta }}_i^{(1)}=X_i.$$ (7)

The final denoised estimate is taken as ${\widehat{\theta }}_i≔{\widehat{\theta }}_i^{(0)}$ for $i=1,\dots ,n$. Notably, this method generalizes Tweedie’s formula in (3), which corresponds to a single-step approximation with $\Delta t=1$. Our multi-step denoising procedure mirrors optimization dynamics, particularly Langevin-type flows over probability spaces, where smaller step sizes yield more controlled updates. This approach allows gradual refinement of the denoised signal, helping mitigate the over-shrinkage bias observed in one-step Tweedie denoising methods.

Learning the latent prior ${\mathit{G}}^{\mathbf{*}}$.

Our stochastic framework also naturally induces a denoising score matching (DSM) objective [Vincent, 2011] for learning the prior $G^{*}$. We propose to further perturb the noisy observations $X_i$ using the forward SDE (5) to generate conditional samples from the kernel $p_{1t}\left(\cdot \mid X_i\right)={\phi }_{{\sum }^{(t)}-\sum }\left(\cdot -X_i\right)$, up to some stopping time $T$. We then match the score of the model’s marginal density to the conditional score induced by this perturbation:

$$\widehat{G}≔\underset{G\in 𝒫\left({\mathbb{R}}^d\right)}{\text{arg}\text{min}}\frac{1}{n}\sum _{i=1}^n{\int }_1^T{\mathbb{E}}_{p_{1t}\left(\cdot \mid X_i\right)}\left[g(t{)}^2{‖\nabla \text{log}{f}_{G,{\sum }^{(t)}}(x)-\nabla \text{log}{p}_{1t}\left(x\mid X_i\right)‖}_{\sum }^2\right]dt,$$ (8)

where $\Vert x{\Vert }_{\sum }^2≔x^{\top }\sum x$ denotes a weighted ${\ell }_2$-norm, extending the likelihood-weighting idea of Song et al. [2021b]. Intuitively, under our stochastic view, the observations at $t=1$ already contain calibrated noise relative to $G^{*}$, and additional perturbation simulates conditional transitions. Matching model scores to these local conditional scores allows global recovery of the latent prior. Unlike standard diffusion models with the DSM objective that directly parameterize the score function with neural networks [Song et al., 2021a], we model the distribution $G$ explicitly in (8). Such $G$-modeling is crucial in deriving new theoretical guarantees for the score estimation risk. Computationally, although the objective 8 is nonconvex, diffusion-type squared loss objectives have been empirically observed to exhibit smoother optimization landscapes and avoid poor local minima [Xu et al., 2024]. This empirical tractability is further supported by theory on the benign optimization behavior of overparameterized neural networks trained over mean-squared losses that can typically satisfy the Polyak–Łojasiewicz condition or behave linearly in the neural tangent kernel regime [Liu et al., 2022].

In summary, we propose a novel integration of score-based diffusion models with the classical EB framework via G-modeling. The rest of the paper is organized as follows: In Section 2, we present a stochastic process framework to model heteroscedastic Gaussian mixtures, where each observation is treated as a marginal realization from a distinct diffusion process with its own score function. Section 3 establishes the theoretical advantages of our approach, showing that G-modeling combined with diffusion yields near-parametric rates for score estimation. We further discuss how to numerically solve the $G$-modeling based score matching objective in Section 4. In Section 5, we present experimental results on both synthetic and real astronomical data, demonstrating substantial improvements in denoising and prior recovery over NPMLE and one-step Tweedie-based methods, supporting the improved computational scalability of our framework in high-dimensional settings.

1.2. Related Work

Score-based models for mixtures.

Recent theory has provided provable guarantees for diffusion and score-based models in learning Gaussian mixtures, yet these results neither handle the heteroscedastic measurement setting we study nor aim at recovering the latent prior itself. Works such as Gatmiry et al. [2024], Chen et al. [2024], Shah et al. [2023] propose diffusion-based algorithms for mixture models under either finite discrete structures or continuous latent mixtures with homoscedastic noise. These approaches typically assume access to an infinite supply of i.i.d. samples and yield slow convergence guarantees of order $O\left((\text{log}n{)}^{-1/p}\right)$ in score or total variation distance, for some exponent $p>1$. Importantly, these works directly parametrize the score function, using either neural networks or piecewise polynomial surrogates, without attempting nonparametric recovery of the latent mixing measure $G^{*}$ or incorporate explicit G-modeling into the objective. Because of this direct score parametrization, these methods cannot be readily applied to our core task of latent recovery and are incompatible with our de-biased multi-step denoising procedure. A notable exception is Ghosh et al. [2025], which presents the first analysis of combining score matching with G-modeling, achieving a near-parametric rate $O\left({\text{log}}^{6}n/n\right)$ for score estimation risk. However, their results are limited to the univariate, homoscedastic setting $(d=1)$. In contrast, our work tackles the more challenging regime of multivariate, heteroscedastic Gaussian location mixtures with arbitrary dimension $d\ge 1$. By coupling G-modeling with a stochastic process perspective, we develop an explicit estimator of the latent prior $G^{*}$ and establish near-parametric rates for score estimation (up to small smoothing), along with finite-sample guarantees for recovering the latent signals.

Image denoising in computer vision.

While our problem and the proposed method may appear superficially related to score-based diffusion approaches for image denoising, the settings and statistical challenges are fundamentally different. Most notably, our task lacks access to clean ground truth data, precluding the use of pre-trained score models, and all inference must be performed from noisy observations alone. In contrast, supervised denoising approaches such as BM3D [Dabov et al., 2006] requires access to a large number of paired noisy-clean images, while self-supervised methods like Noise2Score [Kim and Ye, 2021] rely on one-step Tweedie approximations that ignore the unknown prior $G^{*}$ and leverage only partial model information ${\sum }_i$. These methods are either unrealistic or statistically inefficient in scientific domains where latent structure is present but clean data are unavailable, such as the astronomy measurement data. Even in moderate dimensions (e.g., $d=20$), empirical Bayes is a highly nontrivial statistical task due to the data sparsity: each observation $X_i$ corresponds to a unique latent parameter ${\theta }_i^{*}$ and an individual covariance ${\sum }_i$, yielding just one single noisy sample to infer each ${\theta }_i^{*}$.

2. SDGM: Score-based Diffusion with G-Modeling

In this section, we present our general treatment of the score-based diffusion approach with G-modeling under the heteroscedastic setting. The key idea is still to reinterpret each noisy observation as a time-1 snapshot from a latent stochastic diffusion process. While our homoscedastic formulation serves as a motivating special case, the heteroscedastic setting introduces a crucial distinction: each observation $X_i$ corresponds to a different SDE, despite all sharing a common latent initialization. Specifically, we consider the following family of zero-drift linear SDEs:

$$d{X}_i^{(t)}=g_i(t){\sum }_i^{1/2}d{B}_t,\text{with}{X}_i^{(0)}~G^{*},\text{for}i=1,\dots ,n,$$ (9)

where ${\left\{g_i(t)\right\}}_{i=1}^n\subset \mathbb{R}$ are the time-dependent diffusion coefficients. Without loss of generality, we assume ${\alpha }_i(t)≔{\int }_0^t{g}_i(s{)}^{2}ds$ satisfies ${\alpha }_i(1)=1$ for all $i$. Then the time-dependent covariances are given by ${\sum }_i^{(t)}≔{\alpha }_i(t){\sum }_i$, and the marginal law of $X_i^{(t)}$ becomes

$$X_i^{(t)}~G^{*}*𝒩\left(0,{\sum }_i^{(t)}\right),\text{for}t>0,i=1,\dots ,n.$$ (10)

Thus, each observation $X_i$ in the original mixture model (1) can be viewed as a marginal sample from this diffusion process at $t=1$, i.e., $X_i\stackrel{d}{=}{X}_i^{(1)}~G^{*}*𝒩\left(0,{\sum }_i\right)$ for all $i=1,\dots ,n$.

As in the homoscedastic case, each forward SDE in (9) is coupled with the reverse-time ODE:

$$d{\bar{X}}_i^{(t)}=-\frac{1}{2}{g}_i(t{)}^2\cdot {\sum }_i\cdot {\nabla }_x\text{log}{f}_{G^{*},{\sum }_i^{(t)}}\left({\bar{X}}_i\right)dt,\text{for}i=1,\dots ,n.$$

Given any plug-in estimator $\widehat{G}$ for $G^{*}$, substituting ${\sum }^{(t)}={\sum }_i^{(t)}$ in (7) defines a multi-step score-based denoising procedure that remains consistent with the EB framework. Furthermore, as $\Delta t\to 0$, the dynamics in (7) induces a deterministic transport map from the noisy empirical measure toward a latent measure $\widehat{G}\approx G^{*}$. This perspective aligns with the design principles behind recent optimal transport approaches for empirical Bayes denoising [Soloff et al., 2025, Zhang et al., 2024], mitigating the over-shrinkage limitation of one-step Tweedie-type procedures.

Importantly, in our setting, the SDEs in (9) serve as a belief model for the true data-generating process, and the noisy observations in (1) are treated as marginal realizations from latent diffusion paths. This fundamentally differs from the generative score-based diffusion models like Song et al. [2021a], where SDEs are artificially constructed for the purpose of noise injection and score network training. In an idealized setting where the true latent measure $G^{*}$ is known or accessible via unlimited samples, one could apply the framework of Song et al. [2021b,a] to learn a time-dependent score network $s_{\theta }(\cdot ,t)$, mapping $G^{*}$ to pure noise through an invertible path. However, our setting departs significantly from this setup. First, we observe only a finite collection of noisy samples at forward time $t=1$, each generated by a unique, heteroscedastic SDE rather than homogeneous, artificially perturbed noise. Second, our goal is to explicitly estimate the latent distribution $G^{*}$, as opposed to modeling it implicitly via score functions. Finally, because no data is available at intermediate times $t\in (\mathrm{0,1})$, the standard time-indexed score network $s_{\theta }(\cdot ,t)$ cannot be trained in our setting.

To address these challenges, we propose a new objective that aggregates information across all diffusion paths to estimate $G^{*}$. Our method, which we term score-based diffusion with G-modeling (SDGM), obtains the estimator $\widehat{G}$ by:

$$\widehat{G}\in \underset{G\in 𝒫\left({\mathbb{R}}^d\right)}{\text{arg}\text{min}}{𝒥}_n(G)≔\frac{1}{n}\sum _{i=1}^n{\int }_{t_0}^T{g}_i(t{)}^2{\mathbb{E}}_{x\mid i,t}\left[{‖\nabla \text{log}{q}_G^{(t)}(x)-\nabla \text{log}{q}_n^{(t)}(x)‖}_{{\sum }_i}^2\right]dt,$$ (11)

where, $t_0\ge 1$ is the minimum diffusion time, and to simulate the homoscedastic behavior, we consider the rescaled Gaussian transition kernel: $x\mid i,t~K_t(\cdot \mid i)≔𝒩\left({\left[{\alpha }_i(t){\sum }_i\right]}^{-1/2}{X}_i,(1-\right.\left.\left.{\alpha }_i(t{)}^{-1}\right)I_d\right)$ defined for $t\ge t_0\ge 1$, and define the mixture densities $q_n^{(t)}$ and $q_G^{(t)}$ to be:

$$(\text{empirical})q_n^{(t)}≔\frac{1}{n}\sum _{i=1}^n{K}_t(\cdot \mid i),(G\text{-induced})q_G^{(t)}≔\frac{1}{n}\sum _{i=1}^n{f}_{{\left({\sum }_i^{(t)}\right)}^{-1/2}\#G,I_d},$$

where $T\#G$ denotes the pushforward measure of $G$ under map $T$ such that for every set $A\subseteq {\mathbb{R}}^d$, $(T\#G)(A)=G\left(T^{-1}A\right)$. The defined mixtures satisfy that $\mathbb{E}{q}_n^{(t)}=q_{G^{*}}^{(t)}$. Intuitively, $q_G^{(t)}$ aggregates the global behavior of the latent measure $G$ across all SDEs, while $q_n^{(t)}$ serves as its empirical counterpart based on noisy observations. By matching the score functions of these two mixtures, we guide $\widehat{G}$ toward an accurate approximation of $G^{*}$. In the homoscedastic special case with ${\sum }_i=\sum$ and $g_i=g$, the mixtures simplify to a rescaled version of the objective (11) that is equivalent to the score matching formulation in (8) up to constants.

3. Theoretical Guarantees

In this section, we provide guarantees on the statistical performance for both homoscedastic and heteroscedastic noise models. We primarily study three statistical properties to quantify how well the SDGM estimator $\widehat{G}$ of (11) performs as a plug-in estimator for the true prior $G^{*}$.

The successful convergence of diffusion-based generative models is underlaid by a low score estimation error [Yakovlev and Puchkin, 2025, Bortoli, 2022]. Define the following time-cumulative average Fisher divergence for all path measures $\mu$ and $\nu$:

$${\overline{\mathfrak{F}}}_{\left[t_0,T\right]}(\mu \Vert \nu )≔\frac{1}{n}\sum _{i=1}^n{\int }_{t_0}^T{\mathbb{E}}_{x~{\mu }_i^{(t)}}{‖\nabla \text{log}{\mu }_i^{(t)}-\nabla \text{log}{\nu }_i^{(t)}‖}_2^{2}dt.$$

Then, the specialized metric ${\overline{\mathfrak{F}}}_{\left[t_0,T\right]}\left(f_{G^{*},{\sum }^{(t)}}\Vert f_{\widehat{G},{\sum }^{(t)}}\right)$ measures the time-cumulative risk on how $\nabla \text{log}{f}_{G,{\sum }^{(t)}}$ approximates $\nabla \text{log}{f}_{G^{*},{\sum }^{(t)}}$ along the forward-time SDE path beyond the observable time. It is consistent with the generalization risk concerned in Yakovlev and Puchkin [2025].

For density estimation, $\widehat{G}$ also provides a natural estimate for the mixture marginal law $q_{G^{*}}^{(t)}$ aggregated from the paths of (9). We measure the quality of density estimation by the standard squared Hellinger accuracy, ${\mathfrak{H}}^2(p,q)≔\frac{1}{2}\int (\sqrt{p}-\sqrt{q}{)}^2$ between densities $p$ and $q$. In particular, we are interested in the following squared Hellinger distances:

$$(\text{homoscedastic}){\mathfrak{H}}^2\left(f_{G^{*},{\sum }^{(t)}},f_{\widehat{G},{\sum }^{(t)}}\right),(\text{heteroscedastic}){\mathfrak{H}}^2\left(q_{G^{*}}^{(t)},q_{\widehat{G}}^{(t)}\right),$$

which captures the global fit of the plug-in estimator $\widehat{G}$ for estimating the joint law of ${\left\{X_i^{(t)}\right\}}_{t\ge 1}$.

The third metric of interest we establish is the squared Wasserstein distance (see e.g., Villani [2008]), which is defined for all $\mu ,\nu \in 𝒫\left({\mathbb{R}}^d\right)$ that $W_2^2(\mu ,\nu )≔{\text{inf}}_{\gamma \in \Pi (\mu ,\nu )}{\mathbb{E}}_{(x,y)~\gamma }\left[\Vert x-y{\Vert }_2^2\right]$ with $\Pi (\mu ,\nu )$ denotes the set of all joint measures over $\mu ,\nu$. The distance $W_2^2\left(\widehat{G},G^{*}\right)$ naturally characterizes the direct consistency of $\widehat{G}$ for the true $G^{*}$.

We make the following assumptions on the latent structure and the data generation process (9). In the assumptions, the support conditions are widely considered in literature [Jiang and Zhang, 2009, Saha and Guntuboyina, 2020, Ghosh et al., 2025]. They encompass a wide range of realistic scenarios, including discrete mixtures, compactly supported manifold-supported latent structures, and representative experiment settings in Section 5. Furthermore, well-conditioned assumptions on noise covariances are ubiquitous in prior works [Soloff et al., 2025] and are essential to ensure smoothly-defined score everywhere and to enable information-theoretically meaningful latent recovery.

Assumption 1. The true latent measure in model (1) is compactly supported, i.e., $G^{*}\in 𝒫\left([-M,M{]}^d\right)$ for some finite radius $M>0$. There exists $0<\underset{\_}{\sigma }\le \bar{\sigma }<\infty$ such that the covariances satisfy $\underset{\_}{\sigma }{I}_d⪯{\sum }_i⪯\bar{\sigma }{I}_d$ for all $i=1,\dots ,n$. The diffusion coefficient satisfies ${\alpha }_i(1)=1$ (such that ${\sum }_i^{(1)}={\sum }_i$) and other mild regularity conditions specified in the supplemental.

Our first theorem characterizes the performance of our SDGM framework for the canonical homoscedastic model (4), which is specialized from our guarantee for the heteroscedastic model.

Theorem 1 (Homoscedastic noise). Under Assumption 1 and further assume that ${\sum }_i=\sum ,g_i=g$ for all $i=1,\dots ,n$. Let $\widehat{G}$ minimizes objective (8) constrained over the family $𝒫\left([-M,M{]}^d\right)$ and with parameters $t_0$ and $T$ satisfying $1+\frac{1}{\text{log}n}\le \alpha \left(t_0\right)\le \alpha (T)\le 2$. Then provided $n$ sufficiently large, the following bounds hold with probability at least $1-n^{-2}$:

Score and density estimation:

$$\mathbb{E}{\overline{\mathfrak{F}}}_{\left[t_0,T\right]}\left(f_{G^{\text{*}},{\Sigma }^{\left(t\right)}}\Vert f_{\widehat{G},{\Sigma }^{\left(t\right)}}\right){\lesssim }_{M,d,\left(\underset{\_}{\sigma },\overline{\sigma }\right),g}\frac{1}{n}{(\text{log}n)}^{2d+3},$$ (12)

$$\mathbb{E}{\mathfrak{H}}^2\left(f_{G^{*},{\sum }^{\left(t_0\right)}},f_{\widehat{G},{\sum }^{\left(t_0\right)}}\right){\lesssim }_{M,d,\left(\underset{\_}{\sigma },\bar{\sigma }\right),g}\frac{1}{n}(\text{log}n{)}^{2d+3}.$$ (13)

Deconvolution:

$$\mathbb{E}{W}_2^2\left(G^{*},\widehat{G}\right)\lesssim \frac{1}{\text{log}n}.$$ (14)

Here, we write $f(n)\lesssim g(n)$ if there exists a universal constant $C$ independent of $n$ such that $f(n)\le Cg(n)$, $\forall n$. We write ${\lesssim }_{{\theta }_1,{\theta }_2,\dots }$ to emphasize that the constant $C=C_{\theta }$ also depends on the fixed specification parameters ${\theta }_i$’s. In the following, we discuss implications of this result. Under a latent manifold setting with homoscedastic isotropic noise and neural parameterization of the score function, recent work [Yakovlev and Puchkin, 2025] proves the rate $\mathbb{E}{\overline{\mathfrak{F}}}_{\left[t_0,T\right]}\left(f_{G^{*},{\sum }^{(t)}}\left|\right|f_{\widehat{G},{\sum }^{(t)}}\right)=O\left(n^{-2\beta /\left(4\beta +d_0\right)}\right.\text{polylog}\left.(n)\right)$, where $d_0$ is the intrinsic dimension of the manifold and $\beta$ denotes a smoothness parameter. To our knowledge, this is the sharpest rate known for diffusion-based score estimation under such assumptions. By leveraging explicit $G$-modeling to impose the structure on the estimated scores, our result achieves an improved $O\left(n^{-1}\right)$ leading dependence on sample size $n$ (up to logarithmic factors), under more general assumptions. Moreover, our Hellinger accuracy bound (13), modulo a small distributional shift from $t=1$ to $t_0$, matches the rate achieved in NPMLE literature [Saha and Guntuboyina, 2020]. If the target is to estimate $f_{G^{*},\sum }$ at the observational noise level, an unavoidable smoothing cost due to diffusion arises. In this regime, the task incurs the inevitable minimax rate $O\left(n^{-2/(d+4)}\right)$ that cannot compete with NPMLE [Wibisono et al., 2024]. Finally, in the Gaussian deconvolution setting with fixed noise design, the minimax rate deteriorates to the logarithmic regime $O(1/\text{log}n)$ due to inherent ill-posedness induced by Gaussian smoothing and it is the best rate possibly achieved by both NPMLE and our framework [Soloff et al., 2025, Dedecker and Michel, 2013]. We further discuss limitations of the current theory in Section 6.

Our next theorem generalizes the results to the heteroscedastic settings. The core contents of Theorem 2 parallels that of Theorem 1, but additional care is needed to account for observation-specific noise covariances and SDEs. To make the readers easy to follow, we present the main result informally below and defer detailed technical derivations and proofs of Theorem 1 and 2 to the supplemental material.

Theorem 2 (Heteroscedastic noise; Informal). Under Assumption 1, provided solution $\widehat{G}$ to (11) constrained over the family $𝒫\left([-M,M{]}^d\right)$ with $\left(t_0,T\right)$ satisfying well-behaved smoothing condition, then with high probability,

$$\mathbb{E}{\overline{\mathfrak{F}}}_{\left[t_0,T\right]}\left(q_{G^{*}}^{\left(t_0\right)}\Vert q_{\widehat{G}}^{\left(t_0\right)}\right)\lesssim \frac{(\text{log}n{)}^{2d+3}}{n},\mathbb{E}{\mathfrak{H}}^2\left(q_{G^{*}}^{\left(t_0\right)},q_{\widehat{G}}^{\left(t_0\right)}\right)\lesssim \frac{(\text{log}n{)}^{2d+3}}{n},$$

$$\mathbb{E}{W}_2^2\left(\frac{1}{n}\sum _{i=1}^n{\left({\sum }_i\right)}^{-1/2}\#G^{*},\frac{1}{n}\sum _{i=1}^n{\left({\sum }_i\right)}^{-1/2}\#\widehat{G}\right)\lesssim \frac{1}{\text{log}n}.$$

4. Solution Methods to Solve (11)

The optimization problem in (11) is an infinite dimensional nonconvex problem that is challenging to solve numerically. We discuss practical strategies for approximately computing the SDGM solution.

Strategy one: discretize the support of $\mathit{G}$.

A classical approach to deal with the infinite dimensional variable $G$, as commonly used in NPMLE, is to restrict the class of priors to discrete measures with a finite number of supports: $G=\sum _{j=1}^m{w}_j{\delta }_{{\mu }_j}$, with $\sum _{j=1}^m{w}_j=1,w_j\ge 0$. This reduces the optimization to a finite-dimensional, albeit nonconvex, problem over ${\left\{\left(w_j,{\mu }_j\right)\right\}}_{j=1}^m$.

Strategy two: continuous modeling with normalizing flows.

An alternative approach is to model $G$ continuously using normalizing flows [Rezende and Mohamed, 2015, Dinh et al., 2017, Durkan et al., 2019]. A normalizing flow represents a flexible distribution as the pushforward of a simple base distribution (e.g., standard Gaussian) through an invertible neural network. We parameterize $G$ using a Neural Spline Flow [Durkan et al., 2019] implemented via the nflows library [Durkan et al., 2020], allowing efficient sampling and density evaluation. In our framework, we draw Monte Carlo samples from the flow to estimate $f_G$ and its score during training. Sampling from a continuous prior is theoretically justified by empirical Bayes approximation results: given a continuous $G$, there exists a discrete approximation $G^{\prime }$ that closely matches $f_G$ and its gradient [Saha and Guntuboyina, 2020, Lemma D.3]. Furthermore, continuous modeling enables richer recovery of the underlying structure of $G$, which is crucial for scientific interpretation and denoising applications [Shen and Wu, 2022].

Batched and conditional estimation of ${\mathit{q}}_{\mathit{G}}$.

The full mixture $q_G(x)$ plays a central role in our theoretical guarantees, but evaluating its score ${\nabla }_x\text{log}{q}_G$ is computationally intensive for large $n$. A natural relaxation is to approximate $q_G$ at each iteration using only a mini-batch of observations. In particular, we focus on the case where the mini-batch size $B=1$: at each training step, we align the rescaled conditional score ${\nabla }_x\text{log}{f}_{G,{\sum }_i^{(t)}}(x)$ with the empirical conditional score ${\nabla }_x\text{log}𝒩\left(x;X_i,\left({\alpha }_i(t)-1\right){\sum }_i\right)$. This corresponds to a conditional score matching update, mirroring the classical intuition of denoising score matching [Vincent, 2011], while offering computational simplicity and stability in stochastic optimization. In fact, subsampled estimation for mixture inference has also been used in early works [Celeux et al., 2001, Ihler et al., 2003].

5. Numerical Experiments

In this section, we evaluate the empirical performance of our proposed method on synthetic heteroscedastic datasets and real-world astronomy data. We also investigate whether the empirical scaling behavior aligns with our theoretical guarantees, particularly as we vary the problem dimension. We refer the readers to the supplemental for more implementation details and results.

Baselines and metrics.

We evaluate two variants of our proposed method: SDGM-C, which uses a continuous prior modeled via normalizing flows, and SDGM-D, which uses a discrete prior. We compare our approaches against the state-of-the-art NPMLE solvers for the problem implemented in literature. Specifically, we compare with (1) NPMLE-N, the discretized NPMLE with supports taken fixed at noisy data (thus the objective (2) is optimized over only weights and becomes convex), as advocated in Lashkari and Golland [2007]; we adopt the Newton-based augmented Lagrangian method implemented from [Zhang et al., 2024]. (2) PEM, the partial EM method which is a generalized EM method proposed in Zhang et al. [2022] that used tighter sequential relaxations.

In addition, we consider a hybrid heuristic called SDGM-P to further improve the likelihood value defined in (2). This method first trains the SDGM-C model to learn a continuous prior, from which we sample support points. These support points are then used for a post-processing step via the convex NPMLE-N solver, effectively refining the prior estimate to maximize the likelihood.

We evaluate the performance of different methods over following metrics: (1) Fisher divergence (FD) defined by ${\mathfrak{F}}_n\left(f_G,f_{*}\right)≔\frac{1}{n}\sum _{i=1}^n{‖\nabla \text{log}{f}_{G^{*},{\sum }_i}\left(X_i\right)-\nabla \text{log}{f}_{\widehat{G},{\sum }_i}\left(X_i\right)‖}_2^2$. (2) Negative log-likelihood (NLL) defined by $-\frac{1}{n}\sum _{i=1}^n\text{log}{f}_{\widehat{G},{\sum }_i}\left(X_i\right)$, which serves as a proximal for the KL divergence between $f_{\widehat{G},{\sum }_i}$ and $f_{G^{*},{\sum }_i}$ (3) Squared Wasserstein-2 distance $\left(W_2^2\right)$ between estimated $\widehat{G}$ and $G^{*}$, approximated by the entropy-regularized optimal transport distance [Cuturi, 2013, Genevay et al., 2018]. (4) Mean squared error (MSE) between the denoised and clean data $\frac{1}{n}\sum _{i=1}^n{‖{\widehat{\theta }}_i-{\theta }_i^{*}‖}_2^2$.

Synthetic datasets.

We evaluate our method on two synthetic datasets with low-dimensional manifold structure, both of which are widely used in the NPMLE literature [Saha and Guntuboyina, 2020, Zhang et al., 2024]. Specifically, for $d\ge 2$, we generate clean latent vectors ${\theta }_i^{*}$ that exhibit nontrivial structure only in the first two dimensions, with the remaining dimensions padded with zeros:

• Example 1 (circle): The first two coordinates of each ${\theta }_i^{*}$ are drawn uniformly from a circle of radius 6 centered at the origin.

• Example 2 (uppercase letter “A”): The first two coordinates of each ${\theta }_i^{*}$ are drawn uniformly from one of five line segments that form the shape of the capital letter “A.” These segments connect the points (−4, 6), (−2, 0), (0, 6), (2, 0), and (4, 6), with each segment selected with equal probability 1/5.

We then corrupt each ${\theta }_i^{*}$ with independently generated heteroscedastic Gaussian noise $Z_i~𝒩\left(0,{\sum }_i\right)$, where each ${\sum }_i$ is a diagonal matrix. The diagonal entries of ${\sum }_i$ are sampled independently and uniformly from the interval [(0.8)², (1.2)²].

The results are summarized in Table 1 and Figures 2 and 3, where $\widehat{G}$ of our methods are plotted using the marginal density estimated through geodesic-KDE [Vincent and Bengio, 2002]. As shown in Table 1, SDGM consistently outperforms the NPMLE baselines across all metrics, with the additional post-processing step (SDGM-P) further improving the quality of prior estimation in most cases. The denoising effects of different methods are evident in Figures 2 and 3. Specifically, the estimated signals ${\widehat{\theta }}_i$ from NPMLE and one-step Tweedie denoising appear overly blurred and shrunk, failing to recover the underlying manifold structure. In contrast, SDGM, equipped with multi-step ODE-based denoising, is able to clearly capture and preserve the latent geometry of the data.

Table 1:

Models are trained with dimension $d=16$ and sample size $n=2\times {10}^4$. The ↓ indicates that lower values are better. We report generalization risks evaluated on a separate test dataset of size 10⁴. Denoising of the baseline methods NPMLE-N and PEM is performed via the one-step Tweedie’s formula, while our methods use the ODE solver. Results for stochastic methods are averaged over 5 independent runs.

$G^{*}$	circle				uppercase letter “A”
metric	FD ↓	NLL ↓	${\mathit{W}}_2^2\downarrow$	MSE ↓	FD ↓	NLL ↓	${\mathit{W}}_2^2\downarrow$	MSE ↓
NPMLE-N	46.82	25.88	14.43	3.57	34.88	25.10	12.41	3.04
PEM	4.71	24.89	4.52	1.26	3.78	24.40	4.09	1.35
SDGM-C	1.27±0.04	24.82±0.01	2.07±0.77	1.13±0.01	1.23±0.11	24.33±0.01	1.68±0.31	1.32±0.03
SDGM-D	3.72±0.05	24.88±0.00	3.95±0.03	1.64±0.03	5.26±0.05	24.41±0.00	3.94±0.00	1.73±0.02
SDGM-P	0.68±0.06	24.78±0.00	1.07±0.06	1.29±0.02	0.49±0.02	24.30±0.00	1.25±0.07	1.45±0.02

Figure 2:

Projected plots of the circle experiment onto the first two dimensions. NPMLE-N (columns 2 & 3) performance significantly degrades when the dimension increases from $d=8$ (top row) to $d=16$ (bottom row). Our proposed SDGM-C (columns 4 & 5) stays efficient when $d$ grows.

Figure 3:

Projected plots of the letter $A$ experiment in $d=16$ onto the first two dimensions. NPMLE and PEM (top row) struggled by learning only a noisy latent prior, and causing denoised data to collapse within the region between the two legs of the $A$, illustrating over-shrinkage. The bottom row shows our SDGM variants. Notably, even $\widehat{G}$ is estimated by SDGM, the one-step Tweedie denoising (middle of the bottom row) still over-shrinks ${\theta }_i$. By contrast, the proposed denoising scheme better preserves the $A$ structure, with fewer points confined to the center of the $A$.

Real-world astronomy data.

We further demonstrate our SDGM-C on a challenging real-world scientific problem: estimating the latent distribution of stellar chemical abundances and denoising observations from the Apache Point Observatory Galactic Evolution Experiment (APOGEE) Data Release 14 [Abolfathi et al., 2018, Majewski et al., 2017]. With pre-processing in Ratcliffe et al. [2020], this dataset, comprising $n=\mathrm{27,238}$ stars with $d=19$ chemical abundance measurements (e.g., [Fe/H], [Mg/Fe], [Si/Fe]), features heteroscedastic, element-correlated Gaussian noise for each star. The true underlying abundances ${\theta }_i^{*}$ are expected to reside on a lower-dimensional manifold governed by astrophysical nucleosynthetic processes [Ting et al., 2012, Weinberg et al., 2019]. Accurately revealing the underlying physical structures through the denoised abundances is crucial for reliable studies of Milky Way formation and evolution [Freeman and Bland-Hawthorn, 2002]. As depicted earlier in Figure 1, applying the NPMLE across the full 19-dimensional space is computationally prohibitive and produces unreliable estimates, as evidenced by Soloff et al. [2025], Zhang et al. [2024]. Notably, applying SDGM-C to the APOGEE data (Figure 4) reveals sharpened astrochemically significant structures. Taken together with our theoretical guarantees, these empirical results showcase that SDGM couples superior denoising fidelity with graceful high-dimensional scalability, making it a promising tool for robust inference in complex scientific datasets.

Figure 4:

Left: Noisy APOGEE observations projected onto two 2-D planes. Middle: The denoised data reveals a fusiform-like or banana-like structure in the central region and a thin manifold structure in the upper region. Right: The marginal density of fitted $\widehat{G}$ suggests strong abundance correlation among the chemicals. These results significantly improve the performance of NPMLE and Tweedie denoising in Figure 1.

6. Discussions, Limitations, and Future Work

This work develops a stochastic-process reinterpretation of Gaussian location mixtures that bridges classical empirical Bayes inference and modern diffusion-based modeling. Through this lens, we introduce a diffusion-type $G$-modeling objective and a multi-step generalized Tweedie denoising procedure that jointly estimate the latent measure and recover the underlying signals from heteroscedastic mixtures. Unlike standard score-based diffusion approaches that directly parameterize the score function, our method explicitly models the latent prior distribution $G$, which is crucial for consistent recovery of the latent law. In contrast to traditional NPMLE solvers that suffer from slow convergence and over-shrinkage bias, our framework achieves comparable or improved denoising and prior recovery performance, with better empirical scalability in high-dimensional regimes.

Limitations.

Our theoretical analysis builds on assumptions comparable to those in prior NPMLE studies [Saha and Guntuboyina, 2020, Soloff et al., 2025] and attains a near-parametric rate of $n^{-1}(\text{log}n{)}^{O(d)}$. These results rely on standard regularity conditions, including compact support of the latent measure and well-conditioned covariance matrices, which are common in NPMLE theory but restrict general applicability in heavy-tailed or ill-conditioned settings. Although the rate matches the minimax lower bound in Hellinger distance for Gaussian mixture estimation [Kim and Guntuboyina, 2022], the dependence on the ambient dimension $d$ remains unfavorable and limits its practicality in very high dimensions. This reflects the intrinsic difficulty of estimating a $d$-dimensional latent measure and recovering each ${\theta }_i$ from a single noisy sample $X_i$. Since our experiments show that the score-based $G$-modeling approach is more robust and computationally scalable than existing NPMLE solvers, and diffusion models have been observed to capture low-dimensional data manifolds [Pidstrigach, 2022], a promising future direction is to derive rates depending only on the intrinsic dimension $d_0\ll d$ of the manifold supporting $G^{*}$, following assumptions similar to Yakovlev and Puchkin [2025]. Algorithmically, our current solution methods rely on heuristic sampling and approximations that do not guarantee convergence to global minima of the infinite-dimensional nonconvex objective. The formulation in Section 2 was primarily chosen to ensure theoretical tractability under the heteroscedastic setting and is not yet optimized for large-scale computation. Future work should aim to bridge this gap between theory and implementation by refining both the modeling formulation and algorithmic design, supported by further analysis of optimization and approximation errors within the empirical Bayes framework.

Broader implications.

The line of work suggests that integrating diffusion-based generative modeling with nonparametric empirical Bayes principles can yield a unified approach to statistical deconvolution, manifold recovery, and inverse inference under heteroscedastic noise. Future investigations can further extend this framework toward manifold-aware score estimation and other related directions that enhance its theoretical scope and practical applicability.

NeurIPS Paper Checklist

1. Claims

Question: Do the main claims made in the abstract and introduction accurately reflect the paper’s contributions and scope?

Answer: [Yes]

Justification: Our abstract and introduction set out an overview of our methods, the theoretical guarantees, and empirical improvements that are detailed in Section 2, 3 and 5, respectively. We clearly introduce the scope of study and background context in the first few paragraphs of the introduction.

Guidelines:

• The answer NA means that the abstract and introduction do not include the claims made in the paper.

• The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations. A No or NA answer to this question will not be perceived well by the reviewers.

• The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings.

• It is fine to include aspirational goals as motivation as long as it is clear that these goals are not attained by the paper.

2. Limitations

Question: Does the paper discuss the limitations of the work performed by the authors?

Answer: [Yes]

Justification: We discussed several non-positive implications with our theoretical results following Theorem 1. We also discussed computational challenges of our proposal in Section 4, with efforts to alleviate them.

Guidelines:

• The answer NA means that the paper has no limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.

• The authors are encouraged to create a separate “Limitations” section in their paper.

• The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g., independence assumptions, noiseless settings, model well-specification, asymptotic approximations only holding locally). The authors should reflect on how these assumptions might be violated in practice and what the implications would be.

• The authors should reflect on the scope of the claims made, e.g., if the approach was only tested on a few datasets or with a few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.

• The authors should reflect on the factors that influence the performance of the approach. For example, a facial recognition algorithm may perform poorly when image resolution is low or images are taken in low lighting. Or a speech-to-text system might not be used reliably to provide closed captions for online lectures because it fails to handle technical jargon.

• The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size.

• If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and fairness.

• While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, a worse outcome might be that reviewers discover limitations that aren’t acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an important role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations.

3. Theory assumptions and proofs

Question: For each theoretical result, does the paper provide the full set of assumptions and a complete (and correct) proof?

Answer: [Yes]

Justification: We detailed the most significant assumptions in the main text and presented a proof sketch. We aim to present complete technical derivations in our supplemental.

Guidelines:

• The answer NA means that the paper does not include theoretical results.

• All the theorems, formulas, and proofs in the paper should be numbered and cross-referenced.

• All assumptions should be clearly stated or referenced in the statement of any theorems.

• The proofs can either appear in the main paper or the supplemental material, but if they appear in the supplemental material, the authors are encouraged to provide a short proof sketch to provide intuition.

• Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix or supplemental material.

• Theorems and Lemmas that the proof relies upon should be properly referenced.

4. Experimental result reproducibility

Question: Does the paper fully disclose all the information needed to reproduce the main experimental results of the paper to the extent that it affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided or not)?

Answer: [Yes]

Justification: Our major contribution is a statistical modeling perspective together with an optimization objective that can be implemented with various approaches. We discuss our approaches to approximately solve our proposed model. We detailed datasets, metrics, baselines and model variants in Section 5, and a pointer to the supplemental for implementation details, code, and more results.

Guidelines:

• The answer NA means that the paper does not include experiments.

• If the paper includes experiments, a No answer to this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.

• If the contribution is a dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable.

• Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is a novel architecture, describing the architecture fully might suffice, or if the contribution is a specific model and empirical evaluation, it may be necessary to either make it possible for others to replicate the model with the same dataset, or provide access to the model. In general. releasing code and data is often one good way to accomplish this, but reproducibility can also be provided via detailed instructions for how to replicate the results, access to a hosted model (e.g., in the case of a large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed.

• While NeurIPS does not require releasing code, the conference does require all submissions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example

  1. If the contribution is primarily a new algorithm, the paper should make it clear how to reproduce that algorithm.
  2. If the contribution is primarily a new model architecture, the paper should describe the architecture clearly and fully.
  3. If the contribution is a new model (e.g., a large language model), then there should either be a way to access this model for reproducing the results or a way to reproduce the model (e.g., with an open-source dataset or instructions for how to construct the dataset).
  4. We recognize that reproducibility may be tricky in some cases, in which case authors are welcome to describe the particular way they provide for reproducibility. In the case of closed-source models, it may be that access to the model is limited in some way (e.g., to registered users), but it should be possible for other researchers to have some path to reproducing or verifying the results.

5. Open access to data and code

Question: Does the paper provide open access to the data and code, with sufficient instructions to faithfully reproduce the main experimental results, as described in supplemental material?

Answer: [Yes]

Justification: We detailed the complete synthetic data preparation process in the main text. For the APOGEE data we used, it is publicly available at https://data.sdss.org/sas/dr14/apogee/vac/apogee-rc/cat/. We provided the code together with usages in the supplemental.

Guidelines:

• The answer NA means that paper does not include experiments requiring code.

• Please see the NeurIPS code and data submission guidelines (https://nips.cc/public/guides/CodeSubmissionPolicy) for more details.

• While we encourage the release of code and data, we understand that this might not be possible, so “No” is an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for a new open-source benchmark).

• The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurIPS code and data submission guidelines (https://nips.cc/public/guides/CodeSubmissionPolicy) for more details.

• The authors should provide instructions on data access and preparation, including how to access the raw data, preprocessed data, intermediate data, and generated data, etc.

• The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. If only a subset of experiments are reproducible, they should state which ones are omitted from the script and why.

• At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).

• Providing as much information as possible in supplemental material (appended to the paper) is recommended, but including URLs to data and code is permitted.

6. Experimental setting/details

Question: Does the paper specify all the training and test details (e.g., data splits, hyperparameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?

Answer: [Yes]

Justification: Dataset generation procedures, noise levels, model variants and evaluation protocols are given in Section 5. We will put the architecture and hyperparameters for the continuous normalizing flow model in the supplemental.

Guidelines:

• The answer NA means that the paper does not include experiments.

• The experimental setting should be presented in the core of the paper to a level of detail that is necessary to appreciate the results and make sense of them.

• The full details can be provided either with the code, in appendix, or as supplemental material.

7. Experiment statistical significance

Question: Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?

Answer: [Yes]

Justification: Table 1 reports mean and standard deviation over 5 independent runs for each metric, if the method is run under stochastic optimization methods.

Guidelines:

• The answer NA means that the paper does not include experiments.

• The authors should answer “Yes” if the results are accompanied by error bars, confidence intervals, or statistical significance tests, at least for the experiments that support the main claims of the paper.

• The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall run with given experimental conditions).

• The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.)

• The assumptions made should be given (e.g., Normally distributed errors).

• It should be clear whether the error bar is the standard deviation or the standard error of the mean.

• It is OK to report 1-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have a 96% CI, if the hypothesis of Normality of errors is not verified.

• For asymmetric distributions, the authors should be careful not to show in tables or figures symmetric error bars that would yield results that are out of range (e.g. negative error rates).

• If error bars are reported in tables or plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables in the text.

8. Experiments compute resources

Question: For each experiment, does the paper provide sufficient information on the computer resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?

Answer: [Yes]

Justification: We plan to detail these details together with other implementation details in the supplemental.

Guidelines:

• The answer NA means that the paper does not include experiments.

• The paper should indicate the type of compute workers CPU or GPU, internal cluster, or cloud provider, including relevant memory and storage.

• The paper should provide the amount of compute required for each of the individual experimental runs as well as estimate the total compute.

• The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary or failed experiments that didn’t make it into the paper).

9. Code of ethics

Question: Does the research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics https://neurips.cc/public/EthicsGuidelines?

Answer: [Yes]

Justification: We acknowledge that we have reviewed the Code of ethics.

Guidelines:

• The answer NA means that the authors have not reviewed the NeurIPS Code of Ethics.

• If the authors answer No, they should explain the special circumstances that require a deviation from the Code of Ethics.

• The authors should make sure to preserve anonymity (e.g., if there is a special consideration due to laws or regulations in their jurisdiction).

10. Broader impacts

Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?

Answer: [Yes]

Justification: We discussed the impact of the problem in scientific domains. We demonstrated that our framework help to mitigate challenges in traditional methods, which has the potential to benefit future scientific inference tasks.

Guidelines:

• The answer NA means that there is no societal impact of the work performed.

• If the authors answer NA or No, they should explain why their work has no societal impact or why the paper does not address societal impact.

• Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.

• The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is a direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster.

• The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology.

• If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms to monitor how a system learns from feedback over time, improving the efficiency and accessibility of ML).

11. Safeguards

Question: Does the paper describe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?

Answer: [NA]

Justification: Our work does not release a high-risk model or dataset.

Guidelines:

• The answer NA means that the paper poses no such risks.

• Released models that have a high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere to usage guidelines or restrictions to access the model or implementing safety filters.

• Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.

• We recognize that providing effective safeguards is challenging, and many papers do not require this, but we encourage authors to take this into account and make a best faith effort.

12. Licenses for existing assets

Question: Are the creators or original owners of assets (e.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?

Answer: [Yes]

Justification: We cited the source of APOGEE data and relevant papers. The data is publicly available at https://www.sdss4.org/.

Guidelines:

• The answer NA means that the paper does not use existing assets.

• The authors should cite the original paper that produced the code package or dataset.

• The authors should state which version of the asset is used and, if possible, include a URL.

• The name of the license (e.g., CC-BY 4.0) should be included for each asset.

• For scraped data from a particular source (e.g., website), the copyright and terms of service of that source should be provided.

• If assets are released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, paperswithcode.com/datasets has curated licenses for some datasets. Their licensing guide can help determine the license of a dataset.

• For existing datasets that are re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided.

• If this information is not available online, the authors are encouraged to reach out to the asset’s creators.

13. New assets

Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?

Answer: [NA]

Justification: Our paper does not release any new assets.

Guidelines:

• The answer NA means that the paper does not release new assets.

• Researchers should communicate the details of the dataset/code/model as part of their submissions via structured templates. This includes details about training, license, limitations, etc.

• The paper should discuss whether and how consent was obtained from people whose asset is used.

• At submission time, remember to anonymize your assets (if applicable). You can either create an anonymized URL or include an anonymized zip file.

14. Crowdsourcing and research with human subjects

Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well as details about compensation (if any)?

Answer: [NA]

Justification: The work involves no human-subject data collection.

Guidelines:

• The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.

• Including this information in the supplemental material is fine, but if the main contribution of the paper involves human subjects, then as much detail as possible should be included in the main paper.

• According to the NeurIPS Code of Ethics, workers involved in data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector.

15. Institutional review board (IRB) approvals or equivalent for research with human subjects

Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or an equivalent approval/review based on the requirements of your country or institution) were obtained?

Answer: [NA]

Justification: No human-subject research is conducted.

Guidelines:

• The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.

• Depending on the country in which research is conducted, IRB approval (or equivalent) may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.

• We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere to the NeurIPS Code of Ethics and the guidelines for their institution.

• For initial submissions, do not include any information that would break anonymity (if applicable), such as the institution conducting the review.

16. Declaration of LLM usage

Question: Does the paper describe the usage of LLMs if it is an important, original, or non-standard component of the core methods in this research? Note that if the LLM is used only for writing, editing, or formatting purposes and does not impact the core methodology, scientific rigorousness, or originality of the research, declaration is not required.

Answer: [NA]

Justification: The core original method are developed without any LLMs.

Guidelines:

• The answer NA means that the core method development in this research does not involve LLMs as any important, original, or non-standard components.

• Please refer to our LLM policy (https://neurips.cc/Conferences/2025/LLM) for what should or should not be described.