Denoising Is Stein: Moment Access and a Resource Hierarchy for Learning Hidden Structure in Diffusion Models

Introduction

Diffusion models estimate the score $\nabla\log p_t$ as Gaussian noise progressively obscures a data distribution (Ho et al., 2020; Song et al., 2021). Training provides denoising pairs $(X_0, X_t)$: a clean point and a noisy view. The learner predicts $X_0$ from $X_t$, thereby estimating the score (Vincent, 2011). If the data have a hidden non-Gaussian direction, a successful denoiser must detect it through those pairs. The amount of data needed depends on what the learner can represent and whether it can use repeated views of the same clean point.

Existing theory studies parts of this question under different data models and resource measures. Lower bounds for random features and kernels count features or samples for fixed representations (Ghorbani et al., 2021; Mei et al., 2022); analyses of trained denoisers count fresh pairs for one architecture and optimizer ([5]); statistical-query and cryptographic lower bounds count queries or computation for arbitrary algorithms (Diakonikolas et al., 2017; Song, 2024). These results do not directly compare the cost of learning the same hidden signal across learner classes. We ask:

Which moments of a hidden structure can a denoising learner access, and how do representation size, distinct clean samples, noisy views and the learner’s access to geometry determine the cost of improving on the Gaussian score?

We study two versions of one hidden-component model. Data are $X_0 = UZ + \sigma_0(I - UU^\top)G_0$ with an unknown orthonormal $U \in\mathbb{R}^{d \times k}$, a non-Gaussian latent vector $Z \in\mathbb{R}^k$ and an independent Gaussian $G_0$. In the covariance-visible regime ($\sigma_0 = 0$), the data lie on the subspace, their covariance reveals $U$, and principal components recover it from $k + 1$ points. In the covariance-matched regime, $\operatorname{Cov}(Z) = \sigma_0^2 I_k$, so $\operatorname{Cov}(X_0) = \sigma_0^2 I_d$ carries no information about $U$. The first informative moment has order $k^* \ge3$, the generative exponent of the latent law. This is non-Gaussian component analysis (Diakonikolas et al., 2017; Tan and Vershynin, 2018) seen through a diffusion model.

Main results

Figure 1 shows the observable difference between the two data regimes: a line of clean points in one case, and an isotropic covariance that masks a non-Gaussian direction in the other. Table 1 summarizes the resource costs for each learner. The results follow from a moment-access calculus for the denoising loss (§ 3); Figure 2 compares the leading scales.

Contributions. The Stein–Tweedie identity expresses each residual correlation as a discrepancy between the noisy law and a baseline (Theorem 3.1); on polynomial fields, these correlations recover exactly the corresponding moment differences (Theorem 3.2). Thus a fitted Gaussian baseline removes the moments it matches, shifting the first accessible signal in the covariance-visible mixture from degree two to degree four (Theorem 3.4, Corollary 3.6). Conditioning on the clean point gives an exact clean-data moment query (Lemma 3.3). Averaging fresh noisy views of one point reduces the forward-process variance of the corresponding gradient estimate. Diffusion acts on the hidden coordinate through an Ornstein–Uhlenbeck noise operator; the static signal-to-noise ratio of the best low-degree correlation peaks at the time in Theorem 3.8.

For an orientation-independent dictionary, the covariance-visible mixture has a sharp feature threshold at $d^3/6$, with higher odd-degree steps at $d^J/J!$ (Theorems 4.5, 4.9). For one hidden direction, a one-sided change of measure identifies the height of the first step (Theorem 4.13). The covariance-matched model gives the $d^{k^*-1}$ exponent for dictionaries satisfying the stated $L^4/L^2$ condition (Theorem 4.21). Kernel ridge regression on noisy inputs has the same sample exponent under the conditions of Theorem 5.2; for the exponential kernel this includes every ridge (Lemma 5.3). When covariance reveals the subspace, principal-component projection instead yields fitted-model, score-risk and generated-law distributions independent of ambient dimension (Theorems 4.15, 4.16).

