FastSmallTree: Implementing compact, fast decision trees via Autoresearch

The certified tree is stored as an ordinary sklearn DecisionTreeClassifier in model.estimator_, so it works with sklearn-compatible libraries (plot_tree, export_text, feature_importances_), e.g. dtreeviz:

import dtreeviz  # pip install dtreeviz

model = FastSmallTreeClassifier(
    regularization=0.03)
model.fit(X_train, y_train)

viz = dtreeviz.model(
    model.estimator_, X_train, y_train,
    feature_names=X.columns,
    target_name="diagnosis",
    class_names=["malignant", "benign"])
viz.view()

All 14 come from the GOSDT reference repository, each run at five penalties per leaf (0.1, 0.05, 0.02, 0.01, 0.005): 70 problems. Binary splits is what the solver searches: one per binary feature and one per distinct value of a numeric one, which is why fico_1k and compas_processed hold the problems nobody certifies.

Each development dataset has a stand-in drawn from a TabArena classification dataset by a fixed rule (originals largest first, each taking the smallest unused source with enough rows of each class). A stand-in matches its original's rows and features exactly and its class balance within 0.05, keeping the features with the most mutual information with the label; numeric features are quantile-binned to match the original's number of binary splits, shown in gray. chudi and iris fall short because their sources have too few distinct values. TabArena's HELOC and coupon datasets are excluded, since the development suite already contains them. Builder, manifest and per-problem results: agentic-imodels, baselines/benchmarks/external/.

A. Every returned tree is re-scored from the raw data and checked against its problem's certified optimum, so a solver cannot score by misreporting. v46's points are from the fixed code; a draft certified two wrong trees here (Methods).

B. TabArena-14 has no certified optima, so a certificate there is checked against the best tree any run of any solver returned.

C. One run at a single penalty (0.05) on each of the 14 datasets, so the panel is read for the ordering rather than a close criterion race; several optima at that penalty are a single leaf, which every solver finds at once. Three kinds of baseline failure are charged as no tree: a stop at the 6 GB cap (GOSDT on 8 problems, SPLIT on 8), a run killed after twice its budget (STreeD on 2), and an operating-system kill for outgrowing memory (STreeD on coupon_bar, signal 9 after five minutes). GOSDT runs in a child process that the harness can stop at the cap; STreeD and SPLIT run in-process. gosdt-guesses in exact mode certifies 4 and returns no tree on 8 (5 memory stops, 3 crashes); guided, published and tuned alike, it returns a heuristic tree on 11 and hits the memory cap on the 3 widest datasets, since it builds the exact binarization before guessing thresholds.

D. Run on a Linux server (two Xeon Platinum 8468V, 88 cores), each problem in its own process pinned to its own hardware threads (two for single-threaded solvers, eight for the 8-thread versions), many at once; times are comparable within the panel only. The 4 hours are each solver's time limit, so a run that reaches them returns its best tree. The harness enforces 6 GB on every process: a solver that stops itself at the budget keeps its incumbent, one the harness stops returns no tree. One run per problem, 140 hours of solver time. Two of D's points reflect this harness rather than the search. The shipped model's memory guard counts its own tables, not the process, so on the six widest problems the process passed 6 GB first, the harness stopped it, and it returned no tree on five of them; charged a single leaf there, its criterion (0.348) sits above pygosdt's (0.339) although its search is faster. That is a bug in the guard, not the search, and it is what a 6 GB machine would do; v23, whose guard reads the process, stops itself at 5.99 GB and keeps its tree. STreeD likewise loses three problems it finished on the laptop, where nothing policed its memory (monk_2, car_evaluation and coupon_bar at the smallest penalty, at 6.5 to 7 GB), which is part of why its criterion in D is worse than in B.

gosdt-guesses. In guided mode it reports only its search time, which excludes the threshold guessing where its work is done, so its points sit at the 1 ms floor in every panel. By wall clock, preprocessing included, it takes 3.1 s a problem in C against 9.6 s for v49: about three times faster, not a hundredfold, for a worse tree, no certificate, and no tree on 3 of the 14. Both of its modes were also swept, since a baseline judged at one setting can be judged unfairly. In exact mode its similar-support bound changes nothing measurable (0.489 ± 0.007 s against 0.486 ± 0.018 s). In guided mode the paper's threshold guesser, 40 stumps of depth 1, is the weak part on the visible set: 60 trees of depth 2 improve the criterion from 0.328 to 0.271 at 0.0068 s a problem, while a depth budget of 5 changes nothing. That gain does not transfer: on TabArena-14 the tuned guesser is slightly worse than the published one (0.3505 against 0.3480) and misses 3 trees the default returns. So the post keeps the published configuration as the baseline and shows the tuned points beside it.

