Symbolic Regression via Latent Iterative Refinement

Abstract: Symbolic regression (SR) seeks closed-form mathematical expressions that fit observed data. Neural SR methods amortize the search by training an encoder to map observations directly to expressions in a single pass, but this amortized inference leaves a residual amortization gap between its one-shot prediction and the true posterior. We propose Latent Equation Embedding (LEE), a framework that closes this gap through iterative amortized inference in a functionally grounded latent space. LEE learns a shared latent space Z equipped with three components: an encoder f_theta that jointly embeds symbolic tokens and numerical observations into a single latent vector z; an expression decoder g_expr that reconstructs formulas from z; and an evaluation decoder g_eval that predicts function values from z, explicitly grounding the latent space in functional behavior. At inference, LEE performs iterative refinement by re-encoding decoded expressions jointly with observations, progressively improving the latent estimate. LEE uses the encoder itself as a learned inference optimizer: each re-encoding step implicitly computes the mismatch between the candidate and the data. Because g_eval is differentiable in z, we additionally interleave continuous gradient descent with discrete re-encoding, yielding a hybrid iterative and gradient refinement procedure. On SRBench across three noise levels, against 19 baselines spanning genetic programming, symbolic-neural hybrids, and pre-trained Transformers, LEE produces expressions 2--10x simpler than the strongest accuracy-oriented baselines, including Operon, GP-GOMEA, TPSR, RAG-SR, and GenSR, with complexity 8--11 versus 20--90. These results advance the low-complexity region of the accuracy-complexity Pareto frontier and show graceful degradation as noise increases.

• The paper presents LEE, an encoder-driven iterative framework that refines latent representations to yield concise symbolic expressions.
• It combines discrete iterative updates with gradient-based denoising to achieve robust functional grounding in the latent space.
• Empirical results show LEE produces expressions 2–10 times simpler than leading baselines while maintaining competitive R² accuracy.

Symbolic Regression via Latent Iterative Refinement: Technical Overview

Motivation and Background

Symbolic regression (SR) endeavors to reverse-engineer closed-form equations from observed data, a task relevant for scientific discovery and interpretable machine learning. Traditionally, SR relies on genetic programming (GP) approaches that operate directly in the discrete, combinatorial space of symbolic expressions, typically requiring significant computational resources. Recent neural SR methods have shifted towards amortized inference, mapping observations in a single forward pass to candidate formulas via learned encoders. However, such methods suffer from an inherent "amortization gap"—a discrepancy between the one-shot prediction and the true posterior, particularly pronounced for complex or out-of-distribution data (2605.27245). Attempts to mitigate this gap with contrastive embeddings (e.g., SNIP) or generative latent spaces refined via black-box optimization (e.g., GenSR with CMA-ES) introduce further limitations: lack of functional grounding or inefficiency due to population-based search and disregard for encoder structure.

Latent Equation Embedding (LEE): Model Architecture

LEE introduces a functionally-grounded iterative inference framework in a shared latent space $\mathcal{Z}$ with three primary components: an encoder $f_\theta$ that embeds symbolic tokens and observations; an expression decoder $g_\text{expr}$ for autoregressive formula reconstruction; and an evaluation decoder $g_\text{eval}$, which is differentiable and predicts function values from latents, thereby embedding functional meaning directly into $\mathcal{Z}$. The encoder fuses symbolic and numeric streams via log-compressed embeddings, a Transformer backbone, and mean-pooling, parameterizing a Gaussian latent posterior regularized towards a standard-normal prior.

Figure 1: LEE architecture mapping expressions and data to latent representations, enabling joint reconstruction and evaluation; inference proceeds via iterative refinement in $\mathcal{Z}$.

This explicit grounding imposes the soft equivalence: $\bm{z}_1 \approx \bm{z}_2$ whenever $e_1(\bm{x}) = e_2(\bm{x})$ for all $\bm{x}$, favoring functional similarity over syntactic similarity. Training combines losses for reconstruction, evaluation, latent alignment, iterative refinement via token corruption/denoising, and KL regularization, enforcing a geometry in which movement within latent space corresponds to meaningful functional changes.

