Cross validation while preserving a column (not the target ) distribution

Question

So i'm doing cross validation and then i'm predicting using all the data on a test set ( a hold-out set ). My hold-out set has the same ratio on a column than the train ( seems thats how the test set was generated, a function that sampled it and tried to preserve the ratio for the target classes, and a particular column ) . My local CV is a bit lower than my score on the test set, and i think the problem is stemming from the fact that i'm using stratification only for 'y'.

Can lack of stratification of that feature be the reason of Cv & test scores aren't really close?

And if so how can i perform stratification for the target and a feature! Thanks

Edit : i'm already doing stratification on the target since my data is imbalanced.

Answer 1

One idea would be to combine the two columns (one predictor and the target) and then stratify using the combined column.

Example: say for some observations, the target and the column take the following values: target = [0,1,0] and column = [A,A,B]. Then the combined column could look something like [A0,A1,B0] and could be used for stratification.

I'm obviously assuming that the predictor column is categorical - you might have to take a different approach for a continuous variable.