Technical Summary: Symbolic Density Estimation for Discrete Distributions

Problem Formulation

Discrete probability distributions are fundamental for modeling count data and categorical outcomes. However, the catalog of interpretable, analytically tractable distributions has historically expanded only through case-by-case mathematical derivations. Analysts typically face a trade-off: assuming a specific parametric family risks model misspecification, while flexible nonparametric or neural approaches (e.g., autoregressive networks, normalizing flows) often yield "black-box" implicit models lacking closed-form probability mass functions (PMFs).

The authors address the problem of Symbolic Density Estimation (SDE): automatically recovering a symbolic, closed-form PMF that best fits observed discrete data without prespecifying a distributional family. This task is non-trivial due to the combinatorially large expression space and the strict requirement that any proposed formula must define a valid PMF (non-negative and normalized). Unlike standard regression, discrete distributions involve specific combinatorial constructs like factorials and binomial coefficients.

Methodology: The SDE Framework

The proposed SDE framework operates in the log-probability domain to improve numerical stability and transform multiplicative combinatorial factors into additive components. The method integrates domain-specific structural priors with evolutionary search and a validity-aware inference stage.

1. Operator Set and Search Space

The hypothesis space is induced by a custom operator set $\mathcal{O}$ comprising:

• Basic arithmetic: $+$, $-$, $\times$, $\wedge$.
• Elementary functions: $\log$, $\exp$, $|\cdot|$, $\sin$, $\cos$.
• Log-domain combinatorial primitives: $\log\Gamma(t+1)$ (log-factorial), $\log\binom{n}{k}$, and $\log B(a,b)$ (log-Beta function). These primitives represent factorials, binomial coefficients, and Beta-function terms common in discrete PMFs while maintaining numerical stability for large arguments.

To manage the vast search space (estimated at $>10^{19}$ possibilities for modest complexity), the authors employ two types of constraints:

• Operator Complexity Priors: Two complementary complexity profiles assign soft symbolic costs to operators. One profile favors combinatorial structures (e.g., $\log C$, $\log B$), while the other favors exponential-family motifs (e.g., $\log F$, $\log$). This guides the evolutionary search toward structurally diverse distributions without restricting expressivity.
• Structural Grammar Constraints:
  • Linearity: Arguments to special functions (e.g., $\log\Gamma$, $\log B$) are restricted to affine forms to prevent deeply nested, non-identifiable compositions.
  • Atomicity: Exponentiation is restricted to atomic exponents (single variables or constants) to avoid complex compound structures.

2. Estimation and Weighting

The system estimates the log-PMF from $M$ i.i.d. samples. To handle zero counts and high variance in low-probability regions, a smoothed estimator $\hat{p}(x) = \frac{c(x) + \alpha}{M + \alpha K}$ is used (corresponding to a Dirichlet posterior mean).

• Support Truncation: The fitting domain is restricted to bins with sufficient effective counts ($M\hat{p}(x) \geq \tau$) to mitigate finite-sample noise in the tails.
• Weighted Least Squares: The search targets a weighted reconstruction error on the log-PMF. Weights are derived from the inverse variance of the log-frequency estimates (via the Delta method), down-weighting noisy observations.

3. Search and Inference Procedure

The core search utilizes a multi-start evolutionary strategy:

• Multiple evolutionary runs are executed, each using a different complexity profile to bias the search toward different canonical motifs.
• Candidates are evolved via mutation and crossover, followed by cost-aware selection balancing reconstruction error and symbolic complexity.
• Post-Search Inference: Pooled candidates undergo a three-step validation:
  1. Loss Screening: Discarding expressions with high reconstruction error.
  2. Probabilistic Validity: Ensuring approximate normalization ($|\sum e^{f(x)} - 1| < \epsilon$) and bounded log-mass.
  3. Complexity Control: Rejecting expressions exceeding operator budgets.
• Valid candidates are ranked lexicographically by symbolic complexity and residual error. Parameters are refined via discrete local search starting from identifiable symbolic terms.

4. Extensions

The framework naturally extends to complex structures by augmenting the operator set:

• Finite Mixtures: Introduced via the binary primitive $\text{logaddexp}(u, v) = \log(e^u + e^v)$, allowing nested applications to represent arbitrary mixtures in log-space.
• Zero-Inflated Models: Introduced via a primitive $\log \delta_0(x)$ encoding a point mass at zero, combined with $\text{logaddexp}$ to model the mixture of a zero mass and a baseline distribution.

Key Contributions

1. First Symbolic Framework for Discrete Distributions: SDE is the first framework to integrate PMF-validity verification (non-negativity and normalization) directly into the symbolic discovery process, combining weighted log-PMF fitting with built-in validity checks.
2. SDEBench: The authors contribute a benchmark dataset spanning 14 base discrete distributions (including exponential families, power laws, and mixed models) and their zero-inflated/mixture variants to support systematic evaluation.
3. Empirical Validation: The method is shown to recover classical, generalized, and structured models from finite samples, bridging the gap between rigid parametric selection and black-box methods.

Experimental Results

The authors evaluate SDE on the SDEBench dataset ($M=50,000$) and a real-world single-cell RNA-sequencing (scRNA-seq) dataset.

• Recovery Accuracy: SDE successfully recovers the correct symbolic form (up to algebraic equivalence) for all 14 base distributions in the benchmark, including complex combinatorial structures like the Beta-Binomial and Negative-Hypergeometric. It also correctly identifies nested mixture structures and zero-inflated forms.
• Parameter Precision: Inferred parameters are typically within a few percent of ground truth. For example, in a noisy setting, the method recovered $r=9.999, p=0.700$ for a Negative Binomial with true $r=10, p=0.70$.
• Robustness: The method demonstrates resilience to sampling noise. While complex models require more data (e.g., Beta-Binomial recovery degrades at $M=500$), simpler targets remain recoverable even at small sample sizes.
• Comparative Performance:
  • Against Oracle Parametric Baselines (MoM/MLE): SDE achieves comparable precision without requiring the family to be known a priori.
  • Against Nonparametric/Neural Baselines (KDE, Pyro): SDE provides interpretable closed-form expressions, whereas KDE and Pyro often produce non-smooth, oscillatory fits or boundary artifacts (e.g., near zero for Zero-Inflated Poisson).
  • Against Generic Symbolic Regression (PySR): PySR frequently yields structurally invalid expressions that fail to satisfy PMF constraints, whereas SDE's validity-aware search ensures physical plausibility.
• Real-World Application: On PBMC scRNA-seq data (Gene 4046), SDE discovered a concise, interpretable mixture model combining two geometric-decay components and a constant baseline. This model achieved a lower Mean Squared Error (MSE) on the empirical log-PMF than standard fixed-family models (selected via AIC) and black-box neural baselines.

Significance and Claims

The paper claims that SDE offers a principled alternative to manual parametric modeling and opaque black-box density estimators. By operating in the log-probability domain with a tailored set of primitive operators, the framework can automatically discover interpretable closed-form PMFs that reveal the compositional structure of underlying distributions.

The authors emphasize that the framework's core strength lies in its extensibility: complex structures like finite mixtures and zero-inflated models are recovered via minimal operator extensions without modifying the underlying search procedure, loss function, or constraints. This capability supports the identification of structured, heterogeneous discrete models where interpretability and structural insight are essential.

The work positions SDE as a practical approach for distribution modeling, particularly in scenarios where the true family is unknown or where the data exhibits complex structural variations (e.g., excess zeros or multimodality) that standard parametric families fail to capture. The authors note that while the current formulation focuses on univariate discrete distributions, the approach provides a foundation for future expansion to multivariate and continuous settings.