I am looking to train a machine learning model to solve a regression problem. The targets, y, are not normally distributed. It is perhaps best described as somewhat bimodal. For context, I have 13,000 points in total.

I am looking to have an 80%/20% train/test split, and on the training set, I intend to do 5-fold cross-validation. My concern is that, depending on the seed used to split the data, there can be somewhat significant run-to-run variability because the distribution is sampled randomly and the subsamples do not have a similar distribution as the parent dataset or between other subsamples.

For this reason, I am considering doing stratified sampling with the train/test split as well as during the 5-fold cross-validation. My plan is to set up five strata of the type (-$\infty$,1], (1,2], (3,4], (4,+$\infty$] and have the subsamples contain a similar distribution as the parent dataset.

However, I see very few people in my field do stratified sampling for regression problems and am concerned about unforeseen circumstances. What concerns are there that one should consider when it comes to stratified sampling in machine learning? Of course, I can also run the ML regression several times and average the results, regardless of whether I do stratified sampling or purely random sampling.


In general, variation is a good thing in cross-validation or train/test split, so there's little reason to reduce variability by stratified sampling. I can think of some situations where stratified sampling may make sense though. For example, if your outcome is binary where the proportion of 1 (or 0) is very low. Then effectively you have a very small sample size, and you may want to ensure there are enough 1 (or 0) in each of the folds. Other than that, if your sample size is around 13000 and it is unlikely that your splits would lead to perfect (or near perfect) classification in one of the folds then I don't think you should do stratified sampling.

  • $\begingroup$ Thanks for your comment. I agree about the binary classification example -- I ran into that exact issue and plan to use stratified sampling there. In the case of the regression example, the variation I'm referring to is upon using different random seeds to shuffle the data. Each seed will lead to a slightly different MAE with purely random sampling because the distributions in train/test (and the folds) will differ. The variation in run-to-run performance would be lower with stratified sampling. So, in what way is this kind of variation a good thing? $\endgroup$ May 12 '20 at 3:01
  • 1
    $\begingroup$ The point of cross-validation (rather than just fitting to the whole data in one go) is to improve out-of-sample prediction by reducing overfitting. You don't know what the "out-of-sample" will be, so increasing variation in training, heuristically, helps you reduce overfitting to one particular type of distribution of the data. $\endgroup$
    – Tim Mak
    May 12 '20 at 5:16
  • $\begingroup$ That makes a lot of sense. Indeed, when doing stratified sampling on an 80%/20% train/test split with stratified 5-fold CV, I get essentially the same results as random sampling. When I do stratified sampling on neural network model with 80%/10%/10% train/hold-out val/test (no CV), it does seem to have a minor but noticeable effect compared to random sampling, but I'm not sure if that's artificial or not. $\endgroup$ May 12 '20 at 5:27

My thoughts are first, to perform a discriminate analysis to assign the y's to one of two populations. This may be based on a simple rule (perhaps size based) or something more complex depending on available information. The goal is to remove (or, at least, mitigate) the bimodal nature of the data.

Second, as the assignment is not likely perfect in Step 1, use a robust regression like Least-Absolute-Deviations which is median, not mean-centered (the latter being sensitive to outliers, which in the current context, are possibly misclassified points). The robust regression analysis is applied separately to each identified subpopulation (two in the current case).

Also, optionally consider a power transformation also to induce normality and homogeneity of variance (like Box-Cox) especially if considering non-robust regression analysis in each population. Caution, this may create interpretations issues with the data results.

As such, your concerns over "seed used to split the data, there can be somewhat significant run-to-run variability because the distribution is sampled randomly and the subsamples do not have a similar distribution as the parent dataset or between other subsamples" may be reduced.


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.