How to Get Feature Importance in XGBoost in Python

You’ve chosen XGBoost as your algorithm, and now you’re wondering: “How do I figure out which features are the most important in my model?”

That’s what ‘feature importance’ is all about. It’s a way of finding out which features in your data are doing the heavy lifting when it comes to your model’s predictions.

Understanding which features are important can help you interpret your model better.

Maybe you’ll find a feature you didn’t expect to be important. Or maybe you’ll see that a feature you thought was important isn’t really doing much at all.

In this tutorial, we’ll dive into how to implement this in your code and interpret the results.

How Is Feature Importance Calculated in XGBoost?

Now that you understand what feature importance is, let’s talk about the different types of feature importance you can get with XGBoost.

There are five main types: ‘weight’, ‘gain’, ‘cover’, ’total_gain’, and ’total_cover’.

Let’s break each one down:

1. Weight: This is the simplest type of feature importance. It’s just the number of times a feature is used to split the data across all trees. If a feature is used a lot, it has a high weight.

2. Gain: This is a bit more complex. Gain measures the average improvement in (training set) loss brought by a feature. In other words, it tells you how much a feature helps to make accurate predictions in your training data.

3. Cover: Cover measures the number of data points a given feature affects. This is done via taking the average of the hessian values of all splits the feature is used in.

4. Total_gain: This is similar to gain, but instead of the average, it gives you the total improvement brought by a feature.

5. Total_cover: This is like cover, but it gives you the total number, rather than the average.

So, when should you use each type?

Unfortunately, there’s no one-size-fits-all answer.

I like to compute all of them and then use my best judgement to decide which one is most useful for my particular problem.

Implementing XGBoost Feature Importance in Python

We’ll be using the Red Wine dataset from UCI for this example.

This dataset contains information about various attributes of red wines, such as ‘fixed acidity’, ‘volatile acidity’, ‘citric acid’, etc.

The task is to predict the quality of the wine based on these attributes.

First, you’ll need to import the necessary libraries.

We’ll need pandas for data manipulation, XGBoost for our model, and matplotlib for plotting our feature importances later.

import pandas as pd
from xgboost import XGBRegressor
import matplotlib.pyplot as plt

Next, let’s load our dataset. We’ll use pandas’ read_csv function for this.

df = pd.read_csv('winequality-red.csv', delimiter=';')

Now, we need to separate our features (X) from our target variable (y).

In this case, we’re trying to predict the ‘quality’ column.

So, we’ll drop that column from our dataframe and assign it to y.

Then, we’ll assign the rest of the dataframe to X.

X = df.drop('quality', axis=1)
y = df['quality']

Here I am not splitting the data into training and test sets to keep things brief, so assume that we have already done that.

The native feature importances are calculated using the training set only, as it’s the data used to create the decision trees.

Next, we’ll create our XGBoost model. We’ll use the XGBRegressor for this example.

Everything is valid for XGBClassifier as well.

Let’s use the default hyperparameters.

model = XGBRegressor()

Then, we fit our model to the data.

model.fit(X, y)

Now comes the fun part: getting our feature importances.

XGBoost sklearn API makes this easy with the feature_importances_ attribute.

To make it easier to understand, let’s put everything into a pandas Series where the index is the feature name and the values are the feature importances.

importances = model.feature_importances_
pd.Series(importances, index=X.columns).sort_values()

As you can see, the ‘alcohol’ feature is by far the most important, with the highest feature importance score of 0.358042.

By default, XGBoost will use the ‘gain’ type of feature importance.

Let’s try other types to see if we get different results.

You need to change it in the model initialization using the importance_type parameter.

model = XGBRegressor(importance_type='weight')
model.fit(X, y)

importances = model.feature_importances_
pd.Series(importances, index=X.columns).sort_values()

This time, the ‘fixed acidity’ feature is the most important, with a feature importance score of 0.155749.

This means it’s the one that’s used the most to split the data across all trees.

Be careful, as this can be misleading.

In the case of high cardinality categorical features, the ‘weight’ type of feature importance will be biased towards those features, as they naturally have more splits.

Random Forests suffer from the same problem.

Here is the same example with the ‘cover’ type of feature importance for completion.

model = XGBRegressor(importance_type='cover')

Finally, here are the results for all types (averaged, not doing the total ones to keep it brief):

Visualizing XGBoost Feature Importance

Finally, let’s plot our feature importances to make them easier to understand.

We’ll create a bar plot with our features on the x-axis and their importance on the y-axis.

As we have our feature importances organized in a pandas Dataframe, we can easily plot them using the integrated matplotlib functions.

res.plot.barh(figsize=(1280/96, 720/96), title='Feature importance')

Here I plotted all the types of feature importance for comparison, but you can plot only one if you want.

The biggest downside of these feature importances is that they don’t tell the direction of the relationship, only the strength of it.

So we don’t know if a feature is positively or negatively correlated with the target variable.

In any case, it’s still useful to get a general idea of which features are important.