flash-ansr-v25.0-T8-120M

flash-ansr-v25.0-T8-120M

A 123.6M-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. This checkpoint is the 120M point of the v25.0-T8 series (3M, 20M and 120M parameters, one recipe, one data prior).

Usage

pip install "flash-ansr>=0.17"
flash_ansr install psaegert/flash-ansr-v25.0-T8-120M

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-T8-120M"),
    generation_config=SoftmaxSamplingConfig(draws=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())

fit draws candidate expressions, fits their constants and ranks the refined candidates; it returns a FitResult (every candidate, the ledger, predict(X, rank), rerank(...)) and keeps it as model.result_. draws= on the call sets the search budget per problem. See the documentation for the full surface.

Model

Set Transformer encoder (640-d, 10 heads, 5 ISAB + 5 SAB, 128 inducing points and 128 seeds) over up to 18 input variables, numeric inputs pre-encoded as binary64 bit patterns; Transformer decoder (640-d, 10 layers, 10 heads, RoPE self-attention, RMS norms) over a vocabulary in which constants are spelled as 8 byte tokens, with a float32 output head under mixed precision. 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,500,000 steps at batch size 128 with AdaMuon (peak learning rate 1e-4 after a 15,000-step warm-up, constant to 1,200,000 steps, linear decay to zero) and a z-loss of 0.0001, 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-T8-120M 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 (commit f6d918e), 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},
}