Publications

Securing the Model Context Protocol: Defending LLMs Against Tool Poisoning and Adversarial Attacks
Saeid Jamshidi
Kawser Wazed Nafi
Negar Shahabi
Naser Ezzati-Jivan
Adsorption energies are necessary but not sufficient to identify good catalysts
Alexander Davis
Alexandre AGM Duval
Oleksandr Voznyy
Alex Hern'andez-Garcia
Evaluation and improvement of algorithmic fairness for COVID-19 severity classification using Explainable Artificial Intelligence-based bias mitigation
Shayan Nejadshamsi
Charlene H. Chu
Katherine S. McGilton
Xiaoxiao Li
Charlene Ronquillo
The COVID-19 pandemic has highlighted the growing reliance on machine learning (ML) models for predicting disease severity, which is importa… (voir plus)nt for clinical decision-making and equitable resource allocation. While achieving high predictive accuracy is important, ensuring fairness in the prediction output of these models is equally important to prevent bias-driven disparities in healthcare. This study evaluates fairness in a machine learning-based COVID-19 severity classification model and proposes an Explainable AI (XAI)-based bias mitigation strategy to address sex-related bias. Using data from the Quebec Biobank, we developed an XGBoost-based multi-class classification model. Fairness was assessed using Subset Accuracy Parity Difference (SAPD) and Label-wise Equal Opportunity Difference (LEOD) metrics. Four bias mitigation strategies were implemented and evaluated: Fair Representation Learning, Fair Classifier Using Constraints, Adversarial Debiasing, and our proposed XAI-based method utilizing SHapley Additive exPlanations (SHAP) method for feature importance analysis. The study cohort included 1642 COVID-19 positive older adults (mean age: 77.5), balance equally between males and females. The baseline (unmitigated) classification model achieved 90.68% accuracy but exhibited a 10.11% Subset Accuracy Parity Difference between sexes, indicating a relatively large bias. The introduced XAI-based method demonstrated a better trade-off between model performance and fairness compared to existing bias mitigation methods by identifying sex-sensitive feature interactions and integrating them into the model re-training. Traditional fairness interventions often compromise accuracy to a greater extent. Our XAI-based method achieves the best balance between classification performance and bias, enhancing its clinical applicability. The XAI-driven bias mitigation intervention effectively reduces sex-based disparities in COVID-19 severity prediction without the significant accuracy loss observed in traditional methods. This approach provides a framework for developing fair and accurate clinical decision support systems for older adults, which ensures equitable care in clinical risk stratification and resource allocation.
Neural Coherence : Find higher performance to out-of-distribution tasks from few samples
Mats Richter
A. Chandar
Christopher Pal
Probing the effectiveness of World Models for Spatial Reasoning through Test-time Scaling
Muhammad Jehanzeb Mirza
Wei Lin
Shiqi Yang
A. Chandar
Vision-Language Models (VLMs) remain limited in spatial reasoning tasks that require multi-view understanding and embodied perspective shift… (voir plus)s. Recent approaches such as MindJourney attempt to mitigate this gap through test-time scaling where a world model imagines action-conditioned trajectories and a heuristic verifier selects helpful views from such trajectories. In this work, we systematically examine how such test-time verifiers behave across benchmarks, uncovering both their promise and their pitfalls. Our uncertainty-based analyses show that MindJourney's verifier provides little meaningful calibration, and that random scoring often reduces answer entropy equally well, thus exposing systematic action biases and unreliable reward signals. To mitigate these, we introduce a Verification through Spatial Assertions (ViSA) framework that grounds the test-time reward in verifiable, frame-anchored micro-claims. This principled verifier consistently improves spatial reasoning on the SAT-Real benchmark and corrects trajectory-selection biases through more balanced exploratory behavior. However, on the challenging MMSI-Bench, none of the verifiers, including ours, achieve consistent scaling, suggesting that the current world models form an information bottleneck where imagined views fail to enrich fine-grained reasoning. Together, these findings chart the bad, good, and ugly aspects of test-time verification for world-model-based reasoning. Our code is available at https://github.com/chandar-lab/visa-for-mindjourney.
Foundations of Diffusion Models in General State Spaces: A Self-Contained Introduction
Vincent Pauline
Kirill Neklyudov
A Survey of Bugs in AI-Generated Code
Ruofan Gao
Amjed Tahir
Peng Liang
Teo Susnjak
Developers are widely using AI code-generation models, aiming to increase productivity and efficiency. However, there are also quality conce… (voir plus)rns regarding the AI-generated code. The generated code is produced by models trained on publicly available code, which are known to contain bugs and quality issues. Those issues can cause trust and maintenance challenges during the development process. Several quality issues associated with AI-generated code have been reported, including bugs and defects. However, these findings are often scattered and lack a systematic summary. A comprehensive review is currently lacking to reveal the types and distribution of these errors, possible remediation strategies, as well as their correlation with the specific models. In this paper, we systematically analyze the existing AI-generated code literature to establish an overall understanding of bugs and defects in generated code, providing a reference for future model improvement and quality assessment. We aim to understand the nature and extent of bugs in AI-generated code, and provide a classification of bug types and patterns present in code generated by different models. We also discuss possible fixes and mitigation strategies adopted to eliminate bugs from the generated code.
Towards More Realistic Extraction Attacks: An Adversarial Perspective
Language models are prone to memorizing their training data, making them vulnerable to extraction attacks. While existing research often exa… (voir plus)mines isolated setups, such as a single model or a fixed prompt, real-world adversaries have a considerably larger attack surface due to access to models across various sizes and checkpoints, and repeated prompting. In this paper, we revisit extraction attacks from an adversarial perspective -- with multi-faceted access to the underlying data. We find significant churn in extraction trends, i.e., even unintuitive changes to the prompt, or targeting smaller models and earlier checkpoints, can extract distinct information. By combining multiple attacks, our adversary doubles (
ACCO: Accumulate While You Communicate for Communication-Overlapped Sharded LLM Training
Training LLMs relies on distributed implementations using multiple GPUs to compute gradients in parallel with sharded optimizers. However, s… (voir plus)ynchronizing gradients in data parallel setups introduces communication overhead that grows with the number of workers, limiting parallelization efficiency. Local optimization algorithms reduce communications but incur high memory costs as they prevent optimizer state sharding, hindering scalability. To address this, we propose \textbf{AC}cumulate while \textbf{CO}mmunicate (ACCO), a memory-efficient optimization algorithm for distributed LLM training. By synchronizing delayed gradients while computing new ones, ACCO reduces GPU idle time and supports heterogeneous hardware. To mitigate the convergence issues caused by delayed updates, we introduce a novel technique ensuring training dynamics align with standard distributed optimization. Compared to ZeRO-1, our approach is significantly faster and scales effectively across heterogeneous hardware.
Adaptive Inference-Time Scaling via Cyclic Diffusion Search
Gyubin Lee
Truong Nhat Nguyen Bao
Jaesik Yoon
Dongwoo Lee
Diffusion models have demonstrated strong generative capabilities across domains ranging from image synthesis to complex reasoning tasks. Ho… (voir plus)wever, most inference-time scaling methods rely on fixed denoising schedules, limiting their ability to allocate computation based on instance difficulty or task-specific demands adaptively. We introduce the challenge of adaptive inference-time scaling-dynamically adjusting computational effort during inference-and propose Adaptive Bi-directional Cyclic Diffusion (ABCD), a flexible, search-based inference framework. ABCD refines outputs through bi-directional diffusion cycles while adaptively controlling exploration depth and termination. It comprises three components: Cyclic Diffusion Search, Automatic Exploration-Exploitation Balancing, and Adaptive Thinking Time. Experiments show that ABCD improves performance across diverse tasks while maintaining computational efficiency.
AlignVLM: Bridging Vision and Language Latent Spaces for Multimodal Document Understanding
Juan A. Rodriguez
Chao Wang
Akshay Kalkunte Suresh
Xiangru Jian
Pierre-Andre Noel
Sathwik Tejaswi Madhusudhan
Enamul Hoque
Christopher Pal
Issam H. Laradji
Spandana Gella
Sai Rajeswar
Aligning visual features with language embeddings is a key challenge in vision-language models (VLMs). The performance of such models hinges… (voir plus) on having a good connector that maps visual features generated by a vision encoder to a shared embedding space with the LLM while preserving semantic similarity. Existing connectors, such as multilayer perceptrons (MLPs), lack inductive bias to constrain visual features within the linguistic structure of the LLM's embedding space, making them data-hungry and prone to cross-modal misalignment. In this work, we propose a novel vision-text alignment method, AlignVLM, that maps visual features to a weighted average of LLM text embeddings. Our approach leverages the linguistic priors encoded by the LLM to ensure that visual features are mapped to regions of the space that the LLM can effectively interpret. AlignVLM is particularly effective for document understanding tasks, where visual and textual modalities are highly correlated. Our extensive experiments show that AlignVLM achieves state-of-the-art performance compared to prior alignment methods, with larger gains on document understanding tasks and under low-resource setups. We provide further analysis demonstrating its efficiency and robustness to noise.
Amortized Sampling with Transferable Normalizing Flows
Charlie B. Tan
Leon Klein
Saifuddin Syed
Michael M. Bronstein
Kirill Neklyudov
Efficient equilibrium sampling of molecular conformations remains a core challenge in computational chemistry and statistical inference. Cla… (voir plus)ssical approaches such as molecular dynamics or Markov chain Monte Carlo inherently lack amortization; the computational cost of sampling must be paid in full for each system of interest. The widespread success of generative models has inspired interest towards overcoming this limitation through learning sampling algorithms. Despite performing competitively with conventional methods when trained on a single system, learned samplers have so far demonstrated limited ability to transfer across systems. We demonstrate that deep learning enables the design of scalable and transferable samplers by introducing Prose, a 285 million parameter all-atom transferable normalizing flow trained on a corpus of peptide molecular dynamics trajectories up to 8 residues in length. Prose draws zero-shot uncorrelated proposal samples for arbitrary peptide systems, achieving the previously intractable transferability across sequence length, whilst retaining the efficient likelihood evaluation of normalizing flows. Through extensive empirical evaluation we demonstrate the efficacy of Prose as a proposal for a variety of sampling algorithms, finding a simple importance sampling-based finetuning procedure to achieve competitive performance to established methods such as sequential Monte Carlo. We open-source the Prose codebase, model weights, and training dataset, to further stimulate research into amortized sampling methods and finetuning objectives.