Siamak Ravanbakhsh

Laurence Perreault-Levasseur

Biographie

Siamak Ravanbakhsh est professeur adjoint à l’École d’informatique de l’Université McGill depuis août 2019. Avant de se joindre à McGill et à Mila – Institut québécois d’intelligence artificielle, il a occupé un poste similaire à l’Université de la Colombie-Britannique. De 2015 à 2017, il a été stagiaire postdoctoral au Département d’apprentissage automatique et à l’Institut de robotique de l’Université Carnegie Mellon, et il a obtenu un doctorat de l’Université de l’Alberta. Il s’intéresse aux problèmes de l’apprentissage de la représentation et de l’inférence dans l’IA.

Ses recherches actuelles portent sur le rôle de la symétrie et de l’invariance dans l’apprentissage profond des représentations.

Étudiants actuels

Tara Akhound-Sadegh

Doctorat - McGill

Co-superviseur⋅e :

tara.akhoundsadegh@mila.quebec

Noah El Rimawi-Fine

Maîtrise recherche - McGill

noah.elrimawi-fine@mila.quebec

natalia.espinosa-dice@mila.quebec

Natalia Espinosa Dice

Maîtrise recherche - McGill

daria.goptsii@mail.mcgill.ca

Qianyi Gao

Maîtrise recherche - McGill

qianyi.gao@mila.quebec

Daria Goptsii Goptsii

Maîtrise professionnelle - McGill

Vineet Jain

Doctorat - McGill

jain.vineet@mila.quebec

Oumar Kaba

Doctorat - McGill

Daniel Levy

Doctorat - McGill

daniel.levy@mila.quebec

raphael.michnik@mila.quebec

Xiusi Li

Doctorat - McGill

Maîtrise recherche - McGill

jun-kai.liao@mila.quebec

Raphael Michnik

Stagiaire de recherche - McGill

Khang Ngo

Doctorat - McGill

khang.ngo@mila.quebec

mohammad.pedramfar@mila.quebec

David Nitchi

Maîtrise recherche - McGill

david.nitchi@mila.quebec

Mohammad Pedramfar

Postdoctorat - McGill

Mehran Shakerinava

Doctorat - McGill

mehran.shakerinava@mila.quebec

Publications

The Tilted Sampling Problem: Optimism, Posterior Sampling, and Logarithmic Regret

Mohammad Pedramfar

Many modern learning systems are used not to select a single best action, but to sample diverse high-quality outputs. Examples include entro… (voir plus)py-regularized reinforcement learning, RLHF, and reward-guided generative modeling. Motivated by these settings, we study tilted sampling: an online learning problem in which the goal is to sample from the Gibbs distribution induced by an unknown reward function. We define sampling regret as the temperature-normalized cumulative KL divergence between the learner's distribution and the Gibbs target, corresponding to the suboptimality gap in entropy-regularized reward maximization. Our framework generalizes recent sampling-based and KL-regularized formulations in bandits and reinforcement learning. Our main contribution is a general regret framework that captures the different uncertainty structure of tilted sampling and classical regret minimization. While standard regret scales as

2026-10-14

arXiv (prépublication)

arxiv.org

Bayesian Last Layer for Neural Force Fields

David Nitchi

Daniel T. Levy

Michel Côté

Reliable uncertainty quantification is essential for deploying Machine Learning Interatomic Potentials (MLIPs), also known as Neural Force F… (voir plus)ields, especially when molecular dynamics or materials simulations encounter configurations outside the training distribution. Deep ensembles remain the strongest practical baseline for MLIP uncertainty, but training and storing several copies of a modern pretrained model is often prohibitively expensive. We show that Bayesian Linear Last Layers (BLLs) provide a scalable alternative for MLIPs: a single pretrained backbone supplies atomic features, while exact Bayesian inference over the final force-prediction layer gives predictive uncertainties. BLL is known to underestimate the uncertainties. We provide an in-depth analysis that shows two sources of miscalibration and introduce a simple post-hoc recalibration to address the issue. On MPtrj and rMD17 benchmarks, including both in-distribution tests and increasingly out-of-distribution regimes, BLLs that are recalibrated on in-distribution examples produce uncertainty estimates competitive with ensembles, while using only one base model.

