SynthStats ARIA Milestone 3

Quantitative evaluation, scaling, and public alpha release

ARIA Brief

Given a dugong's age, what lengths should a statistical model predict? We asked a language model to write a program that answers this question from observed data. The language models in this report write PyMC programs: Python descriptions of probability models.

We ask for a program with AI safety in mind: a statistician who opens the dugong program below sees, in a few lines of Python, that it makes length level off with age and uses the same observation-noise scale at every age, and can object to either before relying on its predictions. Before submitting a program, the AI may also request observations and run code of its own.

The agent chooses a query, receives observations and proposes a statistical program. A program that passes the checks and can be scored receives admission credit and a score for its predictions on unseen observations. We combine these in the reward used to update the agent; one separate training comparison also adds a language-model judge's score.

Quantitative evaluationCopy Anchor Link

Two modelling demonstrationsCopy Anchor Link

All eight dugong-growth programs written by the prompted model Qwen3.6-35B-A3B scored above the floor model. Two scored above the simulator-informed reference by more than two Monte Carlo standard errors (the errors from estimating scores with 200 posterior draws) on this one simulated dataset of 20 held-out observations. Single attempts 2 and 4 exceeded the reference by 1.3617 (Monte Carlo standard error 0.0784) and 1.3319 (0.0787). Loop proposal 4 had a positive gain of 0.0465 (0.1162), within two Monte Carlo standard errors of the reference. All six scored disease-spread programs also beat the floor model, but none beat the reference.

Each comparison here concerns one simulated dataset per environment in BoxingGym, a benchmark of scientific environments: the model receives an English description of the task and 80 observations for fitting, while another 20 responses are held out to score its predictions against a simulator-informed reference program and a floor model that uses one predictive distribution at every input. Predictions are scored by held-out log predictive density, defined as the sum of the log predictive densities assigned to unseen responses (higher values within a task mean greater predictive density at those responses).

For each demonstration, we compared four independent attempts with four programs produced in a critic loop: two proposals in each of two rounds. The critic reviewed the proposals and gave feedback for the next round.

Monte Carlo standard errors (MCSEs) use batch means of each draw's delta-method influence on the total held-out log score, in ten consecutive batches of 20 draws; MCSEs for differences combine the two fits in quadrature under the independent-fit approximation.

Scores below are shown as estimate ± MCSE.

Checks before fittingCopy Anchor Link

In PyMC, a single pm.Potential statement adds an arbitrary term to a model's log density, so a program that uses one can raise a reward computed from its likelihood without defining the intended probability model. The first check below rejects such a model before fitting, along with any model using the other constructions in its row; the last two test whether a program conditions on the observations that we supplied:

Programs and predictionsCopy Anchor Link

Dugong growth: circles mark the observations supplied to the model and crosses mark the held-out responses, with predictive medians (lines) and central 90% intervals of predictive draws (bands) for the reference model and single attempt 1. Single attempt 2 is shown in the listing below.

alpha = pm.Normal("alpha", mu=2.0, sigma=0.2)
beta = pm.Normal("beta", mu=1.5, sigma=0.5)
lam = pm.Normal("lam", mu=0.4, sigma=0.5)
mu = alpha - beta * pm.math.abs(lam) ** age_data
pm.Normal("length", mu=mu, sigma=0.25, observed=y)

L_inf = pm.HalfNormal("L_inf", sigma=2.0)
k = pm.HalfNormal("k", sigma=1.0)
t0 = pm.Normal("t0", mu=-0.5, sigma=0.5)
mu = L_inf * (1 - pm.math.exp(-k * (age - t0)))
sigma = pm.HalfNormal("sigma", sigma=0.2)
length_obs = pm.Normal("length_obs", mu=mu, sigma=sigma, observed=length)

import pymc as pm
import numpy as np

