Why Imputation Must Be Part of the Fold
A regression model usually needs usable predictor values to produce predictions. When a predictor is missing, an imputation rule can replace the missing entry with a value such as the training-sample mean or median. But that replacement is itself learned from data. In cross-validation, the held-out fold is meant to stand in for observations not used to fit the prediction procedure. If held-out predictor values help determine the imputation rule, they have influenced the procedure being evaluated.
The remedy is to treat imputation as part of the prediction procedure: estimate its quantities using only the training portion of each fold, then apply that fitted rule to the fold’s held-out observations. The held-out observations may contain observed predictor values, but those values do not get to alter the imputer before predictions are made.
For example, a mean imputer for a numerical predictor replaces missing entries with the mean of the observed values in that predictor among the rows used to fit the imputer. A mode imputer for a categorical predictor replaces missing entries with its most frequent observed category among those rows. These are simple rules, not universal solutions: mean or mode imputation does not recover the unknown value with certainty or account for the uncertainty of imputation. The point here is how to evaluate a chosen rule without letting held-out data help fit it.
The basic setup is the \(K\)-fold cross-validation algorithm: the data are divided into \(K\) folds, and in turn each fold is held out while the remaining folds are used for training. For each turn, fitting includes every data-dependent preprocessing step. The held-out fold is transformed using the quantities learned from the training folds, and the prediction model is fitted using the transformed training observations.
Fold-Specific Fitting and Application
Suppose \(F_k\) is the set of row indices in held-out fold \(k\), and \(T_k\) is the set of indices in the corresponding training folds. For a numerical predictor \(X\), calculate its imputation value from observed \(X\) entries in \(T_k\) only. Use that value to replace missing \(X\) entries in both \(T_k\) and \(F_k\). Do not recalculate it from \(F_k\), even if some \(X\) values there are observed.
Here \(\widehat{\mu}_{X,k}\) is the training-fold mean for predictor \(X\), and the denominator counts the training-fold observations where \(X\) is observed. When \(X_i\) is missing, replace it with \(\widehat{\mu}_{X,k}\). The corresponding categorical rule uses the training-fold mode rather than this mean.
Before applying a particular imputation procedure, its scope should be clear. The examples below concern missing predictor values and a rule chosen in advance. Each fold must have at least one observed training value for a mean or mode imputer to be defined. If a feature is entirely missing in a training fold, the fold-specific mean or mode does not exist; a fallback rule must be specified in advance using information available to the training procedure, or the fold design or feature handling must be reconsidered. A held-out value cannot be used to fill that gap without compromising the held-out evaluation.
Worked Example: A Mean Imputer in One Held-Out Fold
Suppose a hypothetical data set is divided so that rows A, B, and C are in the training portion and rows D and E are held out. A numerical predictor, hours of weekly practice, has values 2, 4, and missing in A, B, and C, and values 10 and missing in D and E. Use a training-fold mean imputer.
Fit using training rows only. The observed training values are 2 and 4 hours. Thus the imputation value is
Replace C’s missing value with 3 hours. Apply the same fitted value to the held-out rows: D retains its observed value of 10 hours, and E’s missing value is replaced with 3 hours. The held-out value of 10 hours does not enter the imputation calculation.
If instead the mean were calculated once from every observed row, it would be \((2+4+10)/3=5.333\) hours. That value uses D’s held-out predictor measurement, so it is not the result of fitting the imputer on the training portion alone. The fact that E’s imputed value might look plausible does not make the evaluation procedure valid.
Why Imputing Once Before Cross-Validation Leaks Information
Leakage occurs when information from held-out observations influences a step that should be fitted using training data only. Imputing the entire data set before making folds is one way this can happen: the imputation value for a predictor can depend on observed values in folds that are later treated as held out. Those values have then influenced the transformed training data and, through them, the fitted prediction model.
The mean calculation shows the source of the problem. Suppose the observed predictor values in a training portion have mean \(\bar{x}_{T}\), with \(n_T\) observed values, and the observed predictor values in its held-out portion have mean \(\bar{x}_{H}\), with \(n_H\) observed values. If the data are pooled before fitting the imputer, the resulting mean is
When \(n_H>0\), this value generally differs from \(\bar{x}_{T}\) unless \(\bar{x}_{H}=\bar{x}_{T}\). More importantly, \(\bar{x}_{\text{all}}\) depends on held-out information even in cases where the two means happen to match. Cross-validation is intended to evaluate the procedure that would be applied to new observations, for which the held-out values would not be available when the imputer is fitted.
Leakage does not guarantee that the reported prediction error will be smaller in every particular data set. Changing the imputation value can help or hurt predictions. The problem is that the estimate no longer evaluates the intended training-only procedure: information unavailable at fitting time has entered the evaluation. A fair assessment requires matching the evaluation process to the way the procedure will be used.
Worked Example: Imputation Across Three Folds
Consider a hypothetical numerical predictor with six rows: row 1 has value 2, row 2 is missing, row 3 has value 4, row 4 has value 8, row 5 is missing, and row 6 has value 10. Use three held-out folds: \(F_1=\{1,2\}\), \(F_2=\{3,4\}\), and \(F_3=\{5,6\}\). In each turn, calculate the mean from observed training-fold values only.
Fold 1. The training rows are 3, 4, 5, and 6. Their observed values are 4, 8, and 10, so the fitted mean is
In the held-out fold, row 1 retains its observed value 2 and row 2 receives the imputed value 7.333.
Fold 2. The training rows are 1, 2, 5, and 6. Their observed values are 2 and 10, giving
Rows 3 and 4 are held out and retain their observed values, 4 and 8. Neither is used to calculate this fold’s imputation value.
Fold 3. The training rows are 1, 2, 3, and 4. Their observed values are 2, 4, and 8, giving
Row 5, which is missing, is held out and receives 4.667; row 6 retains its observed value of 10. The missing values in rows 2 and 5 therefore receive different imputations in different turns because their training folds differ. A single mean calculated before creating the folds would be \((2+4+8+10)/4=6\), using observed values from the held-out folds. It would assign 6 to both missing rows, but that is not the fold-specific training-only procedure.
The Same Rule Applies to Categorical Predictors
The boundary between training and held-out data matters for more than numerical means. For a categorical predictor, a mode imputer uses the most frequent observed category in the training portion. That category is then applied to missing entries in the held-out portion. If multiple categories tie for most frequent, the procedure needs a prespecified tie rule; the held-out fold must not decide the tie.
Worked Example: A Training-Fold Mode
In a hypothetical training fold, a categorical predictor records A, A, and B. The held-out fold records C, C, C, and missing. Use a mode imputer and a rule that selects the alphabetically first category if there is a tie.
Among the training values, A occurs twice and B occurs once, so the fitted mode is A. The missing held-out value is therefore replaced with A. The three observed held-out values of C stay as recorded, but they do not change the fitted mode.
If the mode were calculated from all seven observed values before cross-validation, C would occur three times, A twice, and B once, so the missing value would instead be filled with C. This full-data mode uses held-out information and is not the training-fold imputation rule. As with numerical imputation, the correct evaluation applies the training-fitted rule unchanged to held-out observations.
A Reliable Fold-by-Fold Workflow
For each fold, separate fitting from application. This keeps imputation aligned with the role of the held-out observations and makes the procedure reproducible.
Identify the training rows and held-out rows before calculating any imputation quantities.
Estimate each feature’s mean, median, mode, or other prespecified imputation quantity from observed training values only.
Use the training-estimated quantities to fill missing values in the training and held-out predictors; do not refit on the held-out portion.
The held-out outcomes and predictor values do not participate in fitting the model or the imputer.
Use the model and the already-fitted imputer to process the held-out predictors and form predictions. Repeat the complete process for each fold.
This process is often implemented as a pipeline that connects imputation and model fitting. The important property is not the software label but the order of operations: when a fold is held out, every data-dependent step is fitted without it. If the imputation rule is changed, for example from a mean to a median, that choice must also be part of a prespecified evaluation procedure; choosing it based on held-out performance would introduce a different form of information reuse.
Imputing missing predictor values should also be distinguished from handling missing outcomes. A missing outcome cannot be used as an observed value when calculating a prediction error. The outcome-handling rule must be specified separately; the predictor-imputation procedure described here does not justify filling in held-out outcomes.
Why the Distinction Matters
Cross-validation estimates the performance of a learning procedure, not merely of a final fitted regression model. That procedure includes steps such as calculating imputation values. If those steps are performed differently during validation than they would be when predicting new cases, the resulting estimate can be misleading. Repeating the correct training-only imputation in every fold makes the evaluation match the intended use: learn the imputation rule from available training data, then apply it to new predictor rows.
A common mistake is to impute all rows first and then create cross-validation folds, reasoning that imputation only changes missing entries and therefore cannot affect the evaluation. The correction is that an imputed value is calculated from data, and the whole-data calculation can use observed values from held-out rows. Even when the held-out row has no missing value itself, its observed predictor value may affect the imputation value used for missing entries in the training rows. Fit the imputer inside each fold, then apply it unchanged to that fold’s held-out rows.
Check Your Understanding
Answer each question using the training-fold and held-out distinction.
- A training fold has observed values 3, 5, and 7 for a predictor; its held-out fold has a missing value and an observed value of 15. What value does the training-mean imputer assign to the held-out missing entry, and why is 15 excluded?
- In one fold, the training values for a categorical predictor are B, B, and A, while the held-out values are C, C, and missing. What value does a training-fold mode imputer assign to the missing entry?
- Explain why calculating one mean from all observations before creating cross-validation folds can be leakage, even if the validation error does not decrease.
- What should be done if every training-fold value of a predictor is missing but the held-out fold contains observed values?
- In the three-fold example, why does row 2 receive a different imputed value from row 5?