A pretrained transformer's equation sketches recover a third of exact formulas; repair search reaches 85% (2026)

On 17 September 2026, Peiyi Zheng, Yanming Kang, Hans De Sterck and Giang Tran of the University of Waterloo posted MOSAIC-SR: Transformer-Guided Symbolic Regression for Scientific Equation Recovery. The system recovers closed-form equations from numerical data, and it reports the highest exact-recovery rate on every split of its main benchmark and on each of six further datasets, against every method it compares. On the Easy split of SRSD-Feynman, a 120-equation physics benchmark scored by symbolic equivalence, MOSAIC-SR returns the correct formula for 84.7% of problems, against 63.3% for the strongest competitor the authors re-ran and 60.0% for the published PySR result. 1

The more portable finding sits in the ablation. A pretrained transformer that proposes whole equation sketches recovers the exact formula for 30.0–36.7% of 30 Easy equations across three seeds. Adding Monte Carlo tree search to assign variables lifts that to 30.0–43.3%. Letting the proposals run through a local search that edits operators, variables and subtrees while refitting constants takes the same 30 equations to 83.3–90.0%. 1 The learned prior supplies the starting point; explicit repair supplies the program.

MOSAIC-SR arrives from the symbolic-regression community, and its concern is the closed-form expression: it keeps no library, and no solved equation changes what the model proposes next. It earns a place in this tracker for two reasons. It measures a choice the DreamCoder lineage also faces, namely how much of a program to ask a learned model for and how much to leave to a search allowed to rewrite the model's answer. And it makes a distinction the lineage lives by, between fitting the observations and recovering the program: an expression can trace sampled data almost perfectly over a narrow domain and still carry the wrong algebraic form. 1

What the paper changes

Symbolic regression treats each dataset as a search over expression space, scored by error plus a complexity penalty:

e^{⋆} = ar g e \in E min [i = 1 \sum n ℓ (f_{e} (x_{i}), y_{i}) + λ Ω (e)]

The structure is unknown in advance, so the search chooses the operator tree and the numerical constants together, and the paper notes that exact symbolic regression is hard in general. 1

Two families have attacked that problem. Evolutionary and rule-based searchers such as PySR, Operon, GP-GOMEA and AI Feynman grow expression trees from random starts and can revise structure once they see the data, but as the paper puts it, "they do not amortize structure discovery across problems as useful subexpressions must be rediscovered for every new input table." 1 Pretrained neural generators such as NeSymReS, SNIP-SR and GenSR decode a formula from the observations in one pass and are structurally fragile: an early token error changes everything downstream of it, and constant fitting can hide a wrong structure behind a good local fit.

MOSAIC-SR's bet is to split the task by who knows what. The transformer commits only to the parts no amount of data can determine — which operators sit where — and leaves variable identity and constant values to inference.

Pretraining. A Set Transformer encodes the observations as an unordered set of rows put through a signed-logarithmic transform, and a second encoder reads normalized prefix expressions. The two are aligned with a bidirectional InfoNCE objective, so the numeric embedding learns which expression generated the data. Expressions are normalized with SymPy, and affine transformations are deliberately left out of the training targets, because extra scales and offsets can introduce terms absent from the ground truth. The model holds roughly 17.4M parameters over a 45-token vocabulary of twelve operators. 1
Sketch decoding. An autoregressive decoder predicts a sketch conditioned on the numeric representation: the operator skeleton, with every variable and constant leaf replaced by a placeholder token. Auxiliary heads predict coarse structure — the arity and operator class of the root — so the representation carries the shape of the equation before any leaf is resolved.
Inference. Stochastic decoding yields 48 candidate sketches over three sampling passes; duplicates and grammar violations are dropped. A conditional Monte Carlo tree search assigns the decoded variable tokens to observed inputs and fits constants, turning sketches into complete equations. The best 24 by training error seed 24 independent local searches of 300 mutation attempts each. Every path edits an operator, a variable or a small subtree, then periodically refits constants by Levenberg–Marquardt from logarithmically spaced starts, so scales spanning many orders of magnitude stay reachable. Final selection uses validation error, preferring the expression with fewest operators once a candidate falls below a near-exact threshold, followed by simplification and snapping of near-fit constants to physical constants. 1

Three-stage pipeline diagram: contrastive numeric-symbolic pretraining, sketch generation with placeholder leaves, and an inference stage where symbolic sketches feed parallel constant-fitting and structural-repair loops that emit one winning prediction. — Figure 1 from the paper, which is where the design is easiest to see: the placeholders inside the decoded sketches are what later let the search pick variables and constants from the observed data, and the loop on the right alternates constant fitting with structural repair across 24 parallel search paths. 1

