Gradient boosted trees (spec model)

This is a math/reference specification of gradient boosting using decision trees.

Important caveat:

• Many computations here are written in a straightforward, proof-friendly style rather than as a tuned implementation.

References (classical):

• CART: Breiman, Friedman, Olshen, Stone, "Classification and Regression Trees", 1984.
• Gradient boosting: Friedman, "Greedy Function Approximation: A Gradient Boosting Machine", 2001.
• XGBoost: Chen and Guestrin, "XGBoost: A Scalable Tree Boosting System", 2016.
• LightGBM: Ke et al., "LightGBM: A Highly Efficient Gradient Boosting Decision Tree", 2017.

Squared-error boosting fits trees to target - prediction. The shared mean-squared-error gradient is mseLossGradSpec in NN/Spec/Models/LinearRegression.lean.

Implementation status

No API builder implements this model. NN/Spec/Models/RandomForest.lean reuses its decision trees; no theorem is proved about it.

Tree representation

We represent a decision tree as a small inductive datatype:

• leaf value stores the prediction for that leaf
• split feature threshold left right branches on a single feature

This is kept compact: it is easy to interpret (forward pass) and easy to fit with a simple greedy CART-style algorithm (implemented below).

Note on comparisons: The spec layer’s scalar interface (Context α) gives us a decidable > (via Context.decidableGT) but does not promise a decidable < for every backend. To stay portable, the tree uses the rule:

goRight := (x_feature > threshold)

and goes left otherwise (equivalently, $x_{\mathrm{feature}}\leq\mathrm{threshold}$ for the usual numeric orders).

inductive Spec.TreeNode (α : Type) (nFeatures : ℕ) : ℕ → Type

A regression-tree node for the typed GBDT specification.

• leaf value stores the prediction for that leaf.
• split feature threshold left right branches on a single feature using the rule goRight := (x_feature > threshold).

• leaf {α : Type} {nFeatures depth : ℕ} (value : α) : TreeNode α nFeatures depth
• split {α : Type} {nFeatures depth : ℕ} (feature : Fin nFeatures) (threshold : α) (left right : TreeNode α nFeatures depth) : TreeNode α nFeatures (depth + 1)

@[instance_reducible]

instance Spec.instInhabitedTreeNode {a✝ : Type} [Inhabited a✝] {a✝¹ a✝² : ℕ} : Inhabited (TreeNode a✝ a✝¹ a✝²)

structure Spec.DecisionTreeSpec (α : Type) (nFeatures maxDepth : ℕ) : Type

Decision-tree specification whose type records its feature count and maximum split depth.

Each split consumes one unit of the depth index, so a value of this type cannot contain a path deeper than maxDepth. Its feature index is a Fin nFeatures, ruling out invalid feature access.

• root : TreeNode α nFeatures maxDepth

  Root node with at most maxDepth splits on any path.

@[instance_reducible]

instance Spec.instInhabitedDecisionTreeSpec {a✝ : Type} [Inhabited a✝] {a✝¹ a✝² : ℕ} : Inhabited (DecisionTreeSpec a✝ a✝¹ a✝²)

structure Spec.GradientBoostedTreesSpec (α : Type) (nFeatures nTrees maxDepth : ℕ) : Type

Gradient boosted tree ensemble (regression-style) specification.

The model stores an explicit tensor of trees, a shrinkage parameter, and an initial prediction.

• trees : TorchLean.Tensor (DecisionTreeSpec α nFeatures maxDepth) [nTrees]

  The boosted regression trees, applied in order and summed.

• lr : α

  The shrinkage factor multiplying each tree's contribution.

• initialPrediction : α

  The constant base prediction that the trees correct.

Forward pass

All forward passes in this file are explicit about the feature dimension nFeatures. Tree depth (maxDepth) limits how many splits a tree may contain; it has nothing to do with how many input features exist.

def Spec.decisionTreeForwardSpec {α : Type} [TorchLean.Storage α] [Context α] {maxDepth nFeatures : ℕ} (tree : DecisionTreeSpec α nFeatures maxDepth) (input : TorchLean.Tensor α [nFeatures]) : α

Forward pass for a single decision tree on an input vector of nFeatures features.

def Spec.decisionTreeForwardSpec.traverse {α : Type} [TorchLean.Storage α] [Context α] {nFeatures : ℕ} (input : TorchLean.Tensor α [nFeatures]) {depth : ℕ} (node : TreeNode α nFeatures depth) : α

