dmlc / dmlc/xgboost

RFC: Honest categorical splitting for partition-based categorical features

Open
#12,130 9 comments 1 reaction 0 assignees View on GitHub
Dominant language
C++
Stars
28.8k
Forks
8.9k
Avg merge
1d 12h
Merged PRs (30d)
54

Description

## Summary

Add a boolean training option for partition-based categorical splits that learns the category order on one subsample and evaluates candidate split gains on a disjoint subsample.

The goal is to reduce the extra selection bias that can make target-adaptive categorical splits look artificially better than numeric threshold splits.

This is intended as a simple on/off regularization mode, not another hyperparameter.

## Motivation

For a numeric feature, split search is performed over thresholds on a pre-existing feature order.

For a partition-based categorical feature, the algorithm first computes category-level statistics from the training objective, then sorts categories by those statistics, and finally searches threshold splits along that learned order.

So a categorical partition split is more adaptive than a numeric split in two distinct ways:

1. it searches over thresholds
2. it also learns the ordering being thresholded from the same response-derived information

That extra adaptivity can inflate in-sample gain and make trees prefer categorical splits more often than they should, especially for high-cardinality features with sparse levels.

The proposal is to regularize exactly this extra adaptivity.

## Problem statement

At a node `I`, let the usual second-order gain for a candidate split `s` be:

```text
Gamma_I(s) =
1/2 * (
G_L(s)^2 / (H_L(s) + lambda)
+ G_R(s)^2 / (H_R(s) + lambda)
- G_I^2 / (H_I + lambda)
) - gamma
```

For a fixed candidate class, selecting the split with highest in-sample gain already introduces winner’s-curse optimism.

For partition-based categorical splitting, there is an additional source of optimism because the same sample is used both to:

1. construct the candidate class through the learned category order
2. score splits within that class

A useful decomposition is:

```text
Bias_cat ~= Bias_threshold_search + Bias_adaptive_ordering
```

Numeric threshold splits mainly incur the first term.

Partition-based categorical splits incur both.

The proposal below targets the second term directly.

## Proposed change

Add a boolean parameter, for example:

`honest_categorical_split = false | true`

When `honest_categorical_split = false`, use the current partition-based categorical split procedure.

When `honest_categorical_split = true`, for each partition-based categorical feature at a node:

1. Randomly split the node rows into two disjoint subsets `A` and `B`
2. Use subset `A` only to compute per-category statistics and derive the category order
3. Construct candidate prefix partitions from that order
4. Evaluate candidate split gains only on subset `B`
5. Select the best categorical split using those gains
6. After the split is chosen, recompute the chosen split’s node statistics on all rows in the node for downstream tree growth

In notation:

```text
s_hat_standard = argmax over s in S(Y) of Gamma_Y(s)

s_hat_honest = argmax over s in S(Y_A) of Gamma_Y_B(s)
```

The key idea is that the category order is learned on one sample and scored on another.

## Why this should work

Conditioned on subset `A`, the category order and therefore the candidate family are fixed relative to subset `B`.

From the point of view of `B`, the algorithm is now only searching thresholds on a fixed order.

That makes the gain-evaluation problem closer to the numeric-split case, where the order is exogenous.

In other words, honest categorical splitting converts the category order from endogenous to effectively exogenous with respect to the scoring sample.

This should reduce the extra optimism of partition-based categorical gains and therefore reduce spurious preference for categorical splits over numeric splits.

A compact way to say it is:

> Standard partition-based categorical splitting is biased because the same sample both learns and scores the category order. Honest categorical splitting breaks that feedback loop.

## Why make this on/off instead of another hyperparameter

This proposal is intended as a robust regularization mode, not another tuning knob.

The core issue is structural:

* numeric splits search thresholds on an exogenous order
* partition-based categorical splits search thresholds on an order learned from the same training responses

Honesty addresses that asymmetry directly.

A boolean switch keeps the behavior easy to explain and avoids introducing another loss-scaled or data-scaled parameter that users would have to tune.

## Expected tradeoff

### Expected benefits

* lower optimism for partition-based categorical gains
* less spurious preference for high-cardinality categorical features
* better comparability between categorical and numeric split selection
* more robust default behavior without extra tuning

### Expected costs

* higher variance from using fewer rows to learn the category order
* possibly weaker split discovery when node sample sizes are very small
* some reduction in raw training gain

This is a standard bias-variance tradeoff: less adaptive search, less overfitting, but also less effective sample size for split finding.

## Suggested implementation shape

A minimal API could be:

`honest_categorical_split = false`

Behavior:

* applies only to partition-based categorical splits
* has no effect on one-hot categorical splits
* has no effect on numeric features
* can coexist with `max_cat_threshold`

Implementation details:

* use the node RNG so behavior is deterministic under a fixed seed
* split rows roughly in half into `A` and `B`
* learn ordering on `A`
* score candidate prefix splits on `B`
* after choosing the split, recompute final node statistics on all rows

The recomputation step is important: honesty is used for split selection, not for permanently discarding data in downstream tree growth.

## Why this is a good fit for XGBoost

This proposal targets exactly the overfitting channel unique to partition-based categorical splitting.

It does not change the standard gain formula.

It does not require a new shrinkage constant tied to gradient or Hessian scale.

It does not ask users to tune another categorical-specific penalty.

Instead, it changes the search procedure so that categorical splits are evaluated more fairly relative to numeric splits.

## Evaluation plan

### Experiment 1: null preference test

Generate:

* one informative numeric feature
* one useless high-cardinality categorical feature
* several useless numeric and categorical noise features

Measure:

* probability that the first split selects the useless categorical
* training gain minus fresh-sample gain of the chosen split
* frequency of categorical versus numeric selection

Hypothesis:

* standard partitioning over-selects the useless categorical
* honest categorical splitting reduces that preference gap

### Experiment 2: matched-signal numeric versus categorical

Construct a numeric feature and a categorical feature with the same true predictive strength.

Measure:

* first-split selection frequency
* out-of-sample gain
* validation loss

Hypothesis:

* standard partitioning prefers the categorical too often
* honest splitting brings the selection frequencies closer

### Experiment 3: real datasets with sparse categoricals

Use datasets with a mix of low- and high-cardinality categorical features.

Measure:

* validation loss
* split usage by feature type
* depth distribution of categorical splits
* gap between training gain and fresh-sample gain for chosen categorical splits

## Related issues

* `#6503` Categorical data support
* `#7899` Categorical data support (part 2)
* `#10844` What does `max_cat_threshold` actually control?

Contributor guide

No contributing guide indexed for this repository

Research direction

Start by reading the categorical data discussions in issues #6503, #7899, and #10844, then trace the existing partition-based categorical split path. Define the boolean option and its interaction with one-hot splits, numeric features, and max_cat_threshold. Done means the behavior is implemented with deterministic node-RNG handling, all-row statistic recomputation, and the proposed null, matched-signal, and real-dataset evaluations.

Written by the indexing model from the issue text.

Assessment

Tech stack
cpp
Domain
machine-learning
Issue type
Feature
Difficulty
5/5
Estimated time
Over a week
Activity status
Quiet
Clarity
Mostly clear
Newbie friendliness
35/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.