StratifiedKFold yields slightly different class label distributions (prior to 0.22 different train and test sizes)

[GregoryMorse]

Description

Unfortunately, probably due to the stratification balancing, the train and test size can vary between iteration. This is pretty highly undesirable and it should always rebalance to same sizes. Otherwise combining all the folds and using efficient vectorized operations to do computations in a matrix with an extra dimension is not easily possible.

This honestly feels like a poor stratification algorithm. At worst this is a flat out bug. At best, its a mode of operation, and a Boolean flag should be added such as something like force_exact_division=True. Of course, in some cases it is impossible to balance the classes, but even in those cases, it could still prefer to guarantee equal distribution, rather than guarantee equal class sizes since we are talking about only fractional remainders here - and the class sizes are not going to be perfect percentage wise balanced anyway as it can be impossible for example if only 1 value belongs to a class. This is surprisingly awkward functionality.

Steps/Code to Reproduce

import numpy as np
from sklearn.model_selection import StratifiedKFold
kfold = StratifiedKFold(n_splits=5, shuffle=True, random_state=0)
X, y = np.arange(7590)[:,np.newaxis], np.random.RandomState(0).randint(0, 3, 7590)
for (train, test) in kfold.split(X, y):
    print(len(train), len(test))
kfold = StratifiedKFold(n_splits=5, shuffle=True, random_state=0)
for (train, test) in kfold.split(X, y):
    assert(len(test) == 7590/5)
    assert(len(train) == 7590*4/5)

Expected Results

6072 1518
6072 1518
6072 1518
6072 1518
6072 1518

Actual Results

6071 1519
6071 1519
6072 1518
6072 1518
6074 1516
Traceback (most recent call last):

  File "<ipython-input-106-1c41a73d61b2>", line 8, in <module>
    assert(len(test) == 7590/5)

AssertionError

Seems to always be in order from min(len(train)) to max(len(train))

Versions

Windows-10-10.0.18362-SP0
Python 3.7.4 (default, Aug 9 2019, 18:34:13) [MSC v.1915 64 bit (AMD64)]
NumPy 1.17.3
SciPy 1.3.1
Scikit-Learn 0.21.3

[GregoryMorse]

Reference #10767 where @jnothman mentions the StratifiedKFold algorithm needs a complete redo and #11777 which although a different issue still involves the way it balances classes.

[GregoryMorse]

It is not impossible to vectorize with different sizes, but it adds complications as the extra dimension will not work. Instead the rows would be concatenated, the train and test sizes will need to be tracked so it can be split apart again.

[jnothman]

Fixed in #14704. Please confirm that this issue is fixed in 0.22, and help us identify other related issues that can be closed. Thanks

[GregoryMorse]

@jnothman, excellent sorry I did not search the PRs only the issues. This is perfect if it preserves the train-test size and the class value distribution amongst folds. Allows for some really clever vectorized operations to be done on all at once in certain contexts.

I realized 0.22 was just freshly released. It is indeed working there but since I did not check class label distribution so let us do that now:

import numpy as np
from sklearn.model_selection import StratifiedKFold
kfold = StratifiedKFold(n_splits=5, shuffle=True, random_state=0)
X, y = np.arange(7590)[:,np.newaxis], np.random.RandomState(0).randint(0, 3, 7590)
counts = []
for (train, test) in kfold.split(X, y):
    _, ctrain = np.unique(y[train], return_counts=True)
    _, ctest = np.unique(y[test], return_counts=True)
    counts.append((len(train), len(test), *ctrain, *ctest))
print(counts)
assert(all([counts[0] == x for x in counts[1:]]))

[(6072, 1518, 2031, 2079, 1962, 508, 520, 490), (6072, 1518, 2031, 2079, 1962, 508, 520, 490), (6072, 1518, 2031, 2079, 1962, 508, 520, 490), (6072, 1518, 2031, 2080, 1961, 508, 519, 491), (6072, 1518, 2032, 2079, 1961, 507, 520, 491)]
Traceback (most recent call last):

  File "<ipython-input-1-b0c88890b497>", line 11, in <module>
    assert(all([counts[0] == x for x in counts[1:]]))

