neural_trees.SoftDecisionTree

class neural_trees.SoftDecisionTree(depth: int = 5, max_epochs: int = 40, learning_rate: float = 0.01, batch_size: int = 64, penalty_coef: float = 0.001, device: str = 'cpu', verbose: bool = False, random_state: int | None = None, class_weight=None, growth: str = 'none', growth_init: str = 'random', growth_jitter: float = 0.2, growth_budget: str = 'split', warm_start: bool = False, learn_temperature: bool = False, early_stopping: bool = False, validation_fraction: float = 0.1, n_iter_no_change: int = 10, split_features: int | None = None, init: str = 'random', init_sharpness: float = 10.0, init_prune: bool = True)

Bases: ClassifierMixin, BaseEstimator

Soft Decision Tree Classifier (sklearn-compatible).

A fully differentiable decision tree where each internal node applies a soft (sigmoid) split, allowing end-to-end gradient training.

Parameters:

depthint, default=5

Depth of the tree. The tree has 2^depth leaves.

max_epochsint, default=40

Number of training epochs.

learning_ratefloat, default=0.01

Learning rate for Adam optimizer.

batch_sizeint, default=64

Mini-batch size for training.

penalty_coeffloat, default=1e-3

Regularization coefficient for the entropy penalty on internal nodes. Higher values encourage more balanced splits.

devicestr, default=”cpu”

PyTorch device. “auto” picks CUDA if it is available, then Apple silicon’s MPS, then CPU. Anything else is passed to torch as given, so “cuda:1” works. Resolved once in fit and recorded as device_, so prediction always runs where training did.

verbosebool, default=False

Whether to print training progress.

random_stateint or None, default=None

Seed for model initialization and shuffled mini-batches.

class_weightdict, “balanced” or None, default=None

Weights per class, combined multiplicatively with sample_weight. “balanced” uses n_samples / (n_classes * bincount(y)), which is what an imbalanced target usually needs: without it a rare class contributes so little to the loss that the tree can ignore it entirely.

growth_init{“random”, “residual”, “residual_gate”}, default=”random”

How a new level is initialised under growth="incremental". Every leaf becomes a gate with two children; if the children started identical the new gate’s gradient would be exactly zero and the level could never learn, so the symmetry has to be broken.

• "random" perturbs the two children with Gaussian noise of scale growth_jitter.

• "residual" moves the two children apart along the leaf’s residual error direction (which classes it under-predicts), by the same amount on average, so the split is placed where the leaf is wrong rather than in a random direction.

• "residual_gate" additionally points the new gate at the samples that pull towards the right child, so routing starts informative.

Under growth="per_leaf" the same choice applies to the two children of the leaf being split, which additionally inherit the parent’s distribution so that the split preserves the function up to the perturbation. "uniform" is the pre-0.7 per-leaf behaviour (children start as untrained uniform distributions, identical to each other), kept only for reproducing the comparison. Ignored for growth="none".

growth_budget{“split”, “full”}, default=”split”

"split" divides max_epochs across the growth rounds, so a grown tree costs about what a fixed-depth one does; measured, that leaves it under-trained relative to a tree of its final shape trained from scratch. "full" gives every round the whole budget, at up to depth (incremental) or 2 * depth (per-leaf) times the cost.

growth_jitterfloat, default=0.2

Size of the perturbation applied to new children when the tree is deepened. Must be positive for the new level to learn; 0.0 is allowed only because it reproduces the failure exactly, which is useful for demonstration. Beyond being nonzero the value barely matters (0.05, 0.2 and 0.5 are within half a point of each other).

warm_startbool, default=False

When True, a second call to fit continues from the parameters the first one left, instead of reinitializing. Useful for training in stages, or for extending a run that turned out too short.

This is warm_start rather than partial_fit deliberately. sklearn’s partial_fit contract promises that a model updated on batches approaches one trained on the union, and requires handling classes that were absent from the first call. Neither holds here: the architecture is fixed at the first fit, and mini-batch gradient descent over a second dataset drifts toward that dataset rather than the union. warm_start promises only what is actually delivered, which is continuation.

The label set must not change between calls; a new class would need an output layer this model cannot grow.

growth{“none”, “incremental”, “per_leaf”}, default=”none”

How the tree reaches its shape.

• "none" builds the complete tree of depth depth up front, which is the Frosst & Hinton (2017) formulation.