def Spec.decisionTreeForwardLeadingSpec {α : Type} [TorchLean.Storage α] [Context α] (leading : Shape) {maxDepth nFeatures : ℕ} (tree : DecisionTreeSpec α nFeatures maxDepth) (input : TorchLean.Tensor α (leading.concat [nFeatures])) : TorchLean.Tensor α (leading.concat Shape.scalar)

Apply a decision tree independently at every index of a leading shape.

def Spec.gradientBoostedTreesForwardSpec {α : Type} [TorchLean.Storage α] [Context α] {nTrees maxDepth nFeatures : ℕ} (model : GradientBoostedTreesSpec α nFeatures nTrees maxDepth) (input : TorchLean.Tensor α [nFeatures]) : TorchLean.Tensor α Shape.scalar

Forward pass for a gradient boosted ensemble on a single input.

Starting at initialPrediction, visit trees in index order and update acc := acc + lr * tree_i(x). Shrinkage is applied before each addition; rounded arithmetic does not in general permit factoring it out of the sum.

@[irreducible]

def Spec.gradientBoostedTreesForwardSpec.accumulate_trees {α : Type} [TorchLean.Storage α] [Context α] {nTrees maxDepth nFeatures : ℕ} (model : GradientBoostedTreesSpec α nFeatures nTrees maxDepth) (input : TorchLean.Tensor α [nFeatures]) (i : ℕ) (acc : α) : α

def Spec.gradientBoostedTreesForwardLeadingSpec {α : Type} [TorchLean.Storage α] [Context α] (leading : Shape) {nTrees maxDepth nFeatures : ℕ} (model : GradientBoostedTreesSpec α nFeatures nTrees maxDepth) (input : TorchLean.Tensor α (leading.concat [nFeatures])) : TorchLean.Tensor α (leading.concat Shape.scalar)

Apply a gradient-boosted ensemble independently at every index of a leading shape.

Classical training: CART-style regression trees (MSE)

Tree-based models here are intended as baselines and reference points. For neural models we implement reverse-mode explicitly; for trees we instead provide a classical (non-gradient) training routine.

The code below implements a small greedy CART-like procedure for regression:

• choose the split (feature, threshold) that minimizes SSE(left) + SSE(right) (sum of squared errors around each side’s mean)
• recurse until depth runs out or the split becomes degenerate

This is deterministic by construction:

• thresholds are chosen from the observed feature values
• ties are broken by the “first best” encountered during folding

This implementation prioritizes clarity and determinism over performance.

structure Spec.RegressionExample {α : Type} [TorchLean.Storage α] (nFeatures : ℕ) : Type

A single regression training example: feature vector x and scalar target y.

• x : TorchLean.Tensor α [nFeatures]

  The feature vector.

• y : α

  The scalar regression target.

def Spec.GradientBoostedTrees.Internal.meanOrZero {α : Type} [Context α] (xs : Array α) : α

Mean of an array, with 0 as a convenient default for an empty sample.

def Spec.GradientBoostedTrees.Internal.sse {α : Type} [Context α] (ys : Array α) : α

Sum of squared deviations from the mean (SSE).

def Spec.GradientBoostedTrees.Internal.partitionBySplit {α : Type} [TorchLean.Storage α] [Context α] {nFeatures : ℕ} (feature : Fin nFeatures) (threshold : α) (xs : Array (RegressionExample nFeatures)) : Array (RegressionExample nFeatures) × Array (RegressionExample nFeatures)

Partition samples into (left, right) for a given split.

def Spec.GradientBoostedTrees.Internal.targets {α : Type} [TorchLean.Storage α] {nFeatures : ℕ} (xs : Array (RegressionExample nFeatures)) : Array α

Extract the regression targets from an array of examples.

def Spec.GradientBoostedTrees.Internal.splitScore {α : Type} [TorchLean.Storage α] [Context α] {nFeatures : ℕ} (feature : Fin nFeatures) (threshold : α) (xs : Array (RegressionExample nFeatures)) : Option (α × Array (RegressionExample nFeatures) × Array (RegressionExample nFeatures))

Score a candidate split by sum of squared errors (SSE).

Returns none for degenerate splits (all samples go to one side).

