nems — Neural-guided EML Symbolic Regression

Learned-policy symbolic regression over the EML operator eml(x, y) = exp(x) − ln(y) (Odrzywolek, arXiv:2603.21852). EML together with the constant 1 is functionally complete for the elementary functions, so any target function can be expressed as a binary tree of eml over the leaves {1, x, y}. The complexity K of a target is the number of RPN tokens in its shortest such tree.

This project replaces brute-force enumeration with a distilled transformer policy + best-first search, and studies where that approach helps, where it plateaus, and why. It is a research log as much as a tool — see FINDINGS.md for the citable results and PROGRESS.md for the session-by-session narrative.

Headline results

K regime	targets	status	how
K ≤ 11	e, exp, ln, 0, Sub, x−1, …	✅ solved	distilled transformer breaks the LLM "encoding ceiling"
K = 15	−1, −x, 1/x (negation/inversion)	✅ solved	corpus augmentation lifts a signature-imbalance wall
K = 17	x·y, x/y, x²	✅ solved	exhaustive K=17 corpus + augmentation
K ≥ 19	Plus, Hypot, Avg, trig, √2	⛔ open	needs K≥19 data or discovery-decoding

The strongest model (v6) solves 12/12 paper-known benchmark targets up to K=17, each in ≈18 nodes of search (essentially greedy decoding of the correct RPN). We also exhaustively enumerated all canonical pure-EML trees up to K=17 (27.7M signatures) and used it to verify — not beat — the paper's Table 4 K bounds for Times/Divide/Sqr (all optimal) and confirm Plus has min K > 17.

Two negative/structural findings worth citing:

• General LLMs (Claude Sonnet/Opus) systematically fail to synthesize valid pure-EML trees beyond K≈13 — a clean encoding-ceiling result with a specific failure-mode taxonomy.
• Naive AlphaZero fails at our compute scale (sparse-reward death spiral); the K-frontier is moved by training-data coverage, not by lookahead.

Repository layout

nems/
├── eml/                       # core: AST, RPN, numeric evaluator, canonical signature, compiler
│   ├── tree.py                #   AST: Const | Var | Eml(left, right) | Op(name, args)
│   ├── rpn.py                 #   postfix (RPN) serialization
│   ├── numeric.py             #   numpy-complex128 evaluator, extended-real semantics
│   ├── canonical.py           #   numeric signature on transcendental probes (+ inf-aware variant)
│   ├── compiler.py            #   parser for EML[a,b] / Wolfram-style strings
│   └── _compiler_vendored.py  #   VA00's eml_compiler_v4.py, vendored (MIT — see THIRD_PARTY_LICENSES)
├── policy/                    # policies plugged into the search loop
│   ├── uniform.py             #   uniform baseline
│   ├── transformer.py         #   signature-conditioned transformer policy (v1..v6)
│   ├── transformer_v4.py      #   policy + value head (AlphaZero experiments)
│   └── transformer_v7.py      #   IO-conditioned policy (v7)
├── search/                    # best-first search: anytime, beam-pruned, length-normalized
├── experiments/               # corpus build, training, eval, MCTS, LLM pilots (see below)
├── benchmark/                 # target sets + CLI runner
└── tests/                     # unit + integration tests

Key experiments/ scripts

script	purpose
build_corpus.py / build_corpus_parallel.py	exhaustive DP enumeration of canonical EML trees (serial / multiprocessing)
transformer_v1.py	signature-conditioned decoder trainer (v1–v6, stratified / sig-stratified sampling)
transformer_v7.py	IO-conditioned two-tower trainer (v7)
make_augmented_corpus.py	inject duplicates of rare target trees to fix signature-imbalance
eval_v5_aug.py, eval_v6.py, eval_v7.py	benchmark evals vs prior checkpoints
mcts.py, mcts_batched.py, alphazero_loop_fast.py	AlphaZero pipeline (parked)
refine.py, sqrt2_opus.py, prompt_inspect.py	LLM-policy field tests (encoding-ceiling study)

Quickstart

python3.12 -m venv .venv
.venv/bin/pip install -e '.[dev]'
.venv/bin/pip install numpy            # required by numeric.py (IEEE-754 inf propagation)

.venv/bin/python -m pytest             # run the test suite
.venv/bin/python -m benchmark.runner --tier EASY --timeout 30 --max-depth 5 --max-size 13

Reproducing the corpus + a trained policy:

# 1. enumerate canonical trees up to K=15 (inf-aware), ~minutes on many cores
.venv/bin/python experiments/build_corpus_parallel.py \
    --max-k 15 --inf-aware --workers 32 --out /tmp/corpus_a2_k15_inf.jsonl

# 2. augment with rare targets, then train a signature-conditioned policy
.venv/bin/python experiments/make_augmented_corpus.py \
    --src /tmp/corpus_a2_k15_inf.jsonl --dst /tmp/corpus_aug.jsonl --n-copies 10000
.venv/bin/python experiments/transformer_v1.py \
    --corpus /tmp/corpus_aug.jsonl --epochs 8 --batch-size 512 --out /tmp/policy.pt

Large artifacts (corpora *.jsonl, checkpoints *.pt, logs) are git-ignored by design — regenerate them with the scripts above.

Design decisions

• Numeric-only canonicalization. Trees are hashed by (round(re·1e10), round(im·1e10)) at a few transcendental probe points (γ, A, …), with no syntactic rewriting (eml/canonical.py). An inf-aware variant keeps ±∞ intermediates as their own equivalence classes, which is required to enumerate paper-known forms that pass through extended-reals (e.g. eml(1, 0) = e − ln 0 = +∞).
• Length-normalized priority. priority = −log_prob / K^0.7 — without this, cumulative log-prob over-prefers short trees and search never reaches depth 6+.
• Two-stage numeric check in search_numeric: cheap MSE on 10 points filters ~99% of junk before a full evaluation.
• Anytime + beam pruning. Best tree so far is retained and returned on timeout; the frontier is trimmed to beam_width when it overflows.

Status & limitations

• The K-frontier reached by a trained policy equals the K range present in its training corpus — no more. Paper-timeout targets (Hypot, Avg, √2, trig) have minimal trees at K ≥ 19, absent from our corpora, and remain unsolved.
• The IO-conditioned model (v7) is a proof-of-concept (trained on K ≤ 13); it is more search-efficient on seen targets and is the most promising direction for unbounded training data, but a full K=17 run needs a precomputed-IO corpus.
• AlphaZero is implemented and architecturally complete but parked — it does not beat supervised training at our compute scale.

Sources & attribution

• Paper: A. Odrzywolek, EML: all elementary functions from a single operator, arXiv:2603.21852.
• Reference implementation: VA00/SymbolicRegressionPackage (MIT). eml/_compiler_vendored.py is vendored verbatim; eml/canonical.py and search/ port the numeric-signature and search ideas. See THIRD_PARTY_LICENSES.

License

MIT — see LICENSE. Vendored third-party code retains its original MIT license; see THIRD_PARTY_LICENSES.