Additive Symbolic Regression

Additive Symbolic Regression#

Additive symbolic regression fits a model as a sum of small symbolic expressions, \(f(x) = c + \sum_k \eta_k\, g_k(x)\), where each \(g_k\) is discovered by the existing JAXSR machinery. This is analogous to gradient boosting, except each weak learner is an interpretable symbolic expression rather than a decision tree.

This notebook shows the StagewiseSymbolicRegressor: it repeatedly fits a small expression to the current residual, then (optionally) refits all linear coefficients by least squares.

import numpy as np

from jaxsr.additive import StagewiseSymbolicRegressor

# Additive target: y = 2*x0 + 0.5*x1^2 + noise
rng = np.random.default_rng(0)
X = rng.uniform(-2, 2, size=(300, 2))
y = 2.0 * X[:, 0] + 0.5 * X[:, 1] ** 2 + 0.1 * rng.normal(size=300)

# hold out a test set
n_train = 200
X_train, y_train = X[:n_train], y[:n_train]
X_test, y_test = X[n_train:], y[n_train:]
X_train.shape, X_test.shape

((200, 2), (100, 2))

Fit a stagewise additive model#

Keep max_complexity small so each stage is a simple expression and the ensemble accumulates many interpretable terms. With refit_coefficients=True, all linear weights are re-solved by least squares after each stage.

model = StagewiseSymbolicRegressor(
    n_terms=5,
    learning_rate=0.2,
    max_complexity=4,
    refit_coefficients=True,
)
model.fit(X_train, y_train)
print(model)

StagewiseSymbolicRegressor(
    intercept = 1.075
    terms =
        + 1.001 * (y = 1.998*x0 - 1.07 + 0.5015*x1^2)
        + 2.31 * (y = - 0.003903*x0^2)
        + 1.865 * (y = 0.002939*x1)
        + 1.124 * (y = 0.001106*x0*x1)
        + 6.931 * (y = - 0.0001612*x1^3)
)

Inspect the learned ensemble#

print('intercept:', model.intercept_)
print('coefficients:', model.coefficients_)
print()
for i, expr in enumerate(model.expressions_):
    print(f'term {i}: {expr}')

intercept: 1.0750973224639893
coefficients: [1.0005701780319214, 2.3100991249084473, 1.8648916482925415, 1.1235284805297852, 6.930931568145752]

term 0: y = 1.998*x0 - 1.07 + 0.5015*x1^2
term 1: y = - 0.003903*x0^2
term 2: y = 0.002939*x1
term 3: y = 0.001106*x0*x1
term 4: y = - 0.0001612*x1^3

print('train R^2:', model.score(X_train, y_train))
print('test  R^2:', model.score(X_test, y_test))

train R^2: 0.9982670545578003

test  R^2: 0.9988469481468201

Combined symbolic expression#

to_expression() collapses the ensemble into a single simplified SymPy expression.

model.to_expression()
\[\displaystyle 0.00124225317639448 x_{0} x_{1} + 1.99923764521122 x_{0} - 0.0090172041649631 x_{0}^{2.0} + 0.00548121186382547 x_{1} + 0.501793767439942 x_{1}^{2.0} - 0.00111730957318264 x_{1}^{3.0} + 0.00453644099624739\]

Why additive? Many small terms vs one big expression#

When the signal is a sum of several simple effects, a tiny per-stage budget still recovers it because the ensemble accumulates terms across stages — where a single expression with the same small budget cannot.

from jaxsr import fit_symbolic

rng = np.random.default_rng(2)
Xm = rng.uniform(-1.5, 1.5, size=(500, 4))
ym = (
    1.5 * Xm[:, 0]
    + 2.0 * Xm[:, 1] ** 2
    - 1.0 * Xm[:, 2] ** 3
    + 0.8 * Xm[:, 3]
    + 0.2 * rng.normal(size=500)
)
Xtr, ytr, Xte, yte = Xm[:300], ym[:300], Xm[300:], ym[300:]

stagewise = StagewiseSymbolicRegressor(n_terms=10, max_complexity=1).fit(Xtr, ytr)
single = fit_symbolic(Xtr, ytr, max_terms=1, include_transcendental=False)

print('stagewise (max_complexity=1, 10 stages) test R^2:', round(float(stagewise.score(Xte, yte)), 4))
print('single expression (max_terms=1)         test R^2:', round(float(single.score(Xte, yte)), 4))