That placement is where the paper claims its novelty. Earlier hybrids attach the learned model to a search decision: seeding a genetic population, restricting which tree actions the search may take, or supplying a learned description-length score as the search objective. MOSAIC-SR lets the prior commit to a whole skeleton and then hands that skeleton to a search free to rewrite it. 1

What the evidence measures

The main benchmark is SRSD-Feynman: 120 equations drawn from physics, split 30/40/50 across Easy, Medium and Hard, with sampling ranges set from each equation's physical meaning and logarithmic sampling across orders of magnitude, and 8,000 training, 1,000 validation and 1,000 test points per problem. A variant adds one to three irrelevant input columns. Every method runs under a 3,600-second budget per problem, results are means over five seeds, the test set is used only for final scoring, and the returned expression is re-parsed and re-evaluated outside the method, with invalid expressions and timeouts counted as failures. 1

Two metrics are reported side by side. Accuracy counts problems whose test fit passes a high threshold; symbolic solution rate counts problems where the returned expression is provably equivalent to the ground truth under the symbolic verifier used by SRBench. The two come apart, and that gap is the paper's subject.

Method	Provenance	Easy	Medium	Hard	Easy + dummies	Medium + dummies	Hard + dummies
MOSAIC-SR	re-run here	84.7	55.0	34.8	74.0	54.5	31.2
SR4MDL	re-run here	63.3	43.0	21.2	40.0	14.0	4.8
TPSR	re-run here	47.3	44.0	20.0	32.0	36.0	17.2
PySR	published	60.0	30.0	4.0	20.0	5.0	0.0
uDSR	published	50.0	17.5	4.0	10.0	7.5	0.0
NeSymReS	re-run here	16.7	2.5	3.6	4.0	0.0	0.8

Symbolic solution rate, in percent, on the three SRSD-Feynman splits, clean and with dummy variables. Rows marked "published" are taken from the original SRSD-Feynman study; rows marked "re-run here" were evaluated by these authors under their own pipeline. 1

The dummy-variable columns are the sharpest part of the result. Irrelevant inputs destroy the exact recovery of search-based methods: PySR's Easy rate falls from 60.0% to 20.0%, and the neural generator DGSR drops from 36.7% to 3.3%. MOSAIC-SR loses 10.7, 0.5 and 3.6 points on the three splits, and beats the strongest baseline by 34.0, 18.5 and 14.0 points there. 1

The transfer study shows where the two metrics separate most clearly. On the Nguyen-12 benchmark, MOSAIC-SR and PySR reach identical numerical accuracy, yet MOSAIC-SR recovers the exact expression for 75.0% of problems against PySR's 8.3%; on Strogatz-14 the same comparison runs 64.3% against 42.9%. Across all six additional datasets it has the highest symbolic solution rate, beating the strongest baseline by 12.5 to 33.3 points, and it is 1.3–2.2× faster than PySR and 6.8–24.3× faster than SR4MDL. The directly neural generators stay below 6.7% exact recovery everywhere. 1

Where the gain comes from

The component ablation on 30 clean Easy equations, run over three seeds, separates the three mechanisms.

Configuration	Symbolic rate (seeds 0/1/2)	Accuracy (seeds 0/1/2)
Decoder only	36.7 / 30.0 / 30.0	60.0 / 46.7 / 46.7
+ conditional MCTS	43.3 / 30.0 / 40.0	66.7 / 60.0 / 66.7
Full: 24 paths × 300 mutations	83.3 / 83.3 / 90.0	96.7 / 96.7 / 100.0

Component ablation on the SRSD-Feynman Easy split, percent. "Decoder only" disables both MCTS variable assignment and mutation search. 1

Assigning variables with MCTS helps, and unevenly: it adds 4.0 to 10.0 points of accuracy on every seed, and it rescues exact recovery on one seed, leaving another untouched. Structural repair is the load-bearing part, roughly doubling the symbolic rate on each seed.

Search budget behaves the same way. At 500 mutation steps per path, going from 1 path to 48 raises numerical recovery from 80.0% to 98.9% and exact recovery from 56.7% to 87.8%. The curve flattens late: at 48 paths, 300 steps already reach 98.9% accuracy and 85.6% exactness, and doubling the paths to 48 at 500 steps buys 1.1 points of accuracy and 2.2 points of exactness. 1

