MLflow at Quantexa

Adam_Stein

In the dynamic landscape of Machine Learning (ML) development, the need for effective experiment tracking is paramount. This is especially true as organizations scale their operations. As the complexity of ML projects grows, so does the need for comprehensive tools to efficiently manage experimentation, iterations, and model versions. At Quantexa, we encountered this challenge head-on and sought a robust solution to streamline our ML workflows. We found a big part of the solution to these challenges was to use MLflow. MLflow is a powerful platform designed to simplify the end-to-end machine learning lifecycle. In this blog, we delve into how we leverage MLflow at Quantexa.

What is MLflow?

MLflow is an open-source platform for managing the end-to-end machine learning lifecycle. It's designed to help Data Scientists and Machine Learning Engineers with tracking experiments, packaging code into reproducible runs, and sharing and deploying models. It comes with a graphical user interface which makes it easy to use.

At Quantexa, we primarily use MLflow for experiment tracking and reproducibility. This will be the focus of this post. If you are interested in learning more about how you can use MLflow for model deployment, we recommend you spend time reading collateral on the MLflow Website.

What issues is MLflow solving?

MLflow addresses several critical challenges in our machine learning workflow. This enhances our overall efficiency and effectiveness. These challenges will also be faced by any client building models on top of Quantexa data. Here's how it tackles three of our most key issues:

Reproducibility and traceability

At Quantexa, we have a considerable number of machine learning models. It is critical that the origin of all these models is well-documented and readily available. This is important from a compliance perspective but also for internal purposes. The team and clients should be aware of performance metrics and each model's strengths and weaknesses. The models are also valuable intellectual property for the company, and we need to be able to re-produce results.

MLflow helps us with these issues. When models are trained, we record versions of code, model parameters, and the source of the original data. We can then use this information to perfectly recreate a model if ever needed. Furthermore, we have a record of any experiments which were used in selecting models. This means that a Data Scientist in the future can understand why certain decisions were made during the model prototyping phase.   

Collaborative insight

We want Data Scientists to be able to collaborate easily. MLflow fosters collaboration by providing a centralized repository for tracking and sharing experiments. This is particularly useful when it comes to gaining insight from non-technical stakeholders. Stakeholders can use the MLflow UI to access all the information they need to assess and give feedback on models. This dramatically speeds up the feedback cycle and leads to better models. It also saves time for all involved. MLflow enables data scientists to collaborate better with one another. This encourages a culture of knowledge exchange and innovation within the team.

Data correctness

At Quantexa, we are dealing with very large datasets when running predictions using our models. This scale can make it difficult to assess the quality and correctness of data. Manual sense checks are always performed, but mistakes can be missed. This is particularly true if errors only affect a certain portion of the dataset (e.g. one country).

To avoid errors slipping through, we produce lots of metrics and plots when we receive new data refreshes. We automatically compare this data to what we have seen previously. We log all of these metrics in MLflow. This makes it easy for the whole team to audit the data and spot mistakes such as missing data. Large differences are highlighted automatically even if only in one geographical region or data source. This information can then be easily forwarded to upstream teams to resolve issues. This process dramatically reduces the risk of erroneous predictions from incorrect data.

How do we use it?

In practice, we mainly use MLflow for three different workflows.

Feature generation

A pre-processing step transforms data produced by Quantexa Entity Resolution into a different tabular format. Our models then use this new format for training and prediction. We record aggregate information about data produced during this pre-processing step in the form of plots and CSVs. This aggregate information is compared to previous runs to highlight big discrepancies in easy-to-read files automatically. As mentioned, this reduces the risk of downstream tasks using erroneous data critical for model training and prediction.

Code versions and parameters for this pipeline are recorded. This allows us always to re-create the pre-processing step. It also provides an understanding of the data source we are training our models on.

Here is a list of some of the information we record at this stage:

• Mean, median, and variance of features
• Data drift scores for each feature
• Percentage of null values in a column

A plot produced to visualize the top drifting features across countries. Country and feature names have been omitted.

A plot produced to visualize the percentage of fields that have missing values in the Documents. These fields are used to produce features. Country and field names have been omitted.

Model training

We use MLflow to record model training. This can be to track multiple experiments during the prototyping phase of a project or be used as a reference for models in production. All the required information is recorded so that models can be re-produced if required.

Here is a list of some of the information we record when training models

• Model hyperparameters
• Test metrics such as precision and recall
• Confusion matrix

An ROC curve (receiver operating characteristic curve) produced against the test set during training.

A cumulative gains curve produced against the test set during training.

Staging model evaluation

We often iterate on our production machine learning models to improve their performance. We do this by adding more labeled examples to their training datasets or by adding more features. When we train a new model that we think should replace the existing model, we want to ensure that the new model is superior to the model we are replacing. We can look at the performance of the model across a test set, but sometimes this does not tell the full story. This is particularly true if the amount of labeled data we have is limited.

Before upgrading to any new model, we run an evaluation across the entire unlabeled dataset to see how the models differ in practice. This is important for understanding the real business impact of a change. This evaluation is recorded in MLflow such that the data scientists and any relevant stakeholders can review it.

Here is some of the information we record during the evaluation:

• The average size of differences in model score between the old model and new model
• Individual examples with the biggest differences in the model score between the old model and the new model
• Individual examples with SHAP explainability plots

Explainability plot produced when the model is used for inference across unlabelled data. Feature names have been omitted.

Visualization of the different scores produced by a model currently in production compared to the model selected to replace it.

Conclusion

In the dynamic realm of machine learning, effective experiment tracking is indispensable, especially with growing organizational scale. At Quantexa, MLflow has emerged as our solution, simplifying the ML lifecycle and addressing critical challenges in reproducibility, collaboration, and data correctness. Leveraging MLflow enhances our internal workflows and empowers clients to build models on Quantexa data to navigate complexities seamlessly.