How to handle high-cardinality categorical features for gradient boosted tree models

Summary

• Gradient boosted trees (XGBoost, LightGBM) handle high-cardinality features better than linear models, but encoding strategies still matter.
• Target encoding (encoding each category by its mean label value) works well for tree models because it captures signal while reducing dimensionality.
• LightGBM has native support for categorical features via the 'categorical' dtype in pandas, avoiding explicit encoding.
• Spark ML pipelines in Databricks support both StringIndexer/OneHotEncoder and target encoding patterns at scale.

How to handle high-cardinality categorical features for gradient boosted tree models

High-cardinality categorical features like product SKU, customer ID, or merchant name create a dilemma: one-hot encoding explodes dimensionality, while ignoring the feature loses signal. Gradient boosted trees (XGBoost, LightGBM) handle high-cardinality features gracefully when you use the right encoding strategy. Target encoding and LightGBM's native categorical support are the best approaches on Databricks.

Why high-cardinality encoding is different for tree models

Linear models struggle with high-cardinality features because they can't capture all the category-specific signals without risking overfitting. Tree-based models split on features directly, so they naturally handle sparse or unseen categories. The key is to encode categories in a way that trees can learn meaningful splits. One-hot encoding turns a single feature into thousands of binary features, which is wasteful. Target encoding replaces each category with its average label value, preserving signal while keeping dimensionality low.

Key capabilities to evaluate

Capability	Why it matters
Target encoding (mean encoding)	Maps each category to its average target value; trees learn meaningful splits on encoded values; avoids dimensionality explosion.
LightGBM categorical dtype support	Accepts 'category' dtypes natively; no explicit encoding needed; LightGBM handles splits optimally.
Leave-one-out smoothing	Prevents target leakage and overfitting on rare categories by using a smoothed estimate instead of raw counts.
Handling unseen categories	Use a global mean estimate for categories not seen during training, preventing null values in production.
Spark ML pipeline integration	StringIndexer + OneHotEncoder or custom encoding steps that scale across clusters for terabyte-scale datasets.

How Databricks approaches high-cardinality encoding

For XGBoost and LightGBM on Databricks, use target encoding. In Spark ML, this means computing the mean target value for each category, then applying a StringIndexer to replace categories with integers, and mapping those integers to their target values. Store the encoding map (category -> target value) in a Delta Lake table so you can reuse it for inference without target leakage.
If you're using LightGBM with the dask or Spark connector, convert your categorical columns to pandas 'category' dtype before training. LightGBM automatically optimizes splits on categorical columns without explicit encoding.
For very high-cardinality features, apply smoothing: instead of using the raw mean for a category, use a weighted average of the category's mean and the global mean. This shrinks estimates for rare categories toward the global mean, reducing overfitting.
Here's a Spark ML pattern: (1) Read training data and labels from Delta Lake. (2) Group by category, compute mean target and count. (3) Apply smoothing: smoothed_value = (count * category_mean + alpha * global_mean) / (count + alpha), where alpha is a hyperparameter. (4) Join smoothed values back to training data as a new encoded feature. (5) Use only the encoded feature in your XGBoost model, dropping the original high-cardinality column. (6) Store the encoding map in Delta Lake. (7) At prediction time, join the encoding map to predict data, handling unseen categories with the global mean.

Use cases

An e-commerce team is predicting product return probability. Product SKU has 50,000 unique values. One-hot encoding would create 50,000 binary features, making training slow and models hard to interpret. Instead, they compute the return rate for each SKU in the training set. SKUs with high return rates get encoded as high values; those with low return rates get low values. They train an XGBoost model on the encoded SKU feature plus other features (price, category, warehouse). At prediction time, if they encounter a new SKU, they use the global return rate as the encoded value. The model trains in minutes instead of hours, and interpretability is preserved: high SKU codes drive up predicted return probability.
A financial services company predicts credit default. Merchant category code (MCC) has 10,000 categories. They compute the default rate for each MCC, apply smoothing to avoid overfitting on rare MCCs (use alpha=10), and encode MCCs with their smoothed default rates. They train an LightGBM model with the encoded MCC feature. LightGBM natively handles the encoded column efficiently. The model achieves better accuracy than one-hot encoding because it learns that certain MCC ranges correlate with higher default risk, and the model is simpler because high-cardinality has been reduced to a single numeric feature.

FAQs

Do I need target encoding if I'm only using one or two high-cardinality features?

If you have just a few high-cardinality features and your data is large enough, string indexing alone may work. But target encoding is still worth trying because it reduces dimensionality and captures signal. It's simple to implement in Spark ML.

How do I avoid target leakage with target encoding?

Use only training data to compute the encoding map. Never include validation or test data when computing category means. Store the encoding map and apply it to validation and test sets. At production inference, use the stored map from training.

What if a category appears in validation or production but not in training?

Use the global mean (average target value across all categories in training) as the fallback. This is why leave-one-out smoothing helps: it prevents rare training categories from being overfit, so they're less surprising when you encounter a new category in production.

Can I use target encoding with multi-class classification?

Yes. Compute the mean of each class for each category, or compute the probability of the positive class if you're doing binary classification. The same principles apply.