• "incremental" starts from a single split and deepens one level at a time, keeping a level only if it improves validation loss (Irsoy, Yildiz & Alpaydin, ICPR 2012). depth becomes an upper bound and tree_depth_ reports what was actually kept. Requires a validation split, so validation_fraction applies whether or not early_stopping is on, and the max_epochs budget is divided across rounds rather than spent per round. That last point matters in practice: a budget that trains a fixed tree adequately can leave an incremental one under-trained, so raise max_epochs when switching.

• "per_leaf" splits one leaf at a time, the one carrying the most expected error, so the tree can end up unbalanced and spend depth only where the data needs it. This is the growth rule of İrsoy, Yıldız & Alpaydın (ICPR 2012); level-wise growth was the tractable approximation of it.

  It produces by far the sparsest trees, and wins where a fixed depth over-parameterizes. 3 seeds of 5-fold CV, accuracy and splits kept:

  none / incremental / per_leaf Iris 0.958 / 15 0.931 / 15 0.880 / 7.7 Wine 0.977 / 15 0.981 / 15 0.966 / 5.2 Breast Cancer 0.971 / 15 0.971 / 9.9 0.978 / 4.0 synthetic 20d 0.839 / 63 0.881 / 13 0.885 / 3.7

  On the synthetic problem it reaches better accuracy than a fixed depth-6 tree using 3.7 splits against 63. On Iris it loses, which is why the default is still “none”.

learn_temperaturebool, default=False

Learn a per-node inverse temperature on the gate, so a node can sharpen its split instead of saturating in the flat part of the sigmoid (Frosst & Hinton, 2017). Off by default because the effect is mixed: averaged over 5 seeds of 5-fold CV at depth 4 it moved Iris from 0.900 to 0.928, and cost about 0.6 points on Wine and Breast Cancer.

early_stoppingbool, default=False

Hold out validation_fraction of the training data and stop once validation loss has not improved for n_iter_no_change epochs. The parameters of the best epoch are restored.

validation_fractionfloat, default=0.1

Fraction held out when early_stopping=True.

n_iter_no_changeint, default=10

Epochs without validation improvement before stopping.

init{“random”, “cart”}, default=”random”

How the gates start. "cart" fits a CART tree of the same depth on the training data and starts from a soft copy of it (see from_hard_tree): greedy search places sharp single-feature thresholds, gradient descent then moves them jointly and tilts them. Requires growth="none". Pair it with learn_temperature=True so a gate that needs CART’s sharp threshold can keep it.

init_sharpnessfloat, default=10.0

Gate slope of the CART copy, in standard deviations of the feature. Large values copy CART exactly but give small gradients; small values learn more and blur sharp thresholds. Worth tuning per dataset.

init_prunebool, default=True

With init="cart", prune the CART tree before copying it, with the cost-complexity strength chosen by 3-fold cross-validation, so a small dataset does not start from an overfitted shape. With split_features=1 as well, every gate keeps the single feature CART chose for the whole fit.

split_featuresint or None, default=None

At most this many features per split. A soft split is oblique: its gate weighs every feature, which copies a complex boundary well and reads badly. With split_features=k the tree is first trained as usual, then pruned in rounds that halve the features each gate keeps until k remain, ranking by |weight| * std(feature) (the feature’s share of the gate’s variation, so the choice does not depend on the input scale). After each round the tree is fine-tuned with the pruned weights held at zero; the rounds share max_epochs // 2 extra epochs. split_features=1 gives axis-aligned soft splits, a soft threshold on one feature each, as in CART. None keeps every feature.

Attributes:

classes_ndarray of shape (n_classes,)

The class labels.

n_features_in_int

Number of features seen during fit.

training_history_list of dict

Loss and accuracy per epoch, plus validation loss when early stopping is on.

feature_importances_ndarray of shape (n_features,)

Gate weight magnitudes, weighted by how much probability mass reaches each node on the training data, normalized to sum to 1.

n_iter_int

Epochs actually run.

tree_depth_int

Depth of the fitted tree. Equals depth unless growth=”incremental” stopped earlier.

growth_str

The growth mode actually used. Falls back to “none” when the data is too small to hold out a stratified validation split.

References

İrsoy, O., Yıldız, O. T., & Alpaydın, E. (2012). Soft Decision Trees. ICPR 2012.

Examples

>>> from neural_trees import SoftDecisionTree
>>> from sklearn.datasets import load_iris
>>> X, y = load_iris(return_X_y=True)
>>> sdt = SoftDecisionTree(depth=4, max_epochs=30)
>>> sdt.fit(X, y)
>>> sdt.score(X, y)

explain(X, feature_names=None, max_terms: int = 3, counterfactual: bool = True, feature_bounds=None)