def Spec.bestSplit {α : Type} [TorchLean.Storage α] [Context α] {nFeatures : ℕ} (xs : Array (RegressionExample nFeatures)) : Option (Fin nFeatures × α × Array (RegressionExample nFeatures) × Array (RegressionExample nFeatures) × α)

Find the best split (feature, threshold) by exhaustive search over observed thresholds.

def Spec.leafValue {α : Type} [TorchLean.Storage α] [Context α] {nFeatures : ℕ} (xs : Array (RegressionExample nFeatures)) : α

Leaf prediction value for regression: the mean target.

def Spec.fitRegressionNode {α : Type} [TorchLean.Storage α] [Context α] {nFeatures : ℕ} (depth : ℕ) : Array (RegressionExample nFeatures) → TreeNode α nFeatures depth

Fit a regression tree by greedy CART-style splitting (MSE/SSE), with a depth budget.

depthLeft counts how many splits we are still allowed to make.

def Spec.decisionTreeFitRegressionMseSpec {α : Type} [TorchLean.Storage α] [Context α] {batch maxDepth nFeatures : ℕ} (x : TorchLean.Tensor α [batch, nFeatures]) (y : TorchLean.Tensor α [batch]) : DecisionTreeSpec α nFeatures maxDepth

Fit a regression decision tree from a batched dataset.

Classical training: CART-style classification trees (Gini impurity)

For classification we often want the leaf prediction to be a label (e.g. String or Nat), while split thresholds remain numeric. To avoid forcing labels into the numeric scalar type α, we define a separate classifier tree type parameterized by the label type β.

The training algorithm mirrors the regression case:

• enumerate candidate thresholds from observed feature values
• pick the split that minimizes weighted Gini impurity
• recurse until depth runs out or no improvement is possible
• leaf prediction is the majority class (deterministic tie-breaking via first occurrence)

Why β is separate from α:

• Splits compare numeric features (α), so they need ordering/decidable comparison.
• Leaf predictions are usually discrete (β), and we do not want to pretend that labels form a numeric scalar domain.

PyTorch / sklearn analogies:

• This is closest in spirit to sklearn.tree.DecisionTreeClassifier with criterion="gini", expressed as a small pure spec.
• The boosting semantics (adding many trees sequentially) matches the high-level idea of sklearn.ensemble.GradientBoostingClassifier, but TorchLean does not try to reproduce all of sklearn’s engineering details (regularization knobs, histogram binning, etc.).

inductive Spec.ClassifierTreeNode (α β : Type) (nFeatures : ℕ) : ℕ → Type

A classifier tree node: numeric splits, label-valued leaves.

• leaf {α β : Type} {nFeatures depth : ℕ} (label : β) : ClassifierTreeNode α β nFeatures depth
• split {α β : Type} {nFeatures depth : ℕ} (feature : Fin nFeatures) (threshold : α) (left right : ClassifierTreeNode α β nFeatures depth) : ClassifierTreeNode α β nFeatures (depth + 1)

@[instance_reducible]

instance Spec.instInhabitedClassifierTreeNode {a✝ a✝¹ : Type} [Inhabited a✝¹] {a✝² a✝³ : ℕ} : Inhabited (ClassifierTreeNode a✝ a✝¹ a✝² a✝³)

structure Spec.DecisionTreeClassifierSpec (α β : Type) (nFeatures maxDepth : ℕ) : Type

Specification wrapper for a classification decision tree (numeric splits, label-valued leaves).

• root : ClassifierTreeNode α β nFeatures maxDepth

  Root node with bounded depth and valid feature indices.

@[instance_reducible]

instance Spec.instInhabitedDecisionTreeClassifierSpec {a✝ a✝¹ : Type} [Inhabited a✝¹] {a✝² a✝³ : ℕ} : Inhabited (DecisionTreeClassifierSpec a✝ a✝¹ a✝² a✝³)

def Spec.decisionTreeClassifyForwardSpec {α : Type} [TorchLean.Storage α] [Context α] {β : Type} {maxDepth nFeatures : ℕ} (tree : DecisionTreeClassifierSpec α β nFeatures maxDepth) (input : TorchLean.Tensor α [nFeatures]) : β

Forward pass for a classifier decision tree on an input vector of nFeatures features.

Branching convention:

• go right iff (x[feature] > threshold),
• otherwise go left.

This mirrors the common convention $x_{\mathrm{feature}}\leq\mathrm{threshold}$ goes left and $x_{\mathrm{feature}}>\mathrm{threshold}$ goes right, but avoids needing a decidable < for every Context α backend.

