Organizations: Department of Mathematics, University of Waterloo, Waterloo, Canada · Department of Computer Science, Georgia Tech, Atlanta, USA · University of California, Irvine, Irvine, USA
We introduce AI-Kolmogorov, a novel framework for Symbolic Density Estimation (SymDE). Symbolic regression (SR) has been effectively used to produce interpretable models in standard regression settings but its applicability to density estimation tasks has largely been unexplored. To address the SymDE task we introduce a multi-stage pipeline: (i) problem decomposition through clustering and/or probabilistic graphical model structure learning; (ii) nonparametric density estimation; (iii) support estimation; and finally (iv) SR on the density estimate. We demonstrate the efficacy of AI-Kolmogorov on synthetic mixture models, multivariate normal distributions, and three exotic distributions, two of which are motivated by applications in high-energy physics. We show that AI-Kolmogorov can discover underlying distributions or otherwise provide valuable insight into the mathematical expressions describing them.
Figures & tables
Figure 1: The AI-Kolmogorov pipeline: decomposition (clustering/structure learning), density estimation, support estimation, symbolic regression, and warm-start refinement. Optional workflows are indicated by dashed arrows.
Figure 2: Ablation study on clustering and structure learning. Top: clustering results. Bottom: structure learning results.
Dataset
MSE without decomposition
MSE with decomposition
Gaussian mixture
4.71×10−4
3.57×10−4
4D Gaussian
1.01×10−3
4.24×10−5
Table 1: MSE comparison between direct application of AI-Kolmogorov and augmented AI-Kolmogorov. The MSE with decomposition reflects the MSE after recombining the most accurate expressions from the pareto fronts
Figure 3: Prediction (left) and residuals with respect to the ground truth (right) of the lowest loss expression. The maximum predicted density and absolute residual are 0.140 and 0.003 respectively.
Figure 4: Prediction and regions for local probability mass validation indicated as A & B (left). Residuals with respect to KDE (right) of the lowest loss expression. The maximum predicted density is 3.677, and the maximum absolute residual is 0.640.
Method
Region A
Region B
Empirical
1.4960×10−3
9.3200×10−4
KDE
1.4811×10−3
9.1705×10−4
SR
1.4638×10−3
9.1224×10−4
Table 2: Muon decay local probability mass validation on two 0.02×0.02 square regions.
Figure 5: Prediction and regions for local probability mass validation indicated as A & B (left). Residuals with respect to NSF (right) of the lowest loss expression. The maximum predicted density is 4555.25 and the maximum absolute residual is 1789.14.
Method
Region A
Region B
Empirical
1.0270×10−1
2.3200×10−3
KDE
6.7234×10−2
1.5040×10−2
NSF
1.0974×10−1
1.9397×10−3
SR
1.1493×10−1
2.2609×10−3
Table 3: Dijet local probability mass validation on two 0.02×0.02 square regions.
Figure 6: Mean log-likelihood scores against expression complexity for SR using only MSE loss and SR using MSE, mean negative log-likelihood (NLL), and penalties for negative density predictions (NP) in the loss function. Triangles indicate expressions that made negative predictions on test samples.
Appendix figures & tables39 assets
Supplementary material from the paper’s appendix.
Appendix
Figure 7: Example PGM for distribution p(x1,...,x7)=p(x1)p(x2)p(x3)p(x4∣x1,x2,x3)⋅p(x5∣x1,x3)p(x6∣x4)p(x7∣x4,x5) ( Bishop and Bishop, 2023 ) .
Discrete probability laws underpin statistical modeling, yet the catalog of interpretable distributions has expanded only gradually through centuries of case-by-case mathematical derivations. We introduce symbolic density estimation (SDE), an unsupervised framework that automatically recovers closed-form probability mass functions by composing elementary analytic operations within a structured search space. Our method integrates domain-specific structural priors with evolutionary search and a validity-aware inference stage, and it extends to richer distribution families such as zero inflation and finite mixtures. To support systematic evaluation and future research, we contribute a benchmark dataset spanning a broad collection of commonly used discrete distributions. The proposed algorithm recovers all benchmark families with accurate parameter estimates. A real data application shows that it identifies concise and interpretable mixture models that improve goodness-of-fit over standard models.
Estimating the symbolic or analytical form of probability density functions (PDFs) from observed samples is a fundamental challenge in statistical and computational modelling. This process is critical for deriving interpretable and generalizable relationships characterizing the underlying phenomenon. Traditionally, this estimation depends strongly on domain expertise and prior field-specific knowledge, with experts selecting appropriate functional forms or parametric families based on empirical evidence and theoretical understanding. The coefficients of these forms are then typically determined through parameter estimation. In this paper, we develop a framework for estimating symbolic expressions of unnormalized distributions from observed samples using domain-specific prior knowledge, such as the range of interactions and a predefined set of primitive functions. We integrate deep generative models with symbolic regression (SR), incorporating inductive biases, such as factorizing large distributions, to keep the problem tractable. The deep generative models we examine include likelihood-based models, viz., flow models, and score-based models. Experiments show the effectiveness of the proposed framework for estimating density functions for multivariate toy distributions as well as lattices from computational physics, namely, XY model and φ4 theory. When applied to the renormalization problem in φ4 theory, the proposed framework estimates compact symbolic approximations of the hamiltonian function at different scales directly from samples, yielding expressions that may be challenging to derive using traditional perturbative or analytic approaches in nonperturbative settings.
Symbolic regression (SR) has gained recent traction in AI-driven scientific discovery for learning closed-form physical laws. Yet existing methods are dominated by heuristic search or data-intensive approaches that often assume low-noise regimes and lack principled uncertainty quantification, while fully probabilistic SR formulations remain scarce. We introduce a scalable probabilistic framework for SR, VaSST, based on variational inference. VaSST uses soft symbolic trees, a continuous relaxation of symbolic expression trees in which discrete operator and feature assignments are replaced by probability distributions over allowable components. This transforms combinatorial symbolic search through an astronomically large expression space into efficient gradient-based optimization while preserving a coherent probabilistic interpretation. The learned soft representations induce posterior distributions over symbolic structures, enabling uncertainty quantification across plausible symbolic forms through posterior-aware symbolic model selection. On simulated experiments and the Feynman Symbolic Regression Database, VaSST achieves strong structural recovery and predictive accuracy compared to state-of-the-art competing SR methods.
Somjit Roy, Pritam Dey, Bani K. Mallick
Department of Statistics, Texas A&M University, College Station, Texas, USA