Explain individual predictions.

For each row of X: the predicted class and its probability, the leaf that received most of the sample’s probability mass and that leaf’s class distribution, every gate on the path to that leaf with the direction taken, the probability of that direction and the largest feature terms behind it, a per-feature contribution score aggregated along the path, and the smallest single-feature change that flips the predicted class, verified by re-predicting the changed sample.

The contribution score reads the linear gates directly and is not SHAP. The counterfactual is searched only among the gates on the dominant path and is None when no single-feature change there flips the class. All values are in the model’s input units.

Parameters:

Xarray-like of shape (n_samples, n_features) or (n_features,)

feature_namessequence of str, optional

max_termsint, default=3

Feature terms shown per gate, largest |weight * value| first.

counterfactualbool, default=True

Search for the flipping change; costs one prediction per candidate, so turn it off when explaining many samples.

feature_bounds(lower, upper), optional

Two arrays of length n_features. A counterfactual value is clipped into [lower[j], upper[j]] before it is tried, so the reported change stays inside values the feature can take: pass (0, 1) for one-hot columns and the training range for the rest. Without bounds the search may report a value no real row has.

Returns:

Explanation or list of Explanation

One object per row; a single object if X was one-dimensional. Each has to_text() and to_dict().

fit(X, y, sample_weight=None) → SoftDecisionTree

Fit the Soft Decision Tree.

Parameters:

Xarray-like of shape (n_samples, n_features)

yarray-like of shape (n_samples,)

sample_weightarray-like of shape (n_samples,), default=None

Per-sample weights applied to the loss. Combined multiplicatively with class_weight when both are given. The entropy penalty is left unweighted: it regularizes the shape of the tree, not the fit to any particular sample.

Weighting a sample by k gives the same loss and the same gradient as repeating it k times, but not bit-for-bit the same fit: the repeated dataset is larger, so mini-batches are composed differently and the optimizer follows a different path. This is why check_sample_weight_equivalence_on_dense_data is the one estimator check this class does not pass (62 of 63), and it is not satisfiable by any stochastic mini-batch learner.

Returns:

self

classmethod from_hard_tree(tree, X=None, sharpness: float = 10.0, **params) → SoftDecisionTree

A soft tree that starts as a copy of a fitted scikit-learn DecisionTreeClassifier, ready to be fine-tuned with fit.

Every CART split x[f] <= t becomes the axis-aligned gate sigmoid(sharpness * (x[f] - t) / std(x[f])): right when the feature is above the threshold, as in CART, with a transition about 1 / sharpness standard deviations wide. Every node takes CART’s class distribution, and nodes CART did not split stay leaves, so an unbalanced CART tree keeps its shape.

The point is to start where greedy search is strong (sharp, single-feature thresholds, found exactly) and let gradients do what it cannot (move thresholds jointly, tilt splits towards other features). Call fit(X, y) to fine-tune; the returned estimator has warm_start=True so that fit continues from the copy instead of reinitialising. Without fine-tuning it already predicts.

Parameters:

treefitted sklearn.tree.DecisionTreeClassifier

Xarray-like of shape (n_samples, n_features), optional

Data on the scale the tree was fitted on, used only for the feature standard deviations. Without it the features are assumed standardised (std 1).

sharpnessfloat, default=10.0

Gate slope in units of the feature’s standard deviation. Larger copies CART more exactly at the start but leaves the gates in the flat part of the sigmoid, where gradients are small.

**params

Any other SoftDecisionTree parameter (max_epochs, learning_rate, split_features, …). depth is taken from the tree, growth must stay "none".

Returns:

SoftDecisionTree

get_leaf_distributions() → ndarray

Return the class distribution stored in each leaf node.

Returns:

distributionsndarray of shape (n_leaves, n_classes)

get_split_weights() → List[ndarray]

Return the weight vectors for each internal node’s split.

Returns:

weightslist of ndarray, one per internal node

predict(X) → ndarray

Predict class labels.

Parameters:

Xarray-like of shape (n_samples, n_features)

Returns:

y_predndarray of shape (n_samples,)

predict_proba(X) → ndarray

Predict class probabilities.

Each sample reaches every leaf with some probability, so the returned distribution is the path-probability weighted average of the leaf distributions, P(y | x) = sum_l mu_l(x) Q_l(y). This is why the output is smooth rather than the piecewise constant output of a hard tree.

Parameters:

Xarray-like of shape (n_samples, n_features)

Samples to score. Cast to float32 internally, so any numeric dtype is accepted. Must have the same number of features seen in fit.

Returns:

