K-Fold and Stratified K-Fold Cross-Validation

The Problem with a Single Split

Imagine you train a model and get an accuracy of 92% on your test set. That sounds great, but how confident are you in that single number? What if you had just used a different random_state in your train_test_split function and the accuracy dropped to 85%? Which score is "correct"? The answer is that neither is the full story. Both are just single data points reflecting a particular random split of your data.

To truly understand this variability, let's watch a practical demonstration.

This demonstration shows that to get a trustworthy assessment of a model's ability to generalize to new, unseen data, we need a more systematic approach than a one-off split.

K-Fold Cross-Validation: The Solution

K-fold cross-validation solves this problem by performing multiple train-test splits and averaging the results. The process is straightforward:

1. Split: The entire dataset is partitioned into k equally sized segments, or "folds". A common choice for k is 5 or 10.
2. Iterate: The process then iterates k times. In each iteration:
   • One different fold is held out as the test set.
   • The remaining k-1 folds are combined and used as the training set.
   • The model is trained on the training set and evaluated on the test set, and the performance score is recorded.
3. Aggregate: After k iterations, you will have k performance scores. The final cross-validation score is the average of these k scores. This average provides a much more stable and reliable estimate of the model's performance.

The image below illustrates this process for k=5. In each iteration, a different fold (in yellow) serves as the test set.

For a clear written explanation of this process and its benefits, the following article is excellent.

Implementing K-Fold with Scikit-Learn

Your proficiency with Python and scikit-learn allows us to look under the hood before using the convenient high-level functions. The codebasics channel provides an excellent walkthrough.

As the video shows, while you can write a loop yourself using an iterator object like KFold, scikit-learn provides the cross_val_score function to do all the work in one line. It handles the splitting, looping, training, and scoring for you.

You provide the model, the full dataset (X and y), and the number of folds (cv), and it returns an array of scores, one for each fold.

from sklearn.model_selection import cross_val_score
from sklearn.linear_model import LinearRegression




# Assuming X and y are your features and target
model = LinearRegression()




# Perform 5-fold cross-validation
# For regression, you can specify a scoring metric like 'neg_mean_squared_error'
scores = cross_val_score(model, X, y, cv=5, scoring='neg_mean_squared_error')




# The scores are negative because scikit-learn's convention is to maximize scores.
# So we take the absolute value and often the mean.
import numpy as np
rmse_scores = np.sqrt(-scores)

print(f"Scores for each fold: {rmse_scores}")
print(f"Average RMSE: {rmse_scores.mean():.4f}")
print(f"Standard Deviation of RMSE: {rmse_scores.std():.4f}")

Reporting both the mean and the standard deviation of the scores gives a complete picture: the mean tells you the expected performance, and the standard deviation tells you how much that performance varies between different training sets.

Stratified K-Fold Cross-Validation

Standard K-Fold randomly splits the data. This works well for regression or balanced classification problems. However, for imbalanced classification problems (e.g., 95% non-fraud vs. 5% fraud), a random split could lead to a test fold containing very few, or even zero, instances of the minority class. Training or evaluating on such a fold would be uninformative or even cause errors.

Stratified K-Fold is the solution. It's a variation of k-fold that ensures the proportion of each class is the same in every fold as it is in the original dataset.

Let's turn to the scikit-learn documentation to see the technical difference.

The good news is that scikit-learn often makes the smart choice for you. When you pass an integer for the cv parameter in cross_val_score for a classification model, it automatically uses StratifiedKFold by default. For regression models, it uses KFold. This is a design choice that reflects best practices.

Conclusion

You now have a robust method for evaluating machine learning models that goes far beyond a simple train-test split. This is a cornerstone of reliable model development and selection.

Key Takeaways:

• A single train_test_split yields an unstable performance score that depends on the randomness of the split.
• K-Fold Cross-Validation provides a more reliable performance estimate by training and testing a model on multiple (k) different subsets of the data and averaging the results.
• Stratified K-Fold is a crucial variant for classification tasks. It ensures that the class proportions in each fold mirror the proportions in the overall dataset, which is vital for imbalanced data.
• In scikit-learn, the cross_val_score function is the primary tool for efficiently performing cross-validation.

In our upcoming lessons, we will delve into one of the most important theoretical concepts in machine learning: the bias-variance tradeoff. Cross-validation is not just an evaluation tool; it is also our primary diagnostic tool for understanding whether a model is too simple (high bias) or too complex (high variance). Your understanding of CV will be essential as we learn to diagnose and correct these common modeling issues.