def model(data):
    age = data['age']
    length = data['length']

    with pm.Model() as dugong_model:
        # Priors for the von Bertalanffy growth parameters
        # L_infinity: asymptotic length
        L_inf = pm.HalfNormal('L_inf', sigma=2)

        # k: growth rate
        k = pm.HalfNormal('k', sigma=1)

        # Observation noise
        sigma = pm.HalfNormal('sigma', sigma=0.2)

        # Expected length
        # Using the von Bertalanffy growth function: L(t) = L_inf * (1 - exp(-k * t))
        mu = L_inf * (1 - pm.math.exp(-k * age))

        # Likelihood
        likelihood = pm.Normal('likelihood', mu=mu, sigma=sigma, observed=length)

    return dugong_model

The critic reviewed the first-round dugong programs and recommended "Exploring alternative likelihoods such as Log-Normal if the data is strictly positive and skewed." The resulting Log-Normal proposal had the lowest dugong log predictive density among the model-written programs.

Disease spread: the programs predict infected counts from time in a population of 50, and the plot shows the predictive medians and central 90% intervals for the reference model and single attempt 1. Single attempt 2 is shown in the listing below.

with pm.Model() as result:
        t = pm.Data("t", data["Time"])
        y = pm.Data("y", data["Infected_Count"])
        theta = pm.TruncatedNormal("theta", mu=1, sigma=1, lower=0, upper=2)
        pm.Binomial("Infected_Count", n=50, p=1-pm.math.exp(-theta*t), observed=y)

times = np.array(data['Time'])
infected_counts = np.array(data['Infected_Count'])
N = 50  # Total population size
with pm.Model() as pmodel:
    alpha = pm.Normal('alpha', mu=0, sigma=10)
    beta = pm.Normal('beta', mu=0, sigma=10)
    logit_p = alpha + beta * times
    p = pm.math.sigmoid(logit_p)
    infected_observed = pm.Binomial('infected_observed', n=N, p=p, observed=infected_counts)

import pymc as pm
import numpy as np

def model(data):
    """
    Defines a PyMC model for predicting infected counts in a population of 50.

    Args:
    data (dict): A dictionary containing 'Time' and 'Infected_Count' arrays.

    Returns:
    pm.Model: The defined PyMC model.
    """
    time = np.array(data['Time'])
    infected_counts = np.array(data['Infected_Count'])
    N = 50  # Total population size

    with pm.Model() as m:
        # Priors for the logistic regression coefficients
        # alpha: intercept, beta: slope
        alpha = pm.Normal('alpha', mu=0, sigma=5)
        beta = pm.Normal('beta', mu=0, sigma=5)

        # Compute the probability of infection at each time point
        # Using the logistic function: p = 1 / (1 + exp(-(alpha + beta * t)))
        logit_p = alpha + beta * time
        p = pm.math.invlogit(logit_p)

        # Observed likelihood: Binomial distribution
        # The number of infected individuals is binomially distributed with N trials and probability p
        infected_observed = pm.Binomial(
            'infected_observed',
            n=N,
            p=p,
            observed=infected_counts
        )

    return m

Inspect every program, score and cost, or follow one dugong prediction through its score.

All ten programs in each environmentCopy Anchor Link

Full-screen interactive predictions

Full-screen interactive predictions

Reward designCopy Anchor Link

Asking the same question twice can produce different programs: the AI has a distribution over dugong programs. Each program, in turn, describes a distribution of lengths at a given age, possibly through unnormalised weights that need not sum to one. Our research goal is to learn a distribution over plausible programs, so a scientist can compare different modelling assumptions about the same observations.

For a dugong program, we measure predictive fit on the responses marked by crosses in the figure, which were withheld from the AI. We also give credit for passing the modelling checks, and call a program that passes them admitted. A scorable admitted program earns admission credit plus a predictive-fit term, q, which compares its held-out expected log predictive density (ELPD) with the task's baseline ELPD:

Admission rulesCopy Anchor Link

Matched training comparisonsCopy Anchor Link

GLM-5.3-Flash is the larger model in these comparisons, a tool-using agent trained by reinforcement learning through Tinker, a hosted training service. An attempt, which we call an episode, can include data requests and a program submission.

