Skip to content

Use MSE criterion in base tree regressor - #492

Open
slievens wants to merge 1 commit into
smartcorelib:mainfrom
slievens:use-criterion
Open

slievens wants to merge 1 commit into
smartcorelib:mainfrom
slievens:use-criterion

Conversation

@slievens

Copy link
Copy Markdown
Contributor

This is the first step in making base tree "regressor" be used by all the tree classifiers. Next step is to make the implementation generic over the used Criterion with callers deciding it.

Add a Criterion trait with MSE and Gini implementations in src/tree/criterion.rs. Use MeanStats in BaseTreeRegressor in place of explicit sum/mass bookkeeping.

  • Store node stats in NodeVisitor; remove true/false_child_output
  • Keep scratch stats in SplitWorkspace to avoid new allocations
  • Rank splits by proxy improvement; store the real impurity improvement in split_score
  • Pass the sample count to find_best_cutoff and check min_samples_split at the caller

Fixes #491

Checklist

  • [ x] My branch is up-to-date with main branch.
  • [ x] Everything works and tested on latest stable Rust.
  • [ x] Coverage and Linting have been applied

Current behaviour

BaseTreeRegressor has hardcoded MSE stats.

New expected behaviour

BaseTreeRegressor uses a MSECriterion.

Change logs

Added

  • src/tree/criterion.rs File with the Criterion trait and implementations for MSE and Gini.
  • src/tree/base_tree_regressor.rs Uses MSECriterion.
  • src/tree/mod.rs Declares criterion as pub(mod)

This is the first step in making base tree "regressor" be used by all the
tree classifiers. Next step is to make the implementation generic over the
used Criterion with callers deciding it.

Add a Criterion trait with MSE and Gini implementations in
src/tree/criterion.rs. Use MeanStats in BaseTreeRegressor in place of
explicit sum/mass bookkeeping.

- Store node stats in NodeVisitor; remove true/false_child_output
- Keep scratch stats in SplitWorkspace to avoid new allocations
- Rank splits by proxy improvement; store the real impurity
  improvement in split_score
- Pass the sample count to find_best_cutoff and check
  min_samples_split at the caller
@slievens
slievens requested a review from Mec-iS as a code owner October 14, 2026 10:52
@codecov

codecov Bot commented Oct 14, 2026

Copy link
Copy Markdown

Codecov Report

✅ All modified and coverable lines are covered by tests.
✅ Project coverage is 85.72%. Comparing base (9eaae9e) to head (39f64c1).
⚠️ Report is 200 commits behind head on main.

Additional details and impacted files
@@             Coverage Diff             @@
##             main     #492       +/-   ##
===========================================
+ Coverage   43.97%   85.72%   +41.75%     
===========================================
  Files          85       97       +12     
  Lines        7281     9411     +2130     
===========================================
+ Hits         3202     8068     +4866     
+ Misses       4079     1343     -2736     

☔ View full report in Codecov by Harness.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:
  • ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

@Mec-iS

Mec-iS commented Oct 14, 2026 •

Copy link
Copy Markdown
Collaborator

Thanks @slievens. Centralising the sufficient statistics in one place is a good direction, and recomputing child leaf values from true_stats/false_stats in split_node is cleaner than carrying true_child_output/false_child_output through the visitor. Notes from reading the diff:

Process

  • GiniCriterion/GiniStats (plus roughly 250 lines of tests) are unused in this PR and are kept alive with #[cfg_attr(not(test), expect(dead_code))]. The title and description only cover the regressor. Please move the Gini code to the follow-up PR that wires it into the classifier.
  • Project conventions are to limit comments and not to add TODOs. // TODO: add the entropy and classification-error criteria. should become an issue. The long module-level docs and several inline comments could be trimmed.
  • #[expect(...)] needs Rust 1.81+. Please confirm it fits the crate's MSRV. It goes away if the Gini code is dropped.

Possible behaviour changes in base_tree_regressor.rs

  1. find_best_split now seeds prevx from sorted_by_feature[j][visitor.start_idx]. The old Option handled an empty range, but this indexes unconditionally. It panics if start_idx == end_idx (or start_idx == len). find_random_split guards this case explicitly. Please add the same early return, or show that the callers guarantee a non-empty range.
  2. The n < min_samples_split check moved out of find_best_cutoff. The root now checks it explicitly, but the children rely on true_can_split/false_can_split. Please confirm those conditions include min_samples_split, and add a test where min_samples_split is the binding constraint on a child.
  3. split_score meaning changes. The search functions now store the proxy (sum²/mass terms, without the parent term), and split_node overwrites it with impurity_improvement only when the split is accepted. If a node is evaluated but not split (for example min_samples_leaf rejects it in split_node), it keeps the proxy value. Please check whether split_score is read anywhere else, such as feature importance, and either always store the real improvement or use a separate field for the proxy.
  4. difference() computes parent - part for both sum and mass. With fractional sample weights, the residue can leave mass at a tiny positive value, so sum²/mass is evaluated on noise. The old code had the same subtraction, so this is not a regression. A small relative epsilon on the mass > 0.0 guards, or a test with fractional weights and a near-empty side, would document the behaviour.

