Moksh Jain

prof_pic.jpg

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

selected publications

2025

  1. Recursive Self-Aggregation Unlocks Deep Thinking in Large Language Models
    Siddarth Venkatraman, Vineet Jain, Sarthak Mittal, and 9 more authors
    arXiv preprint arXiv:2509.26626, 2025
  2. 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

2024

  1. 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
  2. Amortizing intractable inference in large language models
    Edward Hu*, Moksh Jain*, Eric Elmoznino, and 4 more authors
    In International Conference on Learning Representations, 2024

2023

  1. Multi-objective gflownets
    Moksh Jain, Sharath Chandra Raparthy, Alex Hernández-Garcı́a, and 4 more authors
    In International Conference on Machine Learning, 2023
  2. 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
  3. GFlowNets for AI-driven scientific discovery
    Moksh Jain, Tristan Deleu, Jason Hartford, and 3 more authors
    Digital Discovery, 2023

2022

  1. Biological Sequence Design with GFlowNets
    Moksh Jain, Emmanuel Bengio, Alex Hernandez-Garcia, and 8 more authors
    In International Conference on Machine Learning, 2022

2021

  1. 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