SymGPT: Training a Transformer to Do Symbolic Regression

The genetic-programming and SISSO engines both start every problem from nothing — a fresh evolutionary search, a fresh feature-space sweep — for each new table of numbers. A transformer changes the economics. Train it once on millions of synthetic (data, formula) pairs, and it learns the prior that maps a cloud of points to the equation most likely to have produced them; a brand-new problem is then answered in a single forward pass. This is symbolic regression as translation — points in, formula out — and it is the idea behind the recent "neural symbolic regression that scales" line of work. This page builds a small one from scratch: a two-million-parameter decoder trained on a stream of random expressions, which then recovers formulas it has never seen, including ones it has to algebraically simplify to get right.

Trees as sequences, formulas as language

To let a language model emit a formula, write the formula as a sentence. Every expression is a tree, and a tree has a unique prefix (Polish) serialization — a preorder walk that is parenthesis-free and reconstructs the tree exactly. $\sin (Cx)$ becomes the token stream sin mul C x; $(e^{-x}{)}^2$ becomes sq exp neg x. The vocabulary is tiny — a dozen operators plus the variable and a single placeholder — so the model is doing next-token prediction over a fourteen-word language whose sentences are equations.

# A formula is a tree; serialize it in PREFIX (Polish) order -> a token sequence,
# parenthesis-free and bijective with the tree, exactly what a decoder emits.
#   vocab (14):  <pad> <bos> <eos>  add sub mul div neg sq sin cos exp  x  C

#   f(x) = sin(C * x)     ->  prefix:  sin mul C x
#   f(x) = (exp(-x))^2    ->  prefix:  sq exp neg x
#   f(x) = x + cos(x)     ->  prefix:  add x cos x

# Every constant collapses to ONE placeholder token 'C'. The network predicts the
# SKELETON (the structure); the actual numbers are fit afterward. So training data is
# infinite and free: sample a random tree, evaluate it at 30 points, emit (points, prefix).

The one move that makes this tractable: every numerical constant collapses to a single placeholder token $C$. The network never has to spell out $3.14159$ digit by digit — it predicts the skeleton, the bare structure of the formula, and the real constants are fit numerically afterward. That decouples the hard combinatorial problem (what is the shape?) from the easy continuous one (what are the numbers?), and it means the training data is effectively infinite: sample a random tree, draw random constants, evaluate it at thirty points, and emit the (points, skeleton) pair. No labels to collect, no dataset to curate — the supervision is generated on the fly.

The model: read the points, write the formula

The architecture is a plain decoder-only transformer with one twist for the input. The thirty $(x_i,y_i)$ points are a set, not a sequence — their order carries no information — so each point is embedded by a small MLP and fed as context with no positional encoding, leaving the model permutation-invariant over the samples. The expression tokens that follow do get positional encodings and a causal mask, so generation is autoregressive: each token attends to all of the point-context and to the tokens already written, and predicts the next one.

class SymGPT(nn.Module):        # a small decoder-only transformer (2.27M params)
    def forward(self, points, tokens):
        # each (x_i, y_i) -> MLP -> d_model, used as PERMUTATION-INVARIANT context
        # (no positional encoding on the point set -- order of the samples is meaningless)
        pe = self.point_mlp(points) + self.type_emb[POINT]        # (B, n_points, d)
        te = self.tok_emb(tokens) + self.pos_emb + self.type_emb[TOKEN]   # (B, L, d)
        x  = cat([pe, te], dim=1)
        # attention mask: points attend among themselves; each expression token attends
        # to ALL points + causally to earlier tokens (it may not peek at the answer ahead)
        x  = self.transformer(x, mask=self._mask(n_points, L))
        return self.head(x)[:, n_points:]        # next-token logits over the 14-word vocab

Trained on a stream of fresh random expressions — 2.27 million parameters, ten thousand steps, about twenty-six minutes on a laptop GPU — the loss falls and token accuracy climbs past one in two. Token accuracy looks modest, but it understates the model badly, because getting a formula right does not require reproducing the training tree token for token. The real test is whether the predicted skeleton, once its constants are fit, reproduces the data.

Does it recover formulas it has never seen?

On held-out expressions the model has never encountered, the median $R^2$ after constant-fitting is $1.0$, with clean recovery on sixteen of twenty-five. And the way it succeeds is the interesting part — it is not memorizing trees, it is finding functions. Handed data from $\cos (x)\cdot x/x$ it returns $\cos (x)$; handed $x-(x^2-x^2)$ it returns $x$. It has learned the function behind the points, not the particular algebra that generated them, which is exactly what recovering a law means.

