How to arrange test dataset distribution for an imbalanced classification problem?

Question

I have a dataset that contains 560 datapoints, and I would like to do binary classification on it. 400 datapoints belong to class 1, and 160 points belong to class 2. In the case of an imbalanced dataset like this, how to arrange the test dataset to get valid performance results? Should I keep the same imbalanced data distribution for the test set which is similar to the distribution of the training data, or arrange it in a way that half of the test points belong to the first class and the remaining half belongs to the second class?

Answer 1

The test set should represents the "real" data distribution your model will tackle once deployed and used in real applications. So the quick answer is yes, the test data should be imbalanced, which comes also as a sort of forced choice for you considering the super small size of you dataset.

You want to keep in mind though that this makes a bit trickier to interpret the resulting metrics.

For example, if your test set is made of 90 instances labeled A and 10 instances labeled B, a model might predict class A for all instances and still have a final accuracy of 90%, correct, but completely unrepresentative of the real behavior of the model. So make sure to compute metrics like precision and recall for each class.

Extra: considering the size of your dataset, you might consider using k-fold cross validation. Split your dataset into k train-test subsets, train k models and average the resulting metrics. The logic being that with such few instances the dataset is almost guaranteed to contain biases of some sort.