TAO Refinement — a walkthrough

Tree-Alternating Optimization (TAO) refines an already fitted SGT in place, without changing its topology. It sweeps internal nodes bottom-up and replaces routing rules that don’t hurt training performance. This notebook shows what it buys and how to control (or turn off) the refinement sgtlearn runs automatically.

What is TAO?

The native ShapeCART trainer grows a tree greedily top-down: a rule chosen high in the tree is never revisited. TAO fixes routing rules after the fact — for each internal node (bottom-up) it re-optimizes the rule against the samples that actually reach it.

  • Topology is preserved. No node is added, removed, or moved; only the rules inside existing nodes change.

  • Training performance never decreases (with the default lambda_=0): a rule is replaced only if it does at least as well as the old one.

[1]:

import numpy as np
from sklearn.datasets import fetch_openml
from sklearn.model_selection import train_test_split
from sgtlearn import SGTClassifier, tao

# Phoneme: 5404 samples, 5 continuous acoustic features, binary target —
# non-monotone per-feature structure, where greedy top-down splits leave room
# for TAO to improve.
X, y = fetch_openml("phoneme", version=1, return_X_y=True, as_frame=False, parser="liac-arff")
X = np.asarray(X, float)
X_tr, X_te, y_tr, y_te = train_test_split(X, y, test_size=0.3, random_state=0, stratify=y)

# Fit WITHOUT TAO (raw native ShapeCART).
model = SGTClassifier(max_depth=4, tao_n_runs=0, random_state=0).fit(X_tr, y_tr)
print(f"native only  train={model.score(X_tr, y_tr):.3f}  test={model.score(X_te, y_te):.3f}")

# Refine the SAME model in place. Pass the SAME (X, y) used to fit — per-sample
# routing is not stored after fit, so TAO needs the data again.
tao.TAO_refine(model, X_tr, y_tr, n_runs=10)
print(f"after TAO    train={model.score(X_tr, y_tr):.3f}  test={model.score(X_te, y_te):.3f}")

native only  train=0.855  test=0.808
after TAO    train=0.871  test=0.813

Train accuracy is equal or higher after refinement (guaranteed with lambda_=0) — here it rises from 0.855 to 0.871. Test accuracy improves too (0.808 → 0.813): TAO fixes greedy top-down routing mistakes, which tends to generalize. The structure is unchanged — same nodes, same depth — only the rules inside them.

Regularizing with lambda_

lambda_ is a per-sample complexity rate. A non-constant rule must beat the constant “send everyone one way” rule by more than lambda_ * n_samples to survive, so raising it collapses weak splits into constants.

To actually see regularization help, we need a model that overfits at lambda_=0. Phoneme at depth 4 barely overfits, so here we switch to a deeper tree (depth 6) on the noisier German credit dataset (1000 rows, 20 features), where the extra depth overfits and lambda_ can rein it back in.

[2]:

from sklearn.datasets import fetch_openml

Xc, yc = fetch_openml("credit-g", version=1, return_X_y=True, as_frame=False, parser="liac-arff")
Xc = np.asarray(Xc, float)
Xc_tr, Xc_te, yc_tr, yc_te = train_test_split(Xc, yc, test_size=0.3, random_state=0, stratify=yc)

# Refit the untouched baseline for each lambda so the runs are independent.
for lam in (0.0, 3e-3, 1e-2, 3e-2, 1e-1):
    m = SGTClassifier(max_depth=6, tao_n_runs=0, random_state=0).fit(Xc_tr, yc_tr)
    tao.TAO_refine(m, Xc_tr, yc_tr, n_runs=10, lambda_=lam)
    print(f"lambda={lam:6.3f}: train={m.score(Xc_tr, yc_tr):.3f}  test={m.score(Xc_te, yc_te):.3f}")

lambda= 0.000: train=0.863  test=0.697
lambda= 0.003: train=0.839  test=0.720
lambda= 0.010: train=0.799  test=0.760
lambda= 0.030: train=0.734  test=0.727
lambda= 0.100: train=0.700  test=0.700

This is a textbook bias–variance tradeoff. At lambda_=0 the tree overfits — train accuracy is 0.863 but test only 0.697. Raising lambda_ collapses weak splits, so train accuracy falls steadily (0.863 → 0.700). Test accuracy first rises — up to 0.760 at lambda_=0.01, as pruning removes overfit splits and closes the train/test gap — then falls once lambda_ is large enough to underfit (0.700, where train ≈ test). Pick the lambda_ that maximizes validation accuracy, not the smallest or largest.

Forests

TAO refines every base tree of a RandomSGForestClassifier / RandomSGForestRegressor, optionally in parallel.

[3]:

from sgtlearn import RandomSGForestClassifier

forest = RandomSGForestClassifier(
    n_estimators=25, max_depth=3, tao_n_runs=0, random_state=0, n_jobs=-1
).fit(X_tr, y_tr)
print(f"forest before  test={forest.score(X_te, y_te):.3f}")

# Refine every tree on the FULL training set (fit-time forest TAO would use each
# tree's bootstrap sample instead). Uses forest.n_jobs by default.
tao.TAO_refine(forest, X_tr, y_tr, n_runs=10)
print(f"forest after   test={forest.score(X_te, y_te):.3f}")

forest before  test=0.805
forest after   test=0.814

How to turn TAO off

TAO is on by default (tao_n_runs=10). Set tao_n_runs=0 in the constructor to skip it entirely at fit time and get raw native ShapeCART.

[4]:

# No TAO — native trainer only.
raw = SGTClassifier(max_depth=4, tao_n_runs=0, random_state=0).fit(X_tr, y_tr)
print(f"tao_n_runs=0  test={raw.score(X_te, y_te):.3f}")

tao_n_runs=0  test=0.808

Fit-time vs. post-hoc

  • Leave it on (tao_n_runs=10) for the default: every model is refined once, no extra code.

  • Fit with TAO off, then ``tao.TAO_refine`` when you want to compare before/after, sweep lambda_ without refitting, or refine a forest on the full training set. Safe to call repeatedly — each call refines further in place.

See the TAO API reference for the full parameter list.