PeopleofSuperintelligence

Researcher

Moksh Jain

Moksh Jain is a machine learning researcher who works on generative flow networks (GFlowNets), a framework for sampling diverse solutions that they apply to scientific discovery and large language models. A graduate of the National Institute of Technology Karnataka, Surathkal, they spent a year at Microsoft Turing compressing language models for Bing and Office, and became second author of the 2021 GFlowNet paper with Yoshua Bengio and others. Their paper Amortizing intractable inference in large language models received an Honorable Mention at ICLR 2024, and in 2025 they were last author of Recursive Self-Aggregation, a test-time scaling method for language models. In 2026 they were first author of the MaD Physics benchmark for scientific agents. They are a PhD student at Mila and Université de Montréal, supervised by Bengio and supported by a FRQNT Doctoral Fellowship.

Timeline

  1. 2021

    The GFlowNet paper

    In June 2021, they are second author of the GFlowNet paper with Emmanuel Bengio, Maksym Korablyov, Doina Precup and Yoshua Bengio, accepted at NeurIPS 2021; it learns a policy that generates objects such as molecular graphs with probability proportional to a reward. Their background includes a bachelor's degree at the National Institute of Technology Karnataka, Surathkal, a year at Microsoft Turing compressing and optimizing large language models for Bing and Office, and a visiting-researcher role with Bengio on uncertainty estimation and drug discovery before joining Mila as a graduate student.

  2. 2022

    Trajectory balance for GFlowNets

    In January 2022, they coauthor Trajectory balance: Improved credit assignment in GFlowNets alongside Emmanuel Bengio, Chen Sun and Yoshua Bengio, for NeurIPS 2022. The paper proposes trajectory balance, a learning objective that is more efficient than the earlier flow matching and detailed balance objectives, and proves that any global minimizer can define a policy that samples exactly from the target distribution.

  3. 2023

    GFlowNets for scientific discovery

    In February 2023, they are first author of GFlowNets for AI-Driven Scientific Discovery with Tristan Deleu, Jason Hartford, Cheng-Hao Liu, Alex Hernandez-Garcia and Yoshua Bengio. The paper presents GFlowNets as a probabilistic machine learning framework for the modeling, hypothesis generation and experimental design stages of the experimental science loop.

  4. 2024

    An ICLR honorable mention

    Their paper Amortizing intractable inference in large language models, written with Edward J. Hu and five coauthors including Yoshua Bengio, receives an Honorable Mention at ICLR 2024 and is presented as an oral. It fine-tunes language models with generative flow networks (GFlowNets) so that they sample from intractable posterior distributions.

  5. 2024–2025

    Diffusion models and red-teaming

    In 2024, they are among the authors of Amortizing intractable inference in diffusion models for vision, language, and control, which appears at NeurIPS 2024 and proves the asymptotic correctness of relative trajectory balance for training a diffusion model to sample from a posterior. In 2025, the ICLR paper Learning diverse attacks on large language models for robust red-teaming and safety tuning, which they coauthor, uses GFlowNet fine-tuning to train an attacker model that generates diverse and effective attack prompts.

  6. September 2025

    Recursive self-aggregation

    In September 2025, they are last author of Recursive Self-Aggregation Unlocks Deep Thinking in Large Language Models, which proposes RSA, a test-time scaling method inspired by evolutionary methods. Each step of RSA refines a population of candidate reasoning chains through aggregation of subsets of them.

  7. May 2026

    The MaD Physics benchmark

    In May 2026, they are first author of MaD Physics, a benchmark that evaluates how well agents make informative measurements and draw conclusions under limits on the quality and quantity of measurements. Their coauthors include Nenad Tomašev, their collaborator on adaptive reasoning and planning during a student-researcher role at Google DeepMind.

  8. As of October 2026

    PhD student at Mila

    As of October 2026, they are a PhD student at Mila and Université de Montréal, supervised by Yoshua Bengio and supported by a FRQNT Doctoral Fellowship. At Mila, they also take part in efforts to accelerate drug discovery and material discovery with AI.

Career

  1. Google DeepMindStudent Researcher

Education

  1. National Institute of Technology KarnatakaBachelor's degree
  2. MilaPhD student
  3. Université de MontréalPhD student

From the biography.

  1. Google had the most people shaping AI every year from 2014 to 2025Includes Google. Our estimates of which organizations employed the most people shaping AI each year from 1955 to 2025, tracing the shift from university labs to companies.
  2. Google has the highest combined h-index among AI organizationsIncludes Google. We estimate that Google holds the highest combined h-index at 8,016, while universities account for 73% of the deduplicated total across AI organizations.
  3. 20% of people at companies who shape AI also hold a university role, led by Google and Stanford UniversityIncludes Google. We estimate that 539 of the 6,766 people with a known current organization (8%) hold both a university role and a company role, including 20% of people at companies and 13% of people at universities.
  4. Google alumni have founded or led 133 other organizations, more than any other AI labIncludes Google. We estimate that alumni of 17 major AI labs went on to found or lead 350 other organizations, including 226 they founded, and we trace 472 career moves directly between 23 labs since 2010.

Colleagues

People named in this biography, and people with dated positions at the same organization in the same years.

Sources and credits

Photo: We have no published photo of Moksh Jain yet. Send us a free-licensed photo.

  1. Moksh Jain mj10.github.io. Checked 21 October 2026.
  2. [2106.04399] Flow Network based Generative Models for Non-Iterative Diverse Candidate Generation arxiv.org. Checked 21 October 2026.
  3. [2201.13259] Trajectory balance: Improved credit assignment in GFlowNets arxiv.org. Checked 21 October 2026.
  4. [2302.00615] GFlowNets for AI-Driven Scientific Discovery arxiv.org. Checked 21 October 2026.
  5. ICLR 2024 Awards iclr.cc. Checked 21 October 2026.
  6. ICLR 2024 Outstanding Paper Awards – ICLR Blog blog.iclr.cc. Checked 21 October 2026.
  7. [2310.04363] Amortizing intractable inference in large language models arxiv.org. Checked 21 October 2026.
  8. [2405.20971] Amortizing intractable inference in diffusion models for vision, language, and control arxiv.org. Checked 21 October 2026.
  9. [2405.18540] Learning diverse attacks on large language models for robust red-teaming and safety tuning arxiv.org. Checked 21 October 2026.
  10. [2509.26626] Recursive Self-Aggregation Unlocks Deep Thinking in Large Language Models arxiv.org. Checked 21 October 2026.
  11. [2605.10820] MaD Physics: Evaluating information seeking under constraints in physical environments arxiv.org. Checked 21 October 2026.
  12. Moksh Jain | Mila mila.quebec. Checked 21 October 2026.