Train and Test Splits

Linear regression is a good first model to study because almost everything it does is visible in a few lines of code. Instead of starting with a definition, this article starts with what happens when you fit the model. You will follow the execution line by line and keep track of what exists in Python after each step. The goal is to understand where the numbers that a model uses come from, and then to see how the train and test splits that feed the model are produced.

Linear regression, also called ordinary least squares or OLS, is the simplest and most classic linear method for regression. Its goal is to choose two numbers, a slope $w$ and an offset $b$, so that the line they define makes predictions that sit as close as possible to the true target values $y$ on the training set. When the data has a single input feature, the prediction is the slope multiplied by the input, plus the offset. Once the two numbers are chosen, the model needs no other information about the training data to make a prediction.

Fitting means choosing $w$ and $b$ so that the error is as small as possible. The error used here is the mean squared error. For each sample, take the difference between the prediction and the true value, square that difference, and add up the squares. Squaring removes the sign of each miss, so a prediction that is too high and one that is too low both count as errors, and large misses count for much more than small ones. The word mean refers to dividing that sum by the number of samples $n$, so the error does not grow just because the dataset grows.

The small program below computes this quantity by hand for two samples. Read it as a template for the calculation that fitting tries to drive down, not as the code used to train a model.

In scikit-learn, the model is created from the LinearRegression class and fitted with the training features and targets. The fit method returns the fitted model, which is why the two steps can be written on one line. After fitting, the slope is in the coef_ attribute and the offset is in the intercept_ attribute.

A fitted model's score method reports $R^{2}$ on whatever data you pass it. $R^{2}$ is a measure of how well the model explains the data. Calling score on the training data tells you how well the model fits the examples it learned from. Calling score on the test data tells you how well the model generalizes to examples it has not seen. Because the score describes whichever set you pass in, the two results must be compared together rather than read alone.

The two numbers are read together. When both scores are low and close to each other, the model is probably underfitting, meaning it is too simple to capture the pattern in the data. When the training score is high and the test score is much lower, that gap is a clear sign of overfitting.

Model complexity is part of the picture. A model with a single feature is simple, so there is little danger of overfitting it. A dataset with many features gives a linear model much more room to bend toward the particular training examples, and the chance of overfitting rises. The source illustrates this with the extended Boston Housing dataset, which has 506 samples and 105 derived features.

ShuffleSplit repeats a random split. A single iteration works in two draws from the same dataset. First, it selects train_size points to form the training set. Then it selects test_size points to form the test set, and those test points are chosen so that none of them is also in the training set for that iteration. The model is trained on the first group and scored on the second. Because the two groups are disjoint within an iteration, no point can be used both to learn and to be judged in the same iteration.

The sizes can be given as whole numbers or as fractions. A whole number is an absolute count of points, while a floating-point value is a fraction of the whole dataset. With the values 0.5 and 0.5, each training set holds half of the rows and each test set holds the other half of the same size.

In the illustration, the five training points and the five test points together account for every point in the dataset, and no point appears in both groups. The next iteration draws again, so a point that was test data in this iteration can appear in a later training set.

Two settings in ShuffleSplit do different jobs. The n_splits argument (called n_iter in the textual description) decides how many times the splitting happens. The train_size and test_size arguments decide how big each split is. Because these are separate, you can ask for many iterations while keeping each split small, or for few iterations with large splits.

The source describes a run of four such iterations on the same ten-point dataset, each drawing its own training and test points. Because each iteration samples again, a point that serves as test data in one iteration can appear in the training set of a later one. This is what gives the method its variety, and it is also why the ten scores in the first example differ from each other. Each iteration produces one score, so the total count of scores equals n_splits.

Most surprises with shuffle-split come from a few predictable places. When a result differs from what you expected, look at the setup lines first, then at what each line returned.

StratifiedShuffleSplit is a stratified variant of ShuffleSplit. Its main use is classification problems, where the source says it can give more reliable results than plain shuffle-split. The name hints at the difference: the word stratified refers to keeping each split organized by the class labels, so that the training and test sets reflect the classes in the data rather than depending only on chance.

When you are splitting data for a classification task, prefer StratifiedShuffleSplit when you want each split to keep the class structure of the full dataset. Plain ShuffleSplit can produce splits whose class mix differs from the whole, because the draw depends only on chance.