I am a Ph.D. student at Mila and Université de Montréal supervised by Yoshua Bengio. Most recently, I was a Student Researcher at Google DeepMind, working with Nenad Tomašev 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 am also involved in various efforts to accelerate the process of drug discovery and material discovery with AI. My research is supported by a FRQNT Doctoral Fellowship.

Previous experience

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.

Selected publications All publications

  1. Moksh Jain, Mehdi Bennani, Johannes Bausch, Yuri Chervonyi, Bogdan Georgiev, Simon Osindero, and Nenad Tomašev
    arXiv preprint arXiv:2605.10820, 2026
  2. Siddarth Venkatraman*, Vineet Jain*, Sarthak Mittal*, Vedant Shah, Johan Obando-Ceron, Yoshua Bengio, Brian R. Bartoldson, Bhavya Kailkhura, Guillaume Lajoie, Glen Berseth, Nikolay Malkin, and Moksh Jain
    arXiv preprint arXiv:2509.26626, 2025
  3. Seanie Lee, Minsu Kim, Lynn Cherif, David Dobre, Juho Lee, Sung Ju Hwang, Kenji Kawaguchi, Gauthier Gidel, Yoshua Bengio, Nikolay Malkin, and Moksh Jain
    International Conference on Learning Representations (ICLR), 2025
  4. Siddarth Venkatraman*, Moksh Jain*, Luca Scimeca*, Minsu Kim*, Marcin Sendera*, Mohsin Hasan, Luke Rowe, Sarthak Mittal, Pablo Lemos, Emmanuel Bengio, Alexandre Adam, Jarrid Rector-Brooks, Yoshua Bengio, Glen Berseth, and Nikolay Malkin
    Advances in Neural Information Processing Systems (NeurIPS), 2024
  5. Edward Hu*, Moksh Jain*, Eric Elmoznino, Younesse Kaddar, Guillaume Lajoie, Yoshua Bengio, and Nikolay Malkin
    International Conference on Learning Representations (ICLR), 2024
  6. Moksh Jain, Sharath Chandra Raparthy, Alex Hernández-Garcı́a, Jarrid Rector-Brooks, Yoshua Bengio, Santiago Miret, and Emmanuel Bengio
    International Conference on Machine Learning (ICML), 2023
  7. Edward J Hu*, Nikolay Malkin*, Moksh Jain, Katie E Everett, Alexandros Graikos, and Yoshua Bengio
    International Conference on Machine Learning (ICML), 2023
  8. Moksh Jain, Tristan Deleu, Jason Hartford, Cheng-Hao Liu, Alex Hernandez-Garcia, and Yoshua Bengio
    Digital Discovery (DD), 2023
  9. Moksh Jain, Emmanuel Bengio, Alex Hernandez-Garcia, Jarrid Rector-Brooks, Bonaventure FP Dossou, Chanakya Ajit Ekbote, Jie Fu, Tianyu Zhang, Michael Kilgour, Dinghuai Zhang, and  others
    International Conference on Machine Learning (ICML), 2022
  10. Emmanuel Bengio, Moksh Jain, Maksym Korablyov, Doina Precup, and Yoshua Bengio
    Advances in Neural Information Processing Systems (NeurIPS), 2021

* equal contribution  ·  † equal senior authorship

Writing All posts