Stratified Splitting of Grouped Datasets Using Optimization
By: João Paulo Figueira
Tags: Gradient Descent, Optimization, Search, Machine Learning, Heuristics

Stratified Splitting of Grouped Datasets Using Optimization

Photo by Scott Webb on Unsplash

One of the most frequent steps on a machine learning pipeline is splitting data into training and validation sets. It is one of the necessary skills all practitioners must master before tackling any problem. The splitting process requires a random shuffle of the data followed by a partition using a preset threshold. On classification variants, you may want to use stratification to ensure the same distribution of classes on both sets. When handling time series data, you might want to skip shuffling and keep the earliest observations on the training set. You might also have grouped data, where you will probably want to keep groups together on either dataset. Finally, you might be lucky enough to meet a mix of the above and not have a handy solution for it, as it happened to me.

This article presents the solution I devised for handling a stratified split of grouped data, requiring that groups be kept together on either dataset.

Background

I have recently worked on a predictive maintenance project targeting a few thousand industrial devices. These devices stream telemetry data over a wireless connection to a backend system, where it is stored and further processed. All devices undergo a regular preventive maintenance schedule, which is not enough to avoid occasional unexpected breakdowns. Due to these breakdowns' disruptive nature, the device manufacturer decided to create a set of predictive maintenance models to try and predict the probability of an untimely device failure before the next maintenance period.

My team got assigned to work on a specific device component and devise a predictive maintenance model to estimate the failure probability for a future time frame. We chose to use a classification model and went through the whole process of setting up the pipeline. When we got to the dataset splitting phase, we decided not to split device data between training and validation sets. Each device data would be either on the former or the latter, never divided between the two. We also decided to do a stratified split, ensuring that both sets had the same proportion of positive observations.

At this moment, we understood that there is no documented way to go about this issue. The Scikit-Learn package implements solutions to split grouped datasets or to perform a stratified split, but not both. Thinking a bit, it makes sense as this is an optimization problem with multiple objectives. You must split the data along group boundaries, ensuring the requested split proportion while keeping the overall stratification the same.

Needless to say that this part of the pipeline was manual, using a random number generator and a few trials. No, I am not proud of it, hence this article.

Problem Statement

Given a labeled and grouped dataset, we want to split it into training and validation sets while keeping the label distribution as close as possible on both and group integrity. After breaking the data into the two datasets, the groups must maintain their integrity, assigned to either set and not split among them. Furthermore, the splitting process should closely respect the imposed splitting proportion and label stratification.

The expected result for this problem is, given an input dataset, the list of groups assigned to each dataset, ensuring that both the train/validation split and the stratification are as close as possible to the specified values.

Solution Research

I started by seeking help from classic optimization packages, such as Google’s OR-Tools. Unfortunately, I found no solution there. The implemented optimization solutions seemed close to what I needed, but still no cigar. I was looking for a solution that I could use from a Python pipeline, so some well-known Java-based optimizers seemed to be out of reach. Why not roll my own?

A few years back, I was involved in developing a vehicle routing optimizer, which led me to learn a lot about discrete optimization algorithms. I summed up my learnings in the article below.

After reviewing all these ideas, I decided to develop a Python-based optimization solution for this problem. I am presenting the code for a proof-of-concept that may later lead to an independent Python package. Let us start by reviewing the problem model.

Problem Model

To solve this problem, we don’t need to handle the whole dataset, just its counts. We only need to determine the number of different groups and the class counts per group. To motivate the problem model discussion, let us see a toy example, where we want to impose a 70/30 split.

The above image depicts a toy problem example, along with a possible solution. Rows contain counts of the ten groups, with the first column containing the total count and the second containing the number of class one samples. The last column is the proposed solution, with the value one representing the assignment to the validation dataset. (Image source: Author)

The image above shows the model and solution for an elementary problem. Each table row contains a group’s description, with the size on the first column, the number of class one samples in the second, and the third indicating the validation set inclusion (the proposed solution). You can see the column totals at the bottom, along with the calculated proportion of class one samples (20.70%).

So how does the proposed solution perform? To check the compliance to the 70/30 split, we need to add all the first column values flagged (value 1) in the third as belonging to the validation set. The partial sum (10960) represents 29.74% of the total, quite close to the required 30%. To check the stratification, we do the same thing, but on the second column. The sum (2328) is 21.24% of the previous count, meaning that we got quite close to the required 20.70%.

It is worth noting that both the problem and solution representations are quite simple. We can represent a solution as a boolean array, which is very compact.

So how did I get to the solution?

Optimization Model

As previously seen, we want to find a solution to the given problem that approaches both the split and the stratification target. Note that the stratification will have multiple targets if you have more than two classes.