Tests and evidence

  • All new tests are unit tests of the criterion. Please add a regression test that fits the pre-PR and post-PR regressor on fixed seeds (both Splitter::Best and Splitter::Random, with and without sample_weight, and with max_features < n_features). It should assert identical split_feature, split_value and predictions. This is what shows the refactor is behaviour-preserving.
  • Please include a before/after fit-time benchmark. The change adds scratch stats to SplitWorkspace and calls through #[inline(always)] methods in the hot loop, so it is probably neutral, but the numbers would confirm it.

Design

  • The Criterion trait is generic, but the regressor hard-codes MSECriterion at every call site. Until a second regressor criterion exists, the trait is speculative. Either make BaseTreeRegressor generic over C: Criterion, or keep the trait minimal and add the generic parameter when it is needed.

Happy to re-review after the retarget and the empty-range fix.

@Mec-iS

Mec-iS commented Oct 14, 2026

Copy link
Copy Markdown
Collaborator

Thanks @slievens, moving the regressor onto a Criterion abstraction is a good direction, and the incremental MeanStats + difference approach (as in sklearn) keeps the maths equivalent to the old gain. I reviewed the diff and have the following points.

Possible behaviour changes (please verify)

  1. min_samples_split was removed from find_best_cutoff and is now checked only for the root (n >= min_samples_split in fit). The child calls (find_best_cutoff(&mut true_visitor, ..., true_count) and the false one) are guarded only by true_can_split / false_can_split. Unless those flags already include min_samples_split, children smaller than min_samples_split will now be split further. Please confirm, or restore the check inside find_best_cutoff.
  2. In find_best_split, workspace.sorted_by_feature[j][visitor.start_idx] is indexed before the loop. find_random_split guards start_idx == end_idx, but this function does not, so an empty range would panic. The old Option logic did not have this problem. Add an early return, or explain why the range can never be empty.
  3. false_stats is computed by parent - true at every level, so floating-point residue accumulates down the false path. A child that should have zero mass can end up with mass of about 1e-16 and a garbage leaf_value, because the mass > 0.0 guard passes. The old code had the same weakness, but it is now centralised, so a tolerance or a direct recompute is cheap to add.
  4. split_score now holds the proxy during the search and is overwritten with impurity_improvement in split. If any other code reads split_score before split succeeds (for example feature importance on a node whose split was rejected), it will see the proxy value. Please check the consumers.

Tests

  • All the new tests are unit tests of criterion.rs. There is no tree-level regression test showing that fitted trees (split features, thresholds, leaf values, feature_importances) are unchanged against the previous implementation. A golden test on a fixed seed with sample weights, Splitter::Random and Splitter::Best, and min_samples_leaf / min_samples_split > 1 would cover the points above.
  • Please add a benchmark or a timing note for fit before and after, since add / difference now run in the hot loop.

Scope and project conventions

  • GiniCriterion, GiniStats and about 15 Gini tests are not used by this PR (hence the expect(dead_code) attributes). This is dead code in the PR, and it adds roughly 450 lines to a PR titled "Use MSE criterion in base tree regressor". Please move them to the follow-up PR that wires up the classifier.
  • #[expect(...)] needs Rust 1.81 or later. Please check it against the crate's MSRV and the CI toolchain.
  • // TODO: add the entropy and classification-error criteria. goes against the project rule of opening an issue instead of leaving TODOs. Please remove it and open an issue.
  • The target branch is main. The project convention is development, so please retarget.
  • GiniStats holds a Vec, so empty_stats() allocates per node. That is fine for now, but it is worth noting before the classifier uses it, given the zero-copy goals.
  • SplitWorkspace hardcodes MeanStats, so the trait is not yet usable generically. Please say whether making the workspace generic over Criterion::Stats is planned for the classifier PR.

Nits

  • The for k in 0..mtry.min(workspace.variables.len()) change is needed to avoid the borrow conflict. A short comment would help, or keep the .take(mtry) form if it compiles.
  • The module doc cites sklearn's _criterion.pyx. Please add the proxy-improvement derivation ((\sum_{side} S^2/W)) briefly to the Criterion docs.

Happy to re-review once 1 and 2 are clarified.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

Make BaseTreeRegressor use a "Criterion"

2 participants