Publications

Autoregressive Boltzmann Generators

Charlie B. Tan

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

Can Computational Reducibility Lead to Transferable Models for Graph Combinatorial Optimization?

Michael Perlmutter

A key challenge in deriving unified neural solvers for combinatorial optimization (CO) is efficient generalization of models between a given… (see more) set of tasks to new tasks not used during the initial training process. To address it, we first establish a new model, which uses a GCON module as a form of expressive message passing together with energy-based unsupervised loss functions. This model achieves high performance (often comparable with state-of-the-art results) across multiple CO tasks when trained individually on each task. We then leverage knowledge from the computational reducibility literature to propose pretraining and fine-tuning strategies that transfer effectively (a) between MVC, MIS and MaxClique, and (b) in a multi-task learning setting that additionally incorporates MaxCut, MDS and graph coloring. Additionally, in a leave-one-out, multi-task learning setting, we observe that pretraining on all but one task almost always leads to faster convergence on the remaining task when fine-tuning while avoiding negative transfer. Our findings indicate that learning common representations across multiple graph CO problems is viable through the use of expressive message passing coupled with pretraining strategies that are informed by the polynomial reduction literature, thereby taking an important step towards enabling the development of foundational models for neural CO. We provide an open-source implementation of our work at https://github.com/semihcanturk/COPT-MT .

2026-03-01

arXiv (preprint)

doi.org

arxiv.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

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

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

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

Hierarchical Retrieval at Scale: Bridging Transparency and Efficiency

Tianyi Chen

Valentina Zantedeschi

Information retrieval is a core component of many intelligent systems as it enables conditioning of outputs on new and large-scale datasets.… (see more) While effective, the standard practice of encoding data into high-dimensional representations for similarity search entails large memory and compute footprints, and also makes it hard to inspect the inner workings of the system. Hierarchical retrieval methods offer an interpretable alternative by organizing data at multiple granular levels, yet do not match the efficiency and performance of flat retrieval approaches. In this paper, we propose ReTreever, a tree-based method that makes hierarchical retrieval viable at scale by directly optimizing its structure for retrieval performance while naturally providing transparency through meaningful semantic groupings. Our method offers the flexibility to balance cost and utility by indexing data using representations from any tree level. We show that ReTreever delivers strong coarse (intermediate levels) and fine representations (terminal level), while achieving the highest retrieval accuracy at the lowest latency among hierarchical methods. These results demonstrate that this family of techniques is viable in practical applications.

2026-03-01

Trustworthy AI @ International Conference on Learning Representations (published)

doi.org

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

LatentLens: Revealing Highly Interpretable Visual Tokens in LLMs

Transforming a large language model (LLM) into a Vision-Language Model (VLM) can be achieved by mapping the visual tokens from a vision enco… (see more)der into the embedding space of an LLM. Intriguingly, this mapping can be as simple as a shallow MLP transformation. To understand why LLMs can so readily process visual tokens, we need interpretability methods that reveal what is encoded in the visual token representations at every layer of LLM processing. In this work, we introduce LatentLens, a novel approach for mapping latent representations to descriptions in natural language. LatentLens works by encoding a large text corpus and storing contextualized token representations for each token in that corpus. Visual token representations are then compared to their contextualized textual representations, with the top-k nearest neighbor representations providing descriptions of the visual token. We evaluate this method on 10 different VLMs, showing that commonly used methods, such as LogitLens, substantially underestimate the interpretability of visual tokens. With LatentLens instead, the majority of visual tokens are interpretable across all studied models and all layers. Qualitatively, we show that the descriptions produced by LatentLens are semantically meaningful and provide more fine-grained interpretations for humans compared to individual tokens. More broadly, our findings contribute new evidence on the alignment between vision and language representations, opening up new directions for analyzing latent representations.

2026-03-01

Re-Align @ International Conference on Learning Representations (published)

doi.org

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

Molecule property prediction with molecular orbitals

Yan Zhang

Khang Ngo

Sékou-Oumar Kaba

Daniel T. Levy

Siamak Ravanbakhsh

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… (see more)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)

openreview.net

Multi-scale Predictive Representations for Goal-conditioned Reinforcement Learning

Valliappan CA

David Meger

Sai Rajeswar

Pietro Mazzaglia

Goal-conditioned reinforcement learning (GCRL) requires agents to learn effective state and goal representations, which represents a challen… (see more)ging problem, especially in high-dimensional vision-based environments, as differences in the observations can be uncorrelated with dynamical distances. Classical deep reinforcement learning techniques often fail to capture the alignment between state and goal spaces, requiring additional representation learning techniques. To address this, we propose

2026-03-01

World Models @ International Conference on Learning Representations (published)

openreview.net

Mila Techaide 2026

Venture Scientist Bootcamp

AI Advantage: Productivity in Public Service

Publications

Mila Techaide 2026

Venture Scientist Bootcamp

AI Advantage: Productivity in Public Service

Popular keywords:

Publications