Publications

Attention with Routed-Memory for Learnable Sparse Control
Qiuhao Zeng
Jerry Huang
Peng Lu
Ruiyi Fang
Gezheng Xu
Zihao Jing
Yufei Cui
Charles X. Ling
Gang Niu
Boyu Wang
Despite advances in long-context inference, large language models (LLMs) remain fundamentally limited by the key-value (KV) caching mechanis… (see more)ms that are necessary for stable computation. Management techniques, such as selective token eviction and pruning, have vastly mitigated the issues that have arisen, but often discard potentially useful information to manage the growing memory requirements of the cache. In this paper, we build upon these approaches to propose Attention with Routed Memory ARM, a novel KV caching structure that introduces a fully differentiable, fixed-size memory system organized as a hierarchical routing structure that learns to select memory slots via Gumbel-Softmax and performs sigmoid-gated updates that softly combine new and stored information, avoiding hard eviction and thereby reducing information loss. By combining this with a policy to dynamically select varying amounts of memory at inference, ARM adapts its accesses for simple contexts and expanding retrieval for inputs that require deeper reasoning, enabling more scalable and effective retrieval on both short and long contexts. Experimental results on standard commonsense and long-context reasoning benchmarks demonstrate that ARM achieves superior performance and efficiency compared to fixed KV-caching approaches, while remaining efficient and scalable in terms of both memory and generation latency.
2025-12-31
International Conference on Machine Learning (Accept (regular))
openreview.net
openreview.net
Bayesian Deep Learning for Remaining Useful Life Estimation via Stein Variational Gradient Descent
Luca Della Libera
Jacopo Andreoli
Davide Dalle Pezze
Mirco Ravanaelli
Gian Antonio Susto
A crucial task in predictive maintenance is estimating the remaining useful life of physical systems. In the last decade, deep learning has … (see more)improved considerably upon traditional model-based and statistical approaches in terms of predictive performance. However, in order to optimally plan maintenance operations, it is also important to quantify the uncertainty inherent to the predictions. This issue can be addressed by turning standard frequentist neural networks into Bayesian neural networks, which are naturally capable of providing confidence intervals around the estimates. Several methods exist for training those models. Researchers have focused mostly on parametric variational inference and sampling-based techniques, which notoriously suffer from limited approximation power and large computational burden, respectively. In this work, we use Stein variational gradient descent, a recently proposed algorithm for approximating intractable distributions that overcomes the drawbacks of the aforementioned techniques. In particular, we show through experimental studies on simulated run-to-failure turbofan engine degradation data that Bayesian deep learning models trained via Stein variational gradient descent consistently outperform with respect to convergence speed and predictive performance both the same models trained via parametric variational inference and their frequentist counterparts trained via backpropagation. Furthermore, we propose a method to enhance performance based on the uncertainty information provided by the Bayesian models. We release the source code at https://github.com/lucadellalib/bdl-rul-svgd.
2025-12-31
IEEE Trans Autom. Sci. Eng. (published)
doi.org
arxiv.org
Beyond Multi-Token Prediction: Pretraining LLMs with Future Summaries
Divyat Mahajan
Sachin Goyal
Badr Youbi Idrissi
Mohammad Pezeshki
Ioannis Mitliagkas
David Lopez-Paz
Kartik Ahuja
Next-token prediction (NTP) has driven the success of large language models (LLMs), but it struggles with long-horizon reasoning, planning, … (see more)and creative writing, with these limitations largely attributed to teacher-forced training. Multi-token prediction (MTP) partially mitigates these issues by predicting several future tokens at once, but it mostly captures short-range dependencies and offers limited improvement. We propose future summary prediction (FSP), which trains an auxiliary head to predict a compact representation of the long-term future, preserving information relevant for long-form generations. We explore two variants of FSP: handcrafted summaries, for example, a bag of words summary of the future of the sequence, and learned summaries, which use embeddings produced by a reverse language model trained from right to left. Large-scale pretraining experiments (3B and 8B-parameter models) demonstrate that FSP provides improvements over both NTP and MTP across math, reasoning, and coding benchmarks.
2025-12-31
International Conference on Learning Representations (Accept (Poster))
doi.org
openreview.net
BRIDGE: Predicting Human Task Completion Time From Model Performance
Fengyuan Liu
Jay Gala
Nilaksh
Dzmitry Bahdanau
Siva Reddy
Hugo Larochelle
Evaluating the real-world capabilities of AI systems requires grounding benchmark performance in human-interpretable measures of task diffic… (see more)ulty. Existing approaches that rely on direct human task completion time annotations are costly, noisy, and difficult to scale across benchmarks. In this work, we propose BRIDGE, a unified psychometric framework that learns the latent difficulty scale from model responses and anchors it to human task completion time. Using a two-parameter logistic Item Response Theory model, we jointly estimate latent task difficulty and model capability from model performance data across multiple benchmarks. We demonstrate that latent task difficulty varies linearly with the logarithm of human completion time, allowing human task completion time to be inferred for new benchmarks from model performance alone. Leveraging this alignment, we forecast frontier model capabilities in terms of human task length and independently reproduce METR's exponential scaling results, with the 50% solvable task horizon doubling approximately every 6 months.
2025-12-31
International Conference on Machine Learning (Accept (regular))
doi.org
openreview.net
Bridging Explainability and Embeddings: BEE Aware of Spuriousness
Cristian Daniel Paduraru
Antonio Barbalau
Radu Filipescu
Andrei Liviu Nicolicioiu
Elena Burceanu
Current methods for detecting spurious correlations rely on data splits or error patterns, leaving many harmful shortcuts invisible when cou… (see more)nterexamples are absent. We introduce BEE (Bridging Explainability and Embeddings), a framework that shifts the focus from model predictions to the weight space and embedding geometry underlying decisions. By analyzing how fine-tuning perturbs pretrained representations, BEE uncovers spurious correlations that remain hidden from conventional evaluation pipelines. We use linear probing as a transparent diagnostic lens, revealing spurious features that not only persist after full fine-tuning but also transfer across diverse state-of-the-art models. Our experiments cover numerous datasets and domains: vision (Waterbirds, CelebA, ImageNet-1k), language (CivilComments, MIMIC-CXR medical notes), and multiple embedding families (CLIP, CLIP-DataComp.XL, mGTE, BLIP2, SigLIP2). BEE consistently exposes spurious correlations: from concepts that slash the ImageNet accuracy by up to 95\%, to clinical shortcuts in MIMIC-CXR notes that induce dangerous false negatives. Together, these results position BEE as a general and principled tool for diagnosing spurious correlations in weight space, enabling principled dataset auditing and more trustworthy foundation models. Our code is publicly available.
2025-12-31
International Conference on Learning Representations (Accept (Poster))
doi.org
openreview.net
Caustics: A Python Package for Accelerated Strong Gravitational Lensing Simulations
Connor Stone
Alexandre Adam
Adam Coogan
M. J. Yantovski-Barth
Andreas Filipp
Landung Setiawan
Cordero Core
Ronan Legin
Charles Wilson
Gabriel Missael Barco
Yashar Hezaveh
Laurence Perreault-Levasseur
2025-12-31
Zenodo (published)
doi.org
arxiv.org
Charge allostatique et dimensions transdiagnostiques des patients admis à l’urgence psychiatrique : un protocole de projet utilisant la Biobanque Signature
Robert‐Paul Juster
Cécile Le Page
Névéna Chuntova
Enzo Cipriani
Clara Morin
Claudia Trudel-Fitzgerald
Marie‐France Marin
Ahmed Jérôme Romain
Steve Geoffrion
Stéphane Guay
Vincent Taschereau‐Dumouchel
Stephane Potvin
Naguib Mcchawar
Caroline Ménard
Eugénie Samson-Daoust
Charles‐Édouard Giguère
Janick Boissonneault
Helen Findlay
Jasmine Boulette
Michelle Lonergan … (see 5 more)
Bahram Armoon
Alain Lesage
Alexandre Hudon
Christophe Longpré-Poirier
Emmanuel Stip
Objectif Les personnes en situation de crise aiguë qui sollicitent les services d’urgence psychiatrique manifestent une détresse psychol… (see more)ogique sévère et présentent probablement des signes physiologiques de stress chronique. Toutefois, comment mesurer au mieux le stress chronique ? Notre étude vise à évaluer la charge allostatique (CA), c’est-à-dire l’ « usure » des systèmes biologiques que nous pouvons mesurer à l’aide de différents marqueurs biologiques appelés biomarqueurs. Méthode Dans ce protocole utilisant les données de la Biobanque Signature, nous souhaitons évaluer les dimensions biologiques et psychosociales liées aux changements des symptômes psychiatriques chez les patient(e)s consultant à l’urgence psychiatrique. L’évaluation sera faite à 4 moments consécutifs : (T1) 1-4 jours après leur admission à l’urgence ; (T2) à la fin de leur hospitalisation (en moyenne 2 semaines plus tard) ; (T3) lors d’un suivi en consultation externe (en moyenne 2 mois plus tard) ; et (T4) environ 1 an plus tard. Nous prévoyons que : 1) la CA à T1 permettra de distinguer différents troubles psychiatriques ; 2) une CA plus élevée permettra de prédire l’évolution des symptômes au fil du temps ; et 3) que la combinaison de facteurs de risque et de protection modifiables (p. ex. les comportements liés à la santé, les expériences psychosociales) renforcera ou atténuera ces associations observées. Nous avons obtenu un financement conséquent des IRSC afin d’analyser les échantillons de cheveux, de salive et de sang, prélevés et conservés par la Biobanque Signature. Cette analyse constituera des mesures à T1 ( N = 1984 ; 40 % de femmes ; âge moyen de 41 ans, avec une fourchette allant de 18 à 81 ans) ainsi qu’à T2 ( n = 685), T3 ( n = 606) et T4 ( n = 370). Résultats Le projet proposé permettra d’indexer la CA chez les patient(e)s et les 149 participant(e)s témoins sans trouble psychiatrique récent. Parmi les sous-échantillons avec des mesures T2, T3 et T4, nous évaluerons simultanément, de manière prospective, les changements de la CA et les trajectoires des symptômes psychiatriques. Conclusion De l’admission à la rémission, l’étude des patient(e)s en urgence psychiatrique est une occasion unique d’évaluer des profils extrêmes qui peuvent contribuer à des approches de médecine de précision. La Signature Biobanque apportera des données uniques sur les signatures biologiques des troubles psychiatriques et leur évolution dans le temps, tout en offrant des possibilités de connaissances uniques et renforcement des capacités pour les années à venir, avec le soutien d’une équipe transdisciplinaire.
2025-12-31
Santé Mentale au Québec (published)
doi.org
Clinicians' ethical considerations on use of AI-enabled technologies for primary prevention of cardiovascular disease in female patients
Ting Wang
Amrita Sandhu
Kyle Vamvakas
Howard Bergman
Roland Grad
Isabelle Vedel
Marie-Pierre Gagnon
Shahram Yousefi
Samira Abbasgholizadeh-Rahimi
Artificial intelligence (AI) is increasingly used in healthcare to support the prevention and management of cardiovascular disease (CVD); ho… (see more)wever, its ethical implications in clinical practice, particularly for female patients, remain insufficiently explored. This study aimed to explore clinicians' perspectives on the ethical use of AI for preventing and managing cardiovascular disease (CVD) in female patients. A qualitative descriptive design was employed using semi-structured interviews with clinicians practicing in Montreal, Canada. Interviews were conducted online, audio-recorded with participants’ consent, and transcribed for analysis. Data were analyzed using deductive thematic analysis informed by ethical domains in the established AI frameworks. Ethical approval was obtained from McGill University’s Research Ethics Board. The study adhered to the Consolidated Criteria for Reporting Qualitative Research (COREQ) guidelines. A final sample of twelve clinicians was interviewed, with each interview lasting approximately 60 minutes. Four key themes emerged: fairness, privacy and security, explainability, and data integrity. Clinicians expressed concerns that AI-enabled technologies may introduce or reinforce biases affecting certain populations, including older adults, individuals with limited digital literacy, and those lacking reliable internet access or access to digital technologies. Participants also raised concerns regarding data integrity, privacy and security, and emphasized the importance of transparent and understandable AI outputs to support clinical decision-making. Ethical considerations are fundamental to the responsible integration of AI in cardiovascular care. Addressing concerns related to fairness, privacy and security, explainability, and data integrity may strengthen clinician trust and support the implementation of AI-enabled technologies in clinical practice. Future research should explore practical approaches to address these concerns and assess how ethically informed AI systems can be implemented effectively in clinical practice.
2025-12-31
FACETS (published)
doi.org
Clinicians' needs and perspectives on use of AI-enabled technologies for primary prevention of cardiovascular disease in female patients
Ting Wang
Amrita Sandhu
Kyle Vamvakas
Howard Bergman
Roland Grad
Isabelle Vedel
Marie-Pierre Gagnon
Shahram Yousefi
Samira Abbasgholizadeh-Rahimi
This study aimed to (1) explore clinicians’ perspectives of cardiovascular disease (CVD) and risk management in female patients and (2) de… (see more)scribe clinicians’ needs and desired features in AI-enabled tools for primary prevention and management of CVD among female patients. This work employed a qualitative description design. We conducted semi-structured interviews with 12 clinicians in Montreal, Canada. We used inductive thematic analysis to interpret the data. Seven themes emerged from the analysis. Three themes were related to the first objective: complexity in clinical decision-making, limitations of CVD risk assessment tools, and resources and health literacy. Four themes were related to the second objective: AI efficiency, multilingual design, electronic medical record integration, and ease of use. Clinicians reported challenges in supporting female patients at higher risk for CVD and expressed concerns about existing decision support tools. They showed openness to AI-enabled tools like Xi-Care and provided input on desired features to ensure their usability and effectiveness. There is a demand to support clinicians in the primary prevention and management of CVD among female patients. AI-enabled tools could effectively address this demand, provided their development prioritizes clinicians' needs and perspectives to ensure safe and effective implementation.
2025-12-31
FACETS (published)
doi.org
CLIP Is Shortsighted: Paying Attention Beyond the First Sentence
Marc-Antoine Lavoie
Anas Mahmoud
Aldo Zaimi
Arsène Fansi Tchango
Steven L. Waslander
CLIP models learn transferable multi-modal features via image-text contrastive learning on internet-scale data. They are widely used in zero… (see more)-shot classification, multi-modal retrieval, text-to-image diffusion, and as image encoders in large vision-language models. However, CLIP’s pretraining is dominated by images paired with short captions, biasing the model toward encoding simple descriptions of salient objects and leading to coarse alignment on complex scenes and dense descriptions. While recent work mitigates this by fine-tuning on small-scale long-caption datasets, we identify an important common bias: both human- and LLM-generated long captions typically begin with a one-sentence summary followed by a detailed description. We show that this acts as a shortcut during training, concentrating attention on the opening sentence and early tokens and weakening alignment over the rest of the caption. To resolve this, we introduce DeBias-CLIP, which removes the summary sentence during training and applies sentence sub-sampling and text token padding to distribute supervision across all token positions. DeBias-CLIP achieves state-of-the-art long-text retrieval, improves short-text retrieval, and is less sensitive to sentence order permutations. It is a drop-in replacement for Long-CLIP with no additional trainable parameters.
2025-12-31
IEEE/CVF Conference on Computer Vision and Pattern Recognition (Accept (Highlight))
doi.org
arxiv.org
A Coin Flip for Safety: LLM Judges Fail to Reliably Measure Adversarial Robustness
Leo Schwinn
Moritz Ladenburger
Tim Beyer
Mehrnaz Mofakhami
Gauthier Gidel
Stephan Günnemann
Automated \enquote{LLM-as-a-Judge} frameworks have become the de facto standard for scalable evaluation across natural language processing. … (see more)For instance, in safety evaluation, these judges are relied upon to evaluate harmfulness in order to benchmark the robustness of safety against adversarial attacks. However, we show that existing validation protocols fail to account for substantial distribution shifts inherent to red-teaming: diverse victim models exhibit distinct generation styles, attacks distort output patterns, and semantic ambiguity varies significantly across jailbreak scenarios. Through a comprehensive audit using 6642 human-verified labels, we reveal that the unpredictable interaction of these shifts often causes judge performance to degrade to near random chance. This stands in stark contrast to the high human agreement reported in prior work. Crucially, we find that many attacks inflate their success rates by exploiting judge insufficiencies rather than eliciting genuinely harmful content. To enable more reliable evaluation, we propose ReliableBench, a benchmark of behaviors that remain more consistently judgeable, and JudgeStressTest, a dataset designed to expose judge failures. (Data in supplement).
2025-12-31
International Conference on Machine Learning (Accept (regular))
doi.org
openreview.net
Compositional Behavioral Semantics and Metrics for State Abstraction in Reinforcement Learning
Yivan Zhang
Ziyan "Ray" Luo
Manuel Baltieri
State abstraction plays a central role in scaling reinforcement learning to complex but structured systems. In studying such systems, a wide… (see more) range of behavioral structures have been proposed in reinforcement learning, including value functions, invariants, bisimulation relations, and behavioral metrics. However, a general principle for determining what structures are provably preserved under state abstraction is still lacking. In this paper, we present a unified framework for defining and analyzing behavioral structures in reinforcement learning. Our framework provides a compositional way to specify behavioral semantics and metrics based on local, one-step descriptions of system dynamics. Using this framework, we establish results showing how behavioral structures can be safely transferred between abstract and concrete systems. We further show how to construct quantitative metrics from logical behavioral semantics with soundness guarantees. Together, these results provide a principled foundation for reasoning about behaviors under state abstraction in reinforcement learning and offer reusable definition and proof principles for a broad class of behavioral structures in reinforcement learning.
2025-12-31
International Conference on Machine Learning (Accept (regular))
openreview.net
openreview.net