Equivalent points. Rows with identical features but different labels can never be separated, so each such group costs at least its cheaper label. Summed over X this is a loss floor m(X), and since a split means at least two leaves, R(X) ≥ min(leaf(X), m(X) + 2λ), where leaf(X) is the best single label plus λ. Every node starts from this bound.

Leaf support. Let pot(S) be the most that relabelling the rows S can change the loss: the sum over the rows of the largest minus the smallest cost for their class (for 0/1 loss, their share of the rows). In an optimal tree every leaf has potential at least λ: deleting a leaf L with less and routing its rows to its sibling saves λ and adds at most pot(L) of loss. So a node with pot(X) < 2λ is a leaf, and (from v4) a split with a child of potential below λ is dropped, here and on every subset, since potential only shrinks. The test loss(X) − m(X) < λ also makes a node a leaf, for any nonnegative costs: a split costs at least m + 2λ and the leaf loss + λ. (An inherited test, 1 − m(X) < λ, assumed losses of at most 1; v14 removed it.)

A consequence every later bound uses. Deleting leaves of potential below λ strictly improves any tree and must stop, and it stops at a tree whose every split leaves both sides with potential at least λ. Such a split is also valid at every ancestor, since potential only grows with the rows. So R(X) is always attained by a tree built from splits valid at X, and a bound needs to hold only for such trees. Several hold for nothing more: the pairwise values, the shape relaxation and the column programme are computed over the valid splits, and a child only inherits its parent's valid splits.

Look-ahead budgets. A split's two children are solved in turn, the first with the budget less the second's lower bound, the second with the budget less the first's optimum; if either proves its lower bound above its share, the split is discarded. Budget deepening (v9) only changes the sequence of budgets, and every call respects (iii).

Threshold exchange. R is monotone under inclusion: restricting an optimal tree of X to S ⊂ X, and deleting emptied leaves, adds neither loss nor leaves, so R(S) ≤ R(X). For consecutive thresholds t < t′ of a column, left(t) ⊇ left(t′) and right(t) ⊆ right(t′), so if lb(right(t)) reaches the leaf value of right(t′), split t′ is no worse than t and t is dropped (symmetrically on the left). The two rules exclude each other on a pair, so no cycle drops both.

For thresholds t < t′ of one column, the two splits disagree only on the rows with t ≤ x < t′, whose potential is |pot(left(t)) − pot(left(t′))| because the left sets are nested. Moving those rows turns any tree rooted at t′ into one rooted at t with the same leaves and at most that much more loss, so |R(t) − R(t′)| is bounded by it. A bound proved for one threshold therefore bounds its whole column; the reference applies this only to neighbouring thresholds.

A tree that splits only on one column partitions the rows into intervals of that column, one per leaf, and any partition into k intervals is realised by a chain of k − 1 splits. So the best such tree is the best segmentation of the sorted rows, each segment charged its best label plus λ: a quadratic dynamic programme over class counts the node already has. When one column's thresholds are the only splits left at a node (after the exclusions above, which are monotone under subsets, so a descendant cannot regain a split), every tree on X uses only that column and the programme gives R(X) exactly, for every node of the chain. Elsewhere the chain is a real tree, used only as an upper bound.

Counting rows for every pair of splits gives the exact cost of the best two-leaf tree on each child of each split, best₂. Three consequences: (a) min(leaf, 2λ + best₂) is achieved by a real tree, so it upper-bounds the child; (b) any tree on a child is a leaf, one split, or has at least three leaves, so R(child) ≥ min(leaf, 2λ + best₂, 3λ); (c) every tree on X with at most three leaves is a leaf or a two-level tree, so the best two-level value over the valid splits, with the leaf, is at most the optimum over trees with at most three leaves (it also covers four-leaf (2,2) trees) and is achieved by a real tree. With a budget below 4λ no tree with four or more leaves fits, so that value is R(X) if within budget, and otherwise R(X) ≥ min(best, 4λ) is the bound to store.

v7 stored instead the smaller of best and the smallest bound over all splits, candidates inside the budget included, which can fall at or below the budget, so a caller could treat an unsolved node as solved: a violation of (iii), fixed in v8. Note that (c) needs the minimum over every split that separates the node's rows, not only the candidates that survive the budget filter; v49 returns to this.

For every child S of every split, R(X) ≥ R(S) ≥ lb(S) by monotonicity, so a node's lower bound can be raised to the largest proven bound of any child, for free. v14 also dropped the leaf test that assumes losses of at most 1 and made a user-supplied upper bound a budget rather than an achievable value.

