Model interpretability roadmap

[miquelduranfrigola]

Adding an interpretability module to ZairaChem

Background

This project is related to @HellenNamulinda's MSc thesis at Makerere University. The thesis is co-supervised by Dr. Joyce Nakatumba-Nabende. At the moment, ZairaChem does not have any explainable AI (XAI) capabilities. The goal of this project is to develop an automated tool for model interpretability that can be incorporated into ZairaChem. While there are many approaches for chemistry, here we will focus on the following:

1. A limited set of molecular descriptors: We need to focus on descriptors that a medicinal chemist would understand. Therefore, we will focus on a set of commonly used descriptors. This post can offer some guidance.
2. XGBoost/CatBoost regression: These are tree-based methods that work well in many scenarios, especially for regression. Automatic hyperparameter tuning can be achieved with Optuna.
3. Shapley analysis: Shapley values work naturally well with tree-based methods. We will focus on this approach to interpretability.
4. We will validate the tool in the context of a dataset donated by Medicines for Malaria Venture (MMV) to Ersilia.

Objectives

1. To develop a standalone Python tool for chemistry specifically oriented to Shapley value analysis of molecular descriptors. The tool is called xai4chem.
2. To apply the tool to the MMV dataset.
3. To incorporate the tool into ZairaChem.

Steps

• Create a Python package structure using a boilerplate for the xai4chem repository. @DhanshreeA can assist with this too. #1
  
• Define molecular descriptors to be used. @miquelduranfrigola can assist with this. #3
• Normalize molecular descriptors. #3
• Implement an XGBoost or CatBoost regressor with hyperparameter optimization. Optuna or FLAML can help.
• Apply Shapley value analysis using the SHAP Python library.
• Train-test split of the MMV dataset.
• Normalization of the outcome variable.
• Train the regressor.
• Predict on the training set.
• Do accuracy analysis on the test set.
• Do explainability analysis on the training set.
• Write a report on the modeling of the MMV dataset.
• Incorporate xai4chem as a ZairaChem dependency. Since ZairaChem is focused on binary classification (for now), xai4chem should be used on a reference dataset calculated in a forward pass to a trained model, taking as outcome variable the classification score or probability.

FAQ

Where do we create issues?

Most issues related to this work should be created in the xai4chem repository. When we reach a point of integration to ZairaChem, we can create issue there correspondingly.

Is there a more comprehensive description of the project available?

Yes. This is part of @HellenNamulinda 's MSc project and she is writing a thesis accordingly. A project proposal document is already available.

[HellenNamulinda]

Progress Updates

• Created a python package structure for xai4chem, Use a Python package boilerplate xai4chem#1
• Created a class to calculate molecular descriptors, Calculate molecular descriptors using datamol xai4chem#5
  For the start, we are using the default 22 features of datamol descriptors.
  Others will be experimented and tracked in the issue; Create a class to calculate molecular descriptors using fit and transform logic xai4chem#3
• Implemented a regressor class with the option of XGBoost or CatBoost. This is for training, evaluating, explaining the model.
  Hyperparameter optimization is done using optuna during training.
  To evaluate the model, the evaluate method is called and metrics such as Mean Squared Error, Root Mean Squared Error, Mean Absolute Error, R-squared Score, and Explained Variance Score are saved in a json file.
  To interprete the model, the explain method method is used. Shapely values are calculated and currently three plots are saved; a summary_plot( barplot showing average of shapely values), a beeswarm plot showing shap values of all the inputs, and a waterfall plot as summary for one input(This shows the shapley values for each feature).
  We will explore mapping shapley values unto the molecular graphs, and this will be tracked in the isuue; Applying Shapley Value Analysis: Visualization Plots and Number of Features to Consider xai4chem#7

Next

• Outline steps to run the tool as python package

[miquelduranfrigola]

Thanks @HellenNamulinda - this is useful.

We did not discuss much about optuna. Does it work as expected for you? Do you get better results than using XGBoost with default parameters?

[HellenNamulinda]

Hello @miquelduranfrigola,
I apologize for the delay in providing this update.

I've compared XGBoost's performance with default parameters against those optimized by Optuna. Surprisingly, the default parameters seem to yield better results, with an R2 of 0.50 compared to around 0.3 achieved with Optuna. I will add these findings in the notebook that I will uipload later today.

It's worth noting that the parameters optimized by Optuna can vary in each study, introducing some uncertainty in the results. I may need to adjust the search space to align more closely with the default parameters.

Additionally, I've observed that training a CatBoost model is consistently slower, which prolongs the optimization process with Optuna.

However, I believe Optuna is still valuable. It's essential to carefully define the search parameters to achieve optimal results.

[miquelduranfrigola]

Thanks @HellenNamulinda , this is useful. I agree we need to use optuna. We'll have to play a bit with the search space, then, and perhaps increase the number of iterations.