Read together, the two ablations say the neural prior in this pipeline is a good initializer and a poor finisher. What converts a promising sketch into the exact law is the part of the system that may keep changing the expression after the network has spoken, which is also the expensive part.

Position in the DreamCoder lineage

Method	Persistent object	Neural search guidance	Abstraction engine	Substrate	Scaling bottleneck
MOSAIC-SR	Fixed 12-operator vocabulary plus a pretrained sketch prior; no cross-task library 1	Pretrained transformer proposes whole sketches; auxiliary heads predict root structure	None; subexpressions found for one equation die with it	Real-valued observations and closed-form equations	Search scale on high-dimensional targets; coverage of the pretraining distribution
DreamCoder	Hierarchical typed λ-calculus DSL plus a trained recognition model 2	Recognition model trained by wake-sleep on replayed and imagined tasks	Bayesian refactoring that extends the language with new abstractions	Synthetic symbolic domains, graphics	Enumerative search explosion
LAPS	DreamCoder library plus a jointly learned search model conditioned on language	Neurally guided search trained jointly with library learning	DreamCoder's library learning, driven by natural-language annotations 3	String editing, image composition, scene reasoning	Needs natural-language annotations; inherits DreamCoder's refactoring cost
Stitch	Compressed corpus of programs plus learned abstractions; no neural component	None; the corpus guides top-down search	Top-down anti-unification over e-graphs, 3–4 orders of magnitude faster than DreamCoder's deductive refactoring 4	Large DSL corpora	Memory on massive corpora; no learned guidance
LILO	Named, documented code library	LLM-guided synthesis prompted with library functions and specifications	Stitch compression plus automatic documentation of new abstractions 5	String editing, scene reasoning, graphics composition	Compression limits on arbitrary code; proposal fidelity

Comparison compiled by this article from the cited papers; MOSAIC-SR makes no claim about library learning.

The transformer occupies the recognition model's slot and plays a different role in it. DreamCoder's recognizer is trained by wake-sleep on the system's own waking solves and self-generated dreams, so the proposal distribution tracks the library as the library grows. MOSAIC-SR's prior is trained offline on synthetic equations and never sees the system solve anything: no solved equation flows back into the weights, and the vocabulary it proposes over is fixed at pretraining time.

The missing half is the library. In the lineage, an abstraction recovered for one task shrinks the search for later tasks, and the recognizer learns over the enlarged vocabulary; the gain compounds. Here each equation is a fresh search over the same twelve operators, and the only thing carried between problems is the pretrained weights. The complaint MOSAIC-SR aims at evolutionary search — that useful subexpressions must be rediscovered for every input table — applies to its own mined structure as well, which does not outlive the problem that produced it.

Where the lineage's inductive bias lives in the DSL, here it lives in the pretraining distribution. That makes the vocabulary a hard ceiling, in the way a DSL is, and the paper reports the familiar consequence at the other end: targets far outside the pretraining distribution need a larger search budget, and recovery gets harder as the number of input variables grows. 1

Verbatim quotes

"However, fast proposal is not the same as equation recovery." — Peiyi Zheng, Yanming Kang, Hans De Sterck and Giang Tran, MOSAIC-SR. 1

"However, they do not amortize structure discovery across problems as useful subexpressions must be rediscovered for every new input table." — the same authors, on search-based symbolic regression. 1

"This supports the use of learned sketches as search priors rather than one-shot predictions." — the same authors, in conclusion. 1

"Despite these gains, recovery becomes more difficult for high dimensional equations because the corresponding search space grows rapidly. Targets far outside the pretraining distribution may also require a larger search budget." — the same authors, on the limits of their own method. 1

The unresolved test

The ablation leaves one question open that the paper's own framing makes sharp. Its prior accumulates nothing, while the lineage's recognizer accumulates everything, and both are trying to reduce the same quantity: search spent per exact program.

A test that would settle the difference: take MOSAIC-SR's pretrained prior, then let it run a wake-sleep loop of its own over its solved equations — mine the subexpressions that recur across the corpus, add those to the operator vocabulary as new primitives, and retrain the prior on the system's own solves and self-generated targets. Score exact recovery against search budget on a held-out split, with the frozen prior as the control. A library that shifts the recovery curve the way the repair search does would show that the repair search was standing in for library learning. A library that instead lowers the price of reaching the same recovery would make the two mechanisms complementary, and would give the lineage its first measurement of how much of DreamCoder's compounding gain a fixed vocabulary and a frozen prior can imitate.

A pretrained transformer's equation sketches recover a third of exact formulas; repair search reaches 85%