Symbolic Regression via Latent Iterative Refinement
arXiv:2605.27245v1 [cs.LG] 26 May 2026
Xieting Chu Georgia Institute of Technology
Sriram Vishwanath Georgia Institute of Technology
Vijay Ganesh Georgia Institute of Technology
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: (1) an encoder fθ that jointly embeds symbolic tokens and numerical observations into a single latent vector z; (2) an expression decoder gexpr that reconstructs formulas from z; and (3) an evaluation decoder geval that predicts function values from z, explicitly grounding the latent space in functional behavior. At inference, LEE performs iterative refinement: zt+1 = fθ gexpr (zt ), D , re-encoding decoded expressions jointly with observations to progressively improve 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 geval is differentiable in z, we additionally interleave continuous gradient descent with the discrete re-encoding, yielding a hybrid iterative + gradient refinement. On SRBench across three noise levels, against 19 baselines spanning GP, symbolic– neural hybrids, and pre-trained Transformers, LEE produces expressions 2–10× simpler than the strongest accuracy-oriented baselines—Operon, GP-GOMEA, TPSR, RAG-SR, and GenSR (complexity 8–11 vs. 20–90)—advancing the lowcomplexity region of the accuracy–complexity Pareto frontier and degrading gracefully as noise grows.
1
Introduction
Symbolic regression (SR) recovers interpretable mathematical expressions from data. Formally, given ∗ observations D = {(xi , yi )}N i=1 , the goal is to find e = arg maxe∈E p(e | D), where E is the space of symbolic expressions—a mixed search space whose skeleton (operators, variables, tree structure) is discrete and combinatorial while the embedded real-valued constants are continuous [19, 21, 4]. The amortization gap and existing remedies. Neural SR methods [1, 7] replace the hours-long combinatorial search of classical genetic programming [19, 4] with an amortized inference model fϕ that maps observations directly to approximate posterior parameters λ(i) ← fϕ (x(i) ) (e.g. the logits of an autoregressive decoder). This one-shot prediction is fast but inherently limited: Marino et al. [14] shows that the gap between the amortized estimate and the true optimum—the amortization gap—grows with problem complexity; E2ESR [7] drops from R2 =0.857 on Feynman to R2 =0.361 on out-of-distribution black-box data. A second line of work introduces a latent space that pairs numerical and symbolic representations: SNIP [15] learns a discriminative pairing via contrastive Preprint.
pre-training (which cannot be searched directly), while GenSR [12] learns a generative latent space via a dual-branch CVAE and refines the prior estimate with CMA-ES [6]. GenSR’s CMA-ES is effective, but it treats the learned model as a black-box optimization objective: it uses only scalar fitness evaluations, discards the encoder’s structural information, scales as O(d2z ), and carries no semantic understanding of why a candidate fits poorly. We close the amortization gap with a different lever: a learned latent search that exploits the encoder’s own forward pass. Our approach: iterative amortized inference. Marino et al. [14] show that the amortization gap can be closed by iterative inference models that learn to optimize approximate posteriors by encoding gradients: (i) (i) (i) λt+1 ← ft ∇λ Lt , λt ; ϕ . (1) LEE instantiates this principle in the symbolic regression setting. Rather than encoding explicit gradients, LEE uses the decoded expression itself as an implicit error signal (an instance of the error-encoding variant of iterative amortized inference). The encoder, when given both a candidate expression êt = gexpr (zt ) and the observations D, can internally attend to the mismatch between the expression’s predictions and the data—effectively computing a learned residual. This yields the LEE update rule: zt+1 = fθ êt , D , êt = gexpr (zt ) (2) Each iteration refines z by re-encoding a decoded expression jointly with observations, closing the amortization gap through the model’s own inference pathway. Contributions. We introduce three design choices that together enable iterative amortized inference for SR: 1. Iterative latent search via encode–decode–reencode (Section 3.5). The update rule (2) uses the model’s own encoder as a learned inference optimizer, maintaining a candidate pool for diversity. Unlike GenSR’s CMA-ES, each step is semantically informed: the encoder attends to both the candidate’s tokens and the observations. 2. Evaluation decoder for functional grounding (Section 3.3). A dedicated decoder geval (z, x) → ŷ forces the latent space to encode what a function does, not just its syntactic form. This creates a latent geometry where proximity reflects functional similarity—a prerequisite for meaningful iterative search. 3. Hybrid iterative + gradient refinement (Section 3.5c). Because geval is differentiable in z, continuous gradient descent can interleave with the discrete re-encoding, yielding a search that is more robust to noise than either alone; its importance is verified by a dedicated ablation (Section 5.3). On the SRBench benchmark suite [9], LEE produces expressions 2–10× simpler than the strongest accuracy-oriented baselines (Operon, GP-GOMEA, TPSR, RAG-SR, GenSR) while advancing the low-complexity region of the accuracy–complexity Pareto frontier.
2
Background and Related Work
Direct search in E. Genetic programming (GP) methods sample and mutate expression trees directly in the discrete skeleton space, delegating constant fitting to an inner numeric optimizer. Contemporary high-performance GP systems such as Operon [2], GP-GOMEA [25], PySR [4], and Bingo [18] drive much of the state-of-the-art on SRBench. GP achieves high accuracy but requires repeated per-dataset evolutionary search, incurring wall-clock costs that scale with both dataset and population sizes. Amortized inference (one-shot). Pre-trained neural SR models amortize the per-dataset search with a single forward pass of an autoregressive decoder: NeSymReS [1] and E2E-SR [7] train a direct mapping fϕ : D → ê. A second family combines amortized components with search-style augmentations, including DSR [17], TPSR [20], uDSR [11], and RAG-SR [26]. These models are fast at inference but incur a residual amortization gap: the one-shot estimate rarely matches the true posterior, and the gap widens on distributions the encoder did not see during pre-training. Latent-space methods. Another line of work learns a latent space that pairs numerical and symbolic representations and uses it either as an initialization or as the explicit search domain. SNIP [15] aligns numerical data and symbolic expressions via contrastive pre-training; the resulting discriminative 2
Latent Space 𝒁
Training
Inference
Symbolic Input 𝓏0
sin x + y ^2 ···
pad
VAE Encoder 𝑓𝜃
𝓏 ℒ𝑎𝑙𝑖𝑔𝑛 𝓏′
Expression Decoder 𝑔𝑒𝑥𝑝𝑟
Evaluation Decoder 𝑔𝑒𝑣𝑎𝑙
𝑦ො𝑗
sin x + y ^2 ···
𝓏1 𝑓𝜃
𝑔𝑒𝑥𝑝𝑟
sin x + y ^2 ···
𝑔𝑒𝑥𝑝𝑟
sin x + y ^2 ···
𝒆ො 𝟏
𝓏2
𝓏𝑇
Data 𝐷
gradient refinement
ℒ𝑟𝑒𝑓𝑖𝑛𝑒 Reconstruct Original
𝑓𝜃
𝒆ො 𝟎
𝑔𝑒𝑥𝑝𝑟
Data 𝐷
ℒ𝑒𝑣𝑎𝑙
𝑥𝑗
sin x + y ^2 ···
𝑓𝜃
𝓏1 iterative refinement
Numeric Input
Corrupted Symbolic Input
iterative refinement
𝓏0
iterative refinement
sin x + y ^2 ··· ℒ𝑒𝑥𝑝𝑟
Data 𝐷
pad
𝓏𝑇
𝑓𝜃
𝒆ො 𝑻
𝓏∗ gradient refinement
sin x + y ^2 ···
Numeric Input
Refine 𝓏𝑇 by solving 𝓏 ∗ = arg 𝑚𝑖𝑛𝑧 ∑(𝑔𝑒𝑣𝑎𝑙 𝑧, 𝑥𝑗 − 𝑦𝑗 )2 ℒ = 𝜆1 ℒ𝑒𝑥𝑝𝑟 + 𝜆2 ℒ𝑒𝑣𝑎𝑙 + 𝜆3 ℒ𝐾𝐿 + 𝜆4 ℒ𝑎𝑙𝑖𝑔𝑛 + 𝜆5 ℒ𝑟𝑒𝑓𝑖𝑛𝑒
Figure 1: LEE architecture. Left (training). The encoder fθ maps an expression and its scatter D to a latent z, from which gexpr reconstructs tokens and geval predicts function values at queries xj ; the five training losses (Sec. 3.4) are labeled at the points where they apply. The dashed path repeats the encode with tokens replaced by pad to define the scatter-only branch used by Lalign ; the boxed inset depicts Lrefine (random token corruption → denoised reconstruction). Middle. Inference trajectory in Z: solid red steps are discrete iterative refinement, the dashed segment is the final gradient refinement to z ∗ . Right (inference). Starting from z0 = fθ (pad, D), each step decodes êt = gexpr (zt ) and re-encodes (êt , D) to zt+1 (Eq. 2); after T steps, gradient refinement on geval produces z ∗ . embedding captures shared structure but is not generative, so it typically seeds a downstream decoder rather than acting as a search domain itself. GenSR [12] instead learns a generative latent space via a dual-branch CVAE, framing SR as maximizing p(Equ. | Num.) through the ELBO log p(F | X) ≥ Eq(z|X,F ) log p(F | X, z) − DKL q(z | X, F ) ∥ p(z | X) , (3) where the posterior branch encodes both expression F and numerical data X into q(z | X, F ) and the prior branch encodes only X into p(z | X). At inference, CMA-ES refines the prior-branch estimate in Z using only scalar fitness feedback, so the search is gradient-free w.r.t. the model and discards the encoder–decoder’s structural information. Iterative amortized inference. Marino et al. [14] propose closing the amortization gap by learning to iteratively refine approximate posteriors (Eq. 1); standard amortized inference is the one-step t=0 special case. The principle has not been applied to symbolic regression because (i) the latent space must be functionally meaningful for a refinement step to translate into a better expression, and (ii) the encoder must be trained to consume its own decoded outputs as input—two conditions that the LEE design explicitly satisfies.
3
Method: Latent Equation Embedding
LEE consists of three jointly trained components sharing a latent space Z ⊂ Rdz (Figure 1): z = fθ (t, D) ∈ Rdz ,
(4)
Expression decoder: t̂ = gexpr (z), Evaluation decoder: ŷ(x) = geval (z, x),
(5) (6)
Encoder:
where t ∈ V L is the symbolic token sequence and D is a set of scatter observations. At inference, t is unavailable; the initial encoding uses only D. 3
3.1
Encoder
The encoder takes two input streams: = Embed(tj ) ∈ Rd . Symbolic stream. Each token tj of the expression is embedded as hsym j Numeric stream. Each observation (xi , yi ) is embedded by a two-layer MLP, producing hnum ∈ Rd . i Coordinates and function values are both log-compressed by ũ = sign(u) · log(1 + |u|) for numerical stability across many orders of magnitude. Non-finite values are handled by learnable special embeddings. Fusion. Symbolic and numeric embeddings are concatenated along the sequence dimension, processed by an L-layer Transformer [23], and masked mean-pooled into a single vector h = MeanPool(Transformer([hsym ; hnum ])). Two linear heads map h to Gaussian parameters µ = Wµ h and log σ 2 = Wσ h, and the latent vector is drawn by reparameterization [8]: z = µ + σ ⊙ ϵ,
ϵ ∼ N (0, I).
(7)
The Gaussian posterior q(z | t, D) = N (µ, diag(σ 2 )) is regularized toward the standard-normal prior by LKL (Sec. 3.4). 3.2
Expression Decoder
The expression decoder autoregressively generates prefix-notation tokens from z. The latent vector is projected into K memory tokens that serve as cross-attention keys: Mexpr = reshape(Wm z) ∈ RK×dexpr .
(8)
A causal Transformer decoder cross-attends to Mexpr , producing logits at each step: pθ (tj | z, t<j ) = softmax Wo · TransDec(Mexpr , t<j ) . 3.3
Evaluation Decoder: Grounding the Latent Space
The evaluation decoder predicts function values at arbitrary query coordinates {qj }M j=1 from z: Meval = reshape(We z) ∈ RK×deval , ŷj = MLP TransDec(Meval , MLP(qj )) . (9) Why an evaluation decoder? Consider two expressions e1 = 2 sin(x) cos(y) and e2 = sin(x + y) + sin(x − y). They are syntactically distant — one is a scalar multiple of a product of two trig atoms, the other is a sum of trig functions applied to compound arguments — but functionally identical by the product-to-sum identity. Without geval , no loss explicitly ties z to function values: symbolic reconstruction penalizes e1 and e2 equally for being decoded as each other, and while the scatter input together with Lalign provides an indirect pressure toward numerical consistency, nothing forces functionally-equivalent expressions to map to nearby z. With geval , the latent space must encode functional behavior: for e1 and e2 to both decode to identical ŷ, they are pushed toward nearby z. This encourages the soft equivalence z1 ≈ z2
e1 (x) = e2 (x) ∀x,
whenever
(10)
shaping a latent geometry in which proximity reflects functional similarity. This is the geometry needed for iterative search: moving z toward a functionally better region tends to decode a functionally better expression. Because geval is differentiable in z, it also provides the gradient signal that enables the continuous refinement mode of Sec. 3.5(b). 3.4
Training Objective
The total loss combines five terms: L = λexpr Lexpr + λeval Leval + λKL LKL + λalign Lalign + λrefine Lrefine . Expression reconstruction (Lexpr ). P − |T1 | j∈T log pθ (tj | z, t<j ).
Cross-entropy over non-padding tokens:
4
(11) Lexpr
=
Evaluation loss (Leval ). Scale-invariant MAE: Leval = |V1fin |
|ŷj −yj | j∈Vfin max(|yj |,1) .
P
Latent regularization (LKL ). Standard VAE KL divergence between the posterior q(z | t, D) produced by the encoder’s (µ, σ) head and a unit Gaussian prior N (0, I), with a small weight λKL to avoid posterior collapse. Cross-modal alignment (Lalign ). At inference, only scatter is available. To bridge this modality gap we run the same encoder fθ twice per training sample—once on the full (tokens, scatter) input, producing a posterior q(z | t, D), and once on scatter alone (symbolic stream filled with [pad]), producing a scatter-only distribution p(z | D)—and align the two via a conditional KL: Lalign = DKL q(z | t, D) sg[p(z | D)] . (12) Here sg[·] denotes the standard stop-gradient operator [22, 5]: during the backward pass it treats its argument as a constant, so the gradient of Lalign flows only into the posterior (q) branch, while the scatter-only (p) branch is held fixed as the target. This matches the KL direction in GenSR’s ELBO (Eq. 3) but is used as a standalone weighted term rather than part of a probabilistic objective. Iterative refinement (Lrefine ). To train the encoder for the iterative regime, we simulate the inferencetime loop during training. Expression tokens are randomly corrupted (drops, swaps, substitutions), producing a noisy expression ẽ. The encoder must map (ẽ, D) to a z that decodes to the original expression: 1 X Lrefine = − log pθ tj | fθ (ẽ, D), t<j . (13) |T | j∈T
This trains the encoder to act as a denoising inference optimizer [24]: given a corrupted expression and the data, it must “correct” the latent representation. At inference, decoded expressions play the role of ẽ—they are imperfect approximations that the encoder refines. 3.5
Iterative Latent Search
The inference procedure instantiates Eq. 2 through three complementary refinement strategies: iterative refinement (discrete re-encoding through fθ , operating over a candidate pool), gradient refinement (continuous descent in Z through the evaluation decoder), and their combination iterative + gradient refinement. These form the core of our method and the basis of the ablation study in Sec. 5.3. Initialization. The initial latent vector encodes only the observations: z0 = fθ (∅, D).
(14)
From z0 , we decode ninit candidate expressions via greedy and temperature-sampled decoding, score each by R2 − α · complexity, and keep the top P as the initial pool Π0 . (i)
(a) Iterative refinement. At each step t, we sample a parent expression êt from Πt with rankweighted probability, and apply the update: (i) (i) (j) (i) zt+1 = fθ êt , D , êt+1 ∼ gexpr zt+1 , j = 1, . . . , nnew . (15) New candidates are scored and merged into Πt , keeping the top P with complexity diversity. Constants in decoded expressions are refined via L-BFGS-B [3]. To prevent pool collapse, we periodically re-sample scatter points from D and decode fresh candidates from a new scatter-only z0 (every 5 batches; Sec. 4). This is the discrete, encoder-driven realization of Eq. 2. (b) Gradient refinement. Because geval (z, x) is differentiable in z, we can directly descend on the latent: zt+1 = zt − η ∇z ∥ geval (zt , X) − y∥22 + λprox ∥zt − zanchor ∥22 . (16) The proximal term keeps zt near the decodable region of Z. Every d steps, we decode the current zt and score the resulting expression; the best expression seen over the trajectory is returned. This mode exploits the functional grounding induced by geval (Sec. 3.3). (c) Iterative + gradient refinement. Iterative refinement excels at global exploration through discrete re-encoding; gradient refinement excels at local fine-tuning through continuous descent. The combined mode alternates between the two: every d iterative refinement steps, we take the current pool champion z ∗ , run k steps of gradient descent (16), decode, and merge the resulting 5
expression back into the pool. A safety fallback uses the held-out validation fold: if the gradient step lowers validation R2 relative to the pool champion, we revert to the pool champion. The test fold is untouched until final reporting. As we show in Sec. 5.3, this hybrid is especially valuable on noisy data, where gradient refinement locally denoises coefficients while iterative refinement maintains structural diversity through the pool.
4
Experimental Setup
Benchmarks. We evaluate on the SRBench benchmark suite [9]: Strogatz (14 ODE systems), Feynman (116 physics equations), and black-box (63 PMLB datasets without known ground truth). Ground-truth benchmarks are run at three target noise levels ϵ ∈ {0, 0.01, 0.1} (Gaussian noise with standard deviation proportional to the target range); black-box is noise-free. Data splits. We adopt SRBench’s canonical 75%/25% train/test partition; the 25% test fold matches SRBench’s protocol exactly, making our test numbers directly comparable to published baselines. Internally, we carve a 20% validation slice from the 75% training portion, so each dataset is 60/15/25 train/val/test overall. R2 on the test fold is the reported accuracy metric; the validation fold is used only for round selection within a trial and for the gradient-fallback decision, and is never observed by the model during search. Complexity is SymPy-simplified [16] node count. Evaluation protocol. For each (dataset, ϵ) pair, we run independent trials with distinct random seeds (data splits and search seeds) and report the mean ± standard deviation across trials. The main results in Table 1 use 10 trials per cell to match SRBench’s published-baseline protocol; the ablation studies in Sec. 5.3 and the appendix sensitivity sweeps use 3 trials to keep the compute footprint manageable. Each trial is the best-of-R=10 rounds of iterative + gradient refinement (Sec. 3.5), with the winning round chosen by highest validation R2 . Aggregation at the dataset-group level uses the mean across datasets within each group. Training data. ∼13.4M synthetic expressions from a stochastic context-free grammar (15 operators, 1–10 variables), paired with 200 scatter points from U(−10, 10)k . Model. Encoder: d=768, 6 layers, 12 heads, dz =512 (≈75M params). Expression decoder: d=512, 8 layers, 8 heads, K=4 memory tokens (≈50M). Evaluation decoder: d=512, 4 layers, 8 heads, K=4 (≈25M). Total: ≈150M parameters. Training: AdamW [13] with cosine decay, single NVIDIA GH200. Inference. Pool size P =16, ninit =32, T =200 iterations per round, nnew =3 per iteration, batch k=5 parents processed together, scatter refresh every 5 batches, L-BFGS-B ramped 100 → 300 steps. 2 Candidates are scored by Rtrain − α · C(e) with α=0.002 and C(e) the SymPy-simplified node count; decoder sampling uses temperature τ =0.7. The pool maintains complexity diversity by keeping at most ⌈P/4⌉ candidates per complexity bucket, so that short and long expressions are both retained. For the combined mode, we insert one gradient segment (k=50 steps, η=5×10−3 , λprox =0.1) every 25 iterative refinement steps. Rounds run 8-way parallel on a single GH200 node. Baselines. We compare against 19 SRBench methods: GP-based (Operon, GP-GOMEA, SBPGP, GPlearn, AFP, AFP-FE, EPLEX, ITEA), symbolic + deep hybrids (DSR, RSRM, MDL, SPL, AIFeynman2), and neural SR (NeurSR, E2ESR, SNIP, TPSR, RAG-SR, GenSR). Baseline numbers are taken from the published GenSR paper [12] and the SRBench 2.0 feather data where applicable, which follow an identical 10-trial, 75/25 split, R2 test-fold protocol. Since LEE completes per-dataset in tens of seconds—well below the SRBench compute budget that bounds the baselines (Appendix J)— our R2 and complexity are hardware-agnostic and directly comparable to the published values.
5
Results
Our empirical study answers three questions: (i) How does LEE compare to existing SR methods across noise levels and benchmark types, both in headline metrics and on the accuracy–complexity Pareto frontier (Sec. 5.1–5.2)? (ii) Which refinement strategies (iterative, gradient, or their combination) are responsible for the results (Sec. 5.3)? (iii) Does the iterative re-encoding update specifically—rather than the backbone, training data, or scoring—drive the accuracy gain over oneshot decoding and CMA-ES on the same checkpoint (Sec. 5.4)? Iterative-convergence behavior in Z and detailed timing data are deferred to Appendix C and Appendix J. 6
Table 1: SRBench results across noise levels. Mean test R2 (↑) and mean simplified complexity (Cmplx ↓). Black-box has no ground truth and is run noise-free. LEE numbers are means over 10 trials, matching SRBench’s published-baseline protocol; per-cell standard deviations are reported in Appendix K for compactness. Bold: best neural/hybrid; underline: best overall. Strogatz (14) ϵ=0.01 R2 C
ϵ=0.1 R2 C
ϵ=0 R2 C
ϵ=0.01 R2 C
Black-box (63)
ϵ=0.1 R2 C
— R2
C
Genetic programming Operon .988 59 .983 82 .938 83 .989 70 .988 88 .985 89 .794 GP-GOMEA .992 36 .978 43 .967 44 .996 35 .997 45 .996 46 .738 SBP-GP .981 712 .981 851 .932 901 .994 489 .995 596 .990 622 .787 GPlearn .769 29 .796 31 .823 26 .881 72 .889 60 .891 49 .539 AFP .925 38 .915 39 .911 44 .959 37 .961 41 .958 41 .633 AFP-FE .944 46 .958 49 .950 51 .981 40 .982 47 .983 49 .640 EPLEX .812 50 .856 53 .882 54 .987 53 .991 54 .990 46 .737 ITEA .792 11 — — .910 15 — — — — — — .629
66 30 634 19 35 36 53 117
Symbolic–neural hybrid DSR .760 16 RSRM .550 13 MDL .990 14 SPL .739 15 AIFeynman2 .646 22
.820 .597 .972 .739 .775
18 14 20 15 32
.809 .555 .969 .772 .317
18 14 20 14 24
.844 15 .878 16 .800 13 .809 13 .917 23 .914 31 .707 13 .713 13 .931 124 .873 155
.878 16 .810 13 .910 31 .711 14 .225 177
.562 .332 .626 .547 .211
10 9 30 13 2240
Pre-trained neural NeurSR .521 E2ESR .534 SNIP .995 TPSR .965 RAG-SR .991 GenSR .992
.518 .503 .984 .980 .987 .994
12 36 29 56 49 20
.505 .515 .919 .971 .969 .977
13 38 39 56 46 20
.396 .857 .985 .992 .993 .987
.382 .771 .992 .984 .985 .989
.123 .361 .334 — — .842
13 61 39 — — 35
.854 8.1 .876 8.9 .880 8.3 .884 9.9 .884 10.1 .824 10.6 .559
9.0
Method
LEE (ours)
5.1
ϵ=0 R2 C
Feynman (116)
11 32 29 56 46 20
13 36 32 57 46 23
.394 .834 .987 .991 .990 .987
13 40 33 64 72 23
14 44 38 67 75 24
Overall Comparison
Table 1 reports mean test R2 and simplified complexity across Strogatz, Feynman, and black-box for three noise levels. Both LEE and baseline numbers follow SRBench’s 10-trial 75%/25% protocol on the same 25% test fold; each LEE trial is a best-of-10-rounds run. LEE consistently occupies the low-complexity corner of the accuracy–complexity trade-off, and its accuracy degrades gracefully with noise on the ground-truth benchmarks, in contrast to several neural methods that sharply collapse. Accuracy. On Strogatz and Feynman, LEE lies within 6–14 R2 points of the top GP methods, which is the cost of the simplicity trade-off discussed below; LEE is not an accuracy-SOTA method. Under noise, however, LEE’s Strogatz R2 rises slightly from 0.854 to 0.880 as ϵ grows from 0 to 0.1, whereas SNIP drops by 0.076 and E2ESR by 0.019 (while producing ∼ 4× larger expressions); we hypothesize that input noise widens the encoder’s posterior, increasing search diversity at no accuracy cost. On black-box (no ground truth, out-of-distribution), LEE reaches R2 =0.559, ahead of one-shot neural methods (SNIP 0.334, E2ESR 0.361) while keeping complexity 4–7× smaller. Complexity. Across all settings, LEE produces the simplest expressions—complexity 8–11 versus 15–70+ for all competing neural methods and most GP baselines—while retaining competitive R2 . This is the Pareto-differentiating property of LEE: where other methods trade accuracy against complexity with larger expressions, LEE advances the low-complexity region of the frontier. Speed. LEE’s inference is fast: one dataset completes in tens of seconds of wallclock time on a single GH200, faster than most GP baselines and within a small constant factor of the fastest one-shot neural methods. A detailed timing comparison is given in Appendix J; because hardware varies substantially across baselines, we caution against over-interpreting absolute numbers. 7
Strogatz (N = 14, = 0.1)
AI Feynman (N = 116, = 0.1) clipped at 300
102
Operon TPSR AFP-FE EPLEX RAG-SR AFP GP-GOMEA SNIP
E2ESR GPlearn
MDL GenSR
AIFeynman2
DSR SPL
101
RSRM NeurSR
LEE (Ours)
1.0
0.9
0.8
0.7
0.6
0.5
Mean R 2 (higher is better )
0.4
clipped at 300
SBP-GP (cplx=622)
Equation complexity (lower is better )
Equation complexity (lower is better )
SBP-GP (cplx=901)
AIFeynman2
102
Operon RAG-SR TPSR AFP-FE GPlearn GP-GOMEA EPLEX SNIPAFP MDL
E2ESR
GenSR DSR
0.3
RSRM
SPL
0.8
0.7
NeurSR
LEE (Ours)
101 1.0
0.9
(a) Strogatz, ϵ=0.1
0.6
0.5
Mean R 2 (higher is better )
0.4
0.3
0.2
(b) Feynman, ϵ=0.1
Figure 2: Pareto frontiers at ϵ=0.1 (test R2 vs. complexity, log-x). LEE sits in the low-complexity corner of the frontier, typically 2–7× simpler than accuracy-comparable methods. The same qualitative picture holds at ϵ=0 and ϵ=0.01 (Appendix L) and on black-box (Fig. 7). Table 2: Search-strategy ablation on Strogatz (mean ± std, each best-of-10-rounds with 200 iterations per round; R = 10, T = 200, 8-way parallel). ϵ=0 Strategy iterative refinement (fθ re-encoding) gradient refinement (∇z geval ) iterative + gradient refinement
5.2
2
0.872 ± 0.012 0.742 ± 0.027 0.854 ± 0.005
R
ϵ=0.1 Cmplx
R
2
Cmplx
8.0 9.2 8.1
0.850 ± 0.024 0.744 ± 0.045 0.880 ± 0.024
9.5 9.6 8.3
Pareto Analysis
Figure 2 visualizes the same tables in the (R2 , complexity) plane. Three observations stand out: (i) Consistency across noise (Appendix L, Fig. 6). LEE’s position on the frontier is essentially invariant as ϵ grows from 0 to 0.1, while several neural methods (SNIP, TPSR, RAG-SR) move up and to the right (more complex, less accurate). (ii) Distinct regime. No other method reaches complexity < 15 at the accuracy LEE achieves; DSR and MDL are the closest competitors but sit at higher complexity or lower R2 . (iii) Black-box robustness (Fig. 7, Appendix L). The black-box Pareto frontier is a staircase of non-dominated points at progressively higher complexity tiers—LEE and DSR both anchor the low-complexity end (cplx ≤ 10, R2 ≈ 0.56), then MDL (30), GP-GOMEA (30), and GenSR (35)—with LEE producing the simplest expressions on the frontier. One-shot neural baselines (SNIP, E2ESR, NeurSR) are dominated in both dimensions, reflecting how the single-pass inference distribution fails to transfer to OOD data. 5.3
Why Combine Iterative and Gradient Refinement: Ablation Study
A central claim of this paper is that the combination of discrete encoder-driven iterative refinement and continuous gradient refinement via the evaluation decoder makes iterative amortized inference effective, particularly under noise. We test this with an ablation on Strogatz that disables each component: Gradient refinement alone is insufficient. Pure gradient refinement through geval (row 2 of Table 2) lags iterative refinement by 13 R2 points at ϵ=0 and 10 points at ϵ=0.1, and exhibits the highest run-to-run variance. The gradient signal moves z toward a training-loss minimum, not toward a well-formed expression in gexpr ’s decodable region, so without re-projection through the encoder, it overfits coefficients and drifts off-manifold. Iterative refinement alone is competitive at ϵ=0 but loses accuracy under noise. Iterative refinement alone achieves the best R2 on clean data (0.872), consistent with the idea that on noisefree problems, the discrete search space is well-structured and continuous refinement adds little. Under ϵ=0.1, however, it drops to 0.850 with std 0.024—a three-point accuracy hit and triple the variance of the combined mode at ϵ=0. 8
Table 3: Same-backbone search comparison on Strogatz at ϵ=0.1. All three procedures use the identical pre-trained LEE checkpoint, scoring, and matched decode budget. (a) and (b) are 3-trial averages; (c) reports the headline 10-trial number from Table 1. Procedure (a) One-shot decode from z0 (b) CMA-ES on z (R2 fitness) (c) LEE iterative + gradient
R2 (↑)
Cmplx (↓)
0.795 ± 0.019 0.848 ± 0.011 0.880 ± 0.024
13.0 ± 1.9 13.8 ± 0.8 8.3 ± 1.4
Combining both is the robust choice. Iterative + gradient refinement is within 0.02 of iterative-only on clean data while having 2.4× lower variance (±0.005 vs. ±0.012), and it surpasses iterative-only on noisy data (0.880 vs. 0.850, a +3% gain), while also producing the simplest expressions (Cmplx 8.3 vs. 9.5). This matches our design intuition (Sec. 3.5c): iterative refinement maintains discrete diversity while gradient segments locally adjust z along the manifold shaped by geval so that a subsequent decode lands on a better expression (whose constants L-BFGS-B then refines). 5.4
Same-Backbone Search Comparison
To attribute LEE’s gains specifically to the iterative re-encoding update—rather than to the backbone, training data, or scoring function—we compare three search procedures on the identical pre-trained LEE checkpoint, identical scoring rule s(e) = R2 − αC(e), and matched per-round decode budget of ninit + T · nnew decodes per round (Sec. I): (a) one-shot decode from z0 = fθ (pad, D) (no search); (b) CMA-ES on z with R2 fitness (GenSR-style; population 24); (c) LEE iterative + gradient refinement (ours). Searching in z at all (a→b) buys +0.05 R2 over the one-shot baseline, confirming that the latent geometry is useful for search beyond the initial estimate. Replacing scalar-fitness CMA-ES with our encoder-driven iterative update (b→c) buys another +0.03 R2 and cuts complexity by ∼40%. We note that CMA-ES is given 4−5× the per-dataset wallclock budget of LEE iterative+gradient (its per-generation L-BFGS-B refinement is heavier than LEE’s incremental pool update), and a longer budget would likely close part of the R2 gap; the simplicity gap, however, is structural rather than budget-bound. The encoder’s structured update therefore does two things that CMA-ES cannot match at any budget: (i) it produces simpler expressions, because the encoder’s training distribution is biased toward simple skeletons; and (ii) it makes each step a single forward pass rather than a population evaluation with covariance updates and per-candidate constant refinement, so progress is incremental and unaffected by population synchronization.
6
Conclusion and Future Work
We presented LEE, a framework that casts symbolic regression as iterative amortized inference in a functionally-grounded latent space. The central equation, zt+1 = fθ (gexpr (zt ), D), uses the model’s own encoder as a learned inference optimizer, closing the amortization gap of one-shot methods like E2ESR and avoiding the black-box search of GenSR’s CMA-ES. A differentiable evaluation decoder further grounds Z in functional behavior, enabling both discrete re-encoding and continuous gradient-based refinement; our ablation (Sec. 5.3) shows that combining the two is essential under noise. On SRBench across three noise levels, LEE occupies a distinctive Pareto position: 2–10× simpler expressions than the strongest accuracy-oriented baselines (Operon, GP-GOMEA, TPSR, RAG-SR, GenSR) with accuracy within 0.10–0.17 R2 of those methods, modest wallclock cost, and graceful out-of-distribution behavior—well suited for scientific discovery where interpretability matters as much as fit. Future work. LEE’s framework admits several natural extensions. Scaling the operator vocabulary and pre-training corpus should narrow the 0.10–0.17 R2 gap to top GP methods (Operon, GPGOMEA) on clean benchmarks; coupling the same backbone with a higher-capacity generative prior (e.g., GenSR-style dual-branch encoding) for z0 , and directly measuring latent distance between canonically-equivalent expressions to quantify the functional-grounding claim of Sec. 3.3, would each tighten the framework further. 9
Acknowledgments and Disclosure of Funding We thank our collaborators and colleagues for helpful discussions. We thank our collaborators and colleagues for helpful discussions.
10
References [1] Luca Biggio, Tommaso Bendinelli, Alexander Neitz, Aurelien Lucchi, and Giambattista Parascandolo. Neural symbolic regression that scales. In International Conference on Machine Learning, pages 936–945. PMLR, 2021. [2] Bogdan Burlacu, Gabriel Kronberger, and Michael Kommenda. Operon C++: an efficient genetic programming framework for symbolic regression. In Proceedings of the 2020 Genetic and Evolutionary Computation Conference Companion, pages 1562–1570, 2020. [3] Richard H Byrd, Peihuang Lu, Jorge Nocedal, and Ciyou Zhu. A limited memory algorithm for bound constrained optimization. SIAM Journal on Scientific Computing, 16(5):1190–1208, 1995. [4] Miles Cranmer. Interpretable machine learning for science with PySR and SymbolicRegression.jl. arXiv preprint arXiv:2305.01582, 2023. [5] Jean-Bastien Grill, Florian Strub, Florent Altché, Corentin Tallec, Pierre H Richemond, Elena Buchatskaya, Carl Doersch, Bernardo Avila Pires, Zhaohan Daniel Guo, Mohammad Gheshlaghi Azar, Bilal Piot, Koray Kavukcuoglu, Remi Munos, and Michal Valko. Bootstrap your own latent: a new approach to self-supervised learning. Advances in Neural Information Processing Systems, 33:21271–21284, 2020. [6] Nikolaus Hansen and Andreas Ostermeier. Completely derandomized self-adaptation in evolution strategies. Evolutionary Computation, 9(2):159–195, 2001. [7] Pierre-Alexandre Kamienny, Stéphane d’Ascoli, Guillaume Lample, and François Charton. End-to-end symbolic regression with Transformers. Advances in Neural Information Processing Systems, 35:10269–10281, 2022. [8] Diederik P Kingma and Max Welling. Auto-encoding variational Bayes. arXiv preprint arXiv:1312.6114, 2013. [9] William La Cava, Bogdan Burlacu, Marco Virgolin, Michael Kommenda, Patryk Orzechowski, Fabrício Olivetti de França, Ying Jin, and Jason H Moore. Contemporary symbolic regression methods and their relative performance. In Advances in Neural Information Processing Systems Datasets and Benchmarks Track, 2021. [10] Guillaume Lample and François Charton. Deep learning for symbolic mathematics. arXiv preprint arXiv:1912.01412, 2019. [11] Mikel Landajuela, Chak Shing Lee, Jiachen Yang, Ruben Glatt, Claudio P Santiago, Ignacio Aravena, Terrell Mundhenk, Garrett Mulcahy, and Brenden K Petersen. A unified framework for Deep Symbolic Regression. Advances in Neural Information Processing Systems, 35: 33985–33998, 2022. [12] Qian Li, Yuxiao Hu, Juncheng Liu, and Yuntian Chen. GenSR: Symbolic regression based in equation generative space. arXiv preprint arXiv:2602.20557, 2026. [13] Ilya Loshchilov and Frank Hutter. Decoupled weight decay regularization. arXiv preprint arXiv:1711.05101, 2017. [14] Joe Marino, Yisong Yue, and Stephan Mandt. Iterative amortized inference. In International Conference on Machine Learning, pages 3403–3412. PMLR, 2018. [15] Kazem Meidani, Parshin Shojaee, Chandan K Reddy, and Amir Barati Farimani. SNIP: Bridging mathematical symbolic and numeric realms with unified pre-training. arXiv preprint arXiv:2310.02227, 2023. [16] Aaron Meurer, Christopher P Smith, Mateusz Paprocki, Ondřej Čertík, Sergey B Kirpichev, Matthew Rocklin, Amit Kumar, Sergiu Ivanov, Jason K Moore, Sartaj Singh, Thilina Rathnayake, Sean Vig, Brian E Granger, Richard P Muller, Francesco Bonazzi, Harsh Gupta, Shivam Vats, Fredrik Johansson, Fabian Pedregosa, Matthew J Curry, Andy R Terrel, Štěpán Roučka, Ashutosh Saboo, Isuru Fernando, Sumith Kulal, Robert Cimrman, and Anthony Scopatz. SymPy: symbolic computing in Python. PeerJ Computer Science, 3:e103, 2017. 11
[17] Brenden K Petersen, Mikel Landajuela, T Nathan Mundhenk, Claudio P Santiago, Soo K Kim, and Joanne T Kim. Deep Symbolic Regression: Recovering mathematical expressions from data via risk-seeking policy gradients. arXiv preprint arXiv:1912.04871, 2019. [18] David L Randall, Tyler S Townsend, Jacob D Hochhalter, and Geoffrey F Bomarito. Bingo: a customizable framework for symbolic regression with genetic programming. In Proceedings of the Genetic and Evolutionary Computation Conference Companion, pages 2282–2288, 2022. [19] Michael Schmidt and Hod Lipson. Distilling free-form natural laws from experimental data. Science, 324(5923):81–85, 2009. [20] Parshin Shojaee, Kazem Meidani, Amir Barati Farimani, and Chandan K Reddy. Transformerbased planning for symbolic regression. Advances in Neural Information Processing Systems, 36:45907–45919, 2023. [21] Silviu-Marian Udrescu and Max Tegmark. AI Feynman: A physics-inspired method for symbolic regression. Science Advances, 6(16):eaay2631, 2020. [22] Aaron Van Den Oord, Oriol Vinyals, and Koray Kavukcuoglu. Neural discrete representation learning. Advances in Neural Information Processing Systems, 30, 2017. [23] Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Łukasz Kaiser, and Illia Polosukhin. Attention is all you need. Advances in Neural Information Processing Systems, 30, 2017. [24] Pascal Vincent, Hugo Larochelle, Yoshua Bengio, and Pierre-Antoine Manzagol. Extracting and composing robust features with denoising autoencoders. In Proceedings of the 25th International Conference on Machine Learning, pages 1096–1103, 2008. [25] Marco Virgolin, Tanja Alderliesten, Cees Witteveen, and Peter A N Bosman. Improving model-based genetic programming for symbolic regression of small expressions. Evolutionary Computation, 29(2):211–237, 2021. [26] Hengzhe Zhang, Qi Chen, Wolfgang Banzhaf, and Mengjie Zhang. RAG-SR: Retrievalaugmented generation for neural symbolic regression. In The Thirteenth International Conference on Learning Representations, 2025.
12
A
Formal Connection to Iterative Amortized Inference
We formalize the connection between LEE’s iterative search and the framework of Marino et al. [14]. Table 4 summarizes the high-level differences against GenSR’s CMA-ES; the rest of this section makes the LEE–Marino correspondence precise. Table 4: Search mechanism comparison. LEE’s iterative update uses the model’s own encoder as the optimizer; each step conditions on the full candidate expression and the data, whereas CMA-ES sees only a scalar fitness. Update rule Signal per step Model awareness Cost scaling Latent codes
GenSR (CMA-ES)
LEE (Iterative Amortized)
zt+1 ← CMA-ES(zt , fitness) Scalar fitness (R2 ) Black-box (ignores encoder) O(d2z ) covariance updates Two separate (zsym , znum )
zt+1 ← fθ (gexpr (zt ), D) Full token sequence + scatter Uses encoder’s learned geometry One encoder forward per step One shared z
Setup. Following the notation of Marino et al. [14], let λ(i) denote the approximate posterior parameters for data example x(i) , and let L(x(i) , λ(i) ; θ) be the ELBO. Standard amortized inference uses a direct mapping (their Eq. 5): λ(i) ← f (x(i) ; ϕ). (A.1) Iterative amortized inference refines this estimate (their Eq. 6): (i) (i) (i) λt+1 ← ft ∇λ Lt , λt ; ϕ . (A.2) LEE as an instance. In LEE, the approximate posterior is parameterized by the latent code z ∈ Rdz (i.e., λ ≡ z), from which expressions are decoded autoregressively. The ELBO analog is: L(D, z; θ) = log pθ (t | z) + log pθ (y | z, X) − DKL (q(z)∥p(z)) . | {z } | {z } | {z } expression fit
evaluation fit
(A.3)
regularization
LEE’s iterative update replaces the gradient ∇z L with the decoded expression êt = gexpr (zt ): zt+1 = fθ êt , D = fθ gexpr (zt ), D; θ . (A.4) This corresponds to the error-encoding variant (their Eq. 14), where the bottom-up mismatch between êt (X) and y and the top-down discrepancy between êt ’s structure and the latent prior are implicitly computed by the encoder’s cross-attention. Key difference. In Marino et al. [14], the error signal is a real-valued vector in the same space as λ. In LEE, the “error signal” is the symbolic expression êt —a discrete, structured object. The encoder fθ performs the nontrivial mapping from this structured input to a continuous update in Z, which is why training with Lrefine (Eq. 13) is essential: it teaches the encoder how to extract useful refinement signals from imperfect expressions.
B
Mechanism Analysis
Why does re-encoding close the amortization gap? The encoder computes fundamentally different functions in scatter-only vs. joint mode. In scatter-only mode (t=0), it solves an ambiguous inverse problem: map finite, noisy observations to a latent code. In joint mode (t>0) it receives both a candidate expression êt and D, and can internally attend to their mismatch—effectively computing a residual: zt+1 = fθ êt , D ≈ fθ êt , êt (X) − y, X , (17) where êt (X) − y is the residual the encoder computes internally via cross-attention between the symbolic and numeric streams. This mirrors the “error encoding” variant of Marino et al. [14] (their Eq. 14), which was shown to approximate higher-order derivatives and converge faster than gradient encoding. The refinement loss Lrefine (Eq. 13) explicitly trains the encoder for this regime: at inference, decoded expressions from the pool play the role of the noisy input ẽ that the encoder has been trained to denoise. 13
Expression simplicity as inductive bias. LEE’s tendency toward simple expressions arises from three compounding effects: (1) the autoregressive decoder has an implicit length bias—shorter token sequences have higher probability under teacher forcing; (2) the pool scoring function s(e) = R2 (e) − α · C(e) explicitly favors parsimony; and (3) the evaluation decoder creates a latent geometry where simple functional forms—more prevalent in the training distribution—occupy larger volumes of Z and are therefore more likely to be decoded. GenSR and GP methods, in contrast, have no inherent simplicity bias and rely on post-hoc complexity penalties.
C
Iterative Convergence Latent trajectory on strogatz_shearflow1
Convergence of iterative refinement on Strogatz (N = 14)
4
0.4
0
5
10 15 20 25 30 Iteration (batches of k = 5 parents)
2 1 0 1
±0.5 across Strogatz datasets LEE iterative refinement (mean) one-shot baseline (t = 0): 0.686
0.2
Iteration (pool update)
0.6
0.0
18 16 14 12 10 8 6 4 2 0
3
0.8
PC 2 (21.5% var)
Best-in-pool R 2 on training fold
1.0
z0 (scatter-only) zT (final, R 2 = 0.925)
2
35
2
(a) R2 convergence over iterations
1
0 1 PC 1 (38.0% var)
2
3
(b) Latent trajectory (strogatz_shearflow1)
Figure 3: Iterative convergence. (a) Best-in-pool R2 (mean ± 0.5σ) across the 14 Strogatz datasets. (b) PCA projection of the latent vectors z produced at successive pool-champion updates. Figure 3 summarizes the convergence behavior of LEE’s iterative refinement on the 14 Strogatz datasets. The pool’s best R2 (panel a) improves rapidly in the first 10–15 batches as re-encoding corrects the initial scatter-only estimate, then refines gradually toward a within-round pool-best of 0.827 (up from a t=0 baseline of 0.686); the full best-of-R pipeline reaches 0.854 in Table 1. Panel (b) visualizes the same refinement in the latent space: successive pool-champion encodings on strogatz_shearflow1 trace a path from z0 (scatter-only initialization) to zT (R2 =0.925), so each re-encoding step translates a discrete improvement in the decoded expression into a measurable move in Z.
D
Latent Space Interpolation
A complementary qualitative test of the latent geometry is whether linear interpolation between two encoded equations decodes to expressions that smoothly bridge them in function space. We encode two ground-truth expressions A and B jointly with their scatter observations to obtain zA = fθ (tokensA , DA ) and zB = fθ (tokensB , DB ); for each t ∈ {0, 1/3, 2/3, 1} we set zt = (1−t)zA + tzB , decode several candidates (1 greedy + 31 samples at τ =0.7), and pick the one whose values most closely match the linear blend (1−t) yA + t yB . Figure 4 shows the result for A: y = x → B: y = x2 . The endpoints round-trip back to their inputs, and the intermediate decoded expressions trace a smooth deformation from a linear ramp through a tanh-modulated bowl into the parabola. This qualitative behavior is consistent with the functional-grounding hypothesis (Sec. 3.3): the evaluation decoder shapes Z so that proximity in latent space reflects functional similarity.
E
VAE Architecture Ablation
To probe the contribution of the VAE encoder with conditional-KL alignment (Sec. 3.4), we compare the full LEE model against a non-VAE variant trained on the same data: deterministic encoder, z-norm penalty in place of KL, no cross-modal alignment loss. We compare on two axes: (i) headline 14
4 3
y
2 1 0 t = 0.00 : 1.0x t = 0.33 : xcos (tanh(x)) t = 0.67 : xtanh(x) t = 1.00 : x2
1 2 2.0
1.5
1.0
0.5
0.0 x
0.5
1.0
1.5
2.0
Figure 4: Latent interpolation, y=x → y=x2 on x ∈ [−2, 2]. Solid curves are decoded expressions at four points zt = (1−t)zA + tzB ; the color gradient (light → dark) encodes t. Endpoints round-trip back to their inputs (t=0 decodes to x, t=1 to x2 ); intermediate latents at t=1/3 and t=2/3 decode to syntactically distinct but functionally smooth interpolants (x cos(tanh x) and x tanh x).
accuracy on Strogatz at ϵ=0.1 (Table 5), and (ii) the four-step latent interpolation introduced in Appendix D (Figure 5). Both probes use the same evaluation protocol and inference hyperparameters; only the model checkpoint differs. Table 5: VAE architecture ablation on Strogatz ϵ=0.1 (mean ± std over 3 trials with failed-toconverge datasets dropped from the R2 and complexity averages). Removing the VAE encoder + KL alignment costs ∼0.30 R2 on Strogatz at ϵ=0.1; on 1–2 of 14 datasets per trial the non-VAE variant fails to return a valid expression at all. R2 (↑)
Cmplx (↓)
Failures / 14
0.582 ± 0.084 0.880 ± 0.024
8.5 ± 0.9 8.3 ± 1.4
1.3 ± 0.6 0
Variant Non-VAE variant (z-norm + no KL align) Full LEE (VAE + conditional KL)
(a) LEE (full model, with VAE)
4 3
1.0
2
0.5
1
y
y
(b) Ablation: no VAE / no KL alignment
1.5
0.0
0 t = 0.00 : 1.0x t = 0.33 : xcos(tanh(x)) t = 0.67 : xtanh(x) t = 1.00 : x2
1 2 2.0
1.5
1.0
0.5
0.0 x
0.5
1.0
1.5
0.566 t = 0.00 : 4.94x x 8.97 t = 0.33 : x(0.046x + 0.5) t = 0.67 : tan(tanh(x(0.305x + 0.239))) t = 1.00 : tan(tanh(xtanh(x)))
0.5
2.0
2.0
1.5
1.0
0.5
0.0 x
0.5
1.0
1.5
2.0
Figure 5: Latent interpolation comparison: with vs. without VAE (y=x → y=x2 on x ∈ [−2, 2]). Solid curves are decoded expressions at t ∈ {0, 1/3, 2/3, 1}; the color gradient encodes t. (a) Full LEE round-trips both endpoints and the intermediates are syntactically distinct but functionally smooth (x cos(tanh x) and x tanh x). (b) The non-VAE variant fails to round-trip (t=1 decodes to a saturating bowl rather than x2 ) and intermediates are ill-conditioned tan(tanh(·)) shapes that bear no resemblance to either A or B. The accuracy gap (Table 5) and the qualitative interpolation breakdown (Figure 5) tell the same story: the VAE-induced latent regularity is essential for a step in Z to correspond to a step in function space, which is the geometric prerequisite for the iterative re-encoding update of Eq. 2 to work as designed. 15
Caveat: this non-VAE variant differs from the full LEE not only in lacking the VAE/KL terms but also in the cross-modal alignment loss and in being trained on a smaller variable vocabulary (5 vs. 10), so the comparison is a proxy for the VAE ablation rather than an exact controlled study; a from-scratch retrain with use_vae=false holding all other components fixed is left for future work.
F
Latent Space Property Decoding
To test whether the encoder’s latent space encodes structurally meaningful properties of the input expression, we ask: for each of seven properties of an expression e, is the property linearly recoverable from z = fθ (e, D) alone—i.e., is there a single direction in Z along which expressions satisfying the property are separated from those that do not? Setup. We sample 5000 expressions uniformly at random from the test split of our pre-training corpus, encode each through the full encoder on freshly generated scatter D (yielding a 5000×512 latent matrix), and assign each example labels for the seven properties listed in Table 6. For each property, we train a logistic regression classifier (a single linear vector with L2 penalty, C=1) on a stratified 4000/1000 train/test split and report test accuracy together with the test AUC for binary targets. The 10-way num_variables classifier is multinomial and is reported as accuracy only. Table 6: Property linear-probe panel. Test accuracy and AUC of a single-vector logistic regression on z ∈ R512 . The True % column reports the proportion of positives in the sampled subset (chance for the 10-class probe is 10%). All binary probes achieve AUC ≥ 0.86, indicating each property has a dedicated linear direction in Z. Property has trig (sin / cos / tan / tanh) has log/exp has sq/cube (x2 or x3 ) has division √ is polynomial (no trig, log/exp, abs, ·) high-dim (nvars ≥ 5) num_variables (10-class)
True %
Test acc
Test AUC
53.7 26.8 33.5 73.7 26.5 31.6 10.0 (chance)
0.831 0.865 0.831 0.886 0.866 0.939 0.727
0.914 0.912 0.861 0.944 0.938 0.982 —
Findings. All six binary properties are linearly recoverable with AUC between 0.86 and 0.98. The strongest axis is high-dim (AUC 0.98, accuracy 94%), followed by has division (AUC 0.94) and is polynomial (AUC 0.94); the weakest is has sq/cube (AUC 0.86), still well above chance. The 10-way variable-count classifier reaches 72.7% accuracy (7.3× chance), comparable to the AUC profile of the binary probes. Each property, therefore, corresponds to a distinct linear direction in Z, and the directions evidently coexist: a 512-dimensional space has ample capacity to host one axis per property without conflict. This supports the design intent of Sec. 3.3—the encoder learns a function-grounded latent space whose principal directions correspond to interpretable structural properties of the underlying expression—and helps explain why a small number of iterative reencoding steps suffices to traverse the space (Sec. 3.5): movement along any single axis carries a structurally meaningful change in the decoded expression.
G
Training Details
Tokenization. Expressions are serialized in prefix (Polish) notation over a vocabulary of 40 tokens: 4 special tokens (PAD, BOS, EOS, UNK), 2 structural tokens, 10 variables (x0 , . . . , x9 ), 15 operators √ (+, −, ×, ÷, sin, cos, tan, tanh, exp, log, , x2 , x3 , abs, neg), and 14 digit tokens for constant encoding. Constants are represented at 3 significant figures in scientific notation as 9-token sequences: [sign, d1 , ., d2 , d3 , e, sign′ , e1 , e2 ]. Grammar and expression sampling. The stochastic context-free grammar (SCFG) follows the protocol of Lample and Charton [10], Kamienny et al. [7]: a binary-tree scaffold is sampled first (with 1 ≤ b ≤ bmax =4 binary operators), then unary operators (u drawn with umax =4) are attached, and finally leaf nodes are filled with variables (uniformly over the allowed set) and numerical constants. Constants are drawn from a mixture: 60% integer [−10, 10], 30% log-uniform over [10−2 , 102 ], and 10% from a small catalogue of physics constants (π, e, etc.). We enforce variable coverage: every 16
declared variable appears at least once in the tree. Trees are re-sampled on syntactic failure (NaN/Inf on the fixed query grid). The final training corpus contains ∼13.4M unique prefix sequences, split 80/10/10 into train/val/test. Pre-training cost. Total training wallclock is ≈ 200 GH200-GPU-hours spread across the five phases of Table 7, on a single node with batch size 256 and mixed-precision (bf16) forward/backward. The dataset is generated offline in ∼8 CPU-hours on 16 cores. This up-front cost is amortized across downstream datasets; for SR workflows that evaluate hundreds of datasets, the break-even point against 1–24-hour-per-dataset GP baselines is in the single digits.
Table 7: Loss weights across the five training phases. Phase
λexpr
λeval
λKL
λalign
λrefine
1 (basic) 2 (+align) 3 (+refine) 4 (freeze dec) 5 (unfreeze)
1.0 1.0 1.0 0 1.0
5.0 5.0 5.0 0 5.0
0.001 0.001 0.001 0.001 0.001
0 2.0 2.0 5.0 2.0
0 0 1.0 0 1.0
Training schedule. Phase 1: 50k steps (basic reconstruction, all parameters). Phase 2: 30k steps (add alignment, encoder focused). Phase 3: 50k steps (add refinement, full model). Phase 4: 30k steps (freeze decoders, alignment-only, encoder learns modality bridging). Phase 5: 40k steps (unfreeze all, co-adaptation). Batch size 256, AdamW (β1 =0.9, β2 =0.999), cosine decay from 3×10−4 to 1×10−5 . Data augmentation. During training, scatter points are randomly sub-sampled (128–200 points per example) and coordinate-rotated for multi-variable expressions. Token corruption for Lrefine : each token is independently dropped (15%), swapped with a random token (10%), or kept (75%). Constant optimization. After decoding, numerical constants in each expression are refined by L-BFGS-B, minimizing MSE on the training split. The budget ramps linearly from 100 to 300 steps over the search iterations. For datasets with > 1000 training points, we randomly subsample 1000 points for each L-BFGS-B call.
H
Architecture Details
Table 8: Architecture hyperparameters. Encoder
Expr. Decoder
Eval. Decoder
Model dim (d) Layers Heads FFN dim Dropout Memory tokens (K) Latent dim (dz )
768 6 12 3072 0.1 — 512
512 8 8 2048 0.1 4 512
512 4 8 2048 0.1 4 512
Parameters
≈75M
≈50M
≈25M
The encoder’s scatter-embedding MLP has a hidden dimension of 256 and uses SiLU activation. The log-compressed coordinate x̃ = sign(x) log(1 + |x|) is additionally divided by a fixed scale of 4 before the MLP, so that |x| ≤ 50 lands roughly in [−1, 1]; function values are not rescaled. The evaluation decoder’s query embedding MLP maps k coordinate dimensions to deval , with a hidden dimension of 256. Both decoders project z into K=4 memory tokens via a linear layer, then use the Transformer decoder cross-attention to these memory tokens. The expression decoder uses causal self-attention; the evaluation decoder uses bidirectional self-attention (queries can attend to all other query positions). 17
I
Inference Procedures
Full inference hyperparameters. Table 9 lists every inference-side hyperparameter and its default value. All numbers reported in Sec. 5 use these settings unless noted otherwise. Table 9: Inference hyperparameters. Name R T P ninit nnew batch k refresh period α τ L-BFGS-B budget L-BFGS-B subsample ngrad η λprox decode period MAX_SEARCH_POINTS ntrials
Value 10 200 16 32 3 5 5 0.002 0.7 100 → 300 1000 50 5×10−3 0.1 25 2000 10
Meaning rounds per trial refinement iterations per round candidate pool size candidates decoded from z0 at round start new candidates decoded per iteration parents processed per batch batches between scatter resamples 2 complexity penalty in scoring Rtrain − αC(e) decoder sampling temperature constant-refinement steps, linearly ramped over R rows used for constant fit if Ntrain >1000 gradient refinement steps per segment gradient refinement learning rate proximal anchor weight iterations between gradient segments row cap for per-candidate R2 scoring independent trials per (dataset, ϵ) for main results (3 for ablations)
Pool initialization. The initial latent z0 = fθ (pad, D) is decoded into ninit = 32 candidates: 1 greedy argmax decode and 31 temperature-τ sampled decodes. Each candidate’s constants are immediately 2 refined via L-BFGS-B on the training fold. Candidates are scored by s(e) = clip(Rtrain (e), −1, 1) − α C(e), with C(e) the SymPy-simplified node count. The top P = 16 form the initial pool Π0 . Parent selection and complexity diversity. At each iteration we sample k = 5 parents from Πt with rank-weighted probabilities pi ∝ 1/(i + 1) (so rank-1 is twice as likely as rank-3). When merging new candidates back, we enforce complexity diversity by bucketing C(e) into 4 bins [0, 5), [5, 10), [10, 20), [20, ∞) and capping the pool at ⌈P/4⌉ = 4 entries per bucket (best by s(e)). This prevents the pool from collapsing onto a single expression family. Scatter refresh. Every 5 batches, we (i) re-sample nscatter = 200 scatter points from the training fold, (ii) re-compute a fresh scatter-only z0 , and (iii) decode 3 new candidates that are added to the pool before the next iteration. This injects exploration when the pool has converged. Gradient refinement (pg mode). Every 25 iterations of iterative refinement, we take the current pool champion z ∗ (its constants fixed), run 50 Adam steps on z minimizing ∥geval (z, Xtrain ) − ytrain ∥22 + λprox ∥z − z ∗ ∥22 , decode from the resulting z, and insert the decoded expression (after L-BFGS-B) back into the pool. The proximal term with λprox = 0.1 prevents the gradient descent from wandering off the decodable manifold. Safety fallback. After pg-mode concludes a round, the reported winner is the pool entry with the highest validation R2 ; if that winner was produced by a gradient segment but has lower validation R2 than the best pool entry from the previous iterative-only step, we revert to the latter. This is the safety fallback of Sec. 3.5(c) and ensures the combined mode is never worse than iterative-only up to selection noise. Seeding and reproducibility. Each trial uses a fresh random seed st = sbase + 1000 t, where sbase varies per trial t ∈ {0, . . . , ntrials − 1} (with ntrials =10 for the main results in Table 1 and ntrials =3 for the ablations). Within a trial, st controls (i) the 60/15/25 train/val/test split, (ii) all decoder sampling, (iii) gradient optimizer initialization, and (iv) scatter subsampling. All trials share the same pre-trained checkpoint. Baseline numbers follow the seed protocol of their source publication. Pool size sensitivity. We sweep the candidate pool size P ∈ {8, 16, 32} on Strogatz at ϵ=0.1, holding every other inference hyperparameter fixed (Table 10). The default P =16 used in the main results sits at a small but consistent sweet spot: P =8 slightly underperforms because the rank-weighted 18
parent sampler depletes diversity too quickly, while P =32 slightly underperforms because lower-rank parents are sampled too rarely to inject fresh exploration.
Table 10: Pool size sensitivity (Strogatz, ϵ=0.1, 3 trials, paper protocol). The default P =16 used in the main results sits at a small but consistent sweet spot between under- and over-sized pools.
J
P
R2 (↑)
Cmplx (↓)
8 16 (default) 32
0.863 ± 0.030 0.880 ± 0.024 0.863 ± 0.026
10.5 ± 0.5 8.3 ± 1.4 10.7 ± 0.9
Timing Details
We move detailed wall clock comparisons here because hardware across methods varies substantially, and absolute seconds are therefore not directly comparable. What follows are the operational timings of our method and the reported timings of the baselines from the published SRBench and GenSR data. LEE: per-dataset wallclock. One full LEE run on a dataset consists of R = 10 rounds of T = 200 iterations of iterative + gradient refinement, executed with W = 8 parallel workers on a single NVIDIA GH200. Per-dataset wallclock is therefore lower-bounded by ⌈R/W ⌉ · τ = 2τ (where τ is the mean per-round time), with measured values closer to 1.25τ under continuous round dispatch. Table 11 reports the measured wallclock averaged over all datasets in each benchmark group and 10 trials.
Table 11: LEE per-dataset wallclock time (seconds, mean over 10 trials, single GH200 node with 16 CPU cores, 8-way round parallelism). Values are the full cost to produce the best-of-10-rounds result for a single dataset. Benchmark Strogatz (14) Feynman (116) Black-box (63) Strogatz, iterative only Strogatz, gradient only
ϵ=0
ϵ=0.01
ϵ=0.1
64.8 ± 1.3 109.2 ± 1.0 79.2 ± 1.1
65.6 ± 1.1 109.8 ± 1.2 —
67.6 ± 1.0 116.8 ± 0.3 —
48.2 5.0
— —
49.1 5.3
Comparison to baselines. Reported times for baselines in Table 1 of the main text vary from ~4 seconds (E2ESR one-shot on Feynman) to 149k seconds (SBP-GP on black-box). LEE’s per-dataset cost is in the tens-of-seconds regime: an order of magnitude faster than most GP methods (Operon, GP-GOMEA, SBP-GP) and within ∼ 30× of the fastest one-shot neural methods (E2ESR, SNIP, NeurSR), which perform no iterative search at all. In absolute terms, a full pass over the 116 Feynman datasets at ϵ=0.1 completes in roughly 116 × 117/60 ≈ 3.3 hours of wallclock on a single node. Where the time goes. Within each round, L-BFGS-B constant refinement dominates (roughly 65–75% of τ ), followed by GPU forward passes through encoder/decoders (20–30%) and expression serialization/SymPy simplification (the remainder). The gradient segments in the combined mode are inexpensive (< 5% of τ ), which is consistent with gradient refinement alone being the fastest mode in Table 11 (∼5 s) but the weakest in accuracy: most of the wallclock cost is in evaluating and refining candidate expressions, which only iterative refinement exercises. Caveats. Baseline times were measured on the hardware reported in their respective source papers; some (e.g. SBP-GP) use server-class CPUs over days, others use single-GPU workstations. Our numbers assume a single-node GH200 + 16 CPU cores. We therefore intentionally omit time from the main-text Table 1 and use it here only as a coarse order-of-magnitude reference. 19
K
LEE Summary Statistics with Standard Deviations
Table 12 reports the per-group LEE mean ± standard deviation for R2 and complexity, omitted from the main Table 1 for compactness. Each row aggregates 10 independent trials with distinct splits and search seeds (Sec. 4); the std reflects trial-to-trial variation. Table 12: LEE per-group summary with standard deviations (mean ± std over 10 trials; same trials underlying Table 1). R2 (↑)
Cmplx (↓)
Strogatz (14) ϵ=0 ϵ=0.01 ϵ=0.1
0.854 ± 0.005 0.876 ± 0.012 0.880 ± 0.024
8.1 ± 1.5 8.9 ± 1.2 8.3 ± 1.4
Feynman (116) ϵ=0 ϵ=0.01 ϵ=0.1
0.884 ± 0.001 0.884 ± 0.005 0.824 ± 0.005
9.9 ± 0.5 10.1 ± 0.3 10.6 ± 0.7
Black-box (63)
0.559 ± 0.004
9.0 ± 0.4
Setting
Standard deviations are small across the board: R2 std ranges from 0.001 to 0.024, with the largest variance on noisy Strogatz (ϵ=0.1, σ=0.024), reflecting the higher sensitivity of small-dataset noise realizations. Complexity std is similarly tight (≤1.5), confirming that the simplicity property of LEE is consistent across trials, not an artifact of a lucky seed.
L
Additional Pareto Frontiers
The main text (Fig. 2) shows the Strogatz and Feynman Pareto frontiers at the hardest noise level ϵ=0.1 for clarity. Figure 6 reproduces the same diagram at ϵ=0 and ϵ=0.01; Fig. 7 adds the black-box frontier. Across every setting, LEE sits in the low-complexity corner.
20
Strogatz (N = 14, = 0)
AI Feynman (N = 116, = 0) clipped at 300
EPLEX E2ESR
GPlearn AIFeynman2
GenSR DSR SPL
MDL
RSRM
ITEA
101
NeurSR
LEE (Ours)
1.0
0.9
0.8
0.7
0.6
Mean R 2 (higher is better )
102
DSR
RSRM
SPL
0.8
0.7
NeurSR
LEE (Ours)
0.9
0.6
Mean R 2 (higher is better )
(b) Feynman, ϵ=0 clipped at 300
E2ESR
AIFeynman2 GPlearn
SPL
RSRM
NeurSR
LEE (Ours)
0.8
0.7
0.6
Mean R 2 (higher is better )
0.5
0.4
clipped at 300
SBP-GP (cplx=596)
DSR
0.9
ITEA
101
AI Feynman (N = 116, = 0.01)
SNIP
1.0
E2ESR MDL
(a) Strogatz, ϵ=0
EPLEX
101
GenSR
Strogatz (N = 14, = 0.01)
AFP
GenSR MDL
GPlearn
Operon TPSR EPLEX RAG-SR AFP-FE AFP GP-GOMEA SNIP
1.0
Operon TPSR RAG-SR AFP-FE GP-GOMEA
AIFeynman2
102
0.5
SBP-GP (cplx=851)
Equation complexity (lower is better )
Equation complexity (lower is better )
102 Operon TPSR RAG-SR AFP-FE GP-GOMEA AFP SNIP
clipped at 300
SBP-GP (cplx=489)
Equation complexity (lower is better )
Equation complexity (lower is better )
SBP-GP (cplx=712)
AIFeynman2
102
Operon RAG-SR TPSR GPlearn EPLEX AFP-FE GP-GOMEA AFP E2ESR SNIP MDL
GenSR DSR RSRM
SPL
0.8
0.7
NeurSR
LEE (Ours)
101
0.5
1.0
0.9
(c) Strogatz, ϵ=0.01
0.6
Mean R 2 (higher is better )
0.5
0.4
(d) Feynman, ϵ=0.01
Figure 6: Pareto frontiers on clean and mildly noisy data.
Black-box (N = 63) clipped AIFeynman2 (cplx=2240) at 250
Equation complexity (lower is better )
SBP-GP (cplx=634)
ITEA
102 Operon
E2ESR
EPLEX
GenSR
GP-GOMEA
SNIP
AFP-FE AFP MDL GPlearn
NeurSR
SPL
101
DSR LEE (Ours)
0.8
0.7
0.6
0.5
RSRM
0.4
Mean R 2 (higher is better )
0.3
0.2
0.1
Figure 7: Pareto frontier on the black-box benchmark (test R2 vs. complexity). LEE anchors the low-complexity end; GenSR anchors the high-accuracy end. Several one-shot neural baselines (SNIP, E2ESR, NeurSR) are dominated in both dimensions.
21