Moksh Jain

I am a Ph.D. student at Mila and Université de Montréal supervised by Yoshua Bengio. I am currently a Student Researcher at Google DeepMind, working with Nenad Tomasev on adaptive reasoning and planning.

My research interests fall broadly in the following directions:

Algorithms for probabilistic inference in high-dimensional structured spaces.
Structured inference with foundation model priors to improve reasoning and decision making under uncertainty.
Efficient experimental design in high-dimensional design spaces.
Developing AI systems to augment and accelerate the process of scientific discovery.

At Mila, I lead and contribute to various efforts to develop novel machine learning approaches in the context of drug discovery and material discovery. My research is supported by a FRQNT Doctoral Fellowship.

Previous experience

Most recently, I was a research intern at Basis, working on a benchmark and approach for language agents that can interactively build models of the environment. In the past, I worked with Jason Hartford at Valence Labs (@ Recursion) on experimental design for gene knockout experiments. Before joining Mila as a graduate student, I was a visiting researcher working with Yoshua Bengio on uncertainty estimation and drug discovery. I also spent a year at Microsoft Turing working on compressing and optimizing large langauge models for deployment across Bing and Office.

I completed my bachelor’s degree at the National Institute of Technology Karnataka, Surathkal. I spent a semester at Microsoft Research supervised by Prateek Jain and Harsha Simhadri, working on efficient machine learning algorithms for resource-constrained settings. I also spent a summer at the Machine Learning group at Leuphana Universitat Lüneburg, supervised by Uwe Dick and Ulf Brefeld working on inverse reinforcement learning.

latest posts

Mar 7, 2023	GFlowNets and Scientific Discovery
Jul 10, 2019	Learning Disentangled Representations

selected publications

2025

ICLR
Learning Diverse Attacks on Large Language Models for Robust Red-Teaming and Safety Tuning

Seanie Lee, Minsu Kim, Lynn Cherif, and 8 more authors

In International Conference on Learning Representations, 2025

Abs arXiv Bib

Red-teaming, or identifying prompts that elicit harmful responses, is a critical step in ensuring the safe and responsible deployment of large language models (LLMs). Developing effective protection against many modes of attack prompts requires discovering diverse attacks. Automated red-teaming typically uses reinforcement learning to fine-tune an attacker language model to generate prompts that elicit undesirable responses from a target LLM, as measured, for example, by an auxiliary toxicity classifier. We show that even with explicit regularization to favor novelty and diversity, existing approaches suffer from mode collapse or fail to generate effective attacks. As a flexible and probabilistically principled alternative, we propose to use GFlowNet fine-tuning, followed by a secondary smoothing phase, to train the attacker model to generate diverse and effective attack prompts. We find that the attacks generated by our method are effective against a wide range of target LLMs, both with and without safety tuning, and transfer well between target LLMs. Finally, we demonstrate that models safety-tuned using a dataset of red-teaming prompts generated by our method are robust to attacks from other RL-based red-teaming approaches.
@inproceedings{lee2025learning, title = {Learning Diverse Attacks on Large Language Models for Robust Red-Teaming and Safety Tuning}, author = {Lee, Seanie and Kim, Minsu and Cherif, Lynn and Dobre, David and Lee, Juho and Hwang, Sung Ju and Kawaguchi, Kenji and Gidel, Gauthier and Bengio, Yoshua and Malkin, Nikolay and Jain, Moksh}, booktitle = {International Conference on Learning Representations}, year = {2025}, }

2024

NeurIPS
Amortizing intractable inference in diffusion models for vision, language, and control

Siddarth Venkatraman*, Moksh Jain*, Luca Scimeca*, and 12 more authors

Advances in Neural Information Processing Systems, 2024

Abs arXiv Bib Code

