Portrait of Moksh Jain is unavailable

Moksh Jain

moksh.jain@mila.quebec
PhD - Université de Montréal
Supervisor
Yoshua Bengio
Research Topics
Active Learning
AI for Science
GFlowNets
Probabilistic Models
Website
Google Scholar
GitHub

Publications

Graph-Based Active Machine Learning Method for Diverse and Novel Antimicrobial Peptides Generation and Selection
Bonaventure F. P. Dossou
Dianbo Liu
Xu Ji
Moksh Jain
Almer M. van der Sloot
Roger Palou
Michael Tyers
Yoshua Bengio
As antibiotic-resistant bacterial strains are rapidly spreading worldwide, infections caused by these strains are emerging as a global crisi… (see more)s causing the death of millions of people every year. Antimicrobial Peptides (AMPs) are one of the candidates to tackle this problem because of their potential diversity, and ability to favorably modulate the host immune response. However, large-scale screening of new AMP candidates is expensive, time-consuming, and now affordable in developing countries, which need the treatments the most. In this work, we propose a novel active machine learning-based framework that statistically minimizes the number of wet-lab experiments needed to design new AMPs, while ensuring a high diversity and novelty of generated AMPs sequences, in multi-rounds of wet-lab AMP screening settings. Combining recurrent neural network models and a graph-based filter (GraphCC), our proposed approach delivers novel and diverse candidates and demonstrates better performances according to our defined metrics.
2022-11-09
Research Square (Research Square) (accepted)
doi.org
arxiv.org
Biological Sequence Design with GFlowNets
Moksh Jain
Emmanuel Bengio
Alex-Hernandez Garcia
Jarrid Rector-Brooks
Bonaventure F. P. Dossou
Chanakya Ekbote
Jie Fu
Tianyu Zhang
Michael Kilgour
Dinghuai Zhang
Lena Simine
Payel Das
Yoshua Bengio
Design of de novo biological sequences with desired properties, like protein and DNA sequences, often involves an active loop with several r… (see more)ounds of molecule ideation and expensive wet-lab evaluations. These experiments can consist of multiple stages, with increasing levels of precision and cost of evaluation, where candidates are filtered. This makes the diversity of proposed candidates a key consideration in the ideation phase. In this work, we propose an active learning algorithm leveraging epistemic uncertainty estimation and the recently proposed GFlowNets as a generator of diverse candidate solutions, with the objective to obtain a diverse batch of useful (as defined by some utility function, for example, the predicted anti-microbial activity of a peptide) and informative candidates after each round. We also propose a scheme to incorporate existing labeled datasets of candidates, in addition to a reward function, to speed up learning in GFlowNets. We present empirical results on several biological sequence design tasks, and we find that our method generates more diverse and novel batches with high scoring candidates compared to existing approaches.
2022-07-18
International Conference on Machine Learning (Accept for Short Presentation)
doi.org
proceedings.mlr.press
Flow Network based Generative Models for Non-Iterative Diverse Candidate Generation
Emmanuel Bengio
Moksh Jain
Maksym Korablyov
Doina Precup
Yoshua Bengio
This paper is about the problem of learning a stochastic policy for generating an object (like a molecular graph) from a sequence of actions… (see more), such that the probability of generating an object is proportional to a given positive reward for that object. Whereas standard return maximization tends to converge to a single return-maximizing sequence, there are cases where we would like to sample a diverse set of high-return solutions. These arise, for example, in black-box function optimization when few rounds are possible, each with large batches of queries, where the batches should be diverse, e.g., in the design of new molecules. One can also see this as a problem of approximately converting an energy function to a generative distribution. While MCMC methods can achieve that, they are expensive and generally only perform local exploration. Instead, training a generative policy amortizes the cost of search during training and yields to fast generation. Using insights from Temporal Difference learning, we propose GFlowNet, based on a view of the generative process as a flow network, making it possible to handle the tricky case where different trajectories can yield the same final state, e.g., there are many ways to sequentially add atoms to generate some molecular graph. We cast the set of trajectories as a flow and convert the flow consistency equations into a learning objective, akin to the casting of the Bellman equations into Temporal Difference methods. We prove that any global minimum of the proposed objectives yields a policy which samples from the desired distribution, and demonstrate the improved performance and diversity of GFlowNet on a simple domain where there are many modes to the reward function, and on a molecule synthesis task.
2021-12-06
Conference on Neural Information Processing Systems (Accept (Poster))
doi.org
openreview.net