AssertionError

So it almost looks perfect. The train-test is divided in equal sizes each time. But the 5th run we see class value counts for train and test differing between 2031, 2079, 1962, 508, 520, 490 and 2032, 2079, 1961, 507, 520, 491. I think the test code above might be worthy of integration as a unit test.

@jnothman is this a failure for the algorithm, or is equal class label group sizes not guaranteed?

[jnothman]

No we cannot guarantee equality. Please see https://github.com/scikit-learn/scikit-learn/blob/5f3c3f0378f2f30f3c4340bd9bf1e211e96d5c3c/sklearn/model_selection/_split.py#L613 for the set of invariances that should hold.

[jnothman]

I'm pretty sure there is no issue here.

The docs could be updated to add the general kfold invariances: the test sets together constitute a partition of the full dataset into n_splits subsets, and each test and train set constitute a partition of the full dataset.

With all of these constraints, I think this is as good as it gets.

[GregoryMorse]

Yes everything looks good here especially from a theoretical data science standpoint. The invariance listed are holding.

However, for practical purposes having equal group sizes could be very beneficial in some tasks. Here is the idea: a flag force_equal_class_sizes=True. It would have some theoretical problems as rows would never be included and some would be included twice but these are outliers and not statistically relevant. As long as this is acceptable and understood. Take the remainders of classes_counts % n_splits. All 0 remainder do not have a problem and are already in the current code fine. The remainders are sorted in descending order. The top half or largest of the remainders will be in the +1 category while the bottom half of the remainders will be in the -1 category. A middle value would be in the 0 category. Then any in the +1 category will randomly use extra values for the n_splits-classes_counts%n_splits hole values. And the -1 category will randomly drop a value for its classes_counts%n_splits values. Doing this strategy minimizes the number of repeated values and dropped values, while allowing this class label preservation property. If it does not then slight modification to the strategy might be needed but the goal of minimizing the spurious values is not so complicated. @jnothman what do you think of the idea? Probably the issue will come up every so often without doing some best-try approximation solution like this and adding the feature flag.

[jnothman]

I'd rather not make things more complicated in this space without very clear justification

[GregoryMorse]

1. This would not change the existing code which is very good right now. The refactoring is excellent and should not be changed as its certainly an improvement over the old algorithm. I cannot think of any good argument for the old code though perhaps it is preserving the class sizes and not the test sizes (and should be also allowed via parameter). Both can be argued as data science theoretically justified depending on what your priority is. But its such a small point that I don't think anyone would ever care, and the old code is the most awkward.

2. Admittedly this does make things complicated. Well the only clear justification is doing vectorized operations so a computation can be done on all train/test folds simultaneously. Because you cannot construct an ndarray or matrix of them without them being equal class sizes. Over enough data, the fact that a very small amount of missing rows or duplicated rows is present is statistically insignificant and the convenience of vectorized operations outweighs this cost. I am doing axis-based calculations on data for example, its the slowest place in my code and I cannot optimize with the current stratification algorithm and am stuck with for-loops or list comprehensions or such.

3. Why not make StratifiedKFold totally comprehensive of all possible desirable details in terms of stable test/train sizes, stable class value sizes. Even a mode parameter could easily allow all variations of functionality.

Mode 1: test/train size equivalent as in v0.22.
Mode 2: class label proportions equivalent (I am not sure if this was true in v0.21?)
Mode 3: test/train size equivalent and class label sizes equivalent. This introduces either missing rows, duplicated rows or a mixture of the two theoretically adding 3 submodes of operation.
Submode 1: minimize missing rows/duplicated rows using some simple algorithm along the lines of what I highlighted.
Submode 2: increase the train/test size to make sure only duplicated rows are added and no rows are missing.
Submode 3: decrease the train/test size to make sure no duplicated rows and only rows dropped.

At least we can at least discuss this here, as I cannot think of any further complication that could ever be justified. Granted my situation might be a narrow use case, but is there harm in trying to include it?