stagewise (max_complexity=1, 10 stages) test R^2: 0.9931
single expression (max_terms=1)         test R^2: 0.3119

Early stopping and saving#

Early stopping holds out a validation split and stops once it stops improving. Models serialize to JSON (they are not picklable because the basis functions use closures).

import tempfile, os

es = StagewiseSymbolicRegressor(
    n_terms=20, max_complexity=3,
    early_stopping=True, validation_fraction=0.25, patience=3, random_state=0,
).fit(X_train, y_train)
print('terms used after early stopping:', es.n_terms_)

path = os.path.join(tempfile.mkdtemp(), 'additive_model.json')
es.save(path)
loaded = StagewiseSymbolicRegressor.load(path)
print('reloaded test R^2:', round(float(loaded.score(X_test, y_test)), 4))
os.remove(path)

terms used after early stopping: 2
reloaded test R^2: 0.9989

Robust and quantile regression (beyond ordinary SR)#

Because each term is fit by gradient boosting (fitting the negative gradient), additive symbolic regression can target losses that least-squares selection cannot. Use refit_coefficients=False for non-squared losses.

First, robustness to outliers with the Huber and absolute-error losses:

# Contaminate 8% of the targets with heavy outliers
rng = np.random.default_rng(1)
Xr = rng.uniform(-2, 2, size=(500, 2))
clean = 2.0 * Xr[:, 0] + 0.5 * Xr[:, 1] ** 2
yr = clean + rng.normal(0, 0.1, 500)
yr[rng.choice(500, 40, replace=False)] += 30.0

for loss in ['squared_error', 'huber', 'absolute_error']:
    m = StagewiseSymbolicRegressor(
        n_terms=8, max_complexity=4, learning_rate=0.5,
        loss=loss, refit_coefficients=False,
    ).fit(Xr, yr)
    mae = float(np.mean(np.abs(np.array(m.predict(Xr)) - clean)))
    print(f'{loss:15s}  MAE vs clean signal = {mae:.3f}')

squared_error    MAE vs clean signal = 2.764

huber            MAE vs clean signal = 0.144

absolute_error   MAE vs clean signal = 0.049

The squared-error fit is dragged toward the outliers; Huber and absolute-error stay close to the true signal.

Next, quantile regression — fit several quantiles to form a prediction band:

from jaxsr.additive import QuantileLoss

rng = np.random.default_rng(2)
Xq = rng.uniform(-2, 2, size=(500, 2))
yq = 2.0 * Xq[:, 0] + 0.5 * Xq[:, 1] ** 2 + rng.normal(0, 1.0, 500)

for q in [0.1, 0.5, 0.9]:
    m = StagewiseSymbolicRegressor(
        n_terms=10, max_complexity=3, learning_rate=0.5,
        loss=QuantileLoss(q), refit_coefficients=False,
    ).fit(Xq, yq)
    coverage = float(np.mean(yq <= np.array(m.predict(Xq))))
    print(f'q={q}:  fraction of points below prediction = {coverage:.3f}')

q=0.1:  fraction of points below prediction = 0.100

q=0.5:  fraction of points below prediction = 0.514

q=0.9:  fraction of points below prediction = 0.890

Structural uncertainty: is the discovered expression stable?#

A single fitted expression can hide that the structure itself is uncertain — several basis sets may explain the data about equally well (common with collinear features). bootstrap_additive refits on bootstrap resamples and reports how often each basis function is selected (a proxy for a posterior inclusion probability), plus a predictive ensemble.

from jaxsr.additive import bootstrap_additive, bootstrap_predict_additive

# Collinear features: x2 is nearly a copy of x0, so the structure is not
# well determined even though the fit is good.
rng = np.random.default_rng(0)
a = rng.normal(0, 1, 400)
c0, c1 = a, 0.9 * a + 0.1 * rng.normal(0, 1, 400)
c2 = 0.8 * a + 0.2 * rng.normal(0, 1, 400)
Xc = np.column_stack([c0, c1, c2])
yc = c0 - c1 + 0.5 * c2 + rng.normal(0, 0.1, 400)

est = StagewiseSymbolicRegressor(n_terms=3, max_complexity=1)
res = bootstrap_additive(est, Xc, yc, n_bootstrap=80, random_state=0)
for name, prob in list(res['inclusion_probabilities'].items())[:6]:
    print(f'{name:8s} selected in {prob:.0%} of resamples')

