Accelerating Drug Discovery: Machine Learning for Protein Crystal Detection
Understanding the 3D structure of biological molecules like proteins is a crucial component of drug discovery. This understanding hinges on a process known as protein crystallization. However, identifying these tiny crystals amongst a myriad of similar-looking molecules can be a time-consuming and resource-draining process.
In this article, we will take a deep dive into this process, its relevance, existing practices, and how using AI-powered models can help reduce the cost of developing medicines by using an efficient computer vision model offering benefits for drug discovery.
The proposed model establishes a new state-of-the-art in protein crystal identification, confidently overcoming notorious limitations of former approaches (most notably, their generalization capability).
- We’re presenting a novel machine learning model that significantly accelerates the important process of protein crystal detection in drug discovery.
- Our approach, leveraging pre-trained models and a focused binary classification, outperforms the state-of-the-art MARCO model by reducing the likelihood of missing crystals from 11.1% to 7.6% on MARCO’s test set. This represents a decrease of over 30% in the original error rate.
- Moreover, we are able to achieve superior performance on new data using as few as 60 images of crystals.
- Important Metrics:
- Crystal Identification Accuracy:
- Old Method (MARCO): 88.9% recall and 93.4% precision on crystals. [*]
- New Method: 90.8% recall and 93.9% precision on crystals (4-class model), and a recall of 92.4% with a precision of 93.4% (binary classification model).
- Overall Accuracy:
- Old Method: 93.5% overall accuracy for 4-class classification, 97.7% for detecting crystals.
- New Method: 94.0% overall accuracy (4-class model) and 98.1% (binary classification model).
- Fine-tuning to new data:
- Old Method: on new “VIS” and “UV” datasets, MARCO achieves crystal detection accuracy of 91.1% and 75.7%, respectively.
- New Method: fine-tuned using as few as 60 images of crystals, our model obtains accuracies of 92.9% and 81.1%, respectively; using 2000 images of crystals, the accuracy increases to 95.7% and 91.9%, respectively.
- Computational Effort:
- Old Method: 260 training epochs across 50 GPUs.
- New Method: 3 + 2 training epochs (4-class model + binary model) on a single GPU.
- Crystal Identification Accuracy:
- This breakthrough has the potential to provide substantial cost and time savings in drug development.
Table of Contents
Protein crystallization is crucial in understanding the 3D structure of macromolecules such as proteins or antibodies. This step is essential in drug discovery and development. Drug discovery requires the identification of hot-spots that could be druggable, as well as performing structure activity relationship (SAR) analysis to improve the potency of hit/lead small molecules.
Protein crystallization is a bottleneck in drug discovery. It’s time-consuming and requires multiple lengthy attempts at various conditions; because of this, identifying crystals is crucial. Reports indicate that leveraging structural knowledge in drug discovery can halve the costs and time required from targeting a molecule to filing for clinical trials, which represents up to 10 million USD. Thus, our enhancements in detecting protein crystals could be a key contributor to increasing the chances of identifying crystals that could then be sent to synchrotrons and would result in electron density maps that lead to a 3D protein structure.
“Every time you miss a protein crystal, because they are so rare, you risk missing on an important biomedical discovery.”- Patrick Charbonneau, Duke University Dept. of Chemistry and Lead Researcher, MARCO initiative.
The Previous Frontier: MARCO
A state-of-the-art model already exists, Machine Recognition of Crystallization Outcomes (MARCO). This model achieves 89% recall on crystals [*], meaning that on the test subset of MARCO, it correctly identifies 89% of the crystals, so 11% go unnoticed.
The original research trained a modified Inception-v3 convolutional neural network (CNN) on a 4-class problem. This model achieved up to 93.5% overall accuracy on the test set but saw a performance reduction on crystals, correctly identifying only 89% of crystal images in a holdout set.
This performance difference is due to the fact that crystals are a minority class in the dataset, and the authors did not report any attempt to balance the classes. However, the major goal of drug discovery is to identify crystals; thus minimizing misclassified crystals is of utmost importance.
Furthermore, the original model was trained across 50 GPUs for 260 epochs, a total GPU time of almost 40 days! This is a significant amount of computational time and expense required to train a classification model.
Finally, accuracy of MARCO decreases when evaluated on data from new labs. Earlier this year, researchers prepared two datasets consisting of crystallization images obtained using different apparatus than the ones used when generating data for the MARCO dataset. The pictures in the datasets were obtained using visible and ultraviolet light, so we call those datasets “VIS” and “UV”, respectively. The 4-class classification accuracy of MARCO on those sets drops to as little as 82% and 49%.
Watch Our Presentation
Delivered at Harvard for the Center for Computational Biomedicine (CCB) seminar series, our team presents our work in detail, emphasizing its impact on drug discovery, particularly in reducing drug development costs and timeframes. The presentation also includes a demonstration of shiny.molstar, a visualization toolkit for large-scale molecular data.
Our novel machine learning model has outperformed MARCO by not only identifying crystals with higher recall (missing just 7.6% of crystals) while retaining the same precision but doing so with a fraction of the computational effort required previously.
Where MARCO required 260 training epochs across 50 high-powered graphics processing units (GPUs) – using an equivalent of 40 GPU-days – our model achieved superior accuracy in just 3 epochs (meaning only 19 GPU-hours – a reduction in the computational cost of 98%). Moreover, the model can be adjusted to generalize to novel data sources with significantly smaller effort. This opens the possibility of automating the identification of crystals at scale!
Our Strategy for Success
Our approach to improving crystal detection is based on a series of strategic steps that have enabled us to overcome the limitations of previous models, specifically the MARCO model. Here’s a detailed insight into our methodology:
Our approach improves upon the previous model by:
- Prioritizing crystal detection over overall accuracy: Instead of striving for high overall accuracy across all categories, as was the case with MARCO, we made a conscious decision to prioritize crystal detection. Our primary objective was to minimize errors in this crucial area, especially considering that an 11% misclassification rate of crystals in drug discovery can lead to significant costs and setbacks.
- Significantly reducing computational time and costs: We utilized pre-trained models and transfer learning to ensure high accuracy while drastically reducing the computational time and costs.
The MARCO dataset originates from a comprehensive collection of images from protein
crystallization trials from 5 organizations. These images represent various crystallization
outcomes and are classified into 4 categories:
- Crystals: Images showing successful crystallization of proteins.
- Clear: Images where the solution remains transparent, indicating an absence of Crystals.
- Precipitate: Images displaying a cloudy or amorphous deposit, meaning the protein has precipitated but not formed crystals.
- Other: A category for images that don’t fit neatly into the aforementioned classes, possibly due to multiple phenomena occurring simultaneously or other unforeseen outcomes.
The “VIS” and “UV” datasets have been curated by researchers from the Collaborative Crystallization Center to evaluate performance of the MARCO model on new data. Images in those datasets come from a different crystallization kit than ones used for capturing images in the MARCO dataset. They are split into the same 4 categories.
Here’s our approach towards achieving this.
- Leveraging Pre-trained Models: We leveraged the power of large pre-trained computer vision models, specifically modern CNN (Convolutional Neural Networks) and Vision Transformer networks. By employing transfer learning, we were able to fine-tune these models on the original 4-class problem. This allowed us to directly compare our results with the MARCO model.
Using pre-trained models and transfer learning has two main advantages. Firstly, the baseline model has already been trained on a large and diverse image dataset, meaning it can be easily fine-tuned to produce a highly accurate model for a specific problem, in this case crystal detection. Secondly, since we only have to fine-tune a small fraction of the model parameters, the computational cost is drastically reduced compared to the state-of-the-art models.
- Focused Binary Classification: Recognizing the need to reduce the high error rate in crystal detection, we shifted our strategy to binary classification, specifically targeting crystals. This involved converting the dataset labels to a binary problem: crystal or not crystal, and further fine-tuning our models trained on the original 4-class problem.
However, given that only 12.8% of images in the dataset contained crystals, the data was heavily imbalanced. To address this, we incorporated a weighted loss function during the training phase. This ensured that misclassified crystals were penalized more heavily compared to images without crystals. Since our models were already fine-tuned to the 4-class dataset, this further fine-tuning was efficient and rapid.
- Low-rank Adaptation of Transformer-based Models: The experimental set-ups vary from one laboratory to another, and gathering hundreds of thousands of images for retraining a model to fit new data well is a costly endeavour. We used Low-Rank Adaptation (LoRA) techniques with our Vision Transformer models to fine-tune them to even small samples of new data, i.e., of “VIS” and “UV” datasets.This technique has been successfully used for fine-tuning pre-trained Large Language Models due to its robustness to overfitting. LoRA ensures that the model does not forget too much of what it has learned before, while giving it the flexibility it needs to adapt to the new data or task.
The new method achieved comparable accuracy to the old method but with a substantially reduced computational effort. More importantly, it was far more efficient at identifying crystals, which is the main goal for advancing drug discovery. Finally, it allows efficient fine-tuning to new data using very few samples, making it very cheap to adapt to crystallization data it has not seen before.
- Our 4-class model achieved an overall accuracy of 94.0%.
- On crystals, our 4-class model achieves a recall of 90.8% and a precision of 93.9%.
Our confusion matrix shows strong classification accuracy across the four classes. The full confusion matrix is shown below:
Receiver Operating Characteristic Curve
The Receiver Operating Characteristic (ROC) curve for our 4-class model is shown in comparison to the original MARCO model below:
- Our model achieves an AUC of 0.993 compared to 0.986 for the MARCO model.
Binary Class Model
Further fine-tuning our 4-class model for binary classification significantly improves the performance on crystal detection:
- Our binary model achieved an overall accuracy of 98.1% (MARCO’s crystal detection accuracy is 97.7%)
- On crystals, our binary model achieves a recall of 92.8% and a precision of 92.6%!
Fine-tuning to new data
Fine-tuning even further to perform crystal detection on the “VIS” and “UV” datasets, our models outperformed MARCO using very little data:
- On the “VIS” dataset, our model needed just 60 images of crystals (supplemented by 60 images from each other category) to achieve an accuracy of 92.9% on the held-out test set (MARCO’s crystal detection accuracy on this set is 90.7%), which increased to an accuracy of 95.7% when using 2000 images of crystals.
- On the “UV” dataset, with just 60 images of crystals it achieved an accuracy of 81.1% (MARCO’s accuracy is 75.9% in this setting), which increased to an accuracy of 91.9% when using 2000 images of crystals.
To delve into the intricacies of machine learning evaluation metrics, explore our comprehensive guide on the Top 10 Machine Learning Evaluation Metrics for Classification.
Benefits of this New Approach
- Enhanced Accuracy and Efficiency: This novel strategy outperforms the current state-of-the-art MARCO model with an increase in performance for protein crystal classification (achieving a 92.4% recall rate compared to 88.9% with the same precision) while drastically reducing the computational load (requiring only 3 epochs of training versus 260 epochs).
- Streamlined Computational Resources: Another benefit of this model is it significantly diminishes the computational resources and time required, facilitating the swift development of precise models applicable to a broader range of scenarios.
- Efficient Fine-tuning to New Data: For experimental set-ups not incorporated into the MARCO dataset, we need just 60 images of crystals to make the model fare better than the MARCO model, and with just 2000 images of crystals, it already achieves crystal detection accuracy of 92-96%.
- Cost and Time Savings in Clinical Development: Incorporating structural knowledge has the potential to cut both expenses and duration required to advance a target molecule to clinical trials. Leveraging this method could provide substantial cost reductions.
This new method is a big step forward as it identifies protein crystals more accurately and efficiently, which can accelerate the drug discovery process. Plus, it can be adjusted to work well with new, different data, making it a flexible tool for future research.
By converting the problem to binary classification and further fine-tuning our models, we observe an increase in crystal detection rate by over 3% – reducing the rate of misclassified crystals by over 30% without increasing the number of false positives!
At Appsilon, we’re committed to solving complex drug discovery problems with cutting-edge machine learning solutions.
[*] The results for the MARCO model given in this article differ from the ones provided by the authors. We obtained those using the MARCO model available online at https://github.com/tensorflow/models/tree/master/research/marco and the MARCO validation dataset available at https://marco.ccr.buffalo.edu/download .
We are Open to Collaboration
We are actively seeking collaborations to enhance our capabilities and provide tailored solutions.
If you have crystal images or datasets, we invite you to share them with us. We could apply our top-notch machine learning model and help you optimize crystal detection. Let’s collaborate to unlock the full potential of your data and explore innovative solutions together.