Let g_a(S) bound every tree with exactly a leaves on S: g₁ = leaf(S) and g₂ = 2λ + best₂(S), both exact; g₃ = 3λ plus the cheapest cell any split peels off as a leaf, since a three-leaf tree has a leaf at depth one; g_a = aλ for a ≥ 4. A tree with four or more leaves has a root split with a and b leaves on its sides, a + b ≥ 4, so the minimum over splits and such (a, b) of g_a(left) + g_b(right) bounds them all. If the best two-level tree is no worse than that and than the leaf, it is optimal, and the node is solved at any budget without expanding a child. Children's bounds rise to min(g₁, g₂, g₃, 4λ) by the same argument.

Counting rows for every triple of splits gives each child its exact best three-leaf and (2,2)-leaf trees. The stage ran only on narrow nodes (at most 14 candidates in 64 words, with a budget under the floor plus 8λ), and v40 turned it off: once the search was compiled, the expansions it saved cost less than its cubic count. For shapes it does not enumerate it uses floors from the cells it sees; four leaves as (1,3), for instance, cost at least 4λ + a cell + the cheapest sub-cell on the other side, and five and six leaves are bounded the same way. These floors are admissible, and v19, which charged six or more leaves just 6λ, is exact.

The finding. v20 charges seven or more leaves 6λ + min(c₃, c₅, c₆), which is not admissible. A (3,4) tree keeps a single sub-cell leaf, whose cost can be below all three (c₆ needs a sub-cell on both sides), and a (4,4) tree of two (2,2) subtrees has no leaf above depth three, so its loss can be zero. Take a child whose label is the parity of three binary features, n_X of the n rows: every cell and sub-cell is half wrong, so all three constants are n_X/4n and the floor reads 6λ + n_X/4n, while the eight-leaf tree on all three features is perfect and costs 8λ. Once n_X > 14λn every term of the stored bound exceeds 8λ, so the node's stored lower bound is above its optimum; an independent brute-force check reproduces it (a bound of 0.25 against an optimum of 0.08). The overestimate enters the memo and, through the superset bound, the parent, where it can prune the optimum at a later budget: a false certificate. The loop's tests used at most three features, too few to build a seven-leaf subtree below a split. It never fired on the benchmark, so the numbers of v20 to v39 stand, but their guarantee does not. The admissible floor is 7λ; the package ships with the stage off and now uses that floor, so enabling it is safe.

On wide nodes with few candidates, v49 counts pairs only between the candidates and all other splits. Splits without pair counts keep their own bounds, and in the shape relaxation every tree rooted at one of them is charged that split's bound, which bounds any tree rooted there. The rule below 4λ looks exposed, since its best two-level value now runs over the candidates only, but its failing branch is never reached with the restricted stage. Candidates are chosen at the budget, so every rejected split's bound exceeds it, and those bounds are part of the shape relaxation's value. Whenever the node cannot be solved within budget, min(leaf, best two-level tree, shape bound) therefore already exceeds the budget, and the node returns earlier with that valid bound. A frame under a budget below 4λ is never cached, and one cached at a larger budget and re-entered at a smaller one keeps its rejected bounds above both. An earlier version of this page reported this branch as a second gap. That was wrong: it examined the branch without the early return in front of it. v49 is exact.

After a short sequential phase the root's candidate splits go to eight workers (processes in v25 and v29, threads from v30). The shared incumbent is always a real tree's value and only decreases, so a split a worker proves above the incumbent it saw is above the final answer; every other split is solved exactly, and the root takes the minimum, as in the sequential search. v32 shares proven lower bounds through a lock-free table keyed by capture mask and row count: bounds only rise, a partly written key has the wrong row count, and every published value was proved for exactly that subproblem. v34 shares solved optima the same way and pushes a better incumbent down each task's stack, re-deriving child budgets as when they were pushed; budgets only tighten, so the conclusions match the sequential search.

One assumption the argument needs. A table slot is written as value, bound, row count, key, then a used flag, and read in the reverse order; the argument needs a reader to see those writes in that order. x86 guarantees it. ARM, which the laptop timings ran on, does not without memory barriers, which the code does not issue, and the table is reused across fits with only its flags and values reset. So on ARM a reader could in principle pair a fresh flag with a bound left by the previous fit, which can be too high for the current problem. Nothing in any check showed it, but the proof does not cover it. A release/acquire pair around the flag, or resetting bounds and keys with the flags, would close it. The shipped model runs one thread and never touches the table.

The greedy seed (v43) installs a subtree's value only when it beats the node's, and every value it writes belongs to the structure it writes, so it only supplies (ii). The repair (v45) solves each internal node of the incumbent with its own value as budget, so by (iii) it ends solved or proven. v46's budgeted pair stage evaluates only the most promising splits under a short cap and derives no bounds from them, and v43c's rule that skips an expensive pair stage leaves bounds uncomputed, not changed. None of this can produce a wrong certificate; under a cap it only changes the incumbent. Drafts of v13, v41 and v46 did certify wrong trees through transcription errors (in v46, an indentation slip that skipped a refilter), all caught by the loop's checks before recording.

