KGMM: A K-means Clustering Approach to Gaussian Mixture Modeling for Score Function Estimation

TL;DR

This work tackles the challenge of estimating the score function, the gradient of the log steady-state density, in complex dynamical systems. It introduces KGMM, a two-stage approach that first uses bisecting K-means clustering with Gaussian Mixture Models to obtain cluster-wise, noise-averaged score targets, and then trains a neural network to interpolate the score across state space. KGMM demonstrates accurate recovery of invariant measures and long-time statistics for potential, chaotic Lorenz-type, and KS systems, while offering substantial computational savings over conventional Denoising Score Matching. The method shows particular strength in moderate dimensions and large datasets, with clear guidelines on hyperparameter choice and limitations due to dimensionality and finite-$\sigma$ bias. Overall, KGMM advances data-driven, reduced-order stochastic modelling by providing a robust, efficient framework for score-function estimation in complex dynamical systems.

Abstract

We propose a hybrid method for accurately estimating the score function, i.e., the gradient of the log steady-state density, using a Gaussian Mixture Model (GMM) in conjunction with a bisecting K-means clustering step. Our approach, which we call KGMM, offers a systematic way to combine statistical density estimation with a neural-network-based interpolation of the score, leveraging the strengths of both. We demonstrate its ability to accurately reconstruct the long-time statistical properties of several paradigmatic systems, including potential systems, chaotic Lorenz-type models, and the Kuramoto-Sivashinsky equation. Numerical experiments show that KGMM yields robust estimates of the score function, even for small values of the covariance amplitude in the GMM, where standard GMM methods tend to fail because of noise amplification. We compare the performance of KGMM against the conventional Denoising Score Matching (DSM) approach, demonstrating that KGMM achieves more faithful reconstruction of the steady-state distribution for low-dimensional systems at a fraction of the computational cost. These accurate estimates allow us to build effective stochastic reduced-order models that reproduce the invariant measures of the target dynamics.

Figures (7)

• Figure 1: Comparison for different values of $\sigma$ between the score function obtained through the standard GMM (orange curve) and the one (blue curve) obtained by interpolating the discrete values of the KGMM score function (yellow points). Note that for small $\sigma$, the standard GMM curve becomes significantly noisier, whereas the KGMM approach preserves a close agreement with the true score (red curve). Each panel's white and black background represents the joint distribution of $(x_\omega, -z_\omega/\sigma),\, \omega \in \{1,\cdots N\}$. Fixing a value of $x$ and computing the expected value of the resulting conditional density yields the value of the yellow points.
• Figure 2: Reduced triad model (Eq. \ref{['reduced']}). Left panel: Comparison between the KGMM-estimated score function and its analytical expression given by Eq. \ref{['eq:reduced_score']}. Center panel: Comparison between the observed steady-state distribution (True, red) and the one obtained from integrating Eq. \ref{['eq:langevin']} using the KGMM score function (KGMM, blue), demonstrating that KGMM correctly reproduces the invariant measure. Right panel: Comparison between sample trajectories obtained by integrating Eq. \ref{['reduced']} (True, red) and Eq. \ref{['eq:langevin']} using the KGMM score function (KGMM, blue). Note that individual trajectories differ due to stochastic realizations.
• Figure 3: Two-dimensional asymmetric potential system.First row, left: The force field of the true score function (top) and the force field of the KGMM-estimated score function (bottom). First row, center: Comparison between the observed univariate PDFs for $x$ (top) and $y$ (bottom) with those obtained by integrating Eq. \ref{['lang']} with the KGMM-estimated score function, showing close agreement in marginal distributions. First row, right: Comparison between the observed bivariate probability density (top) and the reconstructed density using the KGMM-based score function (bottom), confirming reproduction of the joint distribution. Bottom row: Comparison between sample trajectories for $x$ (left) and $y$ (right) obtained by integrating Eq. \ref{['potential_2D']} (True, red) and Eq. \ref{['eq:langevin']} using the KGMM score function (KGMM, blue). Note that individual trajectories may differ due to stochastic realizations.
• Figure 4: Lorenz 63 system.First column: Comparison between the observed univariate PDFs for $x$, $y$, and $z$ (True, red) and those obtained integrating the Langevin equation using the KGMM-estimated score function (KGMM, blue), demonstrating accurate marginal distributions. Second and third columns: Comparison between the observed bivariate PDFs for $(x,y)$, $(x,z)$, and $(y,z)$ (True, left column) and those obtained using the KGMM-based score function (KGMM, right column), showing faithful reproduction of joint statistics despite different short-time trajectory behavior. Bottom row: Comparison between sample trajectories for $x$, $y$, and $z$ obtained by integrating Eq. \ref{['lorenz63']} (True, red) and Eq. \ref{['eq:langevin']} using the KGMM score function (KGMM, blue). Note that individual trajectories may differ due to stochastic realizations.
• Figure 5: Lorenz 96 system.Top left: Comparison between the observed univariate PDF (True, red) and the one obtained integrating the Langevin equation using the KGMM-estimated score function (KGMM, blue), showing excellent agreement. Top center and right: Comparison between the observed bivariate PDFs for $x[k]$-$x[k+1]$ (center) and $x[k]$-$x[k+2]$ (right) (True, top row) and those obtained using the KGMM-based score function (KGMM, bottom row), demonstrating that KGMM captures the non-Gaussian structure of the invariant measure. Bottom panel: Comparison between sample trajectories for $x[k]$ obtained by integrating Eq. \ref{['lorenz96']} (True, red) and Eq. \ref{['eq:langevin']} using the KGMM score function (KGMM, blue). Note that individual trajectories may differ due to stochastic realizations.