2026-05-29

AI4Science @ International Conference on Machine Learning (poster)

Long-Horizon Model-Based Offline Reinforcement Learning Without Conservatism

2026-04-29

International Conference on Machine Learning (accepté)

Long-Horizon Model-Based Offline Reinforcement Learning Without Explicit Conservatism

Popular offline reinforcement learning (RL) methods rely on explicit conservatism, penalizing out-of-dataset actions or restricting… (voir plus) rollout horizons. We question the universality of this principle and revisit a complementary Bayesian perspective for test-time adaptation. By modeling a posterior over world models and training a history-dependent agent to maximize expected return, the Bayesian approach directly addresses epistemic uncertainty without explicit conservatism. We first illustrate in a bandit setting that Bayesianism excels on low-quality datasets where conservatism fails. Scaling to realistic tasks, we find that long-horizon rollouts are essential to control value overestimation once conservatism is removed. We introduce design choices that enable learning from long-horizon rollouts while mitigating compounding model errors, yielding our algorithm, NEUBAY, grounded in the neutral Bayesian principle. On D4RL and NeoRL benchmarks, NEUBAY is competitive with leading conservative algorithms, achieving new state-of-the-art on 7 datasets with rollout horizons of several hundred steps. Finally, we characterize datasets by quality and coverage to identify when NEUBAY is preferable to conservative methods.

2026-04-29

International Conference on Machine Learning (accepté)

proceedings.mlr.press

The Role of Symmetry in Optimizing Overparameterized Networks

Sékou-Oumar Kaba

Overparameterization is central to the success of deep learning, yet the mechanisms by which it improves optimization remain incompletely un… (voir plus)derstood. We analyze weight-space symmetries in neural networks and show that overparameterization introduces additional symmetries that benefit optimization in two distinct ways. First, we prove that these symmetries act as a form of diagonal preconditioning on the Hessian, enabling the existence of better-conditioned minima within each equivalence class of functionally identical solutions. Second, we show that overparameterization increases the probability mass of global minima near typical initializations, making these favorable solutions more reachable. Teacher-student network experiments validate our theoretical predictions: as width increases, the Hessian trace decreases, condition numbers improve, and convergence accelerates. Our analysis provides a unified framework for understanding overparameterization and width growth as a geometric transformation of the loss landscape.

2026-04-27

arXiv (prépublication)

arxiv.org

Molecule property prediction with molecular orbitals

Yan Zhang

Khang Ngo

Sékou-Oumar Kaba

Daniel T. Levy

Aristide Baratin

Kisoo Kwon

MiYoung Jang

Eun Hyun Cho

Sangha Park

Sanghyun Yoo

Young-Seok Kim

Hasup Lee

Molecular orbitals describe the distribution of electrons in a molecule and are frequently used by chemists to understand properties of mole… (voir plus)cules, yet machine learning has neglected them so far. If atom coordinates are obtained through DFT anyway, they can be obtained for free at the same time and are thus a useful source of additional data, particularly when data is scarce We give an introduction to molecular orbitals for a machine learning audience and propose models to process three different representations of them. Experiments on a dataset with experimental properties show that including MOs significantly improves performance and sample efficiency over a pretrained molecular foundation model on this real-world task.

2026-03-01

AI4Mat @ International Conference on Learning Representations (poster)

Inverting Data Transformations via Diffusion Sampling

Jinwoo Kim

Sékou-Oumar Kaba

Jiyun Park

Seunghoon Hong

2026-02-08

ArXiv (prépublication)

arxiv.org

Diffusion Tree Search for Inference Time Adaptation of Material Foundation Models

