How to train Boosted Trees models in TensorFlow

And how to interpret them both locally and globally

Mar 5 · 8 min read
Visualizing the prediction surface of a Boosted Trees model. Gradient boosted trees is an ensemble technique that combines the predictions from several (think 10s, 100s or even 1000s) tree models. Increasing the number of trees will generally improve the quality of fit. Try the full example here.

Training a Boosted Trees Model in TensorFlow

Model Understanding

Understanding individual predictions: Directional feature contributions

DFCs for an instance in the titanic dataset. Informally this can be interpreted as “adult_male being True “contributes” about -0.11 to the final probability, while Deck being B contributes to about +0.08 to the final probability, and so on.”
{sex_female: 0.2, age: 0.05, fare= -0.02, num_siblings_aboard=-0.1, fare: 0.09}
Mean absolute values of top DFCs across the entire evaluation dataset.
fare value vs contributions with a LOWESS fit. Contributions across examples provide more granular information than a single feature importance metric. In general, a higher fare features results in the model pushing predictions closer to 1.0 (increasing chance of survival).

How to: Directional feature contributions in TensorFlow

Contributions for an individual example in red. The shaded blue area shows the distributions of contributions for features across the entire validation set.

Model-level interpretability: Gain-based and permutation feature importances

Correlated variables and other considerations


Additional Resources


TensorFlow is an end-to-end open source platform for machine learning.


Written by

TensorFlow is a fast, flexible, and scalable open-source machine learning library for research and production.


TensorFlow is an end-to-end open source platform for machine learning.