def Spec.decisionTreeClassifyForwardSpec.traverse {α : Type} [TorchLean.Storage α] [Context α] {β : Type} {nFeatures : ℕ} (input : TorchLean.Tensor α [nFeatures]) {depth : ℕ} (node : ClassifierTreeNode α β nFeatures depth) : β

structure Spec.ClassificationExample {α : Type} [TorchLean.Storage α] (nFeatures : ℕ) (β : Type) : Type

A single classification training example: feature vector x and label y.

• x : TorchLean.Tensor α [nFeatures]

  The feature vector.

• y : β

  The class label.

def Spec.GradientBoostedTrees.Internal.countEq {β : Type} [DecidableEq β] (lbl : β) (ys : Array β) : ℕ

Count how many times lbl appears in ys.

def Spec.GradientBoostedTrees.Internal.distinct {β : Type} [DecidableEq β] (ys : Array β) : Array β

Remove repeated labels while preserving their first-occurrence order.

def Spec.majorityLabel {β : Type} [DecidableEq β] [Inhabited β] (ys : Array β) : β

Majority label with deterministic tie-breaking.

If there is a tie, we keep the earlier winner from the fold. This is intentional: it avoids non-determinism and keeps the spec stable across backends.

def Spec.GradientBoostedTrees.Internal.gini {α : Type} [Context α] {β : Type} [DecidableEq β] (ys : Array β) : α

Gini impurity of a multiset of labels.

gini(ys) = 1 - Σ_c p(c)^2 where p(c) is the empirical class frequency.

This is the standard CART impurity used by many tree classifiers.

def Spec.giniWeighted {α : Type} [Context α] {β : Type} [DecidableEq β] (ys : Array β) : α

Weighted Gini impurity: |ys| * gini(ys).

def Spec.classTargets {α : Type} [TorchLean.Storage α] {nFeatures : ℕ} {β : Type} (xs : Array (ClassificationExample nFeatures β)) : Array β

Extract the labels from an array of classification examples.

def Spec.partitionBySplitC {α : Type} [TorchLean.Storage α] [Context α] {nFeatures : ℕ} {β : Type} (feature : Fin nFeatures) (threshold : α) (xs : Array (ClassificationExample nFeatures β)) : Array (ClassificationExample nFeatures β) × Array (ClassificationExample nFeatures β)

Partition classification samples into (left, right) for a candidate split.

def Spec.GradientBoostedTrees.Internal.splitScoreC {α : Type} [TorchLean.Storage α] [Context α] {nFeatures : ℕ} {β : Type} [DecidableEq β] (feature : Fin nFeatures) (threshold : α) (xs : Array (ClassificationExample nFeatures β)) : Option (α × Array (ClassificationExample nFeatures β) × Array (ClassificationExample nFeatures β))

Score a candidate classification split (feature, threshold) by weighted Gini impurity.

Returns none when the split is degenerate (one side is empty); otherwise returns the score and the (left, right) partitions.

def Spec.bestSplitC {α : Type} [TorchLean.Storage α] [Context α] {nFeatures : ℕ} {β : Type} [DecidableEq β] (xs : Array (ClassificationExample nFeatures β)) : Option (Fin nFeatures × α × Array (ClassificationExample nFeatures β) × Array (ClassificationExample nFeatures β) × α)

Find the best classification split (feature, threshold) by exhaustive search.

Thresholds are drawn from the observed feature values in the dataset.

def Spec.fitClassificationNode {α : Type} [TorchLean.Storage α] [Context α] {nFeatures : ℕ} {β : Type} [DecidableEq β] [Inhabited β] (depth : ℕ) : Array (ClassificationExample nFeatures β) → ClassifierTreeNode α β nFeatures depth

Fit a classification tree node by greedy CART-style splitting (Gini impurity).

depthLeft counts how many more splits we are allowed to make.

def Spec.decisionTreeFitClassificationGiniSpec {α : Type} [TorchLean.Storage α] [Context α] {β : Type} [TorchLean.Storage β] [DecidableEq β] [Inhabited β] {batch maxDepth nFeatures : ℕ} (x : TorchLean.Tensor α [batch, nFeatures]) (y : TorchLean.Tensor β [batch]) : DecisionTreeClassifierSpec α β nFeatures maxDepth

Fit a classification decision tree (CART-style) using Gini impurity.