[HellenNamulinda]

Progress Updates

• The steps for using xai4chem as a python package are outlined in the README.md, Update ReadME.md with installation and usage instructions, and code improvements xai4chem#8
• After adjusting the search space, the performance improved to R2 > 0.50. This can be seen in the notebook. Finding better parameters requires time, and hence increasing the number of trials(n_trials) can help to get better params.

Next
From the meeting,

• Rename the train method to fit. So that we have fit, predict, evaluate and explain.
• It would be better to reduce the args for the regressor methods. Forexample, the output folder should be defined when instantiating the regressor class.
• Explore Zero Shot AutoML for mining good hyperparameters.
• Include feature selection. This can start by increasing the number of features(datamol descriptors), Create a class to calculate molecular descriptors using fit and transform logic xai4chem#3 (comment), and then 2D mordred descriptors.

[miquelduranfrigola]

Thanks @HellenNamulinda , all next steps sound good to me.

• For the zero-shot, feel free to explore this: https://microsoft.github.io/FLAML/docs/Use-Cases/Zero-Shot-AutoML/
  Good to know that, with more time, you could get to a competitive solution with optuna.
• For feature selection, feel free to explore this: https://github.com/AutoViML/featurewiz
• 2D mordred is less prioritary if a wider datamol descriptor is available.

[HellenNamulinda]

From the meetings,

• We will continue testing with the 3 types of descriptors that were added. That's; datamol, RDKit and Mordred.. Fingerprints will be explored later.
• While featurewiz was added for feature selection, the final number of features can't be determined by the user. So, the logic will be to specify the k-features. If featurewiz is applied and feafures selected are less than the desired number of features, (say 100), we don't use featurewiz. If the number of features returned by featurewiz are morethan k(100), we then select the top k.

@miquelduranfrigola,
From the experiments using the MMV dataset. with feature selection, mordred gives an r2_score of 0.39 and mae of 11.73, compared to rdkit r2_score of 0.35 and mae of 12.24.
Also in the TDC benchmarks, the performance of the pipeline is promising. With the performance on some datasets placing in top 3 in the TDC leaderboards.
More information in the sldes.

We agreed to experiment and compare performance, rdkit descriptors without feature selection(using all the descriptors).

Note: All the experiments were done done using catboost with default parameters.

Trying zero-shot, XGBoost's performance on the pf_3d7_ic50 data improved from r2_score: 0.64 to r2_score: 0.71.
Am going to finalize testing this and we use zero-shot in place for optuna.

[HellenNamulinda]

From last week,
We agreed to use morgan fingerprints. And this was implemented(ersilia-os/xai4chem@d2ff580).

Also, use Zero-shot AutoML(ersilia-os/xai4chem@f5d5ad8). but FLAML zero-shot only supports XGBoost and not Catboost.

[HellenNamulinda]

To be able to interprete other trained models besides the regression models developed using xai4chem, it was best to have the explain_model as a separate module(independent of the regressor).

With the explain_module, interpretability plots can be generated even for trained classification models.

[HellenNamulinda]

Hello @miquelduranfrigola,
From the meeting, we have reviewed everything that has been implemented in xai4chem.

In our pipeline, we choose features to be any of the three descriptors(Small(datamol), Mid-size (RDKit), and Large (Mordred)) or the count-based morgan fingerprints.
Also, Feature selection automatically selects the relevant k features during training, if the value of k is given.

For interpratability, we are currently saving three interpretability plots; barplot, beeswarm plot and a waterfall plot for the first data sample(this can be generated for other samples).

All the other usuage details are documented in the README.

Some pending concerns
Save test results as csv file(not joblib)
Save shapley values as csv file(not only interpretability plots)
Add interpretation onto chemical structures for morgan fingerprints?

Benchmark
With xai4chem pipeline on PPBR_AZ dataset(Plasma Protein Binding Rate(%)): Using RDKit descriptors, the MAE of 7.618 and r2=0.3401 places 2nd in the leaderboard.

The MMV Data:
We started with a small set(LDH assay: 4816 samples). The performance is ranging between 0.36 and 0.40(r2_score) and 12.10 and 11.53 (MEA).

And this brings us to combining descriptors and fingerprints?
What would you advise on Feature Maps as input features -2D(descriptors and fingerprints). I haven't yet implemented it, but it is something I have started looking at next.

[miquelduranfrigola]

Thanks @HellenNamulinda — very informative.

Let's first close the pending concerns and then we will look into blending or not descriptors and fingerprints.

[HellenNamulinda]

@miquelduranfrigola,
For the test results, we are saving a csv file containing the smiles strings and the model output values.
Where as for the interpretability results csv file, we are saving the descriptors/fingerprints and the shapely values.

This week, I'm working on mapping interpretation(shapely values) unto chemical structures for fingerprint features.