x2       selected in 100% of resamples
x1       selected in 62% of resamples
x0       selected in 50% of resamples
x0^3     selected in 45% of resamples
x1^3     selected in 14% of resamples
x2^2     selected in 11% of resamples

Inclusion probabilities near 0 or 1 mean the structure is identifiable; diffuse values (~0.5) mean the data don’t pin down one expression. The bootstrap predictive interval is then the honest summary:

pi = bootstrap_predict_additive(res['models'], Xc[:5], alpha=0.1)
for i in range(5):
    lo, hi, mid = pi['lower'][i], pi['upper'][i], pi['mean'][i]
    print(f'x[{i}]  mean={mid:+.2f}  90% interval=[{lo:+.2f}, {hi:+.2f}]')

x[0]  mean=-0.00  90% interval=[-0.04, +0.04]
x[1]  mean=-0.01  90% interval=[-0.06, +0.04]
x[2]  mean=+0.37  90% interval=[+0.30, +0.47]
x[3]  mean=+0.03  90% interval=[-0.11, +0.11]
x[4]  mean=-0.24  90% interval=[-0.29, -0.16]

Backfitting: revising terms instead of freezing them#

BackfittingSymbolicRegressor keeps a fixed set of terms and re-discovers each one on the partial residual across sweeps (GAM-style), warm-started from a stagewise fit. For squared error it typically matches stagewise+refit rather than beating it (the joint least-squares refit already optimises the linear combination), so prefer it when you specifically want a fixed-size, revisable decomposition.

from jaxsr.additive import BackfittingSymbolicRegressor

bf = BackfittingSymbolicRegressor(n_terms=4, n_sweeps=6, max_complexity=3).fit(
    X_train, y_train
)
print('backfitting test R^2:', round(float(bf.score(X_test, y_test)), 4))
print('sweeps run:', len(bf.training_history_) - 1)

backfitting test R^2: 0.9988
sweeps run: 1

Recursive basis expansion: reaching compositions (experimental)#

A fixed library can only select from basis functions you supply — it cannot discover a composition like exp(x0*x1) that isn’t in it. RecursiveSymbolicRegressor grows the library along the residual, composing the useful terms with unary functions, products, and ratios each round. On this target it recovers exp(x0*x1) exactly, where a flat library with transcendentals cannot.

from jaxsr import fit_symbolic
from jaxsr.additive import RecursiveSymbolicRegressor

rng = np.random.default_rng(0)
Xc = rng.uniform(-2, 2, size=(400, 2))
yc = np.exp(Xc[:, 0] * Xc[:, 1]) + 0.02 * rng.normal(size=400)
Xtr2, ytr2, Xte2, yte2 = Xc[:200], yc[:200], Xc[200:], yc[200:]

def r2(y, p):
    y, p = np.asarray(y, float), np.asarray(p, float)
    return 1 - np.sum((y - p) ** 2) / np.sum((y - y.mean()) ** 2)

rec = RecursiveSymbolicRegressor(n_expansions=3, max_terms=6, beam_width=25).fit(Xtr2, ytr2)
flat = fit_symbolic(Xtr2, ytr2, max_terms=6, include_transcendental=True)
print('recursive  test R^2:', round(r2(yte2, rec.predict(Xte2)), 4), '->', rec.expression_[:40])
print('flat lib   test R^2:', round(r2(yte2, flat.predict(Xte2)), 4))

/home/user/jaxsr/src/jaxsr/regressor.py:1816: UserWarning: Excluding 4 basis function(s) with non-finite values on the training data; they will not be selected.
  return model.fit(X, y)

recursive  test R^2: 1.0 -> y = 0.9998*exp((x0)*(x1))
flat lib   test R^2: 0.7172

It is an experimental, deterministic FFX-style search: competitive with a strong genetic-programming engine on simple compositional targets, but its cost grows with beam_width/n_expansions/feature count and it will not match PySR/Operon on hard problems. The composed bases are closures, so the model is not serialisable.

Difference from single-expression and backfitting#

  • Single-expression (jaxsr.SymbolicRegressor): one sparse expression over a fixed basis library.

  • Stagewise additive: terms are discovered one at a time on the residual and then frozen.

  • Backfitting additive (BackfittingSymbolicRegressor): a fixed set of terms is revised repeatedly, each conditioned on the others (GAM-style; squared error only). A Bayesian (BART/iBART) variant is future work.

See docs/guides/additive-symbolic-regression.md for the full guide.