Aaron Courville

sarvjeet-singh.ghotra@mila.quebec

Juan Duque

Doctorat - UdeM

juan.duque@mila.quebec

Doctorat - UdeM

Doctorat - UdeM

Uday Kapur

Doctorat - UdeM

udaykaran.kapur@mila.quebec

Amr Khalifa

Doctorat - UdeM

amr-khalifa.alshaykh@mila.quebec

Zhixuan Lin

Doctorat - UdeM

Doctorat - UdeM

Co-superviseur⋅e :

alan.milligan@mila.quebec

waltermayor@gmail.com

Doctorat - UdeM

Superviseur⋅e principal⋅e :

Simon Lacoste-Julien

Doctorat - UdeM

Superviseur⋅e principal⋅e :

Gauthier Gidel

andjela.mladenovic@mila.quebec

mohammed.muqeeth@mila.quebec

Muqeeth Mohammed

Doctorat - UdeM

Doctorat - UdeM

Co-superviseur⋅e :

Rishabh Agarwal

morgane.moss@mila.quebec

Andrei Nicolicioiu

Doctorat - UdeM

andrei.nicolicioiu@gmail.com

Doctorat - UdeM

Johan Samir Obando Ceron

Doctorat - UdeM

Co-superviseur⋅e :

michell.payanoperez@mila.quebec

johan.ceron@mila.quebec

Doctorat - UdeM

Dereck Piché

Doctorat - UdeM

pichedereck@gmail.com

khaled.rouissi@mila.quebec

Khaled Rouissi

Maîtrise recherche - UdeM

Esra'a Saleh

Doctorat - UdeM

Superviseur⋅e principal⋅e :

Glen Berseth

esraa.saleh@mila.quebec

Vedant Shah

Doctorat - UdeM

vedant.shah@mila.quebec

colin.yiyang-lu@mila.quebec

Yusong Wu

Doctorat - UdeM

Superviseur⋅e principal⋅e :

Anna (Cheng-Zhi) Huang

wu.yusong@mila.quebec

Colin Yiyang Lu

Doctorat - UdeM

Sujin yun

Doctorat - UdeM

sujin.yun@mila.quebec

xiaofeng.zhang@mila.quebec

Xiaofeng Zhang

Doctorat - UdeM

Publications

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… (voir plus)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-12-30

OpenReview (publié)

FeaTok: Feature-Loss Tokenization for Image Generation

Zhang Xiaofeng

Samuel Lavoie

Nicolas Beltran-Velez

Jianhao Yuan

Felix Friedrich

Jakob Verbeek

Michal Drozdzal

Adriana Romero-Soriano

Discrete image generators learn the distribution of tokens produced by an image tokenizer. These token representations should be easy for ge… (voir plus)nerators to model, not merely enable faithful image reconstruction. Yet these tokenizers are commonly trained to reconstruct pixels, which prioritizes low-level fidelity and may disrupt representation structure useful for generation. Motivated by feature-predictive self-supervised learning, we introduce feature-loss tokenization (FeaTok), which trains discrete tokens by reconstructing pretrained features rather than pixels. Ablations across SigLIP, SigLIP2, DINOv2, and DINOv3 show better generative performance with FeaTok. In class-conditional generation, FeaTok advances the quality-compute Pareto frontier, reducing gFDr6 at 256x256 by up to 29 percent over the strongest discrete baselines and outperforming the strongest unguided continuous baseline (RAEv2) with 38 times less sampling compute. FeaTok also achieves higher human-preference reward-model scores than prior art without reward-based training and surpasses continuous and discrete text-conditional baselines on three benchmarks with only 500 fine-tuning updates after pre-training.

2026-10-14

arXiv (prépublication)

GitSwarm: Decentralized Compounding Inference

Vedant Shah

Ankur Samanta

Paras Dahal

Mikhail Plekhanov

Carole-Jean Wu

Scott Yih

Remi Munos

Rob Fergus

Jakob Foerster

Ruslan Salakhutdinov

Sanjeev Arora

Jason Weston

Anirudh Goyal

Long-horizon problem solving and scientific research require computation to accumulate across successive attempts. Partial solutions, experi… (voir plus)mental findings, and unsuccessful approaches can inform later work, yet most inference-time computation is organized around individual trajectories or candidates rather than a persistent body of reusable work. We call this paradigm compounding inference: organizing inference-time computation so that intermediate work persists and can be inspected, extended, combined, or challenged by subsequent computation. We instantiate compounding inference in GitSwarm, an asynchronous system where homogeneous agents independently decide how to advance a task while collaborating through structured persistent memory. Agents explore, experiment, verify, refine, and synthesize previous work in a shared, branch-able Git repository. Atomic commits preserve intermediate artifacts, while explicit semantic dependencies record how later contributions build on work across branches. We evaluate GitSwarm on long-horizon problem solving and sustained GPU-backed experimental research. On IMOProofBench-Advanced, GitSwarm solves all 30 problems in one run using GPT-5.5. On ProgramBench, it achieves a

2026-10-03