The optimization process starts by generating an initial approximate solution using a random permutation of the groups.

The Initial Solution

The initial solution is our first take at the problem, a leg up in the search process. To give the optimization process a fast start, we will concern ourselves with the train/validation proportion only and start with a random shuffle. As you can see, the function below uses a random permutation of the group sizes to select an approximation of the required proportion, not taking into account the strata. The procedure stops once the validation set exceeds its required capacity.

The code above generates the initial solution using the train/validation proportion as the sole stopping criterion. (Image source: Author)

Now that we have a solution, we can try to improve it. We will do so by making small changes to the solution, essentially flipping bits in the solution array. But how do we know that we are improving the solution? How do we measure improvements? Enter the cost function.

The Cost Function

The cost function measures the “distance” from the existing solution to the theoretical optimum. We know this optimum from the problem statement and want to generate a solution array that best matches both the split and the strata proportions. The function below calculates the cost associated with a given solution array and is JIT-compiled with Numba for performance.

The optimization process evaluates all solutions using the above cost function, with lower costs equating to better solutions. As you can see, the cost function considers the difference between the requested dataset split proportion and the one from the evaluated solution. Next, it does the same with the difference between the problem’s strata and the solution’s strata proportions.

Now that we have a cost function, how do we generate solutions to evaluate with it? We start by developing an initial solution and then iterate from it.

The Search Process

The search process iteratively generates new solutions based on a previous one and evaluates them. In this particular case, we create a solution by flipping a bit of the solution vector, thereby moving a group to or from the validation set. We must be careful not to create previously generated solutions, lest we send the optimization search process into an infinite loop.

To avoid cycles, we do the same as Ariadne did in her maze; we mark the trodden path. Instead of laying a thread, we use a Python set to store our search space's visited states. But keeping the whole state array for each step would not be memory efficient nor fast. To speed up set inclusion tests, we must first transform the state array into a more compact representation, a string. Python handles string hashes quite efficiently, and we also gain by compacting the state representation. The function below shows how to convert a state array into an ASCII string. Note that the procedure only uses six bits per character and is JIT-compiled using Numba.

This function converts a solution vector into an ASCII string by aggregating every six bits. (Image source: Author)

Incumbent Solution

We always keep handy the best solution so far, the incumbent solution, and use it as the final one when the optimization process ends. Whenever the process generates a better solution with a lower cost, it replaces the incumbent with the better one. This way, the optimizer always has an answer ready to serve, independently of the selected termination criteria.

The optimization process generates new solutions by flipping a single bit of the previously-generated solution, using a specific criterion. Here we will explore two different approaches to this process. The first approach uses a heuristic to generate the new solution and is quite similar to methods used in informed search algorithms. The second approach uses the gradient of the cost function to determine the next best solution and is similar in spirit to gradient descent.

Heuristic Approach

The heuristic-based optimizer's main loop starts by generating the list of possible derived states by flipping all the parent solution bits. It then filters out the already visited states using a set that stores the string-encoded binary state arrays (like a tabu list without tenure). As you can see from the function below, generating new solutions is just a matter of iterating through the bit array, flipping each one, and filtering out the already explored solutions.

The above function generates all the possible moves for a given solution. (Image source: Author)

The procedure then evaluates each of the newly-generated states using the cost function and sorts them accordingly. It's time for the process to select the next candidate for expansion, and it’s where it uses the heuristic. Quite simply, the best candidate is the one with the lowest cost. This approach makes the optimization process look for a better solution closer, in cost space, to the previous one. On discrete optimization problems, it is quite usual to find local minima close together, so this approach makes sense. It goes by the name of “intensification” in specialized literature.

It is important to note that we do not necessarily improve over the incumbent solution by selecting the lowest-cost solution. We are merely picking the lowest-cost solution from the generated list, but this might have a higher cost than our best so far. By doing so, we are engaging in what is called “hill-climbing.” As an analogy, you can think of a skier going up a small hill in the hope of finding a steeper slope to descend on the opposing face.

However, by using this strategy, we risk getting stuck in a local high-altitude valley. The optimization procedure keeps track of how often it goes uphill and, after a predetermined number of failed attempts at reducing the incumbent’s cost, it engages a different behavior. Going back to the skier analogy, imagine that she gets bored trying to find a nice slope and calls her private helicopter to take her to a neighboring and more promising peak. With some luck, she might find a more exciting track.

This strategy is called “diversification,” instead of using the lowest cost, the search procedure will go for the highest one when it senses a valley. This strategy allows the optimization process to search for different paths down to the elusive global minimum.

The function that solves the optimization problem using the heuristic search is below.

The function above searches for the best train/validation split using a heuristic. (Image source: Author)