Learning the covariance-matched direction can be cheaper than fixing a representation in advance. Small-step online SGD on one unit needs $d^{k^*-1}$ steps up to $\log^2 d$ (Corollary 6.2), while the statistical-query bound gives tolerance $d^{-k^*/4+o(1)}$ under its output-norm assumption (Proposition 6.3). Averaging $r$ noisy views of each clean point, one gradient step of the matrix score model recovers the direction from $O((1+(\iota_t-1)/r)d^{k^*/2}\log d)$ clean points for even $k^*$ (Theorem 6.5). The same partial-trace matrix is the curvature created by Gaussian smoothing of a single-unit loss (Proposition 6.7). A learner that uses this curvature, sharpens the direction and fits a scalar readout obtains a fixed score gain with $O(d^{k^*/2}\log d)+O(d)+O(1)$ clean points (Theorem 6.8). For $k^*=4$, this separates it from the restricted learner classes above by one power of $d$, up to logarithms and under their stated hypotheses (Corollary 6.9).

The experiments check the predicted barriers and invariances in the covariance-visible regime. In the covariance-matched regime, the partial trace’s measured overlaps are consistent with the proposed spectral transition within Monte Carlo error. A width-8 network trained with clean-point reuse has a $d^2$-like training-set threshold over $d=64$–$192$, with fixed validation and calibration overhead (§ 7).

Related work

Low-dimensional structure in diffusion models. Chen et al. (2023a) give score approximation and estimation rates for data on a linear subspace, Oko et al. (2023) prove minimax optimality of diffusion estimators over Besov classes, and Li and Yan (2024) show that samplers adapt to unknown intrinsic dimension. Given an accurate score, sampling error is controlled by Chen et al. (2023b) and Benton et al. (2024). These are upper bounds for adaptive estimators; our lower bounds identify what restricted classes lack, and our upper bounds count each resource separately.

Fixed representations and kernels. Random-feature and kernel models in high dimension fit only polynomial components of degree $\ell$ with $d^\ell\ll N$ features or samples (Ghorbani et al., 2019, 2021; Mei et al., 2022), and the spectrum of inner-product kernel matrices in polynomial regimes is described by Misiakiewicz (2022) and Lu and Yau (2022). For denoising score matching, George et al. (2026) derive precise learning curves for random features and George and Macris (2026) extend them to manifold data. Here the inputs are non-Gaussian along a hidden direction, and the barrier applies to every orientation-blind dictionary with arbitrary coefficients. For one hidden direction, it also gives the exact risk height relative to the moment-matched Gaussian denoiser.

Exponents, statistical queries and non-Gaussian component analysis. The information exponent governs online SGD for single-index models (Ben Arous et al., 2021), and the generative exponent governs statistical-query and low-degree algorithms (Damian et al., 2024). Statistical-query lower bounds for non-Gaussian component analysis are due to Diakonikolas et al. (2017) and Diakonikolas et al. (2024), spectral algorithms to Vempala and Xiao (2011) and Tan and Vershynin (2018), and score estimation for such data is cryptographically hard in general (Song, 2024). Bardone et al. [5] show that a denoiser trained online on a mixed cumulant model learns pair-wise statistics at linear cost and the fourth cumulant at cubic cost, defining a diffusion information exponent; Székely et al. (2024) report a quadratic sample threshold for supervised networks in a spiked cumulant model. Our Theorem 3.2 identifies the generative exponent as the relevant one for denoising, and Corollary 3.6 gives a mechanism for the staged learning of Bardone et al. [5].

