Nonlinear Moment Matching

• Nonlinear moment matching is a framework that reproduces curated nonlinear moment objects, such as conditional expectations and transfer-function derivatives, to capture complex dynamics.
• It is applied across diverse domains—including diffusion model distillation, model order reduction, and radiative transfer—to maintain fidelity in nonlinear systems.
• The approach leverages techniques like invariance equations, kernel-based MMD/CMMD losses, and projection methods for efficient, simulation-free optimization.

Nonlinear moment matching denotes a family of constructions in which a nonlinear approximation, reduced model, sampler, or estimator is required to reproduce selected moments, conditional moments, multivariate transfer-function derivatives, or steady-state responses of a reference system. Across recent arXiv literature, the term appears in diffusion-model distillation, nonlinear model order reduction, causal generative modeling, radiative-transfer closures, Monte Carlo variance reduction, and distribution fitting. The unifying theme is that the matched object is not restricted to low-order polynomial moments of a fixed distribution; it may instead be a conditional expectation along a sampling trajectory, the solution of an invariance equation, a multivariate Volterra coefficient, or a nonlinear closure relation (Salimans et al., 2024, Varona et al., 2019, Park, 2020).

1. Conceptual scope and principal formulations

In one major line of work, nonlinear moments are conditional expectations. In multistep diffusion distillation, the target is the conditional expectation of clean data given noisy data along the sampling trajectory, and the framework is explicitly interpreted as a generalized method of moments problem whose moments are “nonlinear moments” rather than fixed polynomial moments (Salimans et al., 2024). In nonlinear model reduction, the moment is the steady-state output response induced by a signal generator, characterized by an invariance equation rather than by transfer-function Taylor coefficients alone (Varona et al., 2019). In causal generative modeling, moment matching is realized through MMD and CMMD losses that act on latent variables and on conditional distributions over graph edges (Park, 2020).

Setting	Matched object	Representative relation
Diffusion distillation	Conditional expectation along trajectory	$$\tilde{L}(\eta)=\frac{1}{2}\mathbb{E}_{g(x_s)}\left\|\mathbb{E}_g[\tilde{x}\mid x_s]-\mathbb{E}_q[x\mid x_s]\right\|^2$$
Nonlinear model reduction	Steady-state map under signal generator	$$\frac{\partial \pi}{\partial \omega}s(\omega)=f(\pi(\omega),\ell(\omega))$$
Causal graph networks	Conditional distribution on graph edges	$$D_{\mathrm{CMMD}(Q(\widehat{Z}\mid Z_{\mathrm{pa}})\,\|\,P(Z\mid Z_{\mathrm{pa}}))}$$
Monte Carlo variance reduction	Moment matching in normal-score space	$$\tilde{Y}^{(1)}(k):=F_Y^{-1}(\mathcal{N}(\tilde{X}^{(1)}(k)))$$
Radiative transfer	Nonlinear closure of higher moments	$$E_{N+1}=\int_{-1}^1 \mu^{N+1}\hat{I}(\mu;E_0,\dots,E_N)\,d\mu$$

A common source of ambiguity is that the word “moment” does not carry a single standardized meaning across these literatures. In some papers it refers to classical statistical moments; in others it denotes steady-state interpolation data, derivatives of multivariate transfer functions, or moment closures in kinetic theory. This suggests that “nonlinear moment matching” is best understood as a structural idea—matching appropriately defined moment objects in a nonlinear setting—rather than as a single algorithmic recipe.

2. Conditional-expectation matching in diffusion models

A recent and prominent use of nonlinear moment matching appears in diffusion-model distillation. The sampling process is written in simplified form as

$$x_t=\alpha_t x+\sigma_t \varepsilon_t,\qquad \varepsilon_t\sim \mathcal{N}(0,I),$$

with a teacher denoiser $g_\theta(x_t,t)$ that ideally learns $\mathbb{E}[x\mid x_t]$ and a student denoiser $g_\eta(x_t,t)$. The central requirement is to match conditional expectations of clean data given noisy data along the trajectory: $$\mathbb{E}_g[\tilde{x}\mid x_s]=\mathbb{E}_q[x\mid x_s].$$ The associated objective is

$$\frac{\partial \pi}{\partial \omega}s(\omega)=f(\pi(\omega),\ell(\omega))$$0

and a practical stop-gradient loss is introduced for training (Salimans et al., 2024).

