Is it wise to stratify the continuous y (target) variable when you split your training and testing data from the total sample in regression setting? Here is the approach in python to do implement stratify the continuous target:

In Python (with the same libraries loaded as in the prior code snippet):

# Create the bins.  My `y` variable has
# 506 observations, and I want 50 bins.

bins = np.linspace(0, 506, 50)

# Save your Y values in a new ndarray,
# broken down by the bins created above.

y_binned = np.digitize(y, bins)

# Pass y_binned to the stratify argument,
# and sklearn will handle the rest

X_train, X_test, y_train, y_test = train_test_split(df, y, test_size=0.3, stratify=y_binned)

There you have it: stratification of a continuous numerical target value. However, I am not confident with this approach although stratification of the binary response variable is very common in classification setting to tackle class imbalance issues. I appreciate your suggestions!


2 Answers 2


You are right, it's not common at all, but what you do makes sense though it may be sensitive to your binning strategy. So, it's also a good idea to plot your binned target variable. This stratification makes more sense when the target variable is so skewed such that sometimes some folds do not get enough samples from some regions of the target variable.


Update: First consider whether splitting the data into training and validation subsets makes the best use of your data for building a predictive model.

Split-Sample Model Validation
Bootstrap optimism corrected - results interpretation

If you still want to proceed with a train/validation split, the proposed strategy is equivalent to simple random sampling: it will sample test_size items with equal probability to create a test set and allocate the remaining 1 - test_size items to the training set.

To begin with, the implementation of your strategy has some obvious weaknesses. Why create bins in the interval [0,n] given n data points?

  • Say you have 506 observations between 0 and 100. It seems wasteful to create 40 empty bins for unobserved values > 100.
  • Say you have 506 observations between 0 and 1000. It seems inefficient to put all observations ≥ 506 in the rightmost bin.

So you'll have to be a bit careful about how you define the bins as scikit-learn throws an error if there are bins with fewer than 2 data points.

Other that this technical requirement, it doesn't matter how you define the bins. With stratified sampling each bin is sampled in proportion to its size, so you sample more frequently from bins with more items, which correspond to higher data density regions. But, conditional on the bin, an item in a "dense" bin with many data points has a smaller chance of being sampled than an item in "sparse" bin. In the end the two probabilities even out and each item has the same probability of being sampled, no matter how dense its region. This is simple random sampling. So in expectation, you'll get the same allocation between training and test subsets you'll get with simple random sampling.

To see this, let's compute the probability that item $i$ is selected as the first item to allocate to the test set.

With simple random sampling: $$ \operatorname{Pr}\left\{\text{Select $Y_i$}\right\} = \frac{1}{n} $$

With stratified strategy: $$ \operatorname{Pr}\left\{\text{Select $Y_i$}\right\} = \operatorname{Pr}\left\{\text{Select $\operatorname{Bin} (Y_i)$}\right\}\operatorname{Pr}\left\{\text{Select $Y_i$ from $\operatorname{Bin}(Y_i)$}\right\} \\ = \frac{\operatorname{Size}\left\{\operatorname{Bin}(Y_i)\right\}}{n} \frac{1}{\operatorname{Size}\left\{\operatorname{Bin} (Y_i)\right\}} = \frac{1}{n} $$

So there is no need for the more complex strategy. Just don't forget to shuffle.

A quick simulation to confirm.

import numpy as np
from sklearn.model_selection import train_test_split


n = 1000

# Let's "label" one item, yi, by assigning it the value 501
# Otherwise, the rest of the values are uniformly distributed between 0 and 500
y = np.random.uniform(low=0, high=500, size=n - 1)
y = np.append(y, 501)

bins = np.linspace(0, 506, 50)
y_binned = np.digitize(y, bins)

nreps = 1000
select_yi = 0

for r in range(nreps):
    y_train, y_test = train_test_split(y, test_size=0.3, shuffle=True,
    select_yi += sum(y_test == 501)

# yi is assigned to the test set about 30% of the time
# as expected since the test size is 30%.
# > 297
  • 5
    $\begingroup$ The bigger question is whether it's OK to use data splitting as a validation strategy. In my experience the minimum sample size for split sample validation to work well when Y is binary is n=20,000. Otherwise there is too much luck involved in the split. For continuous Y you might get away with 7,000. $\endgroup$ May 15 at 21:52

Your Answer

By clicking “Post Your Answer”, you agree to our terms of service, privacy policy and cookie policy

Not the answer you're looking for? Browse other questions tagged or ask your own question.