Select Features for Machine Learning Model with Mutual Information | by Andrew Zhu | Towards Data Science

Machine Learning models are amazing when trained with an appropriate set of training data. ML models described in the textbook and using datasets from Scikit-learn, sample datasets are already carefully curated and effective for targeting ML models.
The problem is, in real life, training data could be a mass and not fit for any models at all. For example, you want to train a model to predict the next month's revenue based on current 100+ customer properties, like age, location, country, active visits #, last time visits time... on and on.
Some of those properties may useful, some could be completely useless, say, customer Name.
No wonder that many ML experts say, pick the right features for an ML is one of the most important steps to take, even important than the ML model itself!
Mutual Information can answer the question: Is there a way to build a measurable connection between a feature and target.
Two benefits to using Mutual Information as feature selector:
The MI is model neutral, which means the solution can be applied to various kinds of ML models.
MI solution is fast.
So, What is Mutual Information?
If you are familiar with the Decision Tree Classifier. It works 100% the same as Information Gain which I described in anther article Understand Decision Tree Classifiers.
Mutual Information measures the entropy drops under the condition of the target value. I found the cleanest explanation to this concept is this formula:
MI(feature;target) = Entropy(feature) - Entropy(feature|target)
The MI score will fall in the range from 0 to 1. The higher value, the closer connection between this feature and the target, suggests that we should put this feature in the training dataset. If the MI score is 0 or very low like 0.01. the low score suggests a weak connection between this feature and the target.
You can write a MI function from scratch on your own, for fun, or use the ready-to-use functions from Scikit-Learn.
I am going to use the Breast Cancer dataset from Scikit-Learn to build a sample ML model with Mutual Information applied. Use the Decision Tree Classifier to train 3 datasets from the cancer data and compare the result to see how MI score will impact the ML model effectiveness.
Train dataset 1, use all features.Train dataset 2, use only features whose MI scores are larger than 0.2Train dataset 3, use only features whose MI scores are less than 0.2
Train dataset 1, use all features.
Train dataset 2, use only features whose MI scores are larger than 0.2
Train dataset 3, use only features whose MI scores are less than 0.2
Step 1. load breast cancer data
from sklearn.datasets import load_breast_cancer as LBCcancer = LBC()X = cancer['data']y = cancer['target']
Step 2. compute MI score
from sklearn.feature_selection import mutual_info_classif as MICmi_score = MIC(X,y)print(mi_score)
You shall see the mi_score array like this:
[0.37032947 0.09670886 0.40294198 0.36009957 0.08427789 0.21318114 0.37337734 0.43985571 0.06456878 0.00276314 0.24866738 0.00189163 0.27600984 0.33955538 0.01503326 0.07603828 0.11825812 0.12879402 0.0096701  0.03802394 0.45151801 0.12293047 0.47595645 0.46426102 0.09558928 0.22647456 0.31469449 0.43696443 0.0971793  0.06735096]
30 numbers represent MI score for 30 features.
Step 3. Build 3 train data and test datasets
Dataset 1 with no MI score curation, names marked with _1. Note that all three datasets will use the same y_train and y_test. So, no need to separate target data.
from sklearn.model_selection import train_test_split as ttsX_train_1,X_test_1,y_train,y_test = tts(    X,y    ,random_state=0    ,stratify=y)
Dataset 2 with features have MI scores larger than 0.2
import numpy as npmi_score_selected_index = np.where(mi_scores >0.2)[0]X_2 = X[:,mi_score_selected_index]X_train_2,X_test_2,y_train,y_test = tts(    X_2,y    ,random_state=0    ,stratify=y)
You will see the shape of X_2 is (569,15) instead of (569,30). That is because 15 features have Mi scores larger than 0.2.
Dataset 3 with features have Mi scores less than 0.2
mi_score_selected_index = np.where(mi_scores < 0.2)[0]X_3 = X[:,mi_score_selected_index]X_train_3,X_test_3,y_train,y_test = tts(    X_3,y    ,random_state=0    ,stratify=y)
Coincidentally, the 0.2 MI score separates 30 features to 15 and 15.
Step 4. Compare the 3 datasets with the Decision Tree classifier
from sklearn.tree import DecisionTreeClassifier as DTCmodel_1 = DTC().fit(X_train_1,y_train)model_2 = DTC().fit(X_train_2,y_train)model_3 = DTC().fit(X_train_3,y_train)score_1 = model_1.score(X_test_1,y_test)score_2 = model_2.score(X_test_2,y_test)score_3 = model_3.score(X_test_3,y_test)print(f"score_1:{score_1}\n score_2:{score_2}\n score_3:{score_3}")
The result:
score_1:0.9300699300699301score_2:0.9370629370629371score_3:0.8251748251748252
See, dataset 2 with 15 features has MI > 0.2 reach to accuracy 0.93 as good as dataset 1, which includes all features. While score_3 is merely 0.82 which is the result of 15 features that have MI score < 0.2.
From this sample, it is clear that the MI score can be used as a signal for feature selection.
In real ML projects, you may want to use the top n features, or top n percentile features instead of using a specified number 0.2 like the sample above. Scikit-Learn also provides many selectors as convenient tools. So that you don’t have to manually calculate MI scores and take the needed features. Here is a sample to select the top 50% of features, other selectors share similar usage.
from sklearn.feature_selection import SelectPercentile as SPselector = SP(percentile=50) # select features with top 50% MI scoresselector.fit(X,y)X_4 = selector.transform(X)X_train_4,X_test_4,y_train,y_test = tts(    X_4,y    ,random_state=0    ,stratify=y)model_4 = DTC().fit(X_train_4,y_train)score_4 = model_4.score(X_test_4,y_test)print(f"score_4:{score_4}")
Information Gain and Mutual Information for Machine Learning
1.13. Feature selection
Understand Decision Tree Classifiers
You can copy and run the below code in the Jupyter notebook or vscode python interactive window.
#%% load cancer datafrom sklearn.datasets import load_breast_cancer as LBCcancer = LBC()X = cancer['data']y = cancer['target']#%% compute MI scoresfrom sklearn.feature_selection import mutual_info_classif as MICmi_scores = MIC(X,y)print(mi_scores)#%% prepare dataset 1from sklearn.model_selection import train_test_split as ttsX_train_1,X_test_1,y_train,y_test = tts(    X,y    ,random_state=0    ,stratify=y)#%% prepare dataset 2, MI > 0.2import numpy as npmi_score_selected_index = np.where(mi_scores >0.2)[0]X_2 = X[:,mi_score_selected_index]X_train_2,X_test_2,y_train,y_test = tts(    X_2,y    ,random_state=0    ,stratify=y)#%% prepare dataset 3, MI <0.2mi_score_selected_index = np.where(mi_scores < 0.2)[0]X_3 = X[:,mi_score_selected_index]X_train_3,X_test_3,y_train,y_test = tts(    X_3,y    ,random_state=0    ,stratify=y)#%% compare results with Decision Tree Classifierfrom sklearn.tree import DecisionTreeClassifier as DTCmodel_1 = DTC().fit(X_train_1,y_train)model_2 = DTC().fit(X_train_2,y_train)model_3 = DTC().fit(X_train_3,y_train)score_1 = model_1.score(X_test_1,y_test)score_2 = model_2.score(X_test_2,y_test)score_3 = model_3.score(X_test_3,y_test)print(f"score_1:{score_1}\n score_2:{score_2}\n score_3:{score_3}")#%% use Scikit-learn feature selectorfrom sklearn.feature_selection import SelectPercentile as SPselector = SP(percentile=50) # select features with top 50% MI scoresselector.fit(X,y)X_4 = selector.transform(X)X_train_4,X_test_4,y_train,y_test = tts(    X_4,y    ,random_state=0    ,stratify=y)model_4 = DTC().fit(X_train_4,y_train)score_4 = model_4.score(X_test_4,y_test)print(f"score_4:{score_4}")