Two algorithmic realizations are described. The first uses alternating optimization with an auxiliary denoising model $$\frac{\partial \pi}{\partial \omega}s(\omega)=f(\pi(\omega),\ell(\omega))$$1 that estimates $$\frac{\partial \pi}{\partial \omega}s(\omega)=f(\pi(\omega),\ell(\omega))$$2 on student-generated samples. The second is a parameter-space moment-matching variant that removes the auxiliary model and yields an objective of the form

$$\frac{\partial \pi}{\partial \omega}s(\omega)=f(\pi(\omega),\ell(\omega))$$3

The paper identifies a connection to Efficient Method of Moments, because the matched quantity becomes the teacher gradient rather than a direct sample-space discriminator (Salimans et al., 2024).

This formulation generalizes one-step distillation to the multistep case. The paper emphasizes that the matching is nonlinear because the conditional expectations encode regression functions of the data given noise, “not just fixed polynomial moments (mean, variance) as in classical method of moments.” Empirically, distilled models using up to 8 sampling steps outperform their one-step versions and also their original teacher models on ImageNet, and the method also shows promising large text-to-image results in image space without autoencoders or upsamplers (Salimans et al., 2024).

3. Invariance equations and steady-state matching in nonlinear model reduction

In nonlinear model reduction, moment matching is formulated through a nonlinear signal generator and an invariance equation. For

$$\frac{\partial \pi}{\partial \omega}s(\omega)=f(\pi(\omega),\ell(\omega))$$4

with signal generator

$$\frac{\partial \pi}{\partial \omega}s(\omega)=f(\pi(\omega),\ell(\omega))$$5

nonlinear moments are defined through the map $$\frac{\partial \pi}{\partial \omega}s(\omega)=f(\pi(\omega),\ell(\omega))$$6 solving the Sylvester-type PDE

$$\frac{\partial \pi}{\partial \omega}s(\omega)=f(\pi(\omega),\ell(\omega))$$7

This equation generalizes the linear Sylvester equation and encodes steady-state matching under inputs generated by the exosystem (Varona et al., 2019).

Direct solution of the PDE is generally intractable, so a simulation-free reduction strategy replaces the nonlinear lifting with a linear projection $$\frac{\partial \pi}{\partial \omega}s(\omega)=f(\pi(\omega),\ell(\omega))$$8 and reduces the PDE to nonlinear algebraic equations,

$$\frac{\partial \pi}{\partial \omega}s(\omega)=f(\pi(\omega),\ell(\omega))$$9

followed by column-wise decomposition and collocation. The resulting sample-wise equations,

$$D_{\mathrm{CMMD}(Q(\widehat{Z}\mid Z_{\mathrm{pa}})\,\|\,P(Z\mid Z_{\mathrm{pa}}))}$$0

are solved with Newton-type methods, optional orthonormalization, and SVD deflation. The method is explicitly described as “simulation-free” because it avoids time simulation of the original full-order model during basis construction (Varona et al., 2019).

The same program was transferred to nonlinear second-order structural systems,

$$D_{\mathrm{CMMD}(Q(\widehat{Z}\mid Z_{\mathrm{pa}})\,\|\,P(Z\mid Z_{\mathrm{pa}}))}$$1

where the nonlinear moment-matching condition becomes a second-order Sylvester-like PDE for the embedding map $$D_{\mathrm{CMMD}(Q(\widehat{Z}\mid Z_{\mathrm{pa}})\,\|\,P(Z\mid Z_{\mathrm{pa}}))}$$2. The practical algorithm again uses linear projection, column-wise decoupling, and time discretization, thereby replacing the original PDE by nonlinear algebraic solves at selected snapshot times (Varona et al., 2019).

A least-squares extension recasts nonlinear moment matching as an optimization problem involving the invariance equation and the steady-state behavior of an error system. In that setting, the nonlinear moment is

$$D_{\mathrm{CMMD}(Q(\widehat{Z}\mid Z_{\mathrm{pa}})\,\|\,P(Z\mid Z_{\mathrm{pa}}))}$$3

and the reduction objective is to minimize a derivative-based discrepancy between full and reduced moments, with an associated worst-case steady-state r.m.s. error bound (Padoan, 2021). A further computational development replaces direct PDE solution by a polynomial Galerkin residual method,

$$D_{\mathrm{CMMD}(Q(\widehat{Z}\mid Z_{\mathrm{pa}})\,\|\,P(Z\mid Z_{\mathrm{pa}}))}$$4

