Stefan Bauer

On Closed-Form Couplings

Qiang Liu

Kirill Neklyudov

Few-step generative modelling is an open challenge for flow models. Rectified flows tackle it by distilling a pre-trained “teacher” into… (see more) a few-step “student”, using strong noise–data couplings supplied by the teacher. For a finite dataset and a Gaussian probability path, the probability-flow vector field induced by the empirical distribution is available in closed form, which would allow us to skip training a teacher model. Surprisingly, these couplings turn out to be poor teachers and significantly reduce the performance of the student. We analyse this phenomenon empirically and theoretically, arguing that it stems from intrinsic ambiguity in the induced couplings caused by the strong sensitivity of terminal states to small initialisation perturbations. Under symmetry assumptions, we further prove that the closed-form probability-flow vector field preserves dataset symmetries and induces invariant Voronoi partitions.

2026-03-02

DeLTa @ International Conference on Learning Representations (poster)

openreview.net

Foundations of Diffusion Models in General State Spaces: A Self-Contained Introduction

Vincent Pauline

Kirill Neklyudov

2025-12-03

ArXiv (preprint)

doi.org

arxiv.org

A scalable gene network model of regulatory dynamics in single cells

Paul Bertin

Joseph D Viviano

Alejandro Tejada-Lapuerta

Weixu Wang

Stefan Bauer

Fabian J. Theis

Yoshua Bengio

2025-03-24

ArXiv (preprint)

doi.org

arxiv.org

Learning Decision Trees as Amortized Structure Inference

Mohammed Mahfoud

Ghait Boukachab

Michał Koziarski

Alex Hernández-García

Stefan Bauer

Yoshua Bengio

Nikolay Malkin

2025-03-04

ICLR.cc/2025/Workshop/FPI (poster)

doi.org

openreview.net

Causal machine learning for single-cell genomics

Alejandro Tejada-Lapuerta

Paul Bertin

Stefan Bauer

Hananeh Aliee

Yoshua Bengio

Fabian J. Theis

2023-10-22

ArXiv (preprint)

doi.org

arxiv.org

Neural Causal Structure Discovery from Interventions

Nan Rosemary Ke

Bernhard Schölkopf

Michael Curtis Mozer

Christopher Pal

Yoshua Bengio

Recent promising results have generated a surge of interest in continuous optimization methods for causal discovery from observational data.… (see more) However, there are theoretical limitations on the identifiability of underlying structures obtained solely from observational data. Interventional data, on the other hand, provides richer information about the underlying data-generating process. Nevertheless, extending and applying methods designed for observational data to include interventions is a challenging problem. To address this issue, we propose a general framework based on neural networks to develop models that incorporate both observational and interventional data. Notably, our method can handle the challenging and realistic scenario where the identity of the intervened upon variable is unknown. We evaluate our proposed approach in the context of graph recovery, both de novo and from a partially-known edge set. Our method achieves strong benchmark results on various structure learning tasks, including structure recovery of synthetic graphs as well as standard graphs from the Bayesian Network Repository.

2023-09-09

TMLR (accepted)

openreview.net

Benchmarking Bayesian Causal Discovery Methods for Downstream Treatment Effect Estimation

2023-06-18

ICML.cc/2023/Workshop/SPIGM (poster)

doi.org

openreview.net

Learning Latent Structural Causal Models

Jithendaraa Subramanian

Yashas Annadani

Ivaxi Sheth

Nan Rosemary Ke

Tristan Deleu

Stefan Bauer

D. Nowrouzezahrai

S Ebrahimi Kahou

Causal learning has long concerned itself with the accurate recovery of underlying causal mechanisms. Such causal modelling enables better e… (see more)xplanations of out-of-distribution data. Prior works on causal learning assume that the high-level causal variables are given. However, in machine learning tasks, one often operates on low-level data like image pixels or high-dimensional vectors. In such settings, the entire Structural Causal Model (SCM) -- structure, parameters, \textit{and} high-level causal variables -- is unobserved and needs to be learnt from low-level data. We treat this problem as Bayesian inference of the latent SCM, given low-level data. For linear Gaussian additive noise SCMs, we present a tractable approximate inference method which performs joint inference over the causal variables, structure and parameters of the latent SCM from random, known interventions. Experiments are performed on synthetic datasets and a causally generated image dataset to demonstrate the efficacy of our approach. We also perform image generation from unseen interventions, thereby verifying out of distribution generalization for the proposed causal model.