Tensor methods, smoothing and reuse. The fourth-moment tensor of non-Gaussian component data is a spiked tensor, for which unfolding reaches $d^{k/2}$ (Montanari and Richard, 2014) and gradient descent or Langevin dynamics from random starts needs $d^{k-1}$ (Ben Arous et al., 2020). Gaussian smoothing reaches the unfolding threshold for tensor PCA (Anandkumar et al., 2017) and the $d^{k/2}$ rate for online single-index learning (Damian et al., 2023); SGD can be viewed as optimizing a smoothed objective (Kleinberg et al., 2018). Reusing batches changes what gradient methods can learn in supervised two-layer networks (Dandi et al., 2024; Lee et al., 2024). A link between one gradient step of a diffusion objective and a spectral method was made for a covariance spike by Shah et al. (2023). Here the conditional Stein identity connects reuse to the denoising loss, and the curvature created by smoothing is the partial-trace statistic. The resulting two-stage learner has explicit costs for clean points, noisy views and arithmetic.

Stein’s method and score matching. Theorem 3.1 is a standard consequence of Stein’s lemma and Tweedie’s formula (Robbins, 1956; Efron, 2011); its supremum over a unit ball of fields is a Stein discrepancy (Gorham and Mackey, 2015), and the equivalence of denoising score matching with score matching is due to Hyvärinen (2005) and Vincent (2011). We use this identity to compare what different denoising learners can extract from the same hidden-component model.

Setting, learners and resources

Diffusion pairs, denoisers and baselines

Fix a time $t>0$ and write $a=a_t=e^{-t}$, $h=h_t=1-e^{-2t}$. For clean data $X_0\sim P_0$ on $\mathbb{R}^d$ with all moments finite and an independent $G\sim N(0,I_d)$, the noisy view is

the time-$t$ marginal of the Ornstein–Uhlenbeck process $dX_t=-X_t\,dt+\sqrt{2}\,dB_t$. Write $P_t$ for its law, $p_t$ for its density, $m_t(x)=\mathbb{E}[X_0\mid X_t=x]$ for the denoiser and $s_t=\nabla\log p_t$ for the score. Tweedie’s formula (Robbins, 1956; Efron, 2011) gives

and every denoiser estimate $\widehat{m}$ defines a score estimate $\widehat{s}=(a_t\widehat{m}-x)/h_t$. Then $\mathbb{E}\lVert\widehat{s}(X_t)-s_t(X_t)\rVert^2=(a_t/h_t)^2\mathbb{E}\lVert\widehat{m}(X_t)-m_t(X_t)\rVert^2$, so score risk and denoiser risk differ by the factor $(a_t/h_t)^2$. Regressing $X_0$ on $X_t$ is denoising score matching (Vincent, 2011): for any $f$,

Baselines and the Stein operator. A baseline is a law $Q_0$ with finite moments, with diffused law $Q_t$, smooth positive density $q_t$ and denoiser $m_t^Q$. Learners in this paper improve on a baseline: they fit $m_t^Q+f$ and are judged by how much of the residual $m_t-m_t^Q$ the correction $f$ captures. The Langevin–Stein operator of $Q_t$ acts on $C^1$ vector fields by $\mathcal{T}_{Q_t}\varphi=\operatorname{div}\varphi+\langle\nabla\log q_t,\varphi\rangle$; for $Q_t=N(\bar{m},C)$, $\mathcal{T}_{Q_t}\varphi=\operatorname{div}\varphi-\langle C^{-1}(x-\bar{m}),\varphi\rangle$.

Hermite notation. $\mathrm{He}_j$ is the probabilists’ Hermite polynomial and $h_j = \mathrm{He}_j/\sqrt{j!}$ its orthonormal version under $N(0,1)$. For $x \in\mathbb{R}^d$, $\mathrm{He}_j(x)$ also denotes the degree-$j$ Hermite tensor, the symmetric $j$-tensor-valued polynomial with $\mathrm{He}_j(x)[w^{\otimes j}] = \mathrm{He}_j(w^\mathsf{T}x)$ for unit $w$. Under $N(0,I)$ it satisfies the isometry $\mathbb{E}\,\mathrm{He}_j(x)[A]\,\mathrm{He}_j(x)[B] = \sum_{\pi}\langle A,\pi B\rangle$, summed over the $j!$ permutations of tensor slots. For even $k^* = 2p+2$ the partial trace is the matrix

Two data regimes