In the last decades, the capacity to generate large amounts of data in science and engineering applications has been growing steadily. Meanwhile, the progress in machine learning has turned it into a suitable tool to process and utilise the available data. Nonetheless, many relevant scientific and engineering problems present challenges where current machine learning methods cannot yet efficiently leverage the available data and resources. For example, in scientific discovery, we are often faced with the problem of exploring very large, high-dimensional spaces, where querying a high fidelity, black-box objective function is very expensive. Progress in machine learning methods that can efficiently tackle such problems would help accelerate currently crucial areas such as drug and materials discovery. In this paper, we propose the use of GFlowNets for multi-fidelity active learning, where multiple approximations of the black-box function are available at lower fidelity and cost. GFlowNets are recently proposed methods for amortised probabilistic inference that have proven efficient for exploring large, high-dimensional spaces and can hence be practical in the multi-fidelity setting too. Here, we describe our algorithm for multi-fidelity active learning with GFlowNets and evaluate its performance in both well-studied synthetic tasks and practically relevant applications of molecular discovery. Our results show that multi-fidelity active learning with GFlowNets can efficiently leverage the availability of multiple oracles with different costs and fidelities to accelerate scientific discovery and engineering design.
@article{venkatraman2024amortizing, title = {Amortizing intractable inference in diffusion models for vision, language, and control}, author = {Venkatraman*, Siddarth and Jain*, Moksh and Scimeca*, Luca and Kim*, Minsu and Sendera*, Marcin and Hasan, Mohsin and Rowe, Luke and Mittal, Sarthak and Lemos, Pablo and Bengio, Emmanuel and Adam, Alexandre and Rector-Brooks, Jarrid and Bengio, Yoshua and Berseth, Glen and Malkin, Nikolay}, journal = {Advances in Neural Information Processing Systems}, volume = {37}, year = {2024}, }
ICLR
Amortizing intractable inference in large language models

Edward Hu*, Moksh Jain*, Eric Elmoznino, and 4 more authors

In International Conference on Learning Representations, 2024

Abs arXiv Bib Code

Autoregressive large language models (LLMs) compress knowledge from their training data through next-token conditional distributions. This limits tractable querying of this knowledge to start-to-end autoregressive sampling. However, many tasks of interest—including sequence continuation, infilling, and other forms of constrained generation—involve sampling from intractable posterior distributions. We address this limitation by using amortized Bayesian inference to sample from these intractable posteriors. Such amortization is algorithmically achieved by fine-tuning LLMs via diversity-seeking reinforcement learning algorithms: generative flow networks (GFlowNets). We empirically demonstrate that this distribution-matching paradigm of LLM fine-tuning can serve as an effective alternative to maximum-likelihood training and reward-maximizing policy optimization. As an important application, we interpret chain-of-thought reasoning as a latent variable modeling problem and demonstrate that our approach enables data-efficient adaptation of LLMs to tasks that require multi-step rationalization and tool use.
@inproceedings{hu2024amortizing, title = {Amortizing intractable inference in large language models}, author = {Hu*, Edward and Jain*, Moksh and Elmoznino, Eric and Kaddar, Younesse and Lajoie, Guillaume and Bengio, Yoshua and Malkin, Nikolay}, booktitle = {International Conference on Learning Representations}, year = {2024}, }

2023

ICML
Multi-objective gflownets

Moksh Jain, Sharath Chandra Raparthy, Alex Hernández-Garcı́a, and 4 more authors

In International Conference on Machine Learning, 2023

Abs arXiv Bib Code

