Stratify on regression

I have worked in classification problems, and stratified cross-validation is one of the most useful and simple techniques I've found. In that case, what it means is to build a training and validation set that have the same prorportions of classes of the target variable.

I am wondering if such an strategy exists in regression. A simple approach would be to split the data in quartiles or deciles and make sure that the proportions of training and validation instances in the respective quartiles and deciles are the same.

The question is: is there a standard way to do this? if so, is there an implementation in sklearn?

1 Answer

I have done this before and didn't find a default implementation - the StratifiedKFold and RepeatedStratifiedKFold are only documented to work with classes.

The way I ended up doing it was not quite as you are thinking with quartiles/deciles, but rather using a histogram (it matched my needs).

NumPy has a nice method for doing this, with many different formulas for computing the bin sizes that you can play around with to be match your data i.e. if it is normally distributed or not. I can't post the entire method code, but here is the gist:

samples_per_bin, bins, = np.histogram(data, bins='doane')    # Doane's method worked best for me

min_bin_size = samples_per_bin.min()

# compute the maximum batch size possible, using all samples from the bin with lowest population
n_bins = len(samples_per_bin)
max_batch = min_samples_single_bin * n_bins


I then put the data into a Pandas DataFrame and added a column indicating which bin each data point was in - finally doing something like this to perform the sampling from each of the bins:

df.groupby('bin_name', group_keys=False).apply(
lambda x: x.sample(n_per_group, replace=True))


Obviously you can allow duplicates in a batch or not by changing the replace argument. It might be necessary if you want a larger batch size while forcing stratification. It is definitely a limitation of my approach, which you might be able to overcome by using quartiles etc as you suggest.