How to Split And Resample Imbalanced Dataset Into Train, Validation and Test

Question

I want to understand how to split the imbalanced data set with a binary target variable where 87% of the samples are negative and 13% of the samples are positive. Now, I know that you should always split the data into train and test set before doing any kind of resampling to avoid info leakage but how what strategy gives better results - random train_test_split or StratifiedShuffleSplit? The general sentiment is that the latter seems to be outperforming the former in case of an imbalanced data set. Btw, how does StratifiedShuffleSplit work?

Also, how to factor in the cross-validation set and how to process it?

Should it also contain the same number of positive and negative samples as per the sampling was done on the training set?

Answer 1

Grzegorz's answer gets to the root of the problem - if the future data this model will be used to make predictions on will have the same distribution, then stratification by class % makes sense.

One thing I wanted to add is I typically use the normal train_test_split function and just pass the class labels to its stratify parameter like so:

train_test_split(X, y, random_state=0, stratify=y, shuffle=True)

This will both shuffle the dataset and match the %s of classes in the result of train_test_split.

Answer 2

It all depends on what is your objective. For instance, if you deploy a model where you expect that the test data (e.g. streaming data) will have the same pos/neg distribution then it makes sense to stratify your splitting. In that way, you bias your learning procedure to label more as negative. On the other hand, if you are told that more weight is given to detect the sparse positive examples then you may want to try an equal distribution.

Answer 3

For a large dataset, non-stratified splitting is usually okay. You have to consider what is the probability of significant differences in the balance between test and training sets. For me, 13% of thousands of records is probably no issue but 13% of a hundred records is likely to be an issue.

Marc's answer explains how the stratification works. StratifiedShuffleSplit just makes this easier in the context of cross-validation by creating a generator to split it multiple times. You can also pass it the random_state so that the splits remain consistent in each fold across iterations, which is important for hyper-parameter tuning. Check out the example in the documentation.