The fit-only and judge-assisted comparisons sampled their untrained episodes separately and used different verifier versions.

For any one program, admission credit is larger than any gain or loss from q.

Training methods for the main comparisonsCopy Anchor Link

Budgeted observation acquisition and three arms of eight updatesCopy Anchor Link

The small agent is Qwen3.5-0.8B with tools. We first trained it to imitate positively rewarded episodes from other language models, which we call teachers. After this supervised warm start, we trained it with Subtrajectory Balance (SubTB) in two arms, at learning rates of 1 × 10⁻⁵ and 1 × 10⁻⁶, and with group relative policy optimization (GRPO) in a third, each over eight updates on 220 training packages in 81 scenario groups.

In each episode, we gave the agent a query budget to spend on selecting observations (in the common buzzard task quoted below, survey routes) before writing a program. During training, an episode's reward adds an admission component, for passing the program checks, to a fit term that measures how well the program predicts relative to the task's baseline. All three arms trained on the agent's sequences of tool actions and their rewards.

The Swiss breeding-bird survey records whether Common Buzzard was detected on the first visit of each route. Site covariates describe elevation in metres, forest cover in percent, and route length in kilometres. The analyst wants to learn how these covariates relate to first-visit detection and predict whether the species is detected on other routes. Replicate structure and imperfect detection are not modelled.

Held-out comparison with the warm startCopy Anchor Link

Predicting coral sizeCopy Anchor Link

In one GLM-5.3-Flash training task, we wanted to predict the maximum diameter of brain coral at simulated reef sites from environmental observations, in a setting where the AI could request measurements and run code before submitting a statistical program to predict other sites. The coral prompt also contained outcome summaries computed from all 48 candidate sites.

The submitted coral program specifies its regression and its likelihood in ordinary PyMC.

Scale of the task collection and the training workloadCopy Anchor Link

The coral task and the common buzzard survey belong to a collection of 323 PyMC task packages that covers both simulated processes and species surveys, and when we split the collection, we kept related tasks together in scenario groups: 220 packages in 81 groups trained the models, and the other 103 packages in 39 groups were held out for evaluation. Each task comes as a package containing the prompt and data, rows withheld for scoring, and the scores of a reference program:

On the held-out tasks, we compared the larger comparator, GLM-5.3-Flash, and the small agent, each against the policy it started from.

The larger comparator trained on its own attempts, such as the coral episode. For the small agent's warm start, we used positively rewarded teacher episodes, including the tool actions that led to a submitted program:

We also collected programs written by teacher language models, in two separate collections:

Supporting measurementsCopy Anchor Link

Public alpha releaseCopy Anchor Link

The release provides executable modelling examples with their observations and predictions, evaluation tools, and shared task, trace and reward schemas. The reference dugong program is one of its Python files: the three commands below download the observations and fitting script, fit the reference to the visible rows and compute its held-out score.

Fit the reference programCopy Anchor Link

The fitting script selects the visible rows of the dugong observations (those supplied to the model) and fits the reference program to them, then prints the first posterior draw, the mean of its posterior parameters and its held-out score. The demonstration was scored under PyMC 5.21.1 and ArviZ 0.23.4. In an empty directory, the commands below use curl and uv to download the observations and the script, and then run the fit.

curl -fsS https://synthstats.org/data/dugongs/dugongs-observations.csv -o observations.csv
curl -fsS https://synthstats.org/examples/fit-dugong.py -o fit-dugong.py
uv run fit-dugong.py

UAI 2026 paper on likelihood hackingCopy Anchor Link

In Likelihood hacking in probabilistic program synthesis, published in the UAI 2026 proceedings, we trained a language model to write PyMC programs with the log marginal likelihood as its reward. In SynthStats, the strict PyMC checks before fitting refuse pm.Potential, with which a program can add an arbitrary term to its log density and so inflate a likelihood reward.

TeamCopy Anchor Link

SynthStats is the work of Jacek Karwowski, Younesse Kaddar, Daniella Ye, Esmeralda S. Whitammer and Sam Staton.

Research results · Supporting results