and solves the resulting nonlinear algebraic system with Newton iteration on the monomial coefficients (Doebeli et al., 2024).

4. Volterra-series, quadratic-bilinear, and parametric formulations

For quadratic-bilinear systems, nonlinear moment matching is expressed through multivariate transfer functions and Volterra kernels. A representative quadratic-bilinear descriptor system is

$$D_{\mathrm{CMMD}(Q(\widehat{Z}\mid Z_{\mathrm{pa}})\,\|\,P(Z\mid Z_{\mathrm{pa}}))}$$5

and the reduction objective is to construct projection matrices $$D_{\mathrm{CMMD}(Q(\widehat{Z}\mid Z_{\mathrm{pa}})\,\|\,P(Z\mid Z_{\mathrm{pa}}))}$$6 such that reduced multivariate transfer functions match the original ones at selected interpolation points (Khattak et al., 2021). In this setting, moment matching means Hermite interpolation of $$D_{\mathrm{CMMD}(Q(\widehat{Z}\mid Z_{\mathrm{pa}})\,\|\,P(Z\mid Z_{\mathrm{pa}}))}$$7 and their derivatives, and a greedy framework is built around a posteriori error bounds

$$D_{\mathrm{CMMD}(Q(\widehat{Z}\mid Z_{\mathrm{pa}})\,\|\,P(Z\mid Z_{\mathrm{pa}}))}$$8

to select interpolation points adaptively (Khattak et al., 2021).

A related projection-based method for quadratic-bilinear systems distinguishes symmetric, triangular, and regular forms of multivariate transfer functions. The regular form is emphasized because it avoids the combinatorial complexity of permutation symmetrization and enables matching not only the first two but also the first three multivariate transfer functions. For example,

$$D_{\mathrm{CMMD}(Q(\widehat{Z}\mid Z_{\mathrm{pa}})\,\|\,P(Z\mid Z_{\mathrm{pa}}))}$$9

$$\tilde{Y}^{(1)}(k):=F_Y^{-1}(\mathcal{N}(\tilde{X}^{(1)}(k)))$$0

and the basis construction is expanded recursively to enforce higher-order multivariate moment matching (Asif et al., 2019).

For nonlinear parametric systems, the same invariance structure is lifted to include parameters. If

$$\tilde{Y}^{(1)}(k):=F_Y^{-1}(\mathcal{N}(\tilde{X}^{(1)}(k)))$$1

and the signal generator is

$$\tilde{Y}^{(1)}(k):=F_Y^{-1}(\mathcal{N}(\tilde{X}^{(1)}(k)))$$2

then the parametric moment is defined by

$$\tilde{Y}^{(1)}(k):=F_Y^{-1}(\mathcal{N}(\tilde{X}^{(1)}(k)))$$3

A data-driven approximation uses basis functions in $$\tilde{Y}^{(1)}(k):=F_Y^{-1}(\mathcal{N}(\tilde{X}^{(1)}(k)))$$4,

$$\tilde{Y}^{(1)}(k):=F_Y^{-1}(\mathcal{N}(\tilde{X}^{(1)}(k)))$$5

with coefficients obtained by least squares from measured steady-state responses. The resulting reduced model preserves the moment approximation across the parameter range and is designed to preserve asymptotic stability and dissipativity in the linear case, while the nonlinear construction controls asymptotic stability through the feedback mapping $$\tilde{Y}^{(1)}(k):=F_Y^{-1}(\mathcal{N}(\tilde{X}^{(1)}(k)))$$6 (Zhang et al., 12 Jun 2025).

A further development for MIMO polynomial nonlinear systems uses formal power-series decomposition of the center-manifold PDE. The expansion

$$\tilde{Y}^{(1)}(k):=F_Y^{-1}(\mathcal{N}(\tilde{X}^{(1)}(k)))$$7

turns nonlinear moment matching up to degree $$\tilde{Y}^{(1)}(k):=F_Y^{-1}(\mathcal{N}(\tilde{X}^{(1)}(k)))$$8 into a recursive sequence of Sylvester equations. The paper derives lower bounds for the reduced-model order, shows that in the MIMO case the lower bound can be strictly less than the number of matched moments, and shows that the bound is affected by the ratio of input and output channels (Huang et al., 19 Aug 2025).

5. Conditional-distribution matching, distribution fitting, and uncertainty propagation

