Aaron Courville

sarvjeet-singh.ghotra@mila.quebec

Juan Duque

PhD - Université de Montréal

juan.duque@mila.quebec

PhD - Université de Montréal

udaykaran.kapur@mila.quebec

Arian Hosseini

PhD - Université de Montréal

Uday Kapur

PhD - Université de Montréal

Amr Khalifa

PhD - Université de Montréal

amr-khalifa.alshaykh@mila.quebec

samuel.lavoie@mila.quebec

Samuel Lavoie

PhD - Université de Montréal

Zhixuan Lin

PhD - Université de Montréal

PhD - Université de Montréal

Principal supervisor :

Simon Lacoste-Julien

alan.milligan@mila.quebec

PhD - Université de Montréal

Principal supervisor :

Gauthier Gidel

andjela.mladenovic@mila.quebec

mohammed.muqeeth@mila.quebec

Muqeeth Mohammed

PhD - Université de Montréal

PhD - Université de Montréal

Co-supervisor :

Rishabh Agarwal

morgane.moss@mila.quebec

Andrei Nicolicioiu

PhD - Université de Montréal

andrei.nicolicioiu@gmail.com

michell.payanoperez@mila.quebec

Michael Noukhovitch

PhD - Université de Montréal

Johan Samir Obando Ceron

PhD - Université de Montréal

Co-supervisor :

Pablo Samuel Castro

johan.ceron@mila.quebec

Collaborating researcher - Université de Montréal

Dereck Piché

PhD - Université de Montréal

pichedereck@gmail.com

khaled.rouissi@mila.quebec

Khaled Rouissi

Master's Research - Université de Montréal

Esra'a Saleh

PhD - Université de Montréal

Principal supervisor :

Glen Berseth

esraa.saleh@mila.quebec

Vedant Shah

PhD - Université de Montréal

vedant.shah@mila.quebec

PhD - Université de Montréal

tsirigoc@mila.quebec

Yusong Wu

PhD - Université de Montréal

Principal supervisor :

Anna (Cheng-Zhi) Huang

wu.yusong@mila.quebec

Sujin yun

PhD - Université de Montréal

sujin.yun@mila.quebec

xiaofeng.zhang@mila.quebec

Xiaofeng Zhang

PhD - Université de Montréal

Publications

Data Selection Through Iterative Self-Filtering for Vision-Language Settings

Andrei Liviu Nicolicioiu

Sarvjeet Singh Ghotra

Morgane M. Moss

The availability of large amounts of clean data is paramount to training neural networks. However, at large scales, manual oversight is impr… (see more)actical, resulting in sizeable datasets that can be very noisy. Attempts to mitigate this obstacle to producing performant vision-language models have so far involved heuristics, curated reference datasets, and using pre-trained models. Here we propose a novel, bootstrapped method in which a CLIP model is trained on an evolving, self-selected dataset. This evolving dataset constitutes a balance of filtered, highly probable clean samples as well as diverse samples from the entire distribution. Our proposed Self-Filtering method iterates between training the model and selecting a subsequently improved data mixture. Training on vision-language datasets filtered by the proposed approach improves downstream performance without the need for additional data or pre-trained models.

2026-06-23

Transactions on Machine Learning Research (accepted)

Representation Learning Enables Scalable Multitask Deep Reinforcement Learning

Johan Obando-Ceron

Scaling reinforcement learning (RL) to diverse multitask settings remains a central challenge. While recent advances in model-based RL achie… (see more)ve strong performance, they rely on planning and complex training pipelines, making it unclear which components are essential for scalability. We revisit this question and argue that the primary driver of scalable multitask RL is not model-based control, but \emph{representation learning}. In particular, we show that combining predictive, model-based representations with high-capacity value function approximation is sufficient to achieve strong performance, even without planning. We evaluate a simple model-free algorithm, MR.Q, coupled with auxiliary predictive objectives into a scalable actor-critic architecture. This approach outperforms a recent world-model-based method and a range of deep RL baselines across a diverse suite of multitask continuous control tasks, while significantly reducing computational overhead and improving wall-clock efficiency. We observe consistent improvements with increased model capacity and show through ablations that predictive representation learning is critical for performance.

2026-06-03

arXiv (preprint)

One More Time: Revisiting Neural Quantum States from a Reinforcement Learning Perspective

Juan Agustin Duque

Sergio García-Heredia

Vinicius Hernandes

Eliška Greplová

Thomas Spriggs

Anna Dawid

Neural quantum states (NQS) provides a flexible and scalable framework for approximating quantum many-body wavefunctions. Among NQS paramete… (see more)rizations, autoregressive models are especially attractive because they enable exact, independent sampling from the Born distribution, avoiding the autocorrelation and mixing issues of Markov chain methods. Yet their optimization remains comparatively underexplored: Adam is a scalable method but ignores function space geometry, while stochastic reconfiguration is principled but costly and numerically fragile in large models. To address this gap, we show that variational energy minimization can be viewed as an advantage policy-gradient problem over the Born distribution, motivating trust-region optimization for NQS training. We introduce \emph{Proximal Wavefunction Optimization} (PWO), a trust-region algorithm that clips probability-ratio changes in the amplitude channel and wrapped phase increments in the phase channel. PWO avoids explicit matrix inversion, reuses samples across inner updates, and preserves the scalability of first-order optimization. Across Ising, Heisenberg, and frustrated

