Moksh Jain

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.
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 approaches for 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.
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
Mar 7, 2023 | GFlowNets and Scientific Discovery |
---|---|
Jul 10, 2019 | Learning Disentangled Representations |