A distinct strand of nonlinear moment matching operates directly on distributions and conditional distributions. In causal inference, a generative conditional-moment-matching graph-neural-network applies MMD to the latent joint and CMMD to graph edges. The combined loss is

$$\tilde{Y}^{(1)}(k):=F_Y^{-1}(\mathcal{N}(\tilde{X}^{(1)}(k)))$$9

with $$E_{N+1}=\int_{-1}^1 \mu^{N+1}\hat{I}(\mu;E_0,\dots,E_N)\,d\mu$$0 so that conditional moment matching is prioritized. The paper emphasizes that MMD and CMMD are nonparametric kernel-based discrepancies that can match “all moments,” including higher-order nonlinear moments, and that the edge-wise construction supports out-of-sample interventional sampling (Park, 2020).

For phase-type distributions, the moment-matching problem is cast as high-dimensional unconstrained optimization. The classical constrained objective

$$E_{N+1}=\int_{-1}^1 \mu^{N+1}\hat{I}(\mu;E_0,\dots,E_N)\,d\mu$$1

is replaced by a smooth re-parametrization into unconstrained variables $$E_{N+1}=\int_{-1}^1 \mu^{N+1}\hat{I}(\mu;E_0,\dots,E_N)\,d\mu$$2, with

$$E_{N+1}=\int_{-1}^1 \mu^{N+1}\hat{I}(\mu;E_0,\dots,E_N)\,d\mu$$3

This allows fitting as many as 20 moments to phase-type distributions with as many as 100 phases, and the same framework can incorporate additional differentiable targets such as CDF values (Sherzer et al., 26 May 2025).

For Monte Carlo variance reduction, the paper on asymptotic universal moment matching properties of normal distributions shows that classical linear moment matching asymptotically reduces variance for all integrands if and only if the underlying distribution is normal. To extend the guarantee to any continuous distribution, it proposes nonlinear moment matching through a quantile transform. If

$$E_{N+1}=\int_{-1}^1 \mu^{N+1}\hat{I}(\mu;E_0,\dots,E_N)\,d\mu$$4

then linear moment matching is applied in normal-score space, followed by transformation back: $$E_{N+1}=\int_{-1}^1 \mu^{N+1}\hat{I}(\mu;E_0,\dots,E_N)\,d\mu$$5 The resulting estimators enjoy asymptotic variance reduction for any continuous input distribution (Liu, 5 Aug 2025).

In Gaussian-mixture uncertainty propagation, moment matching means preserving the mean and covariance of a split mixand under nonlinear propagation. If a parent Gaussian $$E_{N+1}=\int_{-1}^1 \mu^{N+1}\hat{I}(\mu;E_0,\dots,E_N)\,d\mu$$6 is split along direction $$E_{N+1}=\int_{-1}^1 \mu^{N+1}\hat{I}(\mu;E_0,\dots,E_N)\,d\mu$$7, the child means and covariances are chosen as

$$E_{N+1}=\int_{-1}^1 \mu^{N+1}\hat{I}(\mu;E_0,\dots,E_N)\,d\mu$$8

with

$$E_{N+1}=\int_{-1}^1 \mu^{N+1}\hat{I}(\mu;E_0,\dots,E_N)\,d\mu$$9

The paper then proposes heuristics for selecting the splitting direction based on uncertainty, higher-order nonlinearity, and whitening-based natural scaling (Kulik et al., 2024).

6. Nonlinear closure models in radiative transfer

In radiative transfer, nonlinear moment matching appears as closure of truncated moment systems by ansatz functions that depend on low-order moments. For slab geometry, a new nonlinear moment model uses the $$x_t=\alpha_t x+\sigma_t \varepsilon_t,\qquad \varepsilon_t\sim \mathcal{N}(0,I),$$0 ansatz as the weight function,

$$x_t=\alpha_t x+\sigma_t \varepsilon_t,\qquad \varepsilon_t\sim \mathcal{N}(0,I),$$1

with

$$x_t=\alpha_t x+\sigma_t \varepsilon_t,\qquad \varepsilon_t\sim \mathcal{N}(0,I),$$2

The full ansatz is

$$x_t=\alpha_t x+\sigma_t \varepsilon_t,\qquad \varepsilon_t\sim \mathcal{N}(0,I),$$3

where $$x_t=\alpha_t x+\sigma_t \varepsilon_t,\qquad \varepsilon_t\sim \mathcal{N}(0,I),$$4 and the coefficients are determined by enforcing moment constraints. The closed moment is then

