From 4a52706e471a042179b6c3c94559a7f270716352 Mon Sep 17 00:00:00 2001 From: Maciej Myszkiewicz Date: Thu, 6 Aug 2026 15:20:57 +0200 Subject: [PATCH] fix(trees): make DecisionTree::fit deterministic MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Repeated `fit` calls on identical data could return differently shaped trees. On the iris dataset, 60 fits in one process produced 4 distinct trees, differing in leaf count (9, 10, 11 and 12) and therefore in their predictions. Class frequencies were carried in a `HashMap`, and `HashMap` iteration order is randomised per instance. Every consumer of those frequencies depended on the order: * `find_modal_class` folded with `best_freq > freq`, so on a tie it kept the *last* maximum it happened to visit — the opposite of the "first class found" its doc comment promises. * `gini_impurity` and `entropy` sum the frequencies in `f32`, and floating-point addition is not associative, so the same frequencies summed in a different order produce slightly different scores. Those score differences flip the `score < best_score` comparison in `TreeNode::fit` at near-ties, which selects a different split and changes the shape of the whole subtree. Frequencies are now `Vec` indexed by position in a sorted class list built once per fit, so the iteration order is pinned. The tree this produces is one of the trees the previous code already produced, just consistently rather than at random. As a side effect this is also faster: the innermost split-search loop ran `right_class_freq.get_mut(class)` and `left_class_freq.entry(class.clone())` for every (sample, feature) pair, i.e. two hash lookups plus a label clone, which for `String` labels meant an allocation. Indexing a `Vec` removes all three. n=2000 d=20 k=3 707.2 ms -> 271.5 ms (2.6x) n=8000 d=20 k=3 2570.1 ms -> 860.9 ms (3.0x) n=8000 d=20 k=10 2198.7 ms -> 861.1 ms (2.6x) n=16000 d=50 k=5 14652.5 ms -> 4494.8 ms (3.3x) Co-Authored-By: Claude Opus 5 (1M context) Claude-Session: https://claude.ai/code/session_01L4uUK6s1CWFwx6u7Rx28go --- .../src/decision_trees/algorithm.rs | 176 ++++++++++++++---- 1 file changed, 143 insertions(+), 33 deletions(-) diff --git a/algorithms/linfa-trees/src/decision_trees/algorithm.rs b/algorithms/linfa-trees/src/decision_trees/algorithm.rs index c507b0845..f6b1d41a6 100644 --- a/algorithms/linfa-trees/src/decision_trees/algorithm.rs +++ b/algorithms/linfa-trees/src/decision_trees/algorithm.rs @@ -1,11 +1,11 @@ //! Linear decision trees //! use std::cmp::Ordering; -use std::collections::{HashMap, HashSet, VecDeque}; +use std::collections::{HashSet, VecDeque}; use std::hash::{Hash, Hasher}; use linfa::dataset::AsSingleTargets; -use ndarray::{Array1, ArrayBase, Axis, Data, Ix1, Ix2}; +use ndarray::{Array1, ArrayBase, ArrayView1, Axis, Data, Ix1, Ix2}; use super::NodeIter; use super::Tikz; @@ -74,6 +74,67 @@ impl RowMask { } } +/// A stable, sorted view of the label space of a dataset. +/// +/// Class frequencies are carried as `Vec` indexed by position in `classes` +/// rather than as `HashMap`. `HashMap` iteration order is randomised per +/// instance, and every consumer of the frequencies depends on that order: +/// `find_modal_class` for its tie-break, and `gini_impurity`/`entropy` because +/// `f32` addition is not associative, so the same frequencies summed in a +/// different order give slightly different scores. Those differences flip the +/// `score < best_score` comparison in `TreeNode::fit` at near-ties, which changes +/// the chosen split — two `fit` calls on identical data could return differently +/// shaped trees. Indexing by a sorted class list pins the order, and also keeps +/// two hash lookups and a label clone out of the innermost split-search loop. +struct ClassIndex { + /// Labels occurring in the dataset, sorted and deduplicated. + classes: Vec, + /// For each observation, the position of its label in `classes`. + of_sample: Vec, +} + +impl ClassIndex { + fn build(targets: &ArrayView1) -> Self { + let mut classes: Vec = targets.iter().cloned().collect(); + classes.sort_unstable(); + classes.dedup(); + + let of_sample = targets + .iter() + .map(|label| { + classes + .binary_search(label) + .expect("every label was collected into the class list") + }) + .collect(); + + ClassIndex { classes, of_sample } + } + + fn n_classes(&self) -> usize { + self.classes.len() + } + + /// Weighted class frequencies over the observations selected by `mask`. + /// + /// Equivalent to `DatasetBase::label_frequencies_with_mask`, except that the + /// result is indexed by class position and classes absent from this subset are + /// present with a frequency of zero rather than being omitted. That makes no + /// difference downstream: a zero frequency contributes nothing to either + /// impurity measure and can never be the modal class of a non-empty node. + fn frequencies(&self, data: &DatasetBase, mask: &RowMask) -> Vec { + let mut freqs = vec![0.0; self.n_classes()]; + + for (i, class) in self.of_sample.iter().enumerate() { + if *mask.mask.get(i).unwrap_or(&true) { + freqs[*class] += data.weight_for(i); + } + } + + freqs + } +} + /// Sorted values of observations with indices (always for a particular feature) struct SortedIndex<'a, F: Float> { feature_name: &'a str, @@ -203,14 +264,13 @@ impl TreeNode { mask: &RowMask, hyperparameters: &DecisionTreeValidParams, sorted_indices: &[SortedIndex], + class_index: &ClassIndex, depth: usize, ) -> Result { // compute weighted frequencies for target classes - let parent_class_freq = data.label_frequencies_with_mask(&mask.mask); + let parent_class_freq = class_index.frequencies(data, mask); // set our prediction for this subset to the modal class - let prediction = find_modal_class(&parent_class_freq); - // get targets from dataset - let target = data.as_single_targets(); + let prediction = class_index.classes[find_modal_class(&parent_class_freq)].clone(); // return empty leaf when we don't have enough samples or the maximal depth is reached if (mask.nsamples as f32) < hyperparameters.min_weight_split() @@ -228,11 +288,11 @@ impl TreeNode { // Iterate over all features for (feature_idx, sorted_index) in sorted_indices.iter().enumerate() { let mut right_class_freq = parent_class_freq.clone(); - let mut left_class_freq = HashMap::new(); + let mut left_class_freq = vec![0.0; class_index.n_classes()]; // We keep a running total of the aggregate weight in the right split - // to avoid having to sum over the hash map - let total_weight = parent_class_freq.values().sum::(); + // to avoid having to sum over the frequencies + let total_weight = parent_class_freq.iter().sum::(); let mut weight_on_right_side = total_weight; let mut weight_on_left_side = 0.0; @@ -254,19 +314,19 @@ impl TreeNode { } // Target and weight of the current observation - let sample_class = &target[presorted_index]; + let sample_class = class_index.of_sample[presorted_index]; let sample_weight = data.weight_for(presorted_index); // Move the observation from the right subtree to the left subtree // Decrement the weight on the class for this sample on the right // side by the weight of this sample - *right_class_freq.get_mut(sample_class).unwrap() -= sample_weight; + right_class_freq[sample_class] -= sample_weight; weight_on_right_side -= sample_weight; // Increment the weight on the class for this sample on the // right side by the weight of this sample - *left_class_freq.entry(sample_class.clone()).or_insert(0.0) += sample_weight; + left_class_freq[sample_class] += sample_weight; weight_on_left_side += sample_weight; // Continue if the next value is equal, so that equal values end up in the same subtree @@ -362,6 +422,7 @@ impl TreeNode { &left_mask, hyperparameters, sorted_indices, + class_index, depth + 1, )?)) } else { @@ -374,6 +435,7 @@ impl TreeNode { &right_mask, hyperparameters, sorted_indices, + class_index, depth + 1, )?)) } else { @@ -537,7 +599,10 @@ where }) .collect(); - let mut root_node = TreeNode::fit(dataset, &all_idxs, self, &sorted_indices, 0)?; + let class_index = ClassIndex::build(&dataset.as_single_targets()); + + let mut root_node = + TreeNode::fit(dataset, &all_idxs, self, &sorted_indices, &class_index, 0)?; root_node.prune(); Ok(DecisionTree { @@ -649,34 +714,31 @@ fn make_prediction( /// Finds the most frequent class for a hash map of frequencies. If two /// classes have the same weight then the first class found with that /// frequency is returned. -fn find_modal_class(class_freq: &HashMap) -> L { - // TODO: Refactor this with fold_first - - let val = class_freq +fn find_modal_class(class_freq: &[f32]) -> usize { + class_freq .iter() + .enumerate() .fold(None, |acc, (idx, freq)| match acc { None => Some((idx, freq)), Some((_best_idx, best_freq)) => { - if best_freq > freq { + if best_freq >= freq { acc } else { Some((idx, freq)) } } }) - .unwrap() - .0; - - (*val).clone() + .expect("a node always has at least one class") + .0 } /// Given the class frequencies calculates the gini impurity of the subset. -fn gini_impurity(class_freq: &HashMap) -> f32 { - let n_samples = class_freq.values().sum::(); +fn gini_impurity(class_freq: &[f32]) -> f32 { + let n_samples = class_freq.iter().sum::(); assert!(n_samples > 0.0); let purity = class_freq - .values() + .iter() .map(|x| x / n_samples) .map(|x| x * x) .sum::(); @@ -685,12 +747,12 @@ fn gini_impurity(class_freq: &HashMap) -> f32 { } /// Given the class frequencies calculates the entropy of the subset. -fn entropy(class_freq: &HashMap) -> f32 { - let n_samples = class_freq.values().sum::(); +fn entropy(class_freq: &[f32]) -> f32 { + let n_samples = class_freq.iter().sum::(); assert!(n_samples > 0.0); class_freq - .values() + .iter() .map(|x| x / n_samples) .map(|x| if x > 0.0 { -x * x.log2() } else { 0.0 }) .sum() @@ -725,14 +787,62 @@ mod tests { let row_mask = RowMask::all(labels.len()); let dataset: DatasetBase<(), Array1> = DatasetBase::new((), labels); - let class_freq = dataset.label_frequencies_with_mask(&row_mask.mask); + let class_index = ClassIndex::build(&dataset.as_single_targets()); + let class_freq = class_index.frequencies(&dataset, &row_mask); + + assert_eq!(class_freq, vec![6.0, 2.0]); + assert_eq!(class_index.classes[find_modal_class(&class_freq)], 0); + } + + #[test] + fn fit_is_deterministic() { + // Class frequencies used to live in a `HashMap`, whose iteration order is + // randomised per instance. That order leaked into the split scores through + // the non-associative `f32` sums in `gini_impurity`/`entropy`, so repeated + // `fit` calls on identical data could return differently shaped trees. + let data = array![ + [0.1, 1.0], + [0.2, 0.9], + [0.3, 1.1], + [1.1, 0.1], + [1.2, 0.2], + [1.3, 0.15], + [2.1, 2.0], + [2.2, 2.1], + [2.3, 1.9], + ]; + let targets = array![0, 0, 0, 1, 1, 1, 2, 2, 2]; + let dataset = Dataset::new(data, targets); + + let reference = DecisionTree::params().fit(&dataset).unwrap(); + let reference: Vec<_> = reference + .iter_nodes() + .map(|node| (node.split(), node.prediction())) + .collect(); + + for _ in 0..16 { + let tree = DecisionTree::params().fit(&dataset).unwrap(); + let shape: Vec<_> = tree + .iter_nodes() + .map(|node| (node.split(), node.prediction())) + .collect(); - assert_eq!(find_modal_class(&class_freq), 0); + assert_eq!(shape, reference); + } + } + + #[test] + fn modal_class_ties_pick_the_first_class() { + // documented behaviour: "If two classes have the same weight then the + // first class found with that frequency is returned" + assert_eq!(find_modal_class(&[2.0, 2.0, 1.0]), 0); + assert_eq!(find_modal_class(&[1.0, 2.0, 2.0]), 1); + assert_eq!(find_modal_class(&[0.0, 0.0, 3.0]), 2); } #[test] fn gini_impurity_example() { - let class_freq = vec![(0, 6.0), (1, 2.0), (2, 0.0)].into_iter().collect(); + let class_freq = [6.0, 2.0, 0.0]; // Class 0 occurs 75% of the time // Class 1 occurs 25% of the time @@ -743,7 +853,7 @@ mod tests { #[test] fn entropy_example() { - let class_freq = vec![(0, 6.0), (1, 2.0), (2, 0.0)].into_iter().collect(); + let class_freq = [6.0, 2.0, 0.0]; // Class 0 occurs 75% of the time // Class 1 occurs 25% of the time @@ -752,7 +862,7 @@ mod tests { assert_abs_diff_eq!(entropy(&class_freq), 0.81127, epsilon = 1e-5); // If split is perfect then entropy is zero - let perfect_class_freq = vec![(0, 8.0), (1, 0.0), (2, 0.0)].into_iter().collect(); + let perfect_class_freq = [8.0, 0.0, 0.0]; assert_abs_diff_eq!(entropy(&perfect_class_freq), 0.0, epsilon = 1e-5); }