2026-05-29

AI4Physics @ International Conference on Machine Learning (published)

Why Open Source? A Game-Theoretic Analysis of the AI Race

Andjela Mladenovic

Gauthier Gidel

In recent years, with the advancement of frontier AI, we have observed certain dynamics in open-sourcing and closed-sourcing decisions. We p… (see more)ropose a game-theoretic model to analyze these dynamics in the current landscape of the AI race. Our model builds on an R&D race framework under a winner-takes-all setting, and it accounts for the cases where the players' actions can be either discrete or continuous (i.e., partial open-sourcing, such as open weights). We show that determining the existence of a discrete pure non-trivial Nash equilibrium is NP-hard in general but that we can transform the discrete Nash existence computation into a MIP (Mixed-Integer Programming) problem, making it tractable for small instances using a standard MIP solver. Next, we show the existence and tractability of pure Nash equilibria in the continuous version of our problem, leveraging standard convex analysis results, and constructing an equivalent MIP formulation. Throughout this work, we leverage both our main technical results as well as surrounding technical analysis, to derive socially relevant insights that we believe can serve both to understand already existing decisions and dynamics and to potentially inform new policies.

2026-04-16

arXiv (preprint)

A Mechanistic Analysis of Looped Reasoning Language Models

Hugh Blayney

Álvaro Arroyo

Johan Obando-Ceron

Pablo Samuel Castro

Michael M. Bronstein

Xiaowen Dong

Reasoning has become a central capability in large language models. Recent research has shown that reasoning performance can be improved by … (see more)looping an LLM's layers in the latent dimension, resulting in looped reasoning language models. Despite promising results, few works have investigated how their internal dynamics differ from those of standard feedforward models. In this paper, we conduct a mechanistic analysis of the latent states in looped language models, focusing in particular on how the stages of inference observed in feedforward models compare to those observed in looped ones. To this end, we analyze cyclic recurrence and show that for many of the studied models each layer in the cycle converges to a distinct fixed point; consequently, the recurrent block follows a consistent cyclic trajectory in the latent space. We provide evidence that as these fixed points are reached, attention-head behavior stabilizes, leading to constant behavior across recurrences. Empirically, we discover that recurrent blocks learn stages of inference that closely mirror those of feedforward models, repeating these stages in depth with each iteration. We study how recurrent block size, input injection, and normalization influence the emergence and stability of these cyclic fixed points. We believe these findings help translate mechanistic insights into practical guidance for architectural design.

2026-04-12

arXiv (preprint)

Stable Deep Reinforcement Learning via Isotropic Gaussian Representations

Johan Obando-Ceron

Deep reinforcement learning systems often suffer from unstable training dynamics due to non-stationarity, where learning objectives and data… (see more) distributions evolve over time. We show that under non-stationary targets, isotropic Gaussian embeddings are provably advantageous. In particular, they induce stable tracking of time-varying targets for linear readouts, achieve maximal entropy under a fixed variance budget, and encourage a balanced use of all representational dimensions--all of which enable agents to be more adaptive and stable. Building on this insight, we propose the use of Sketched Isotropic Gaussian Regularization for shaping representations toward an isotropic Gaussian distribution during training. We demonstrate empirically, over a variety of domains, that this simple and computationally inexpensive method improves performance under non-stationarity while reducing representation collapse, neuron dormancy, and training instability.

2026-02-21

arXiv (unknown)

Generative Adversarial Post-Training Mitigates Reward Hacking in Live Human-AI Music Interaction

Yusong Wu

Stephen Brade

Aleksandra Teng Ma

Tia-Jane Fowler

Enning Yang

Berker Banar

Natasha Jaques

Cheng-Zhi Anna Huang

Most applications of generative AI involve a sequential interaction in which a person inputs a prompt and waits for a response, and where re… (see more)action time and adaptivity are not important factors. In contrast, live jamming is a collaborative interaction that requires real-time coordination and adaptation without access to the other player’s future moves, while preserving diversity to sustain a creative flow. Reinforcement learning post-training enables effective adaptation through on-policy interaction, yet it often reduces output diversity by exploiting coherence-based rewards. This collapse, known as ``reward hacking'', affects many RL post-training pipelines, but is especially harmful in live jamming, where musical creativity relies on dynamic variation and mutual responsiveness. In this paper, we propose a novel adversarial training method on policy-generated trajectories to mitigate reward hacking in RL post-training for melody-to-chord accompaniment. A co-evolving discriminator separates policy trajectories from the data distribution, while the policy maximizes the discriminator output in addition to coherence rewards to prevent collapse to trivial outputs. We evaluate accompaniment quality and output diversity in simulation with both fixed test melodies and learned melody agents, and we conduct a user study with the model deployed in a real-time interactive system with expert musicians. Quantitative evaluation and user feedback demonstrate improved output diversity, harmonic coherence, adaptation speed and user agency. Our results demonstrate a simple yet effective method to mitigate reward hacking in RL post-training of generative sequence models.