probandarray of shape (n_samples, n_classes)

Class probabilities in the order of self.classes_. Each row sums to 1.

save_onnx(path: str, feature_names=None, opset: int = 17) → None

Write to_onnx() to a file.

set_fit_request(*, sample_weight: bool | None | str = '$UNCHANGED$') → SoftDecisionTree

Configure whether metadata should be requested to be passed to the fit method.

Note that this method is only relevant when this estimator is used as a sub-estimator within a meta-estimator and metadata routing is enabled with enable_metadata_routing=True (see sklearn.set_config()). Please check the User Guide on how the routing mechanism works.

The options for each parameter are:

• True: metadata is requested, and passed to fit if provided. The request is ignored if metadata is not provided.

• False: metadata is not requested and the meta-estimator will not pass it to fit.

• None: metadata is not requested, and the meta-estimator will raise an error if the user provides it.

• str: metadata should be passed to the meta-estimator with this given alias instead of the original name.

The default (sklearn.utils.metadata_routing.UNCHANGED) retains the existing request. This allows you to change the request for some parameters and not others.

Added in version 1.3.

Parameters:

sample_weightstr, True, False, or None, default=sklearn.utils.metadata_routing.UNCHANGED

Metadata routing for sample_weight parameter in fit.

Returns:

selfobject

The updated object.

set_score_request(*, sample_weight: bool | None | str = '$UNCHANGED$') → SoftDecisionTree

Configure whether metadata should be requested to be passed to the score method.

Note that this method is only relevant when this estimator is used as a sub-estimator within a meta-estimator and metadata routing is enabled with enable_metadata_routing=True (see sklearn.set_config()). Please check the User Guide on how the routing mechanism works.

The options for each parameter are:

• True: metadata is requested, and passed to score if provided. The request is ignored if metadata is not provided.

• False: metadata is not requested and the meta-estimator will not pass it to score.

• None: metadata is not requested, and the meta-estimator will raise an error if the user provides it.

• str: metadata should be passed to the meta-estimator with this given alias instead of the original name.

The default (sklearn.utils.metadata_routing.UNCHANGED) retains the existing request. This allows you to change the request for some parameters and not others.

Added in version 1.3.

Parameters:

sample_weightstr, True, False, or None, default=sklearn.utils.metadata_routing.UNCHANGED

Metadata routing for sample_weight parameter in score.

Returns:

selfobject

The updated object.

to_hard_tree(rule: str = 'gate')

Export the trained tree with its gates read as hard decisions.

Each internal node’s gate is sigmoid(beta * (w . x + b)); the sign of w . x + b is the decision it has settled on, and beta only sharpens it. Taking that sign and routing each sample down one path gives a plain numpy model with readable rules and no PyTorch in the prediction path.

This is a different model, not a re-encoding: a mixture over leaves is not a single path, and the two disagree on samples that sit near a split. Measure the agreement on held-out data before relying on it.

Parameters:

rule{“gate”, “leaf”, “contribution”}, default=”gate”

How a sample reaches a leaf. "gate" takes the sign of each gate on the way down, depth dot products per sample, and is the routing the printed rules describe. "leaf" sends the sample to the leaf with the largest arrival probability, which can differ from the greedy walk when an early gate is nearly balanced and a later one is decisive. "contribution" sends it to the leaf that contributes most to the soft mixture’s winning class. Both alternatives evaluate every gate, so they cost more per sample, and under them the printed thresholds describe the gates, not the routing. Walking down by the larger child probability is not a separate rule: a sigmoid exceeds one half exactly when its argument is positive, so it is "gate". Measured agreement and cost per rule are in the documentation; the default stays "gate" because no alternative was better on both axes.

Returns:

HardDecisionTree

Examples

>>> hard = sdt.to_hard_tree()
>>> (hard.predict(X_test) == sdt.predict(X_test)).mean()
>>> print(hard.export_text(feature_names=feature_names))

to_numpy(feature_names=None)

Export the fitted tree as a NumpySoftTree.

Unlike to_hard_tree(), this is the same model: the mixture over leaves is kept and predictions agree with the torch model to float32 precision, without PyTorch in the prediction path. The object serialises to JSON and back, which is the intended way to ship a fitted tree to a service that does not install torch, or to freeze one for audit.

to_onnx(feature_names=None, opset: int = 17)

The fitted tree as an onnx.ModelProto (pip install neural-trees[onnx]).

Standard operators only, so it runs wherever ONNX runs, without torch or this library, and predicts the same thing as predict() to float32 precision. Input X is float32 (n, n_features).