Innovations in Virtual Drug Screening

Natalia Vassilieva, Product Management Director | June 16, 2021

Given the investment of money and time required to develop new drugs, companies are increasingly recognizing the value of artificial intelligence drug discovery. Billions of dollars and decades of research are spent on the development of new drugs. Virtual screening has emerged as an important tool, reducing both cost and time.

Drug discovery has seen substantial innovation in recent years, but challenges and areas for improvement remain. The Wafer-Scale Engine (WSE) from Cerebras Systems accelerates drug discovery in the next era of deep learning.

How Drug Discovery Works

The drug target is typically a human molecule (often a protein) that is known to be associated with a particular disease. This can be used to identify the most promising small-molecule compounds, which can bind to the target protein and modulate its function.

Virtual Drug Screening

Virtual screening takes place in the early discovery stage, with the goal of discovering a drug target.

Traditionally, identifying the most promising compounds was done in vitro in a lab. This method can screen only thousands of compounds at a time, which is a tiny fraction of the compounds that can be screened using AI. Virtual screening uses computers and software to do screening in silico, enabling a much faster and broader process with billions of compounds in existing libraries of synthesizable chemical compounds.

Instead of physically testing every compound, virtual screening relies on algorithms to predict which ones will be most likely to bind to the desired drug target. This computational technique shrinks the time and cost of developing new therapeutics by focusing efforts on those compounds that show the most promise of producing effective drugs. These select candidates are then synthesized and put through expensive physical lab trials.

The use of artificial intelligence in the drug discovery process has already led to great leaps forward in time and cost efficiency, but significant limitations remain — if the AI has to rely on traditional GPUs.

Challenges of AI Drug Discovery

Deep neural networks have shown great promise in predicting binding activity between a potential drug compound and a target protein. Rather than conducting computationally expensive simulations, a good deep learning model is able to screen for predicted hits very quickly. However, on existing GPUs it can still take months to identify the most likely drugs among billions of compounds, slowing research and limiting the impact of virtual drug screening.

The neural networks used for drug screening have interesting properties. It is common to pre-compute embeddings for both compounds and proteins offline, and then use small, fully connected models to predict drug-target binding based on these embeddings. With good embeddings, you don’t need large models for this task. Each layer of the model is often relatively small, which means that the main computational component is matrix-matrix multiplication, with relatively small weight matrices. This works against traditional graphics processor architectures which have to contend with low utilization even when minibatch sizes are large. These models are too small to be efficiently accelerated by GPUs. Increasing the batch size drives up only one dimension of the input matrices, and because the remaining two dimensions remain small, the adjustment is insufficient to achieve efficient GPU utilization. As a result, it can take months to run predictions for large libraries of compounds. Dependent on off-chip memory, the graphics processor spends the bulk of its time reading and writing from memory, rather than doing useful compute work. The consequence is months of waiting for the computations to finish.

Luckily, Cerebras Systems has developed a solution with the WSE, making it possible to accelerate virtual drug screening.

Innovations in AI Drug Screening

In collaboration with industry and academic researchers, Cerebras has demonstrated vast improvements in virtual screening on its CS-1 and CS-2 systems. In one instance, the CS-2 reduced screening time for a very large library of compounds from 183 days to 3.5 days, accelerating research by more than 50 X and shrinking time to answer by six months.

The Cerebras WSE architecture brings together up to 850,000 AI-optimized cores on a single chip. Each core has its own dedicated high-speed memory and is connected to its neighboring cores with a low-latency on-chip fabric. For small neural networks like those used in drug screening, the Cerebras compiler maps multiple copies of the small models onto the WSE, while simultaneously and efficiently distributing computations for every model copy across exactly the right number of processing elements for that model. This enables the WSE to train and run inference for multiple models simultaneously, at extremely high throughput. On a CS-2 powered by a WSE, an AI researcher can run hundreds of copies of the virtual drug screening model, while simultaneously achieving massive inference acceleration.

In virtual drug screening, faster inference means a faster search through the unimaginably vast number of potential drug-target combinations. Faster inference therefore means a faster pursuit of potential cures. In the case mentioned above, our partners were able to screen their library of compounds six months faster with the CS-2 than they were with their existing cluster of graphics processing units.