Feature Selection¶
When working with tabular machine learning datasets it is quite common to collect as many columns as possible, and extract many features from those columns. This can lead to very large feature spaces, which can cause issues when fitting certain kinds of ML models, as they may use irrelevant or noisy features leading to poor test time performance. Additionally, large feature spaces can increase training times, and make it difficult to build interpretable models as there are too many features to reason over. For these reasons the field of feature selection developed, which aims to find the relevant subset of features for a given prediction task from a large set of possible features.
Tribuo v4.3 introduces feature selection algorithms for classification problems which score feature quality (called relevancy), and minimize redundant information in the feature set, using measures from information theory. It adds the core interfaces and classes for working with selected feature sets, along with classification specific implementations. In the future we may add feature selection algorithms for the other supported prediction types in Tribuo.
This tutorial will cover applying feature selection algorithms to a dataset, showing how reducing the feature space affects classification performance in terms of training time, accuracy and model size.
Setup¶
As usual we'll add some jars to the classpath and import some classes from Tribuo and the JDK. We're going to use the feature selection and SGD classification jars.
%jars tribuo-classification-sgd-4.3.0-jar-with-dependencies.jar
%jars tribuo-classification-fs-4.3.0-jar-with-dependencies.jar
import java.nio.file.Files;
import java.nio.file.Paths;
import com.oracle.labs.mlrg.olcut.provenance.ProvenanceUtil;
import org.tribuo.*;
import org.tribuo.dataset.SelectedFeatureDataset;
import org.tribuo.datasource.IDXDataSource;
import org.tribuo.classification.*;
import org.tribuo.classification.evaluation.*;
import org.tribuo.classification.fs.*;
import org.tribuo.classification.sgd.fm.*;
import org.tribuo.classification.sgd.objectives.LogMulticlass;
import org.tribuo.math.optimisers.AdaGrad;
import org.tribuo.transform.*;
import org.tribuo.transform.transformations.LinearScalingTransformation;
import org.tribuo.util.Util;
We'll also need some data to work with, so we'll load in the MNIST train and test sets.
We'll use Tribuo's built in IDXDataSource to read them, same as the configuration tutorial. If you've already downloaded MNIST then you can skip this step.
First download the training data:
wget http://yann.lecun.com/exdb/mnist/train-images-idx3-ubyte.gz; wget http://yann.lecun.com/exdb/mnist/train-labels-idx1-ubyte.gz
Then the test data:
wget http://yann.lecun.com/exdb/mnist/t10k-images-idx3-ubyte.gz; wget http://yann.lecun.com/exdb/mnist/t10k-labels-idx1-ubyte.gz
Tribuo's IDX loader natively reads gzipped files so you don't need to unzip them. Tribuo doesn't natively understand the 2d pixel arrangement, so the feature names from the IDXDataSource are just the integers 000 through 783, with leading zero padding to make it up to 3 digits.
var labelFactory = new LabelFactory();
var mnistTrainSource = new IDXDataSource<>(Paths.get("train-images-idx3-ubyte.gz"),Paths.get("train-labels-idx1-ubyte.gz"),labelFactory);
var mnistTestSource = new IDXDataSource<>(Paths.get("t10k-images-idx3-ubyte.gz"),Paths.get("t10k-labels-idx1-ubyte.gz"),labelFactory);
var mnistTrain = new MutableDataset<>(mnistTrainSource);
var mnistTest = new MutableDataset<>(mnistTestSource);
System.out.println(String.format("Training data size = %d, number of features = %d, number of classes = %d",mnistTrain.size(),mnistTrain.getFeatureMap().size(),mnistTrain.getOutputInfo().size()));
System.out.println(String.format("Testing data size = %d, number of features = %d, number of classes = %d",mnistTest.size(),mnistTest.getFeatureMap().size(),mnistTest.getOutputInfo().size()));
Building a baseline¶
First we'll train a model using all the features. We're going to use a small factorization machine as it's a high quality predictor, but still fairly time consuming to train. We're also going to wrap it in a transform trainer, as the factorization machine is sensitive to feature ranges and so is much happier if all the features are in the range [0,1).
var fmTrainer = new FMClassificationTrainer(new LogMulticlass(), // Loss function
new AdaGrad(0.1,0.1), // Gradient optimiser
3, // Number of training epochs
30000, // Logging interval
Trainer.DEFAULT_SEED, // RNG seed
5, // Factor size
0.1 // Factor initialisation variance
);
var transformations = new TransformationMap(List.of(new LinearScalingTransformation(0,1)));
var trainer = new TransformTrainer(fmTrainer,transformations);
var fmStartTime = System.currentTimeMillis();
var fmModel = trainer.train(mnistTrain);
var fmEndTime = System.currentTimeMillis();
System.out.println("Training factorization machine on 783 features took " + Util.formatDuration(fmStartTime,fmEndTime));
We can measure the accuracy of the full feature set:
var fmEvaluation = labelFactory.getEvaluator().evaluate(fmModel,mnistTest);
fmEvaluation.toString();
So we get 96% accuracy, which is pretty good for a simple non-convolutional model on MNIST, but it took a little while to train, and we depend on all 783 features, so need ((783 5) 10) parameters for the feature embeddings and 783 * 10 parameters for the linear part of the model. Let's try to reduce the model complexity a little bit by selecting only the most important features.
Simple feature selection¶
We'll start by selecting features based on their mutual information with the class label. This measures how predictive each feature value is of the label on its own. We're going to use at most 100 features for the experiments from this point, so we'll tell each feature selection algorithm to stop after picking 100 features. Each of the algorithms we'll discuss in this tutorial needs the data to be discretised, and so we'll split each feature into 5 equal width bins. We use the MIM algorithm to select features:
var mim = new MIM(100,5);
var mimSelectStartTime = System.currentTimeMillis();
var mimSet = mim.select(mnistTrain);
var mimSelectEndTime = System.currentTimeMillis();
System.out.println("Selecting the top 100 features with MIM took " + Util.formatDuration(mimSelectStartTime,mimSelectEndTime));
That was pretty quick, but it's also only a single calculation per feature. The select(Dataset<Label>) call returns a SelectedFeatureSet object which contains the feature names, scores, if the set is ordered or not (and all current implementations are), and the provenance for the selection procedure. The provenance can be serialized to a file or a byte array in the same way that other Tribuo objects can be. Let's briefly look at the provenance recorded in the feature set:
var setProvenance = mimSet.getProvenance();
System.out.println(ProvenanceUtil.formattedProvenanceString(setProvenance));
We can see the feature selection algorithm is recorded, the number of features that were selected, the number of discretisation bins, and the number of computation threads. We can also see the information from the dataset that the features were selected from, showing the usual information like the location on disk, hashes and timestamps. This provenance is recorded in the SelectedFeatureDataset used to train the models, and is in turn captured in the ModelProvenance for each of the models trained on those sets.
Now let's actually look at the selected features:
System.out.println("MIM feature set: " + mimSet.featureNames());
Most of the features are roughly in the middle of the 28x28 pixel grid in MNIST. This seems pretty sensible, there's little information in the pixels around the edges as they aren't used very often.
To actually use the feature set we build a SelectedFeatureDataset which copies an existing dataset (in this case the MutableDataset containing the MNIST training set), and removes all the unselected features. It has a parameter which lets the user control how many of the features in the selected feature set are used, which makes it easier to optimize over the feature set returned by a selection algorithm.
So now we'll build a smaller dataset with only the top 100 features in, and see how that affects the factorization machine.
var mimData = new SelectedFeatureDataset(mnistTrain, mimSet);
var mimStartTime = System.currentTimeMillis();
var mimModel = trainer.train(mimData);
var mimEndTime = System.currentTimeMillis();
System.out.println("Training factorization machine on 100 features took " + Util.formatDuration(mimStartTime,mimEndTime));
That took roughly a third of the time of training the full model, and is a seventh the size of the full model, but we'd expect some reduction in accuracy. So let's check that:
var mimEvaluation = labelFactory.getEvaluator().evaluate(mimModel,mnistTest);
mimEvaluation.toString();
The smaller model gets roughly 91% accuracy, and has a much lower F1 for the classes "4" and "9". Those digits can look quite similar, and when you throw away many of the pixels you might miss the crucial ones which distinguish those digits.
The MIM algorithm can be quite prone to this, as the features it selects ignore redundancy and complementarity. Redundancy is when two features provide the same information, only one of those features is necessary to convey that information to the classifier. Complementarity is when two features combine to provide more information than the sum of their individual informations. This can occur when the information in one feature is conditional on the presence of another one.
So let's look at using a more complicated algorithm, JMI. As before we build a JMI object, and ask it to select 100 features. The JMI implementation in Tribuo is multi-threaded, so we'll also tell it to use 4 compute threads.
var jmi = new JMI(100,5,4);
var jmiSelectStartTime = System.currentTimeMillis();
var jmiSet = jmi.select(mnistTrain);
var jmiSelectEndTime = System.currentTimeMillis();
System.out.println("Selecting the top 100 features with JMI took " + Util.formatDuration(jmiSelectStartTime,jmiSelectEndTime));
System.out.println("JMI feature set: " + jmiSet.featureNames());
The first feature selected by JMI is the same as the one from MIM, because they both pick the feature with the highest mutual information. After that point the feature sets diverge, as JMI is selecting features which combine well with the already selected features, in terms of maximising relevancy and complementarity while minimising redundancy, whereas MIM selects features solely based on relevancy. Let's see how that changes the performance of the factorization machine:
var jmiData = new SelectedFeatureDataset(mnistTrain, jmiSet);
var jmiStartTime = System.currentTimeMillis();
var jmiModel = trainer.train(jmiData);
var jmiEndTime = System.currentTimeMillis();
System.out.println("Training factorization machine on 100 features took " + Util.formatDuration(jmiStartTime,jmiEndTime));
As before, training on 100 features takes about a third of the time as the full model, and its still got the same size benefit. Let's look at the accuracy:
var jmiEvaluation = labelFactory.getEvaluator().evaluate(jmiModel,mnistTest);
jmiEvaluation.toString();
Performance is a little better, and it's about 1 point of F1 better at predicting "9"s. We could further tune the performance by optimizing over the number of features to find a model which retained most of the performance while still shrinking in size (which usually makes deployment simpler & prediction faster).
We've been using Tribuo's sparsity features implicitly during this exploration, as each time we passed in the full test dataset containing all 784 features. Tribuo's models discard all feature names that weren't present at training time, so each model automatically discarded all the features that weren't in the 100 features it was trained on. This makes it much easier to do feature selection experiments than other systems where you'd need to explicitly subselect the test time dataset and repack things into the test time arrays/tensors.
Feature selection for data exploration¶
We've looked at using feature selection to improve training time and reduce model complexity, but it can also be used without building a model as part of a data exploration procedure. We've selected the most important 100 features for making predictions on the MNIST data, which for MNIST isn't particularly useful as we collect all the pixel values at once. However if we're working on tabular problems where each column of data has a particular cost of collection or storage, or has different privacy/governance requirements, then discovering which features we need to collect can either make our data pipelines much faster/cheaper to implement. We can also try to discover why these features are relevant for predictions using knowledge of the problem domain, and this might provide some insight into the prediction problem itself.
Conclusions¶
We discussed the process of feature selection, and ran a few algorithms on MNIST showing the tradeoffs between feature set size, accuracy and training speed. We also looked at the provenance captured by feature selection algorithms, showing how it records all the relevant information in the feature set object.