VaSST: Variational Inference for Symbolic Regression using Soft Symbolic Trees
| Authors | Somjit Roy et al. |
| Year | 2026 |
| Field | AI / ML |
| arXiv | 2602.23561 |
| Download | |
| Categories | stat.ME, cs.LG, cs.SC, stat.CO |
Abstract
Symbolic regression has recently gained traction in AI-driven scientific discovery, aiming to recover explicit closed-form expressions from data that reveal underlying physical laws. Despite recent advances, existing methods remain dominated by heuristic search algorithms or data-intensive approaches that assume low-noise regimes and lack principled uncertainty quantification. Fully probabilistic formulations are scarce, and existing Markov chain Monte Carlo-based Bayesian methods often struggle to efficiently explore the highly multimodal combinatorial space of symbolic expressions. We introduce VaSST, a scalable probabilistic framework for symbolic regression based on variational inference. VaSST employs a continuous relaxation of symbolic expression trees, termed soft symbolic trees, where discrete operator and feature assignments are replaced by soft distributions over allowable components. This relaxation transforms the combinatorial search over an astronomically large symbolic space into an efficient gradient-based optimization problem while preserving a coherent probabilistic interpretation. The learned soft representations induce posterior distributions over symbolic structures, enabling principled uncertainty quantification. Across simulated experiments and Feynman Symbolic Regression Database within SRBench, VaSST achieves superior performance in both structural recovery and predictive accuracy compared to state-of-the-art symbolic regression methods.
Engineering Breakdown
Plain English
VaSST tackles symbolic regression—the problem of recovering explicit mathematical equations from experimental data—by proposing a fully probabilistic approach based on variational inference. Existing methods rely on heuristic search or require clean, low-noise data and fail to quantify uncertainty in discovered equations. The authors introduce soft symbolic trees, a continuous relaxation of discrete expression trees that enables efficient gradient-based optimization, allowing the method to scale to complex symbolic spaces while providing Bayesian uncertainty estimates. This addresses a critical gap in scientific discovery where practitioners need both accurate equation recovery and principled confidence intervals on model parameters.
Core Technical Contribution
The core innovation is the continuous relaxation of symbolic expression trees into differentiable soft trees, enabling variational inference where prior discrete symbolic regression methods relied on combinatorial search or MCMC sampling. Instead of explicitly enumerating or randomly sampling expressions, VaSST represents each node in the expression tree as a continuous distribution over possible operations (addition, multiplication, sin, exp, etc.) and continuously relaxes variable bindings, making the entire space amenable to efficient gradient-based optimization. This allows the framework to jointly learn both the symbolic structure and the continuous parameters (coefficients) in a principled Bayesian way, leveraging evidence lower bounds (ELBO) as the objective. The approach fundamentally differs from prior work by providing a scalable, fully differentiable probabilistic model rather than heuristic or sampling-based alternatives.
How It Works
VaSST encodes a symbolic expression as a tree where each node is represented as a continuous probability distribution over available operations (arithmetic, trigonometric, polynomial, etc.). During variational inference, the framework maintains these soft distributions and backpropagates gradients through the entire relaxed tree structure to maximize the ELBO, which balances data fit (likelihood) against model complexity (KL divergence to priors on structure and parameters). The continuous relaxation allows operations like matrix multiplications to aggregate contributions from all possible node types weighted by their posterior probabilities, rather than committing to discrete choices early. In the forward pass, the model computes approximate expression outputs by marginalizing over the continuous distributions; in the backward pass, gradients flow to both the distribution parameters (controlling which operations are likely) and the continuous coefficients. At inference time, the learned soft tree is discretized by selecting the maximum-probability operation at each node, yielding a final closed-form equation with uncertainty estimates derived from the posterior variance.
Production Impact
For teams building scientific discovery pipelines or automated model discovery systems, VaSST offers a critical advantage: Bayesian uncertainty quantification on discovered equations, not just point estimates. Instead of running 100 heuristic restarts to validate equation robustness, practitioners can use posterior variance to assess confidence in discovered terms and coefficients, reducing wasted validation time. The gradient-based optimization is orders of magnitude faster than MCMC methods on high-dimensional expression spaces, enabling real-time or near-real-time equation discovery on noisy experimental datasets. Integration requires: (1) reformulating data into a tensor-friendly format with specified noise levels, (2) choosing appropriate priors over expression complexity and operation types, and (3) a differentiable constraint satisfaction layer if domain-specific equations must be enforced. Trade-offs include modest computational overhead during training (variational inference requires ELBO computation at each step) and sensitivity to prior specification—poor priors can bias discovered equations—but the method remains tractable even on datasets with 1000+ samples and expression trees of moderate depth.
Limitations and When Not to Use This
VaSST assumes that the underlying data is generated from a symbolic expression and that stochasticity is additive Gaussian noise—it will struggle with data containing structural outliers, missing causal variables, or multi-modal relationships where multiple equations explain different regimes. The continuous relaxation introduces approximation error that grows with tree depth and operation diversity; very deep trees or spaces with hundreds of possible operations may suffer from soft-tree approximation bias. The method requires careful prior specification over both expression complexity and operation types; if practitioners use overly restrictive or misaligned priors, the variational posterior may converge to suboptimal equations (similar to how priors can bias Bayesian inference generally). Finally, the discretization step from soft trees to final equations is a greedy arg-max operation, which may not recover the globally optimal equation even if the variational objective was well-optimized—post-hoc refinement or ensemble methods are sometimes needed.
Research Context
VaSST builds on a long tradition of Bayesian symbolic regression (grounded in work by Udrescu, Tegmark, and others on AI for science) and extends recent advances in continuous relaxations of discrete structures (similar to techniques in neural architecture search and differentiable programming). It directly addresses limitations of prior probabilistic approaches like Bayesian symbolic regression via MCMC, which scale poorly; the paper implicitly benchmarks against heuristic methods like genetic programming and data-driven neural-symbolic hybrids. The variational inference framework opens new research directions: combining VaSST with active learning for adaptive data acquisition, extending the approach to multi-objective discovery (balancing accuracy against equation sparsity), and scaling soft trees to higher-order expression spaces. This work reinforces the trend toward differentiable, gradient-based alternatives to combinatorial search in AI for science.
:::tip Subscribe Get weekly breakdowns of papers like this in AI Letters - the newsletter for engineers building production AI systems. :::