held-out recovery  (fresh formulas, never seen in training):
  true:  mul div cos x x x       ->  pred:  cos(x)        R2 = 1.0000   <- cancelled  x/x
  true:  sub x sub mul x x sq x  ->  pred:  x             R2 = 1.0000   <- saw  x^2 - x^2 = 0
  true:  sq exp neg x            ->  pred:  exp(-x)^2     R2 = 1.0000
  true:  sin add sub x x sq x    ->  pred:  sin(x^2)      R2 = 1.0000   <- dropped  x - x = 0
  true:  exp sin x               ->  pred:  exp(sin(x))   R2 = 1.0000
  ...
  median R2 = 1.0000        recovered (R2 > 0.999):  16 / 25

benchmark vs the genetic-programming engine   (20 in-distribution problems):
  SymGPT  greedy (one forward pass)   recovered   3/20        20 ms / problem
  SymGPT  beam-14                     recovered  11/20      1075 ms / problem
  GP      pop=500  gens=40            recovered  13/20      1436 ms / problem

Amortization versus search: the benchmark

Run the transformer head-to-head against the from-scratch genetic-programming engine on the same twenty problems and the result is a draw, not a rout. The GP search is a strong baseline: on these small, in-distribution problems it edges recovery, thirteen against eleven. The transformer matches it through beam search at a comparable wall-clock, and its single greedy forward pass is sixty-five times faster but trades away most of the recovery. At this toy scale, neither dominates. What the numbers do not show — but the timeouts behind them do — is the axis where the transformer wins: its inference cost is flat, independent of how hard the formula is, while a search's cost balloons with problem difficulty. Push the GP toward harder targets and it needs far larger populations and generations to keep up, while the trained model answers in the same fixed pass. Amortization is the whole point: you pay the training cost once and buy constant-time inference forever.

Better together: cross-seeding the search

The benchmark frames the two engines as rivals, but they are more useful as collaborators, and the fix costs nothing at training time. The transformer produces, in one pass, a few high-quality skeleton guesses; the genetic search produces a whole population of expressions it evaluates and mostly discards. Feed each to the other — seed the genetic population with the transformer's beam predictions so the search starts near good structure instead of at random, and pool the search's Pareto front back into the transformer's candidate set so the best expression overall is returned. It is pure inference-time composition, neural-guided genetic programming in the lineage of Deep Symbolic Optimization.

The warm start earns its keep exactly when search is expensive. Starve the genetic engine to a small population and a few generations and it recovers 13 of 20 held-out formulas alone; seed it with the transformer's guesses and it recovers 15, while the transformer by itself — fast but limited — gets 8. One case shows the mechanism cleanly: on a formula where the transformer stalls at $R^2=0.90$ and the starved search stalls at $0.84$, the search refining the transformer's imperfect guess lands exactly. Structure from the transformer, refinement from the search, and the hybrid beats either alone. The benefit fades once you can afford a large search — which finds the answer unaided — so the composition matters most under a tight compute budget, which is the regime that motivated the amortized model to begin with.

What makes it hard, and what real systems add

• The grammar must contain the answer. This model was trained on a nine-operator grammar with no $\log$ or ​; hand it data from a logarithm and it cannot emit the token, so it approximates and fails cleanly. As with every symbolic method, the operator set is the hypothesis class — production models (Kamienny, Biggio) train on far richer grammars and multiple input variables.
• Generalization is in-distribution. The model is only as good as the distribution it was trained on. Expressions deeper or stranger than its training rungs are out-of-distribution, and recovery degrades — the failure mode a search does not have, since a search explores rather than recalls.
• Beam width trades speed for recovery. Greedy decoding commits to the highest-probability token at every step and often paints itself into a corner; a beam keeps several candidate skeletons alive and fits constants to each. That one knob moved recovery from three-in-twenty to eleven-in-twenty here.
• Skeleton then constants. The split — transformer predicts structure, least-squares fits the numbers — is what real systems (SymbolicGPT, NeSymReS) do too. The alternative, generating digits of constants as tokens, is far harder to train and less accurate than handing the numbers to an optimizer.

The point, for a computational scientist

The genetic search and SISSO give you a formula for one table; the transformer gives you a function from tables to formulas. That is what you want when the fitting is inside a loop — screening thousands of candidate materials, tagging every trajectory in a simulation with the closed form of its transient, distilling a controller's learned policy on the fly — anywhere the per-problem cost of a fresh search is the bottleneck. It is the amortized member of the symbolic-regression canon: search, compressed sensing, and now learning, three routes from a column of numbers to an equation, each strongest where the others are weak.

Reproduce it

The engine lives in scripts/symgpt/: data.py generates the synthetic (points, prefix) stream and proves the tree↔sequence round-trip exact; model.py is the decoder with its point-embedding prefix and hand-built attention mask; train.py streams fresh batches to a saved checkpoint; infer.py beam-decodes a skeleton and fits its constants by least squares; and benchmark.py runs the head-to-head against the genetic-programming engine. Retrain on a wider grammar, add input variables, or swap greedy for a wider beam and watch recovery move.