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 Research Intern at Basis. I am interested in developing probabilistic machine learning algorithms that enable reasoning to discover new knowledge through a virtuous cycle of modelling and experimentation. I am interested in applications of these algorithms to accelerate the process of making new scientific discoveries. At Mila, I lead various efforts to develop novel machine learning approaches in the context of drug discovery.

Previously, I was a research intern with Jason Hartford at Valence Labs (@Recursion) working 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. To work on my thesis, I spent a semester at Microsoft Research supervised by Prateek Jain and Harsha Simhadri as a part of the EdgeML project, working on efficient machine learning algorithms for resource-constrained environments. I also spent a summer at the Machine Learning group at Leuphana Universitat Lüneburg, supervised by Uwe Dick and Ulf Brefeld.

latest posts

selected publications

2025

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