As you can see, on each iteration, the function generates a list of possible solutions and uses only one as the next move. For problems with lots of groups, this approach will become too slow for practical use.

Gradient Approach

We can try to select the best expansion by following the cost function gradient as an alternative to the previous approach. By looking at the cost function, we can quickly derive its gradient. The code below calculates the cost function gradient as a vector. Note that the gradient vector merely contains the difference in counts along each dimension of the cost.

The function above calculates the cost function gradient. (Image source: Author)

We can now use this vector to search for the best modification to the solution. Remember that we are merely moving groups between the train and validation sets, so we want to use the cost gradient to determine which move will most impact the cost value. Note that the first component of the cost vector (line 14 above) is positive when the validation set count is below its expected value. To improve the cost, we should add another group to the validation set. The same is also true for the strata components. If the optimization process gets a cost gradient vector with positive elements, it should select one group from the training set and move it to the validation set. On the other hand, if all cost gradient components have negative values, we should remove a group from the validation set and put it back on the training set. The decision for all other cases in between will depend on vector similarity, as explained in the following.

How should the optimization procedure choose the best group to transfer between datasets? Here we must use vector similarity to match the cost gradient vector to each vector from the problem matrix, and the chosen metric was the cosine similarity.

The cosine similarity is the ratio between the internal product of two vectors and the multiplication of their norms. (Image source: Author)

The cosine similarity derives from the definition of the inner product between two vectors and is maximal when both vectors point in the same direction and minimal when opposite. The code below implements the cosine similarity calculation with the adaptation to our problem, reversing the validation group vectors' sign.

By reversing the validation vectors' sign (line 11 above), the algorithm ensures that these will be closer to a cost gradient that points in the reverse direction of a move from the training dataset.

Now we can use the cosine similarity to generate the next move. The code below contains the function that produces the next best solution based on the cosine similarity criterion.

The solution generation function computes a vector of cosine similarity scores calculated between each of the group vectors from the problem matrix and the calculated cost gradient vector. The next best solution corresponds to flipping the bit of the problem vector with the highest score. Note that for a diversification move, we use the lowest score.

The function above searches for the best train/validation split using the cost gradient. (Image source: Author)

The Code

All the code for this article is available on the accompanying GitHub repository. It lives in a single Python file that you can run from the command line or from within your favorite IDE. There is a set of standalone functions and two classes, each implementing one of the optimization approaches described above. Both solvers continuously loop until they meet a stopping criterion. Presently, they use as stopping criteria a minimum cost value or a maximum number of iterations without improvement.

Search Solver

This class’ solver method expands the current solution into all permissible states. It then ranks them according to the cost value and selects the lowest for the next expansion. This solver can be very slow to converge to a solution as the expansion process may take a long time to run on numerous groups' problems. We can circumvent this disadvantage by artificially limiting the number of expansions while adding a bit of stochasticity (not done here).

Gradient Solver

The gradient solver uses the cost gradient of each solution to select the next best move. Instead of expanding all allowable solutions, this solver goes “straight to the point” by choosing the best solution to develop. As previously explained, this solver uses the cosine similarity to find the best match between the cost gradient vector and the group vectors. This design speeds this solver up considerably, making it outperform the search solver on almost every problem instance.

Results

To compare both methods' performance, I ran a set of experiments varying the number of groups and measuring the runtimes. For each group size, I executed the code for both methods ten times and took the average in seconds. Besides the number of groups, I kept all other parameters constant.

The following chart displays how both optimization methods perform, time-wise, as a function of the number of groups in the problem definition. As you can see, the heuristic optimizer performance is super-linear, while the gradient appears linear.

The chart displays the time performance comparison of both optimization methods as a function of the number of groups. (Image source: Author)

Note that these results should only give you a hint on both methods' performance, as the optimization quality (measured as the final cost value) varied a lot. A better performance comparison would probably entail the restriction of having both methods report results within a narrow interval. Nevertheless, I will always prefer to use the gradient method over the heuristic.

Conclusion

In this article, we have looked at the problem of splitting a grouped machine learning dataset into train and validation while respecting both the required split proportion and the inherent label stratification. We have looked at possible solutions from the standard Python libraries, and when none were apparent, we decided to roll our own. We have framed this problem as a discrete optimization process and devised two possible solutions for it. The first solution derives from the informed search using heuristics, while the second uses the cost function gradient to minimize the solution cost. The gradient-based approach displays much better execution performance than the heuristic one.

Future Work

As I previously stated, the work you see here is still a proof-of-concept. I have shown how to solve the stratified split of a grouped dataset using a discrete optimization approach. The next step will be to convert this code into a full-fledged Python package with improved performance using Cython compilation.

Another possible expansion of this idea would be to enable stratified K-Fold splitting.

Resources

GitHub Repository