⚡Flash-ANSR:
Fast Amortized Neural Symbolic Regression

Flash-ANSR is a library for amortized neural symbolic regression: load a pretrained model, call fit(X, y), and recover a symbolic expression for your tabular data, or train your own model. It is built for fast, ready-to-use inference.

Publications

• Saegert & Köthe 2026, Breaking the Simplification Bottleneck in Amortized Neural Symbolic Regression (ICML 2026) https://arxiv.org/abs/2602.08885

Usage

Requires Python >= 3.12.

pip install flash-ansr
flash_ansr install psaegert/flash-ansr-v25.0-T7-3M   # the reference checkpoint (see "Models")

import torch
import numpy as np
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

# Import flash_ansr
from flash_ansr import (
  FlashANSR,
  SoftmaxSamplingConfig,
)

# The installed checkpoint directory
from flash_ansr import get_path
CHECKPOINT = get_path("models", "psaegert/flash-ansr-v25.0-T7-3M")

# Load the model (KV-cache, auto-batching and static decoding are on by default; see "Inference speed")
model = FlashANSR.load(
  directory=CHECKPOINT,
  generation_config=SoftmaxSamplingConfig(choices=1024),
  # Candidate ranking (default): log10(FVU) + 1e-2 per bit of the refined expression's description
  # length. Alternatives: ranking_mode="weighted" with ranking_weights={"n_nodes": 0.05}, or
  # ranking_mode="pareto" with ranking_metrics=("fvu", "n_nodes").
  ranking_mode="mdl",
).to(device)

# Define data: a small synthetic example, y = 2 * x + sin(3 * x)
X = np.linspace(-5, 5, 100).reshape(-1, 1)
y = 2 * X[:, 0] + np.sin(3 * X[:, 0])

# Fit the model to the data
model.fit(X, y, verbose=True)

# Show the best expression
print(model.get_expression())

# Predict with the best expression
y_pred = model.predict(X)

Get all candidates at once (infer): instead of fit + read-back, call model.infer(X, y), which returns an InferenceResult carrying the best Candidate, the score-sorted refined candidates, and the full CandidateLedger (the generation pool joined with the refined survivors, each classified FIT_OK / FIT_FAILED / INVALID).

result = model.infer(X, y)
print(result.best.expression_infix, result.best.fvu)  # best refined candidate
for c in result.candidates:                            # score-sorted survivors
    print(c.score, c.expression_infix)
print(len(result.ledger))                              # all candidates considered

Explore more in the Demo Notebook.

Train your own: see the training guide.

Models

flash_ansr install psaegert/flash-ansr-v25.0-T7-3M

Every catalog that srbf evaluates on is held out of the training data by canonical form (6,660 expressions across 29 catalogs).

Inference speed

Several inference-speed features are enabled by default and designed to be quality-neutral, so the quickstart above already runs in the fast regime. The speed-relevant settings live on the generation config:

from flash_ansr import SoftmaxSamplingConfig

config = SoftmaxSamplingConfig(
  choices=1024,        # number of candidate expressions to sample
  use_cache=True,      # KV cache (default)
  batch_size='auto',   # candidate-budget-adaptive chunking (default)
  static_decode=None,  # auto for capable models (default)
)

Constant refinement runs in parallel; control it via FlashANSR.load(..., refiner_workers=N, persistent_refine_pool=True). By default (refiner_workers=None) the pool uses every available CPU core, which oversubscribes shared machines; pass an explicit integer to cap it (0 disables multiprocessing).

To opt out of these defaults:

SoftmaxSamplingConfig(choices=1024, use_cache=False, batch_size=128, static_decode=False)

Candidate ranking. Three modes, one sort: ranking_mode="mdl" (default; log10(FVU) plus mdl_strength decades per bit of the refined expression's description length), "weighted" (ranking_weights over n_nodes, n_constants, n_constant_placeholders, n_typed_literals, mdl, neg_log_prob) and "pareto" (the non-dominated front over ranking_metrics, ordered by ranking_tie_break). Each knob belongs to one mode and raises under another. The pre-0.14 ranking is ranking_mode="weighted", ranking_weights={"n_nodes": 0.05}.

Overview

SRSD/FastSRB Results Results on the SRSD/FastSRB benchmark [Matsubara et al. 2022], [Martinek 2025] Left: Validation Numeric Recovery Rate (vNRR) as a function of inference time (log scale). FLASH-ANSR models (shades of blue) scale monotonically with compute, with the 120M model partially surpassing the PySR baseline (red). Baselines NeSymReS [Biggio et al. 2021] and E2E [Kamienny et al. 2022] fail to generalize to the benchmark. Right: Expression Length Ratio (predicted vs ground truth) versus compute. We observe a parsimony inversion: while PySR [Cranmer 2023] increases complexity to minimize error over time, FLASH-ANSR converges toward simpler, more canonical expressions as the sampling budget increases. Shaded regions denote 95% confidence intervals.

Training The Flash-ANSR training pipeline. Following the established standard encoder-decoder paradigm, our framework integrates SimpliPy (top center) into the loop for synchronous simplification of on-the-fly generated training expressions.

Architecture Flash-ANSR model architecture. The Set Transformer [Lee et al. 2019] encoder ingests a variable-sized set of input-output pairs and produces a fixed-size latent representation via Induced Set Attention Blocks (ISAB) and Set Attention Blocks (SAB). The Transformer decoder [Vaswani et al. 2017], [Xiong et al. 2020] autoregressively generates a symbolic expression token-by-token, attending to the encoded dataset at each step.

Related projects

• SimpliPy: the expression simplification engine integrated into the Flash-ANSR training loop.
• symbolic-data: the model-agnostic symbolic-regression data layer (catalogs, ProblemSource, holdouts) that feeds Flash-ANSR training. It is an unconditional runtime dependency and the backbone of the training loop.
• srbf: the companion symbolic-regression evaluation and benchmarking framework (engine, model adapters, benchmarks, metrics), developed alongside Flash-ANSR.

Citation

@inproceedings{saegert2026breakingsimplificationbottleneckamortized,
  title   = {Breaking the Simplification Bottleneck in Amortized Neural Symbolic Regression},
  author  = {Paul Saegert and Ullrich Köthe},
  booktitle = {Proceedings of the 43rd International Conference on Machine Learning (ICML)},
  year    = {2026},
  eprint  = {2602.08885},
  archivePrefix =  {arXiv},
  primaryClass  = {cs.LG},
  url     = {https://arxiv.org/abs/2602.08885},
}

% Optionally
@mastersthesis{flash-ansr2024-thesis,
  author  = {Paul Saegert},
  title   = {Flash Amortized Neural Symbolic Regression},
  school  = {Heidelberg University},
  year    = {2025},
  url     = {https://github.com/psaegert/flash-ansr-thesis}
}
@software{flash-ansr2024,
  author  = {Paul Saegert},
  title   = {Flash Amortized Neural Symbolic Regression},
  year    = {2024},
  publisher   = {GitHub},
  version = {0.14.0},
  url     = {https://github.com/psaegert/flash-ansr}
}