2022-10-23

ArXiv (preprint)

doi.org

openreview.net

On the Generalization and Adaption Performance of Causal Models

Nan Rosemary Ke

2022-07-19

ICML.cc/2022/Workshop/SCIS (poster)

doi.org

openreview.net

Bayesian Structure Learning with Generative Flow Networks

In Bayesian structure learning, we are interested in inferring a distribution over the directed acyclic graph (DAG) structure of Bayesian ne… (see more)tworks, from data. Defining such a distribution is very challenging, due to the combinatorially large sample space, and approximations based on MCMC are often required. Recently, a novel class of probabilistic models, called Generative Flow Networks (GFlowNets), have been introduced as a general framework for generative modeling of discrete and composite objects, such as graphs. In this work, we propose to use a GFlowNet as an alternative to MCMC for approximating the posterior distribution over the structure of Bayesian networks, given a dataset of observations. Generating a sample DAG from this approximate distribution is viewed as a sequential decision problem, where the graph is constructed one edge at a time, based on learned transition probabilities. Through evaluation on both simulated and real data, we show that our approach, called DAG-GFlowNet, provides an accurate approximation of the posterior over DAGs, and it compares favorably against other methods based on MCMC or variational inference.

2022-05-19

Conference on Uncertainty in Artificial Intelligence (poster)

doi.org

proceedings.mlr.press

From Points to Functions: Infinite-dimensional Representations in Diffusion Models

Sarthak Mittal

Guillaume Lajoie

Stefan Bauer

Arash Mehrjou

Diffusion-based generative models learn to iteratively transfer unstructured noise to a complex target distribution as opposed to Generative… (see more) Adversarial Networks (GANs) or the decoder of Variational Autoencoders (VAEs) which produce samples from the target distribution in a single step. Thus, in diffusion models every sample is naturally connected to a random trajectory which is a solution to a learned stochastic differential equation (SDE). Generative models are only concerned with the final state of this trajectory that delivers samples from the desired distribution. Abstreiter et. al showed that these stochastic trajectories can be seen as continuous filters that wash out information along the way. Consequently, it is reasonable to ask if there is an intermediate time step at which the preserved information is optimal for a given downstream task. In this work, we show that a combination of information content from different time steps gives a strictly better representation for the downstream task. We introduce an attention and recurrence based modules that ``learn to mix'' information content of various time-steps such that the resultant representation leads to superior performance in downstream tasks.

2022-03-27

ICLR.cc/2022/Workshop/DGM4HSD (poster)

doi.org

openreview.net

Systematic Evaluation of Causal Discovery in Visual Model Based Reinforcement Learning

Nan Rosemary Ke

Danilo Rezende

Michael Mozer

Christopher Pal

Inducing causal relationships from observations is a classic problem in machine learning. Most work in causality starts from the premise tha… (see more)t the causal variables themselves are observed. However, for AI agents such as robots trying to make sense of their environment, the only observables are low-level variables like pixels in images. To generalize well, an agent must induce high-level variables, particularly those which are causal or are affected by causal variables. A central goal for AI and causality is thus the joint discovery of abstract representations and causal structure. However, we note that existing environments for studying causal induction are poorly suited for this objective because they have complicated task-specific causal graphs which are impossible to manipulate parametrically (e.g., number of nodes, sparsity, causal chain length, etc.). In this work, our goal is to facilitate research in learning representations of high-level variables as well as causal structures among them. In order to systematically probe the ability of methods to identify these variables and structures, we design a suite of benchmarking RL environments. We evaluate various representation learning algorithms from the literature and find that explicitly incorporating structure and modularity in models can help causal induction in model-based reinforcement learning.

2021-10-10

Datasets and Benchmarks Track @ Neural Information Processing Systems (accepted)

doi.org

openreview.net

Mila on Udemy

AI Policy Fellowship Publications

Mila Ventures Launchpad

Stefan Bauer

Publications

Mila on Udemy

AI Policy Fellowship Publications

Mila Ventures Launchpad

Popular keywords:

Stefan Bauer

Publications