Publications

Autoregressive Boltzmann Generators
Danyal REHMAN
Charlie B. Tan
Yoshua Bengio
Joey Bose
Alexander Tong
Efficient sampling of molecular systems at thermodynamic equilibrium is a hallmark challenge in statistical physics. This challenge has driv… (see more)en the development of Boltzmann Generators (BGs), which allow rapid generation of uncorrelated equilibrium samples by combining a generative model with exact likelihoods and an importance sampling correction. However, modern BGs predominantly rely on normalizing flows (NFs), which either suffer from limited expressivity due to strict invertibility constraints (discrete time) or computationally expensive likelihoods (continuous time). In this paper, we propose Autoregressive Boltzmann Generators (ArBG), a novel autoregressive modelling framework that overcomes these limitations by departing from the flow-based BG paradigm. ArBG circumvents the topological constraints of flows and enables sequential inference-time interventions, while offering enhanced scalability by leveraging architectures effective in Large Language Models. We empirically demonstrate that ArBG leads to significant improvements over flow-based models across all benchmarks, but particularly in larger peptide systems such as the 10-residue Chignolin. Furthermore, we introduce Robin, a 132 million parameter transferable model trained with the ArBG framework which improves over the previous state-of-the-art, reducing the zero-shot energy error,
2026-03-01
GEM @ International Conference on Learning Representations (published)
openreview.net
openreview.net
Beyond Reward Maximization: Evaluating the Diversity of Trajectories in Reinforcement Learning with Temporal Vendi Score
Tom Stanic
Marco Jiralerspong
Xiaofeng Zhang
Danilo Vucetic
Gauthier Gidel
In domains such as scientific discovery and automated design using reinforcement learning (RL), the final task of an agent should extend bey… (see more)ond maximising a single scalar reward; it requires identifying diverse sets of high-quality trajectories to uncover distinct solutions that can provide novel insights on how to solve the problems of interest and transfer robustly from simulation to the real world. However, the RL literature currently lacks a holistic, domain-agnostic standard for measuring trajectory diversity. Existing metrics have been developed to improve exploration at training time but not to evaluate and compare diversity induced by different agents, rendering cross-method comparisons inconsistent and challenging. To address this, we introduce the Temporal Vendi Score (TVS), a novel metric designed to evaluate the diversity of an RL agent by computing the entropy of the eigenvalues' similarity matrix of sampled trajectories. Unlike previous approaches, our metric captures the behavioural diversity of trajectories by accounting for both the sequential nature of state visitations and the temporal structure of the underlying MDP, rather than relying on order-agnostic state comparisons. We validate the TVS on simple environments where we can control the number of different ways a problem can be solved, demonstrating that it provides a more robust, semantically meaningful ranking of diversity than standard baselines. We then show that our metric can scale to a high-dimensional, continuous environment.
2026-03-01
LLA @ International Conference on Learning Representations (poster)
openreview.net
openreview.net
Code and Data Repository for An Exact Framework for Solving the Space-Time Dependent TSP
Isaac Rudich
Manuel López-Ibáñez
Michael Romer
Quentin Cappart
Louis-Martin Rousseau
The software and data in this repository are a snapshot of the software and data that were used in the research reported on in the paper An … (see more)Exact Framework for Solving the Space-Time Dependent TSP by Isaac Rudich, Manuel López-Ibáñez, Michael Römer, Quentin Cappart, and Louis-Martin Rousseau.
2026-03-01
INFORMS Journal on Computing (published)
doi.org
A Comparative Study of Molecular Dynamics Approaches for Simulating Ionic Conductivity in Solid Lithium Electrolytes
Dounia Shaaban Kabakibo
Félix Therrien
Yoshua Bengio
Michel Côté
Hongyu Guo
Homin Shin
Alex Hernández-García
Accurate prediction of ionic conductivity is critical for the design of highperformance solid-state electrolytes in next-generation batterie… (see more)s. We benchmark molecular dynamics (MD) approaches for computing ionic conductivity in 21 lithium solid electrolytes for which experimental ionic conductivity has been previously reported in the literature. Specifically, we compare simulations driven by density functional theory (DFT) and by universal machine-learning interatomic potentials (uMLIPs), namely a MACE foundation model. Our results suggest comparable performance between DFT and MACE, with MACE requiring only a fraction of the computational cost. The framework developed here is designed to enable systematic comparisons with additional uMLIPs and fine-tuned models in future work.
2026-03-01
AI4Mat @ International Conference on Learning Representations (poster)
openreview.net
openreview.net
A Decision-Theoretic Formalisation of Steganography With Applications to LLM Monitoring
Usman Anwar
Julianna Piskorz
David D. Baek
David Africa
Jim Weatherall
Max Tegmark
Christian Schroeder de Witt
Mihaela van der Schaar
David Krueger
Large language models are beginning to show steganographic capabilities. Such capabilities could allow misaligned models to evade oversight … (see more)mechanisms. Yet principled methods to detect and quantify such behaviours are lacking. Classical definitions of steganography, and detection methods based on them, require a known reference distribution of non-steganographic signals. For the case of steganographic reasoning in LLMs, knowing such a reference distribution is not feasible; this renders these approaches inapplicable. We propose an alternative, **decision-theoretic view of steganography**. Our central insight is that steganography creates an asymmetry in usable information between agents who can and cannot decode the hidden content (present within a steganographic signal), and this otherwise latent asymmetry can be inferred from the agents’ observable actions. To formalise this perspective, we introduce generalised
2026-03-01
Trustworthy AI @ International Conference on Learning Representations (published)
doi.org
openreview.net
DELTA-CROSSCODER: ROBUST CROSSCODER IN NARROW FINE-TUNING REGIMES
Aly M. Kassem
Thomas Jiralerspong
Negar Rostamzadeh
Golnoosh Farnadi
Model diffing methods aim to identify how fine-tuning changes a model's internal representations. Crosscoders approach this by learning shar… (see more)ed dictionaries of interpretable latent directions between base and fine-tuned models. However, existing formulations struggle with narrow fine-tuning, where behavioral changes are localized and asymmetric. We introduce Delta-Crosscoder, which combines Dual-K BatchTopK sparsity with a delta-based loss prioritizing directions that change between models, plus an implicit contrastive signal from paired activations on matched inputs. Evaluated across synthetic false facts, emergent misalignment, subliminal learning, and taboo word games (Gemma, LLaMA, Qwen; 1B–7B parameters), Delta-Crosscoder reliably isolates latent directions causally responsible for fine-tuned behaviors and enables effective mitigation, substantially outperforming baselines. Our results demonstrate that narrow fine-tuning induces distinctive, recoverable latent shifts and that crosscoder methods remain powerful tools for model diffing.
2026-03-01
Trustworthy AI @ International Conference on Learning Representations (published)
openreview.net
openreview.net
DiffuMamba: High-Throughput Diffusion LMs with Mamba Backbone
Vaibhav Singh
Oleksiy Ostapenko
Pierre-Andre Noel
Eugene Belilovsky
Torsten Scholak
Diffusion language models (DLMs) have emerged as a promising alternative to autoregressive (AR) generation, yet their reliance on Transforme… (see more)r backbones limits inference efficiency due to quadratic attention or KV-cache overhead. We introduce DiffuMamba, a masked diffusion language model built on a bidirectional Mamba backbone that combines the diffusion objective with linear-time sequence modeling, and DiffuMamba-H, a hybrid variant with interleaved attention. Across scales up to 1.3B parameters, our models match Transformer-based diffusion in downstream performance while achieving up to 8.2× and 4.3× higher inference throughput, respectively, on long sequences. We further present a systematic analysis of inference efficiency across modern DLM variants, combining asymptotic complexity with empirical measurements. Notably, cache-efficient block diffusion with Mamba mixers emerges as the only strategy that scales linearly with sequence length and achieves the strongest performance across all baselines, suggesting a promising direction for future diffusion-based generation systems.
2026-03-01
MM_Intelligence @ International Conference on Learning Representations (poster)
doi.org
openreview.net
Efficient Refusal Ablation in LLM through Optimal Transport
Geraldin Nanfack
Eugene Belilovsky
Elvis Dohmatob
Safety-aligned language models refuse harmful requests through learned refusal behaviors encoded in their internal representations. Recent a… (see more)ctivation-based jailbreaking methods circumvent these safety mechanisms by applying orthogonal projections to remove refusal directions, but these approaches treat refusal as a one-dimensional phenomenon and ignore the rich distributional structure of model activations. We introduce a principled framework based on optimal transport theory that transforms the entire distribution of harmful activations to match harmless ones. By combining PCA with closed-form Gaussian optimal transport, we achieve efficient computation in high-dimensional representation spaces while preserving essential geometric structure. Across six models (Llama-2, Llama-3.1, Qwen-2.5; 7B-32B parameters), our method achieves up to 11% higher attack success rates than state-of-the-art baselines while maintaining comparable perplexity, demonstrating superior preservation of model capabilities. Critically, we discover that layer-selective intervention (applying optimal transport to 1-2 carefully chosen layers at approximately 40-60% network depth) substantially outperforms full-network interventions, revealing that refusal mechanisms may be localized rather than distributed. Our analysis provides new insights into the geometric structure of safety representations and suggests that current alignment methods may be vulnerable to distributional attacks beyond simple direction removal.
2026-03-01
Trustworthy AI @ International Conference on Learning Representations (published)
doi.org
openreview.net
Hierarchical Procedural Meta-Reasoning for Generalizable Multimodal Agents
Yao Fu
Shengyi Qian
Pierluca D'Oro
Fanyi Xiao
Honglak Lee
Joseph Tighe
Manchen Wang
While multimodal agents can achieve strong performance through fine-tuning, their ability to generalize remains limited in complex real-worl… (see more)d tasks such as mobile navigation, where diverse applications, frequent system changes, and customized workflows are common. We argue that a fundamental bottleneck lies in whether an agent possesses sufficient task-specific procedural knowledge to accomplish a given goal. In practice, due to the limited or outdated knowledge of the agent, the procedural steps it generates can be hallucinated and misaligned with the environment during execution. However, better procedural knowledge can be provided by the general capabilities of large language models, or obtained from additional external resources such as web search when necessary. Based on this view, we propose Procedure-Aware Multimodal Agent with Meta Reasoning, a framework that explicitly represents task knowledge as natural-language procedures and trains a procedure-aware grounded agent to condition its actions on this knowledge. By learning to leverage procedural knowledge from different sources, our approach enables robust and reliable generalization with reduced procedural hallucination across tasks, applications, interface versions, and multi-app workflows, achieving substantial improvements on challenging Android benchmarks.
2026-03-01
Reliable Autonomy @ International Conference on Learning Representations (poster)
openreview.net
openreview.net
Latent Personality Alignment: Improving Harmlessness Without Mentioning Harms
Linh Le
David Williams-King
Mohamed Amine Merzouk
Aton Kamanda
Adam Oberman
Current adversarial robustness methods for large language models require extensive datasets of harmful prompts (thousands to hundreds of tho… (see more)usands of examples), yet remain vulnerable to novel attack vectors and distributional shifts. We propose Latent Personality Alignment (LPA), a sample-efficient defense that achieves robustness by training models on abstract personality traits rather than specific harmful behaviors. Using fewer than 100 trait statements and latent adversarial training, LPA achieves comparable attack success rates to methods trained on 150k+ examples, while maintaining superior utility. Critically, LPA generalizes better to unseen attack distributions, reducing misclassification rates by 2.6x compared to baseline across six harm benchmarks—without ever seeing harmful examples during training. Our results demonstrate that personality-based alignment offers a principled approach to building robust defenses with minimal cost.
2026-03-01
Trustworthy AI @ International Conference on Learning Representations (published)
openreview.net
openreview.net
Mechanics of Bias and Reasoning: Interpreting the Impact of Chain-of-Thought Prompting on Gender Bias in LLMs
Edie Pearman
Sophia Osborne
Mira Kandlikar-Bloch
Mina Arzaghi
Florian Carichon
Golnoosh Farnadi
Large language models (LLMs) are increasingly deployed in socially sensitive settings despite substantial documentation that they encode gen… (see more)der biases. Chain-of-Thought (CoT) prompting has been proposed as an approach for bias mitigation. However, existing evaluations primarily focus on changes in LLM benchmark performance, providing limited insight into whether apparent bias reductions reflect meaningful changes in a model's internal mechanisms. In this work, we present an investigation of how CoT prompting affects gender bias in LLMs, combining benchmark-based evaluation with mechanistic interpretability techniques, and qualitative analysis of reasoning outputs. Our results confirm a stereotypical bias present in LLM outputs across benchmarks, showing that CoT prompting does not consistently reduce the bias gap. While mechanistic analyses reveal clusters of attention heads whose biased behavior is lessened with CoT, gender bias information remains pervasive throughout hidden representations, indicating any improvements from CoT are superficial and fail to transform internal processing of gender bias. A closer inspection of the reasoning chains themselves shows poor quality CoT by which the models dissociate, hallucinate, and evade the present task rather than meaningfully engage with prompt material.
2026-03-01
AFAA @ International Conference on Learning Representations (oral)
openreview.net
openreview.net
Modular Memory is the Key to Continual Learning Agents
Vaggelis Dorovatas
Malte Schwerin
Andrew D. Bagdanov
Lucas Caccia
Antonio Carta
Laurent Charlin
Barbara Hammer
Tyler L. Hayes
Timm Hess
Christopher Kanan
Dhireesha Kudithipudi
Xialei Liu
Vincenzo Lomonaco
Jorge Mendez-Mendez
Darshan Patil
Ameya Prabhu
Elisa Ricci
Tinne Tuytelaars
Gido M. van de Ven
Liyuan Wang … (see 4 more)
Joost van de Weijer
Jonghyun Choi
Martin Mundt
Rahaf Aljundi
Foundation models have transformed machine learning through large-scale pretraining and increased test-time compute. Despite surpassing huma… (see more)n performance in several domains, these models remain fundamentally limited in continuous operation, experience accumulation, and personalization, capabilities that are central to adaptive intelligence. While continual learning research has long targeted these goals, its historical focus on in-weight learning (IWL), i.e., updating a single model's parameters to absorb new knowledge, has rendered catastrophic forgetting a persistent challenge. Our position is that combining the strengths of In-Weight Learning (IWL) and the newly emerged capabilities of In-Context Learning (ICL) through the design of modular memory is the missing piece for continual adaptation at scale. We outline a conceptual framework for modular memory-centric architectures that leverage ICL for rapid adaptation and knowledge accumulation, and IWL for stable updates to model capabilities, charting a practical roadmap toward continually learning agents.
2026-03-01
arXiv (preprint)
doi.org
arxiv.org