Moksh Jain

prof_pic.jpg

I am a Ph.D. student at Mila and Université de Montréal supervised by Yoshua Bengio. I am interested in developing machine learning algorithms for effective experimental design, incorporating tools from probabilistic inference and modern deep learning. 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.

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

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