Skip to main content

Decision Trees and PyTorch

PyTorch provides tensors, automatic differentiation, neural-network modules, and optimizers. It does not provide a maintained high-level CART classifier or regressor comparable to scikit-learn's decision-tree estimators.

Why the previous “PyTorch decision tree” was removed​

The former example recursively enumerated observed feature values with PyTorch tensor operations. It illustrated a small split-search exercise, but presented that exercise as a full framework estimator. It omitted important behavior expected from a maintained tree implementation:

  • efficient candidate ordering and stopping rules;
  • sample and class weights;
  • explicit missing-value and categorical-feature behavior;
  • cost-complexity pruning and a cross-validated pruning path;
  • regression, multi-output, sparse-input, and fitted-schema handling;
  • stable serialization, inspection, and estimator validation.

It also gained no benefit from autograd: choosing a feature and threshold is a discrete greedy search, not a gradient update. Moving the tensors to a GPU would not turn that Python recursion into an efficient production tree learner.

Use scikit-learn for a classical single tree on tabular data. Select depth, leaf-size, feature-subset, and pruning parameters within training folds. Evaluate the final tree once on the isolated test set, and inspect both its rules and its subgroup errors before treating it as an interpretable model.

PyTorch can still own representation learning. For example, a neural encoder can turn images into embeddings. Freeze the selected encoder, generate embeddings without gradient tracking, and train the downstream tree using training entities only. Keep group-aware or time-aware split boundaries consistent across the raw examples and their embeddings.

If the desired model is a differentiable or “soft” tree from a research paper, use a maintained implementation that states its objective and limitations, and call it a differentiable tree rather than implying it is standard CART. Validate it against a conventional tree baseline.

What to retain​

  • PyTorch tensor syntax alone does not make an algorithm a maintained PyTorch estimator.
  • Classical tree training does not need backpropagation.
  • A faithful tree implementation needs pruning, weighting, schema, and edge-case semantics—not only Gini calculations.
  • The scikit-learn page in this section is the executable reference for a single decision tree.