Symbolic Density Estimation: A Decompositional Approach

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.

• The paper introduces AI-Kolmogorov, a novel pipeline that leverages decomposition and symbolic regression for interpretable density estimation.
• It decomposes complex distributions through additive clustering and multiplicative factorization, significantly reducing mean squared error and enhancing clarity.
• Empirical results on synthetic and high-energy physics data validate the approach, achieving low residual errors and adherence to normalization constraints.

Symbolic Density Estimation via Decomposition: The AI-Kolmogorov Framework

Introduction

"Symbolic Density Estimation: A Decompositional Approach" (2603.27955) addresses the unsolved problem of interpretable density estimation through general symbolic functional forms. While symbolic regression (SR) is established for supervised learning, there is a fundamental gap in extending such explicit model discovery processes to unsupervised density estimation. The work introduces AI-Kolmogorov, a pipeline that constructs symbolic approximations of continuous probability density functions by leveraging decomposition, nonparametric surrogates, support estimation, and multi-population evolutionary SR. The motivation is rooted in the need for interpretable probabilistic models beyond the black-box deep generative paradigm and the inflexibility of parametric statistics, especially in high-dimensional, structured or "exotic" scientific distributions.

Framework Overview

AI-Kolmogorov is a modular, multi-stage pipeline combining algorithmic decomposition and symbolic search:

• Decomposition: Optionally splits the target distribution by clustering (additive partitioning) or probabilistic graphical model structure learning (multiplicative factorization).
• Nonparametric Density Estimation: A KDE or normalizing flow (Neural Spline Flow, NSF) forms the empirical surrogate, producing continuous targets for SR.
• Support Estimation: Rejects low-density/sparse regions via level-set trimming and convex hulls; boundary corrections use the reflection trick.
• Symbolic Regression: Employs PySR's multi-population genetic programming to fit expressions against density surrogates, balancing parsimony and fit, subject to soft constraints for non-negativity and normalization.
• Warm-Start Refinement: Improves candidate solutions via targeted numerical optimization.

  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.

This design enables both expressivity and interpretability, operating as a general but constrained equation discovery system for densities.

Technical Details of Decomposition

Additive and Multiplicative Decomposition

Additive decomposition (via clustering) separates multi-modal densities into unimodal components, aligning each with a tractable SR sub-problem. Multiplicative decomposition (via structure learning/PGM inference) exploits conditional independence by partitioning variables, reducing the symbolic search dimensionality. Empirically, both strategies enhance accuracy and interpretability relative to monolithic approaches.

Figure 2: Additive decompositon: Results on the Gaussian mixture dataset. Prediction (left) and residuals (center left) by AI-Kolmogorov. Prediction (center right) and residuals (right) by clustering-augmented AI-Kolmogorov.

For a four-dimensional product distribution, PGM-based partitioning provides orders-of-magnitude improvement in MSE versus direct, non-decomposed symbolic regression, yielding expressions that reflect the canonical exponential-of-quadratic Gaussian form in each independent subspace.

Symbolic Regression for Density Discovery

Given the nonparametric density surrogate on restricted support, candidate symbolic expressions are trained to minimize MSE subject to a parsimony penalty. Operator sets are informed by a superset of the ground truth expression structure, except in blind scenarios. The final output is a Pareto front over fit/complexity, with simple polynomials capturing bulk structure and complex expressions reproducing fine multi-modal or oscillatory features.

Rastrigin Probability Density Example

Deployment on the Rastrigin benchmark illustrates the pipeline's capacity to identify both the coarse quadratic bowl structure and superposed oscillatory modes. Maximal absolute density residuals on the test regime are on the order of $10^{-3}$, and symbolic outputs closely recapitulate the expression's ground truth sinusoidal nonlinearity.

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.

High-Energy Physics Applications

Non-standard densities, such as those arising in muon decay or dijet production, exhibit sharp supports and multi-scale features. AI-Kolmogorov, combined with neural flow-based surrogates and convex hull-supported SR, produces candidate distributions with empirical local probability mass estimates aligning with observed frequencies to within $10^{-4}$–$10^{-3}$. While expressions often approximate but do not exactly recover ground truth rational forms, their structure provides valuable scientific interpretability. On challenging LHC-simulated datasets, the best models qualitatively reproduce density concentration and support geometry, though residuals reflect the challenge of discontinuous or highly-peaked densities.

Figure 4: Prediction and regions for local probability mass validation indicated as A and 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.

Empirical Results and Loss Function Design

Extensive ablations show that both additive and multiplicative decomposition significantly reduce MSE while improving model parsimony. Importantly, loss functions using solely raw-sample log-likelihoods or naive MSEs can result in invalid models (negative density, improper normalization); inclusion of normalized surrogate targets and explicit negative prediction penalties is essential for validity. On synthetic and physics datasets, the best AI-Kolmogorov models robustly match empirical sample masses in validation regions, demonstrating reliable fit.*

Figure 5: 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.

Multiple symbolic forms are surfaced, enabling domain experts to select or aggregate over interpretable subsets on the Pareto front.

Practical and Theoretical Implications

AI-Kolmogorov establishes a practical mechanism for interpretable, analytic density discovery in continuous, multi-modal, and multivariate regimes traditionally considered intractable for SR. The results indicate clear gains for scientific modeling, especially in physics-driven disciplines requiring closed-form probabilistic posteriors or explicit uncertainty quantification. The approach is modular and generalizable, supporting surrogate-based learning, domain-aware decomposition, and constraint-soft SR.

Open theoretical directions include SR with direct normalization and non-negativity enforcement, agentic active search over the density support, and probabilistic program induction under physical or causal constraints. The decoupling of density surrogate generation from the symbolic fit establishes a flexible platform for the future integration of LLM-based equation priors, hierarchical decomposition heuristics, and fit diagnostics tailored to heavy-tailed or sparse data.

Conclusion

This work operationalizes symbolic regression for unsupervised density estimation by systematically leveraging decomposition, nonparametric modeling, and support-aware analytic equation search. The result is a class of frameworks that robustly bridges the gap between black-box density estimation and rigid parametric statistics, returning candidate probability laws of direct scientific value. AI-Kolmogorov's empirical performance, across both synthetic multimodal and structurally complex physical distributions, demonstrates that interpretable form discovery for high-dimensional probabilistic laws is tractable given modular decomposition and well-calibrated symbolic search. Extensions to hierarchical, causal, and agentic frameworks for SymDE and integration with domain-specific inductive biases are compelling avenues for further work.