2025-12-31

Neural Information Processing Systems (Accept (poster))

The Expressive Limits of Diagonal SSMs for State-Tracking

State-Space Models (SSMs) have recently been shown to achieve strong empirical performance on a variety of long-range sequence modeling task… (voir plus)s while remaining efficient and highly-parallelizable. However, the theoretical understanding of their expressive power remains limited. In this work, we study the expressivity of input-Dependent Complex-valued Diagonal (DCD) State-Space Models (SSMs) on sequential state-tracking tasks for abstract groups. It is easy to show that a single DCD SSM layer with a universal decoder can track any Abelian group at finite precision by decomposing it into a product of cyclic groups. We show that this is tight by proving that such a model cannot track any non-Abelian group at finite precision. We further establish the expressivity of multi-layer DCD SSMs. We show that a

2025-12-31

International Conference on Learning Representations (Accept (Poster))

Beyond Scalar Rewards: An Axiomatic Framework for Lexicographic MDPs

Mehran Shakerinava

Adam Oberman

Recent work has formalized the reward hypothesis through the lens of expected utility theory, by interpreting reward as utility. Hausner's f… (voir plus)oundational work showed that dropping the continuity axiom leads to a generalization of expected utility theory where utilities are lexicographically ordered vectors of arbitrary dimension. In this paper, we extend this result by identifying a simple and practical condition under which preferences cannot be represented by scalar rewards, necessitating a 2-dimensional reward function. We provide a full characterization of such reward functions, as well as the general d-dimensional case, in Markov Decision Processes (MDPs) under a memorylessness assumption on preferences. Furthermore, we show that optimal policies in this setting retain many desirable properties of their scalar-reward counterparts, while in the Constrained MDP (CMDP) setting -- another common multiobjective setting -- they do not.

2025-12-02

Conference on Neural Information Processing Systems (Accept (spotlight))

Diffusion Tree Sampling: Scalable inference-time alignment of diffusion models

Adapting a pretrained diffusion model to new objectives at inference time remains an open problem in generative modeling. Existing steering … (voir plus)methods suffer from inaccurate value estimation, especially at high noise levels, which biases guidance. Moreover, information from past runs is not reused to improve sample quality, resulting in inefficient use of compute. Inspired by the success of Monte Carlo Tree Search, we address these limitations by casting inference-time alignment as a search problem that reuses past computations. We introduce a tree-based approach that samples from the reward-aligned target density by propagating terminal rewards back through the diffusion chain and iteratively refining value estimates with each additional generation. Our proposed method, Diffusion Tree Sampling (DTS), produces asymptotically exact samples from the target distribution in the limit of infinite rollouts, and its greedy variant, Diffusion Tree Search (DTS

2025-12-02

Neural Information Processing Systems (Accept (poster))

Energy Loss Functions for Physical Systems

Sékou-Oumar Kaba

Kusha Sareen

Daniel Levy

Effectively leveraging prior knowledge of a system's physics is crucial for applications of machine learning to scientific domains. Previous… (voir plus) approaches mostly focused on incorporating physical insights at the architectural level. In this paper, we propose a framework to leverage physical information directly into the loss function for prediction and generative modeling tasks on systems like molecules and spins. We derive energy loss functions assuming that each data sample is in thermal equilibrium with respect to an approximate energy landscape. By using the reverse KL divergence with a Boltzmann distribution around the data, we obtain the loss as an energy difference between the data and the model predictions. This perspective also recasts traditional objectives like MSE as energy-based, but with a physically meaningless energy. In contrast, our formulation yields physically grounded loss functions with gradients that better align with valid configurations, while being architecture-agnostic and computationally efficient. The energy loss functions also inherently respect physical symmetries. We demonstrate our approach on molecular generation and spin ground-state prediction and report significant improvements over baselines.

2025-12-02

Conference on Neural Information Processing Systems (Accept (poster))