Iterative Refinement and Search in Latent Space

Inference commences by encoding observations to yield an initial $\bm{z}_0 = f_\theta(\varnothing, \mathcal{D})$; candidate expressions are decoded and scored by $f_\theta$0, pooled for diversity. LEE's core update is:

$f_\theta$1

iteratively refining the latent vector through re-encoding decoded expressions with observations, thereby leveraging the encoder as a learned inference optimizer. Importantly, $f_\theta$2 enables gradient-based refinement on $f_\theta$3, subject to a proximal penalty to anchor solutions within the decodable manifold. Hybrid iterative and gradient refinement is shown to be essential, synergistically combining global exploration and robust local denoising, particularly under noise.

Figure 2: $f_\theta$4 convergence over iterations, demonstrating rapid early gains and gradual improvement within the candidate pool.

Figure 3: Latent space interpolation between $f_\theta$5 and $f_\theta$6 yields syntactically distinct but functionally smooth intermediates, confirming functional grounding.

Empirical Results and Pareto Frontier

LEE is evaluated on SRBench (Strogatz, Feynman, black-box datasets) across multiple noise levels, following canonical 75/25 train/test splits and reporting metrics over repeated trials. In comparison with 19 baselines spanning GP, symbolic-neural hybrids, and neural SRs, LEE consistently generates expressions 2–10 times simpler than the strongest accuracy-oriented baselines (Operon, GP-GOMEA, TPSR, RAG-SR, GenSR), with complexity in the range of 8–11 versus 20–90, while maintaining competitive accuracy ($f_\theta$7 within 0.06–0.14 points of GP SOTA).

Figure 4: Pareto frontier on the black-box benchmark; LEE anchors the low-complexity regime, outperforming neural baselines in both accuracy and simplicity.

The robustness of LEE under noise is notable, with accuracy degrading gracefully as noise increases and outperforming one-shot neural baselines in out-of-distribution scenarios. Ablations demonstrate the necessity of combining discrete iterative refinement and continuous gradient descent: iterative-only search suffers performance loss under noise, and gradient-only search drifts off-manifold, whereas their combination yields the best accuracy and lowest variance.

Technical Analysis and Mechanism

The iterative amortized inference approach implemented in LEE aligns with the error-encoding variant described by Marino et al., wherein residuals between candidate expressions and observed data are implicitly computed and encoded. Ablation studies reveal that the encoder-driven update yields substantially simpler expressions due to inductive biases—autoregressive decoder length bias, parsimony-favoring scoring, and functional latent geometry shaped by $f_\theta$8. Qualitative interpolation and linear probe experiments confirm that $f_\theta$9 encodes meaningful functional and structural properties, facilitating efficient search.

Figure 5: Latent interpolation comparison; the VAE-enabled latent space exhibits smooth functional transitions, in contrast to instability and failure in non-VAE variants.

Implications and Future Directions

LEE advances the interpretability-accuracy trade-off in symbolic regression, suitable for scientific contexts where parsimony and functional insight are critical. The model's speed—tens of seconds per dataset on a GH200 node—further augments its practical utility relative to GP methods. Potential avenues for research include: scaling operator vocabularies and pre-training corpora to narrow the accuracy gap with GP SOTA, integrating higher-capacity generative priors for improved latent initialization, and deepening the quantification of functional equivalence in latent space. The general principle of encoder-driven iterative refinement with functional grounding is likely extensible to other symbolic or structured data domains.

Conclusion

Latent Equation Embedding (LEE) operationalizes iterative amortized inference in a functionally-grounded latent space as a tractable and effective approach to symbolic regression. Through structured encoder-driven updates and gradient refinement enabled by evaluation decoders, LEE produces concise expressions with competitive accuracy. Empirical evidence and technical analysis underscore the critical role of functional grounding in latent space, establishing LEE as a robust framework for interpretable modeling and scientific discovery.