In many applications of machine learning, like drug discovery and material design, the goal is to generate candidates that simultaneously maximize a set of objectives. As these objectives are often conflicting, there is no single candidate that simultaneously maximizes all objectives, but rather a set of Pareto-optimal candidates where one objective cannot be improved without worsening another. Moreover, in practice, these objectives are often under-specified, making the diversity of candidates a key consideration. The existing multi-objective optimization methods focus predominantly on covering the Pareto front, failing to capture diversity in the space of candidates. Motivated by the success of GFlowNets for generation of diverse candidates in a single objective setting, in this paper we consider Multi-Objective GFlowNets (MOGFNs). MOGFNs consist of a novel Conditional GFlowNet which models a family of single-objective sub-problems derived by decomposing the multi-objective optimization problem. Our work is the first to empirically demonstrate conditional GFlowNets. Through a series of experiments on synthetic and benchmark tasks, we empirically demonstrate that MOGFNs outperform existing methods in terms of Hypervolume, R2-distance and candidate diversity. We also demonstrate the effectiveness of MOGFNs over existing methods in active learning settings. Finally, we supplement our empirical results with a careful analysis of each component of MOGFNs.
@inproceedings{jain2023multi, title = {Multi-objective gflownets}, author = {Jain, Moksh and Raparthy, Sharath Chandra and Hern{\'a}ndez-Garc{\i}́a, Alex and Rector-Brooks, Jarrid and Bengio, Yoshua and Miret, Santiago and Bengio, Emmanuel}, booktitle = {International Conference on Machine Learning}, pages = {14631--14653}, year = {2023}, }
ICML
GFlowNet-EM for learning compositional latent variable models

Edward J Hu*, Nikolay Malkin*, Moksh Jain, and 3 more authors

In International Conference on Machine Learning, 2023

Abs arXiv Bib Code

Latent variable models (LVMs) with discrete compositional latents are an important but challenging setting due to a combinatorially large number of possible configurations of the latents. A key tradeoff in modeling the posteriors over latents is between expressivity and tractable optimization. For algorithms based on expectation-maximization (EM), the E-step is often intractable without restrictive approximations to the posterior. We propose the use of GFlowNets, algorithms for sampling from an unnormalized density by learning a stochastic policy for sequential construction of samples, for this intractable E-step. By training GFlowNets to sample from the posterior over latents, we take advantage of their strengths as amortized variational inference algorithms for complex distributions over discrete structures. Our approach, GFlowNet-EM, enables the training of expressive LVMs with discrete compositional latents, as shown by experiments on non-context-free grammar induction and on images using discrete variational autoencoders (VAEs) without conditional independence enforced in the encoder.
@inproceedings{hu2023gflownet, title = {GFlowNet-EM for learning compositional latent variable models}, author = {Hu*, Edward J and Malkin*, Nikolay and Jain, Moksh and Everett, Katie E and Graikos, Alexandros and Bengio, Yoshua}, booktitle = {International Conference on Machine Learning}, pages = {13528--13549}, year = {2023}, }
DD
GFlowNets for AI-driven scientific discovery

Moksh Jain, Tristan Deleu, Jason Hartford, and 3 more authors

Digital Discovery, 2023

Abs arXiv Bib

Tackling the most pressing problems for humanity, such as the climate crisis and the threat of global pandemics, requires accelerating the pace of scientific discovery. While science has traditionally relied on trial and error and even serendipity to a large extent, the last few decades have seen a surge of data-driven scientific discoveries. However, in order to truly leverage large-scale data sets and high-throughput experimental setups, machine learning methods will need to be further improved and better integrated in the scientific discovery pipeline. A key challenge for current machine learning methods in this context is the efficient exploration of very large search spaces, which requires techniques for estimating reducible (epistemic) uncertainty and generating sets of diverse and informative experiments to perform. This motivated a new probabilistic machine learning framework called GFlowNets, which can be applied in the modeling, hypotheses generation and experimental design stages of the experimental science loop. GFlowNets learn to sample from a distribution given indirectly by a reward function corresponding to an unnormalized probability, which enables sampling diverse, high-reward candidates. GFlowNets can also be used to form efficient and amortized Bayesian posterior estimators for causal models conditioned on the already acquired experimental data. Having such posterior models can then provide estimators of epistemic uncertainty and information gain that can drive an experimental design policy. Altogether, here we will argue that GFlowNets can become a valuable tool for AI-driven scientific discovery, especially in scenarios of very large candidate spaces where we have access to cheap but inaccurate measurements or to expensive but accurate measurements. This is a common setting in the context of drug and material discovery, which we use as examples throughout the paper.
@article{jain2023gflownets, title = {GFlowNets for AI-driven scientific discovery}, author = {Jain, Moksh and Deleu, Tristan and Hartford, Jason and Liu, Cheng-Hao and Hernandez-Garcia, Alex and Bengio, Yoshua}, journal = {Digital Discovery}, volume = {2}, number = {3}, pages = {557--577}, year = {2023}, publisher = {Royal Society of Chemistry}, }

2022

ICML
Biological Sequence Design with GFlowNets

Moksh Jain, Emmanuel Bengio, Alex Hernandez-Garcia, and 8 more authors

In International Conference on Machine Learning, 2022

Abs arXiv Bib Code

Design of de novo biological sequences with desired properties, like protein and DNA sequences, often involves an active loop with several rounds 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.
@inproceedings{jain2022biological, title = {Biological Sequence Design with GFlowNets}, author = {Jain, Moksh and Bengio, Emmanuel and Hernandez-Garcia, Alex and Rector-Brooks, Jarrid and Dossou, Bonaventure FP and Ekbote, Chanakya Ajit and Fu, Jie and Zhang, Tianyu and Kilgour, Michael and Zhang, Dinghuai and others}, booktitle = {International Conference on Machine Learning}, pages = {9786--9801}, year = {2022}, }

2021

NeurIPS
Flow network based generative models for non-iterative diverse candidate generation

Emmanuel Bengio, Moksh Jain, Maksym Korablyov, and 2 more authors

Advances in Neural Information Processing Systems, 2021

Abs arXiv Bib Code

This paper is about the problem of learning a stochastic policy for generating an object (like a molecular graph) from a sequence of actions, 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.
@article{bengio2021flow, title = {Flow network based generative models for non-iterative diverse candidate generation}, author = {Bengio, Emmanuel and Jain, Moksh and Korablyov, Maksym and Precup, Doina and Bengio, Yoshua}, journal = {Advances in Neural Information Processing Systems}, volume = {34}, pages = {27381--27394}, year = {2021}, }