The gosdt-guesses recipe replaces the equivalent-points floor with a reference model's loss plus λ, and caps the depth. The reference loss bounds nothing, since a tree can beat the reference on any subset, so optima are pruned wherever the reference is poor (on monk_2 the regret reached 0.12 to 0.23), and the depth cap removes trees outright. Both are approximations by design.

The proofs above cover rules; this section checks the code that applies them, on the path the package takes. fit runs the compiled engine on one thread, so the parallel phase, the shared table and thread adoption are never entered, and the depth-3 stage is off. Thirteen places in what remains mark a node solved or raise a lower bound, each resting on a rule proved above:

1. Creating a node. lb = min(leaf, floor + 2λ), ub = leaf; solved at once only for a single row, for leaf − floor < λ, or for potential < 2λ. Inherited bounds.
2. No split left. No split leaves both sides with potential at least λ, so the node is a leaf. Leaf support.
3. Everything above the budget. min(leaf, best two-level tree, shape bound) exceeds the budget, so the bound rises to it. Shape relaxation, over every separating split.
4. The single-column chain. When every separating split is on one column, the segmentation value is optimal and every chain node is solved. Segmentation; a child only sees its parent's separating splits, so no column is regained.
5. Installing the best two-level tree. Each child gets min(leaf, 2λ + best₂) and the split achieving it. Real trees, so only (ii).
6. The shape rule. Best two-level tree no worse than the shape bound and the leaf: solved at that value.
7. The rule below 4λ. Solved at best when it fits the budget, else bound min(best, 4λ); best ranges over every separating split.
8. Children in the candidate loop. Solved when ub is within 1e-10 of lb, or when ub fits a budget below 3λ, since ub is exact over one and two leaves and three or more cost at least 3λ. Pairwise stage.
9. Deepening. Each child's budget is the bound less its sibling's lower bound, or less its optimum once that is known; a child above its budget prunes the split. Look-ahead.
10. The epilogue. Solved at best when it fits; else bound max(min(best, cheapest pruned split), largest child bound). Superset bound.
11. The flush. An interruption writes each open frame's incumbent to its node, never a bound or a flag, and the search restarts from a consistent memo.
12. Re-arming a cached frame. Reuses only quantities that depend on the node alone; a frame whose pair stage did not run is never cached.
13. The driver and extraction. A time, memory or depth stop sets optimal_ to False and warns, so a stopped search never certifies. Extraction rebuilds exactly the trees the writes above recorded.

Tolerance. A prune requires exceeding by more than 1e-10 and a solve accepts equality within it. Most conclusions rest on one such comparison, so the certified objective is within 1e-10 of the minimum; the threshold-exchange rule can chain its comparisons along a column, so the worst case is a column's number of thresholds times 1e-10. The returned tree scores no worse than the certificate. With uniform costs two trees differ by a multiple of 1/n plus a multiple of λ, so this matters only if λ matches a multiple of 1/n to ten decimal places; its purpose is to treat ties as ties.

Inputs. The proofs need a nonnegative penalty and nonnegative, finite costs, and nothing more, so nonzero diagonals, unequal errors and balance are covered. The package now rejects anything else; before, a negative penalty would have broken the 4λ floor.

What the check fixed. Four defects, none reached by any benchmark. (1) The single-column chain wrote each node's value before creating the next node's children, so a store that filled mid-chain could leave a node whose value no stored tree achieves, a failure of (ii); the chain now creates every node first. A test that grows a four-node store through dozens of reallocations passes on the old code too, since the failure also needs a memo hit on an inner chain node, so this fix rests on the argument. (2) Candidates were ordered by a rounded key with 1e-9 slack while the loop stops at the first candidate above the bound, so one could be skipped; the order is now exact. (3) The input validation above. (4) With the depth-3 stage forced on, extraction read a stale copy of a node store that had just grown: a crash, not a wrong tree; it now re-reads the store. The new tests enumerate every tree on up to six binary features with two or three classes, random nonnegative cost matrices, and a store forced to grow from four nodes: tripwires for the proofs, not substitutes.

Latent gap: the version ran as an exact solver and certified only known optima on the benchmark, but the proof of its guarantee has the hole described in the named collapsible. A refactor's "no new rule" is its whole argument; the loop's exhaustive checks caught three such drafts (v13, v41, v46), and the gap above is the one they could not reach. The parallel versions from v32 on also rest on the store-ordering assumption described under parallel certification.