2025-12-31

International Conference on Learning Representations (Accept (Poster))

The Intricate Dance of Prompt Complexity, Quality, Diversity and Consistency in T2I Models

Zhang Xiaofeng

Michal Drozdzal

Adriana Romero-Soriano

Text-to-image (T2I) models offer great potential for creating virtually limitless synthetic data, a valuable resource compared to fixed and … (see more)finite real datasets. Previous works evaluate the utility of synthetic data from T2I models on three key desiderata: quality, diversity, and consistency. While prompt engineering is the primary means of interacting with T2I models, the systematic impact of prompt complexity on these critical utility axes remains underexplored. In this paper, we first conduct synthetic experiments to motivate the difficulty of generalization w.r.t. prompt complexity and explain the observed difficulty with theoretical derivations. Then, we introduce a new evaluation framework that can compare the utility of real data and synthetic data, and present a comprehensive analysis of how prompt complexity influences the utility of synthetic data generated by commonly used T2I models. We conduct our study across diverse datasets, including CC12M, ImageNet-1k, and DCI, and evaluate different inference-time intervention methods. Our synthetic experiments show that generalizing to more general conditions is harder than the other way round, since the former needs an estimated likelihood that is not learned by diffusion models. Our large-scale empirical experiments reveal that increasing prompt complexity results in lower conditional diversity and prompt consistency, while reducing the synthetic-to-real distribution shift, which aligns with the synthetic experiments. Moreover, current inference-time interventions can augment the diversity of the generations at the expense of moving outside the support of real data. Among those interventions, prompt expansion, by deliberately using a pre-trained language model as a likelihood estimator, consistently achieves the highest performance in both image diversity and aesthetics, even higher than that of real data. Combining advanced guidance interventions with prompt expansion results in the most appealing utility trade-offs of synthetic data.

2025-12-31

International Conference on Learning Representations (Accept (Poster))

The Markovian Thinker

Milad Aghajohari

Kamran Chitsaz

Amirhossein Kazemnejad

Reasoning LLMs suffer from quadratic compute growth as their context length increases, making reinforcement learning with verifiable rewards… (see more) (RLVR) and test-time scaling prohibitively expensive. Prior work has tried to lighten the computational burden by shortening reasoning traces through pruning, summarization, or multi-stage training, but these methods remain bound to quadratic costs. We introduce Delethink, a thinking algorithm that realizes the Markovian Thinking Paradigm. Instead of producing one long monolithic reasoning trace, Delethink thinks in a sequence of chunks, the Delethink trace. Each chunk continues reasoning by referring only to a fixed number of prior tokens, which functions as a Markovian state sufficient for progressing reasoning, while deleting the rest. This preserves continuity without carrying the quadratic baggage. As a result, compute scales linearly and peak memory remains constant. In experiments, we show that Delethink can be applied directly to off-the-shelf reasoning models ranging from

2025-12-31

International Conference on Learning Representations (Accept (Poster))

A Comedy of Estimators: On KL Regularization in RL Training of LLMs

Vedant Shah

Johan Obando-Ceron

Vineet Jain

Brian Bartoldson

Bhavya Kailkhura

Nikolay Malkin

The reasoning performance of large language models (LLMs) can be substantially improved by training them with reinforcement learning (RL). T… (see more)he RL objective for LLM training involves a regularization term, which is the reverse Kullback-Leibler (KL) divergence between the trained policy and the reference policy. Since computing the KL divergence exactly is intractable, various estimators are used in practice to estimate it from on-policy samples. Despite its wide adoption, including in several open-source libraries, there is no systematic study analyzing the numerous ways of incorporating KL estimators in the objective and their effect on the downstream performance of RL-trained models. Recent works show that prevailing practices for incorporating KL regularization do not provide correct gradients for stated objectives, creating a discrepancy between the objective and its implementation. In this paper, we further analyze these practices and study the gradients of several estimators configurations, revealing how design choices shape gradient bias. We substantiate these findings with empirical observations by RL fine-tuning \texttt{Qwen2.5-7B}, \texttt{Llama-3.1-8B-Instruct} and \texttt{Qwen3-4B-Instruct-2507} with different configurations and evaluating their performance on both in- and out-of-distribution tasks. Through our analysis, we observe that, in on-policy settings: (1) estimator configurations with biased gradients can result in training instabilities; and (2) using estimator configurations resulting in unbiased gradients leads to better performance on in-domain as well as out-of-domain tasks. We also investigate the performance resulting from different KL configurations in off-policy settings and observe that KL regularization can help stabilize off-policy RL training resulting from asynchronous setups.

2025-12-25

ArXiv (preprint)

Shape of Thought: When Distribution Matters More than Correctness in Reasoning Tasks

Abhranil Chandra

Ayush Agrawal

Arian Hosseini

Sebastian Fischmeister

Rishabh Agarwal

Navin Goyal