The label vector has exactly one element per input row. Its length is part of the type, so fitting cannot silently discard labels or invent missing ones.

def Spec.gbtMseLossSpec {α : Type} [TorchLean.Storage α] [Context α] {batch nTrees maxDepth nFeatures : ℕ} (model : GradientBoostedTreesSpec α nFeatures nTrees maxDepth) (input : TorchLean.Tensor α [batch, nFeatures]) (target : TorchLean.Tensor α [batch]) (h : batch ≠ 0) : TorchLean.Tensor α Shape.scalar

Mean squared error (MSE) loss for regression, reduced to a scalar by averaging over the batch.

def Spec.computeResidualsSpec {α : Type} [TorchLean.Storage α] [Context α] {batch nTrees maxDepth nFeatures : ℕ} (model : GradientBoostedTreesSpec α nFeatures nTrees maxDepth) (input : TorchLean.Tensor α [batch, nFeatures]) (target : TorchLean.Tensor α [batch]) : TorchLean.Tensor α [batch]

Residual computation for gradient boosting.

For squared-error regression, the residual is target - prediction.

def Spec.gradientBoostedTreesTrainStepSpec {α : Type} [TorchLean.Storage α] [Context α] {batch nTrees maxDepth nFeatures : ℕ} (model : GradientBoostedTreesSpec α nFeatures nTrees maxDepth) (input : TorchLean.Tensor α [batch, nFeatures]) (target : TorchLean.Tensor α [batch]) (newTree : DecisionTreeSpec α nFeatures maxDepth) (h : batch ≠ 0) : TorchLean.Tensor α Shape.scalar × GradientBoostedTreesSpec α nFeatures (nTrees + 1) maxDepth

One gradient-boosting "add a tree" step, given a pre-fit newTree.

This returns the loss before the update and the model with newTree appended. gradientBoostedTreesTrainStepFitSpec also fits the new tree to the current residuals.

Gradient boosting: a "fit-one-more-tree" step

The original gradientBoostedTreesTrainStepSpec expects a pre-fit newTree. For a more complete baseline, we also provide a deterministic step that fits that tree to the residuals.

def Spec.gradientBoostedTreesTrainStepFitSpec {α : Type} [TorchLean.Storage α] [Context α] {batch nTrees maxDepth nFeatures : ℕ} (model : GradientBoostedTreesSpec α nFeatures nTrees maxDepth) (input : TorchLean.Tensor α [batch, nFeatures]) (target : TorchLean.Tensor α [batch]) (h : batch ≠ 0) : TorchLean.Tensor α Shape.scalar × GradientBoostedTreesSpec α nFeatures (nTrees + 1) maxDepth

Fit a new tree to residuals and append it to the ensemble.

def Spec.GradientBoostedTrees.Internal.incrFeature {α : Type} [TorchLean.Storage α] [Context α] {nFeatures : ℕ} (acc : TorchLean.Tensor α [nFeatures]) (feature : Fin nFeatures) : TorchLean.Tensor α [nFeatures]

Increment a single feature counter by 1 inside a length-nFeatures vector.

This is used by the split-count feature-importance computation below.

def Spec.treeFeatureCounts {α : Type} [TorchLean.Storage α] [Context α] {nFeatures depth : ℕ} : TreeNode α nFeatures depth → TorchLean.Tensor α [nFeatures] → TorchLean.Tensor α [nFeatures]

Count how many times each feature index appears in split nodes of a tree.

This mirrors a very common "split count" importance heuristic.

def Spec.computeFeatureImportanceSpec {α : Type} [TorchLean.Storage α] [Context α] {nTrees maxDepth nFeatures : ℕ} (model : GradientBoostedTreesSpec α nFeatures nTrees maxDepth) : TorchLean.Tensor α [nFeatures]

Simple split-count feature importance for an ensemble.

This mirrors the common "how often was a feature used in a split?" heuristic. It is not the same as gain-based importance in XGBoost/LightGBM, but it is deterministic and easy to interpret.

@[irreducible]

def Spec.computeFeatureImportanceSpec.accumulate_importance {α : Type} [TorchLean.Storage α] [Context α] {nTrees maxDepth nFeatures : ℕ} (model : GradientBoostedTreesSpec α nFeatures nTrees maxDepth) (i : ℕ) (acc : TorchLean.Tensor α [nFeatures]) : TorchLean.Tensor α [nFeatures]