arXiv (prépublication)

V-Simba: Unleashing the Architectural Potential of RL in Visual Continuous Control

Donghu Kim

Youngdo Lee

Hojoon Lee

Johan Obando-Ceron

Byungkun Lee

Jaegul Choo

Clare Lyle

Improving sample efficiency remains a core challenge in reinforcement learning (RL), especially in real-world settings like robotics, where … (voir plus)data collection is costly. This challenge is pronounced in visual RL, where high-dimensional inputs often obscure learning signals. While prior work in visual RL has focused on algorithmic solutions, such as better dynamics models or exploration strategies, recent advances in state-based RL show that architectural design alone can lead to significant gains in sample efficiency. This raises an important question: Can these architectural principles transfer to visual RL? In response, we introduce V-Simba, a simple yet effective visual RL architecture inspired by the Simba architecture from state-based RL. Built on top of Soft Actor-Critic (SAC) with data augmentation, V-Simba modifies the architecture by adding normalization layers to stabilize training and using pointwise convolutions to reduce computation. Despite its simplicity, V-Simba matches or outperforms the state-of-the-art methods across the DMC, Adroit, and Meta-World benchmarks, while being more computationally efficient than DrQ-v2. We make our code publicly available at https://github.com/DAVIAN-Robotics/V-Simba.

2026-08-07

arXiv (prépublication)

On-Policy Self-Distillation with Sampled Demonstrations Reduces Output Diversity

Andrei Liviu Nicolicioiu

Mohammad Pezeshki

On-policy self-distillation achieves strong pass@1 accuracy by using a single model as both teacher and student, with the teacher conditione… (voir plus)d on a correct demonstration to provide dense token-level feedback. We show that this could come at a hidden cost: rollout diversity decreases and pass@k curves flatten (i.e., generating more rollouts fails to improve accuracy). We trace this to compounding biases in the design of self-distillation with sampled demonstrations. The teacher scores each student rollout while conditioned on a sampled correct rollout, channeling its feedback through the model's own biases. We theoretically analyze the optimal self-distillation policy and show that it tilts the base distribution by a pointwise conditional mutual information score between the student's rollout and the correct rollout used as context. Unlike the ideal optimal on-policy reinforcement learning (RL), which preserves probability ratios among equally correct rollouts, self-distillation can amplify existing probability gaps, concentrating mass on already-dominant modes. On a controlled graph path-finding task and science question-answering benchmarks, self-distilled models match or exceed RL on average performance but exhibit substantially lower \textit{functional} and \textit{semantic} diversity, failing on out-of-distribution settings that require diverse strategies.

2026-08-04

Sci-FM @ Conference on Language Modeling (poster)

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… (voir plus)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 (accepté)

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… (voir plus)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 (publié)

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… (voir plus)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 (prépublication)

A Mechanistic Analysis of Looped Reasoning Language Models

Hugh Blayney

Álvaro Arroyo

Johan Obando-Ceron

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 … (voir plus)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 (prépublication)

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

Abhranil Chandra

Ayush Agrawal

Arian Hosseini

Sebastian Fischmeister

Rishabh Agarwal

Navin Goyal

2026-03-04

RSI @ International Conference on Learning Representations (poster)

The Markovian Thinker: Architecture-Agnostic Linear Scaling of Reasoning

Milad Aghajohari

Kamran Chitsaz

Amirhossein Kazemnejad

Reinforcement learning (RL) has recently become a strong recipe for training reasoning LLMs that produce long chains of thought (LongCoT). Y… (voir plus)et the standard RL"thinking environment", where the state is the prompt plus all prior reasoning tokens, makes the state unbounded and forces attention-based policies to pay quadratic compute as thoughts lengthen. We revisit the environment itself. We propose Markovian Thinking, a paradigm in which the policy advances reasoning while conditioning on a constant-size state, decoupling thinking length from context size. As an immediate consequence this yields linear compute with constant memory. We instantiate this idea with Delethink, an RL environment that structures reasoning into fixed-size chunks. Within each chunk, the model thinks as usual; at the boundary, the environment resets the context and reinitializes the prompt with a short carryover. Through RL, the policy learns to write a textual state near the end of each chunk sufficient for seamless continuation of reasoning after reset. Trained in this environment, an R1-Distill 1.5B model reasons in 8K-token chunks yet thinks up to 24K tokens, matching or surpassing LongCoT-RL trained with a 24K budget. With test-time scaling, Delethink continues to improve where LongCoT plateaus. The effect of linear compute is substantial: we empirically estimate at 96K average thinking length LongCoT-RL costs 27 H100-months vs. 7 for Delethink. Analysis at RL initialization shows off-the-shelf reasoning models (1.5B-120B) often sample Markovian traces zero-shot across diverse benchmarks, providing positive samples that make RL effective at scale. Our results show that redesigning the thinking environment is a powerful lever: it enables very long reasoning without quadratic overhead and opens a path toward efficient, scalable reasoning LLMs.

2026-01-25

International Conference on Learning Representations (poster)

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… (voir plus)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))