$$x_t=\alpha_t x+\sigma_t \varepsilon_t,\qquad \varepsilon_t\sim \mathcal{N}(0,I),$$5

Because the weight depends on $$x_t=\alpha_t x+\sigma_t \varepsilon_t,\qquad \varepsilon_t\sim \mathcal{N}(0,I),$$6, the closure is nonlinear and adapts to anisotropy (Fan et al., 2018).

That model combines the primary idea of the $$x_t=\alpha_t x+\sigma_t \varepsilon_t,\qquad \varepsilon_t\sim \mathcal{N}(0,I),$$7 model with the $$x_t=\alpha_t x+\sigma_t \varepsilon_t,\qquad \varepsilon_t\sim \mathcal{N}(0,I),$$8 ansatz from the $$x_t=\alpha_t x+\sigma_t \varepsilon_t,\qquad \varepsilon_t\sim \mathcal{N}(0,I),$$9 family, retaining explicit closures while improving the approximation of anisotropic distributions. For the three-moment case, the paper proves hyperbolicity for all realizable moments and analyzes the characteristic structure of the Riemann problem in detail (Fan et al., 2018).

The three-dimensional extension uses a weighted polynomial expansion around the 3D $g_\theta(x_t,t)$0 ansatz,

$g_\theta(x_t,t)$1

and approximates moments of order $g_\theta(x_t,t)$2 by integrating this ansatz. A hyperbolic regularization based on a modified projection yields a 3D HMPN model with global hyperbolicity, rotational invariance, physical wave speeds, spectral accuracy, and correct higher-order Eddington approximation (Li et al., 2020).

These radiative-transfer models clarify a domain-specific meaning of nonlinear moment matching: the matched objects are moments of the kinetic density, but the closure is nonlinear because the basis or weight itself depends on the lower-order moments. The result is a closure that is explicit, anisotropy-aware, and mathematically structured.

7. Related constructions, adjacent methods, and recurring misunderstandings

Several neighboring literatures illuminate what nonlinear moment matching is, and what it is not. One recurrent misunderstanding is to identify moment matching with matching only mean and covariance. That description fits some Gaussian-mixture splitting methods, but it is too narrow for diffusion distillation, causal CMMD, or nonlinear model reduction. In the diffusion setting, the matched objects are trajectory-wise conditional expectations; in causal graph networks, CMMD is used precisely because mean/variance matching is insufficient in highly nonlinear settings (Salimans et al., 2024, Park, 2020).

A second misunderstanding is to treat nonlinear moment matching as synonymous with a single closure mechanism. In fact, nearby methods employ distinct principles. The derivative-matching closure technique for stochastic differential equations approximates each unclosed moment by a separable monomial in lower-order moments,

$g_\theta(x_t,t)$3

with exponents chosen so that values and first two time derivatives match at deterministic initial conditions. This is a moment-closure method for polynomial and trigonometric stochastic systems, not an invariance-PDE reduction method, but it shares the central idea of replacing inaccessible nonlinear moment quantities by structured surrogates derived from matching conditions (Ghusinga et al., 2017).

Likewise, Prony’s method solves the nonlinear equations of the connected-moments expansion by converting

$g_\theta(x_t,t)$4

into a polynomial root-finding problem. This is a nonlinear moment-matching problem in the classical sense of fitting exponential sums to moment data, and it provides a useful reminder that “nonlinear” may refer either to the matched system or to the equations defining the fit (Fernández, 2011).

Finally, exact moment propagation methods sit adjacent to, rather than inside, most moment-matching frameworks. The Moment-based Kalman Filter computes exact moments of transformed random variables through nonlinear process and measurement models, including non-Gaussian and correlated cases, and uses those moments in a Kalman-style recursion. Its objective is exact propagation rather than reduced-model interpolation or conditional-distribution matching, but it exemplifies the same broader preoccupation with preserving informative moment structure under nonlinear maps (Shimizu et al., 2023).

Overall, the literature presents nonlinear moment matching as a versatile framework rather than a single canon. Depending on the application, it can mean matching conditional expectations along a diffusion trajectory, matching steady-state responses through an invariance PDE, matching multivariate Volterra derivatives, matching conditional distributions with CMMD, fitting many prescribed moments of a phase-type distribution, or constructing nonlinear kinetic closures. The breadth of these uses is not terminological drift alone; it reflects a common strategy of encoding fidelity requirements through moment objects that remain meaningful after linear theory ceases to be adequate.