Symbolic Density Estimation: A Decompositional Approach

Angelo Rajendram; Xieting Chu; Vijay Ganesh; Max Fieg; Aishik Ghosh

cs.LGMar 30, 2026

Symbolic Density Estimation: A Decompositional Approach

Authors: Angelo Rajendram, Xieting Chu, Vijay Ganesh, Max Fieg, Aishik Ghosh

Organizations: Department of Mathematics, University of Waterloo, Waterloo, Canada · Department of Computer Science, Georgia Tech, Atlanta, USA · University of California, Irvine, Irvine, USA

Abstract

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.

Panel of Figure 2 — 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\times 10^{-4}$	$\bm{3.57\times 10^{-4}}$
4D Gaussian	$1.01\times 10^{-3}$	$\bm{4.24\times 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\times 10^{-3}$	$9.3200\times 10^{-4}$
KDE	$1.4811\times 10^{-3}$	$9.1705\times 10^{-4}$
SR	$1.4638\times 10^{-3}$	$9.1224\times 10^{-4}$

Table 2: Muon decay local probability mass validation on two

0.02\times 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\times 10^{-1}$	$2.3200\times 10^{-3}$
KDE	$6.7234\times 10^{-2}$	$1.5040\times 10^{-2}$
NSF	$1.0974\times 10^{-1}$	$1.9397\times 10^{-3}$
SR	$1.1493\times 10^{-1}$	$2.2609\times 10^{-3}$

Table 3: Dijet local probability mass validation on two

0.02\times 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(x_{1},...,x_{7})=p(x_{1})p(x_{2})p(x_{3})p(x_{4}|x_{1},x_{2},x_{3})$ $\cdot\;p(x_{5}|x_{1},x_{3})p(x_{6}|x_{4})p(x_{7}|x_{4},x_{5})$ ( Bishop and Bishop, 2023 ) .

Figure 8: Convex hull support estimates. Left: true support. Right: convex hull estimate ( Cui, 2021 ) .

Figure 9: Reflection trick in one dimension.

Figure 10: DBSCAN on a 2D Gaussian mixture. Core points have $\epsilon$ -neighborhoods with at least minPts samples ( Jang and Jiang, 2019 ) .

Figure 11: Kernel Density Estimation in 1D ( Chen, 2017 )

Hyper-parameter	Value	Notes
Kernel Density Estimation
number of samples	450000
bandwidth	0.101	KDE bandwidth selected using cross-validation.
PySR
binary operators	[+, -, *, /]
unary operators	[exp, log, pow2, pow3]	Allowed unary operators.

Appendix

Table 4: Hyper-parameters for symbolic regression on the Gaussian mixture dataset

Figure 12: Pareto front of recovered expressions for the Gaussian mixture dataset.

Raw Complexity	Simplified Expression
1	$0.13$
2	$0.14$
3	$0.13$
5	$0.14$
6	$0.24-0.0023x_{2}^{2}$
7	$e^{-0.23\left(x_{1}+x_{2}\right)^{2}}$

Appendix

Table 5: Expressions simplified with sympy alongside their raw complexity score from PySR for the Gaussian mixture dataset.

Hyper-parameter	Value	Notes
DBSCAN Clustering
eps	5	Maximum Neighborhood radius.
min samples	10	Neighborhood minimum size.
Kernel Density Estimation
number of samples	225000
bandwidth	0.0923	KDE bandwidth selected using cross-validation.

Appendix

Table 6: Hyper-parameters for symbolic regression on each cluster

Figure 13: Pareto front of recovered expressions for the Gaussian cluster 1 dataset.

Raw Complexity	Simplified Expression
1	$0.15$
2	$0.15$
6	$\dfrac{x_{1}}{x_{1}^{2}+6.6}$
7	$1.1e^{-16\left(-0.25x_{2}-1\right)^{2}}$
8	$\dfrac{1}{16\left(0.25x_{2}+1\right)^{2}+1}$
9	$1.5e^{-64\left(0.12x_{1}-0.12x_{2}-1\right)^{2}}$

Appendix

Table 7: Expressions simplified with sympy alongside their raw complexity score for the Gaussian cluster 1 dataset.

Figure 14: Pareto front of recovered expressions for the Gaussian cluster 2 dataset.

Raw Complexity	Simplified Expression
1	$0.15$
4	$0.15$
6	$\dfrac{0.13}{e^{x_{1}}+0.61}$
7	$e^{-8.5\left(0.25x_{1}+1\right)^{2}}$
9	$e^{-6.9\left(0.062x_{1}x_{2}+1\right)^{2}}$
10	$e^{-660\left(0.062x_{1}x_{2}+1\right)^{6}}$

Appendix

Table 8: Expressions simplified with sympy for the Gaussian Cluster 2 dataset.

Hyper-parameter	Value	Notes
Kernel Density Estimation
number of samples	225000
bandwidth	0.0857	KDE bandwidth selected using cross-validation.
PySR
binary operators	[+, -, *, /]
unary operators	[exp, log, pow2, pow3]	Allowed unary operators.

Appendix

Table 9: Hyper-parameters for symbolic regression on the 4D Gaussian dataset

Figure 15: Pareto front of recovered expressions for the 4 dimensional Gaussian dataset.

Raw Complexity	Simplified Expression
1	$4.0\times 10^{-4}$
2	$3.9\times 10^{-4}$
3	$4.0\times 10^{-4}$
4	$3.9\times 10^{-4}$
5	$4.0\times 10^{-4}$
6	$\dfrac{4.1\times 10^{-4}}{e^{x_{3}}+0.49}$

Appendix

Table 10: Expressions simplified with sympy alongside their raw complexity score from PySR for the Gaussian 4D dataset.

Hyper-parameter	Value	Notes
pgmpy Structure Learning
CI test	Pearson correlation	Conditional independence test.
Kernel Density Estimation
number of samples	225000
bandwidth	0.0923	KDE bandwidth selected using cross-validation.
PySR

Appendix

Table 11: Hyper-parameters for symbolic regression on each independent component

Figure 16: Pareto front of recovered expressions for the 4 dimensional Gaussian marginal ( $x_{1}$ , $x_{2}$ ) dataset.

Raw Complexity	Simplified Expression
1	$0.15$
3	$0.15$
5	$0.15$
6	$\dfrac{x_{1}}{x_{1}^{2}+6.5}$
7	$1.1\,e^{-16\left(0.25x_{2}+1\right)^{2}}$
8	$\dfrac{1.1}{16\left(0.25x_{2}+1\right)^{2}+1.1}$

Appendix

Table 12: Expressions simplified with sympy alongside their raw complexity score from PySR for the Gaussian 4D marginal (

x_{1}

x_{2}

Figure 17: Pareto front of recovered expressions for the 4 dimensional Gaussian marginal ( $x_{3}$ , $x_{4}$ ) dataset.

Raw Complexity	Simplified Expression
1	$0.15$
2	$0.15$
4	$0.14$
6	$\dfrac{0.1}{e^{x_{3}}+0.49}$
7	$e^{-8.3\left(1-0.25x_{4}\right)^{2}}$
8	$e^{-8.6\left(0.25x_{4}-1\right)^{2}}$

Appendix

Table 13: Expressions simplified with sympy alongside their raw complexity score from PySR for the Gaussian 4D Marginal (

x_{3}

x_{4}

) dataset.

Hyper-parameter	Value	Notes
Kernel Density Estimation
number of samples	900000
bandwidth	0.0254	KDE bandwidth selected using cross-validation.
PySR
binary operators	[+, -, *, /]
unary operators	[exp, log, pow2, cos, sin]	Allowed unary operators.

Appendix

Table 14: Hyper-parameters for symbolic regression on the Rastrigin dataset

Figure 18: Pareto front of recovered expressions for the Rastrigin dataset.

Raw Complexity	Simplified Expression
1	$3.6$
4	$x_{1}^{2}+2.6$
5	$x_{1}^{2}+2.6$
6	$x_{1}^{2}x_{2}^{2}+2.5$
7	$x_{1}^{2}+x_{2}^{2}+1.5$
9	$1.2x_{1}^{2}+1.2x_{2}^{2}+1$

Appendix

Table 15: Expressions simplified with sympy alongside their raw complexity score from PySR for the Rastrigin dataset.

Hyper-parameter	Value	Notes
Kernel Density Estimation
number of samples	450000
bandwidth	0.0114	KDE bandwidth selected using cross-validation.
PySR
binary operators	[+, -, *, /]
unary operators	[exp, log, pow2, pow3]	Allowed unary operators.

Appendix

Table 16: Hyper-parameters for symbolic regression on the Muon Decay dataset

Figure 19: Pareto front of recovered expressions for the muon decay dataset.

Raw Complexity	Simplified Expression
1	$4$
3	$4.5-x_{2}$
4	$9.3\left(1-0.48x_{2}\right)^{3}$
5	$6.9\left(1-0.53x_{2}^{2}\right)^{3}$
6	$8.4\left(0.34x_{1}+0.34x_{2}-1\right)^{2}$
7	$8.6-\left(x_{1}+e^{x_{2}}\right)^{2}$

Appendix

Table 17: Expressions simplified with sympy alongside their raw complexity score from PySR for the Muon decay dataset.

Hyper-parameter	Value	Notes
Neural Spline Flow
K	16	Number of autoregressive rational quadratic spline flow blocks.
hidden units	128	Width of the neural networks inside each autoregressive spline flow.
hidden layers	2	Depth of the neural networks inside each autoregressive spline flow
number of samples	90000
epochs	1390	Number of training epochs

Appendix

Table 18: Hyper-parameters for symbolic regression on the dijet dataset (using loss SR MSE + NLL + NP)

Figure 20: Pareto front of recovered expressions for the dijet dataset (SR MSE + NLL + NP).

Raw Complexity	Simplified Expression
1	$0.88$
2	$e^{x_{2}}$
3	$\dfrac{2.5}{x_{1}}$
5	$-4.6\times 10^{-4}+\dfrac{2.5}{x_{1}}$
6	$\dfrac{0.0027}{\left(x_{1}+0.0076\right)^{3}}$
7	$e^{\dfrac{0.029}{\left(x_{1}+0.058\right)^{2}}}$

Appendix

Table 19: Expressions simplified with sympy alongside their raw complexity score from PySR for the dijet dataset (SR MSE + NLL + NP).

Hyper-parameter	Value	Notes
Neural Spline Flow
K	16	Number of autoregressive rational quadratic spline flow blocks.
hidden units	128	Width of the neural networks inside each autoregressive spline flow.
hidden layers	2	Depth of the neural networks inside each autoregressive spline flow
number of samples	90000
epochs	1390	Number of training epochs

Appendix

Table 20: Hyper-parameters for symbolic regression on the dijet dataset (SR MSE + NP)

Figure 21: Pareto front of recovered expressions for the dijet dataset (SR MSE + NP).

Raw Complexity	Simplified Expression
1	$0.82$
2	$0.82$
3	$\dfrac{2.5}{x_{1}}$
4	$\dfrac{2.6}{x_{1}}$
5	$\dfrac{2.5}{x_{1}}$
6	$\dfrac{0.0027}{\left(x_{1}+0.008\right)^{3}}$

Appendix

Table 21: Expressions simplified with sympy alongside their raw complexity score for the dijet dataset recovered using SR MSE + NP.

Hyper-parameter	Value	Notes
Neural Spline Flow
K	16	Number of autoregressive rational quadratic spline flow blocks.
hidden units	128	Width of the neural networks inside each autoregressive spline flow.
hidden layers	2	Depth of the neural networks inside each autoregressive spline flow
number of samples	90000
epochs	1390	Number of training epochs

Appendix

Table 22: Hyper-parameters for symbolic regression on the dijet dataset (using loss SR MSE)

Figure 22: Pareto front of recovered expressions for the dijet dataset (SR MSE).

Raw Complexity	Simplified Expression
1	$0.82$
2	$0.82$
3	$\frac{2.5}{x_{1}}$
4	$\frac{2.7}{x_{1}}$
5	$-7.6+\frac{2.9}{x_{1}}$
6	$\frac{0.0027}{\left(x_{1}+0.0078\right)^{3}}$

Appendix

Table 23: Expressions simplified with sympy alongside their raw complexity score from PySR for the dijet dataset (SR MSE).

Hyper-parameter	Value	Notes
Neural Spline Flow
K	16	Number of autoregressive rational quadratic spline flow blocks.
hidden units	128	Width of the neural networks inside each autoregressive spline flow.
hidden layers	2	Depth of the neural networks inside each autoregressive spline flow
number of samples	90000
epochs	1390	Number of training epochs

Appendix

Table 24: Hyper-parameters for symbolic regression on the dijet dataset (using loss SR NLL + NP)

Figure 23: Pareto front of recovered expressions for the dijet dataset (SR NLL + NP).

Raw Complexity	Simplified Expression
1	$x_{1}$
2	$e^{x_{2}}$
4	$e^{\left(\log\left(x_{1}\right)\right)^{3}}$
5	$x_{1}^{68}$
6	$1.3\cdot 10^{35}e^{-\frac{850}{x_{1}}}$
7	$1.3\cdot 10^{33}e^{-\frac{910}{x_{1}^{3}}}$

Appendix

Table 25: Expressions simplified with sympy alongside their raw complexity score from PySR for the dijet dataset (SR NLL + NP).

Density Estimation

Density estimation learns probability distributions from data to measure likelihood, detect anomalies, quantify uncertainty, and generate samples.

See all 66 papers

Symbolic Density Estimation for Discrete Distributions
May 20, 2026Ziwen Liu, Meng LiDensity Estimation
Rice University Houston, TX
Symbolic Density Estimators for Unnormalized Distributions
Oct 7, 2026Vikas Kanaujia, Riyansha Singh, Shashank Sharma +1Density Estimation Symbolic Regression
Department of Electrical Engineering IIT Kanpur · Department of Computer Science and Engineering IIT Kanpur · Department of Physics IIT Kanpur +1
VaSST: Variational Inference for Symbolic Regression using Soft Symbolic Trees
Feb 27, 2026Somjit Roy, Pritam Dey, Bani K. MallickUncertainty Quantification Symbolic Regression
Department of Statistics, Texas A&M University, College Station, Texas, USA

May 20, 2026J/K move · Enter open · S save

Symbolic Density Estimation for Discrete Distributions

Ziwen Liu, Meng Li

Rice University Houston, TX

Density Estimation

Open paper AskPDF arXiv

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.

Abstract

Density Estimation

Explore similar work

Symbolic Density Estimation for Discrete Distributions

Symbolic Density Estimators for Unnormalized Distributions

VaSST: Variational Inference for Symbolic Regression using Soft Symbolic Trees