Jithendaraa Subramanian

I am currently an AI resident at the Toyota Research Institute, where I work on developing multimodal models for applications in material science. I completed my M.Sc. at McGill University and Mila supervised by Derek Nowrouzezahrai and Samira Ebrahimi Kahou. My research interests span probabilistic inference, multimodal generative models, Large Language Models (LLMs), and causality. I am also passionate about writing GPU-efficient code.

During my master’s, I worked on problems at the intersection of causal discovery and representation learning, some of them involving GFlowNets. My thesis was on learning a joint distribution over causal variables, structures, and mechanisms from low-level data. In 2024, I was a visiting researcher at ServiceNow Research, where I worked on change point detection in temporal causal models, and benchmarking text-conditioned forecasting capabilities of LLMs. I spent summer 2023 as a research scientist intern at Amazon training large-scale models for long-term revenue forecasting.

Before my graduate studies, I completed my undergrad from NIT Trichy, India. As a research intern with the RoboTutor Team at Carnegie Mellon University, under Jack Mostow, I applied reinforcement learning to develop a personalized Intelligent Tutoring System for underprivileged students in Tanzania. I also explored machine learning for security during a research internship at UC Berkeley with Dawn Song and Joseph Near.

News

Feb 10, 2025	Co-organizing the Multimodal Learning for Materials Science (MM4Mat) workshop at CVPR 2025 🔬🧪
Nov 11, 2024	I’m joining the Accelerated Material Design and Discovery team(AMDD) team at Toyota Research Institute where I’ll be working on building multimodal models for applications in material science.
Apr 22, 2024	Started as a Visiting Researcher at ServiceNow Research with Valentina Zantedeschi and Alexandre Drouin.
Sep 21, 2023	Joint Bayesian Inference of Graphical Structure and Parameters with a Single Generative Flow Network is accepted at NeurIPS 2023.
Jun 5, 2023	I am joining Amazon as a research scientist intern for the summer to work on large-scale models for long-term revenue forecasting.
Nov 28, 2022	Bayesian Learning of Causal Structure and Mechanisms with GFlowNets and Variational Bayes is accepted at the 3rd Workshop on Graphs and more Complex structures for Learning and Reasoning (GCLR) at AAAI 2023.
Jun 12, 2022	Latent Variable Models for Bayesian Causal Discovery has been accepted at the ICML 2022 Workshop on Spurious Correlations, Invariance, and Stability.
Jan 6, 2022	I am starting as the lead TA for ECSE 343 Numerical Methods for Engineering.
Sep 1, 2021	Enrolled as a Masters Thesis Computer Science student at Mcgill University/MILA in Montreal, Quebec
Jun 1, 2021	Graduated with a Bachelors in Production Engineering from NIT, Trichy
Nov 15, 2020	Joined as a research intern at MILA/École de Technologie Supérieure advised by Prof. Samira Ebrahimi Kahou and Prof. Christian Desrosiers.
Apr 1, 2020	I am joining as a research intern at Carnegie Mellon University advised by Prof. Jack Mostow to work on The RoboTutor Project using reinforcement learning to personalize intelligent tutoring systems like RoboTutor, to improve literacy rates amoung young students in Tanzania.
Sep 1, 2019	I am joining as a research intern in Prof. Dawn Song’s lab at UC Berkeley to work on enforcing privacy policies for machine learning programs advised by Prof. Joseph Near and Lun Wang.

Selected publications

NeurIPS

Joint Bayesian Inference of Graphical Structure and Parameters with a Single Generative Flow Network

Tristan Deleu, Mizu Nishikawa-Toomey, Jithendaraa Subramanian, and 3 more authors

Advances in Neural Information Processing Systems, 2023

Bib

@article{deleu2023joint,
  title = {Joint Bayesian Inference of Graphical Structure and Parameters with a Single Generative Flow Network},
  author = {Deleu, Tristan and Nishikawa-Toomey, Mizu and Subramanian, Jithendaraa and Malkin, Nikolay and Charlin, Laurent and Bengio, Yoshua},
  journal = {Advances in Neural Information Processing Systems},
  year = {2023},
  selected = true,
}

Under review,
AISTATS 2024

Bayesian learning of Causal Structure and Mechanisms with GFlowNets and Variational Bayes

Mizu Nishikawa-Toomey*, Tristan Deleu*, Jithendaraa Subramanian, and 2 more authors

AAAI Workshop Graphs and More Complex Structures for Learning and Reasoning, 2023

Bib

@article{nishikawa2022vbg,
  title = {{Bayesian learning of Causal Structure and Mechanisms with GFlowNets and Variational Bayes}},
  author = {Nishikawa-Toomey*, Mizu and Deleu*, Tristan and Subramanian, Jithendaraa and Bengio, Yoshua and Charlin, Laurent},
  journal = {AAAI Workshop Graphs and More Complex Structures for Learning and Reasoning},
  year = {2023},
  selected = true,
}

Under review,
ICLR 2024

Learning Latent Structural Causal Models

Jithendaraa Subramanian, Yashas Annadani, Ivaxi Sheth, and 5 more authors

arXiv preprint arXiv:2210.13583, 2022

Bib

@article{subramanian2022learning,
  title = {Learning Latent Structural Causal Models},
  author = {Subramanian, Jithendaraa and Annadani, Yashas and Sheth, Ivaxi and Ke, Nan Rosemary and Deleu, Tristan and Bauer, Stefan and Nowrouzezahrai, Derek and Kahou, Samira Ebrahimi},
  journal = {arXiv preprint arXiv:2210.13583},
  year = {2022},
  selected = true,
}

EDM

Deep Reinforcement Learning to Simulate, Train, and Evaluate Instructional Sequencing Policies

J. Subramanian, and J. Mostow

RL4ED workshop at EDM conference, 2021

Bib PDF

@article{step_rl4education,
  title = {Deep Reinforcement Learning to Simulate, Train, and Evaluate Instructional Sequencing Policies},
  author = {Subramanian, J. and Mostow, J.},
  year = {2021},
  journal = {RL4ED workshop at EDM conference},
  selected = true
}