flash-ansr-v25.0-T7-3M

A 3.5M-parameter Flash-ANSR model for amortized neural symbolic regression: it reads a set of input-output pairs and proposes symbolic expressions for them, with numeric constants predicted at full precision and refined by a downstream optimizer. The reference checkpoint for flash-ansr 0.14.

Usage

pip install "flash-ansr>=0.14,<0.15"
flash_ansr install psaegert/flash-ansr-v25.0-T7-3M

import numpy as np
import torch
from flash_ansr import FlashANSR, SoftmaxSamplingConfig, get_path

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model = FlashANSR.load(
    directory=get_path("models", "psaegert/flash-ansr-v25.0-T7-3M"),
    generation_config=SoftmaxSamplingConfig(choices=1024),
).to(device)

X = np.linspace(-5, 5, 100).reshape(-1, 1)
y = 2 * X[:, 0] + np.sin(3 * X[:, 0])
model.fit(X, y)
print(model.get_expression())

Every call to fit / infer samples candidate expressions, fits their constants and ranks the refined candidates; model.infer(X, y) returns all of them at once. See the documentation for the auxiliary verbs (predict_y, predict_complexity, predict_constants, score_outliers) this checkpoint supports.

Model

Set Transformer encoder (192-d, 3 heads, 1 ISAB + 1 SAB, 128 inducing points and seeds) over up to 18 input variables, numeric inputs pre-encoded as binary64 bit patterns; Transformer decoder (192-d, 3 layers, 3 heads, RoPE self-attention) over a vocabulary in which constants are spelled as 8 byte tokens. The checkpoint carries the optional-conditioning null memory and the per-point outlier head. The configuration is in model.yaml; the tokenizer in tokenizer.yaml.

Training

Trained for 1,000,000 steps at batch size 128 on expressions sampled on the fly from symbolic-data with the prior in catalog_train.yaml and the task mixture in dataset_train.yaml (noise mixture with outliers, complexity conditioning and prediction, held-out point prediction, constant masking and infilling), and simplified in the training loop by SimpliPy with the acj-5-4-llm rule set. Every catalog evaluated by srbf is held out by canonical form: 6,660 expressions across 29 catalogs. The exact run recipe is configs/v25.0-T7 in the flash-ansr repository; train.yaml, dataset_train.yaml, dataset_val.yaml and catalog_train.yaml here are the resolved copies.

Software: flash-ansr 0.14.0, simplipy 0.14.6, symbolic-data 0.18.0.

Publications

• Saegert & Köthe 2026, Breaking the Simplification Bottleneck in Amortized Neural Symbolic Regression (ICML 2026), arXiv:2602.08885

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},
}