PAI Lab
Machine Learning Department, MBZUAI · Abu Dhabi, UAE
PAI Lab (Provable and Actionable Interpretability) is a research group in the Machine Learning Department at MBZUAI, led by Lijie Hu. We study how to explain the behavior of machine learning models, especially large language and multimodal models, in ways that come with guarantees and that people can act on: auditing and steering models, editing and unlearning behavior, and supporting scientific discovery.
We are recruiting PhD students for Fall 2027. See Join Us.
Research
More on our research →- Theoretical Foundations of Usable XAI
- We study stability, faithfulness, controllability, consistency, influence, and evaluation criteria for actionable explanations.
- Useful XAI in Large Models
- We develop mechanistic analyses, representation interventions, model editing, unlearning, and steering methods that make model behavior actionable.
- Systems for XAI
- We build benchmarks, pipelines, tooling, and service patterns for repeatable explanation workflows and interpretable model monitoring.
- Interactive XAI
- We design human-in-the-loop interfaces, feedback protocols, evaluation tasks, and decision workflows that turn interpretability into action.
- XAI for Science
- We apply PAI methods to protein language models, healthcare, autonomous systems, material science, and other discovery settings.
Selected Publications
All publications →-
EAP-GP: Mitigating Saturation Effect in Gradient-based Automated Circuit Identification.
Abstract
Understanding the internal mechanisms of transformer-based language models remains challenging. Mechanistic interpretability based on circuit discovery aims to reverse engineer neural networks by analyzing their internal processes at the level of computational subgraphs. In this paper, we revisit existing gradient-based circuit identification methods and find that their performance is either affected by the zero-gradient problem or saturation effects, where edge attribution scores become insensitive to input changes, resulting in noisy and unreliable attribution evaluations for circuit components. To address the saturation effect, we propose Edge Attribution Patching with GradPath (EAP-GP), EAP-GP introduces an integration path, starting from the input and adaptively following the direction of the difference between the gradients of corrupted and clean inputs to avoid the saturated region. This approach enhances attribution reliability and improves the faithfulness of circuit identification. We evaluate EAP-GP on 6 datasets using GPT-2 Small, GPT-2 Medium, and GPT-2 XL. Experimental results demonstrate that EAP-GP outperforms existing methods in circuit faithfulness, achieving improvements up to 17.7%. Comparisons with manually annotated ground-truth circuits demonstrate that EAP-GP achieves precision and recall comparable to or better than previous approaches, highlighting its effectiveness in identifying accurate circuits.
-
Semi-supervised Concept Bottleneck Models.
Abstract
Concept Bottleneck Models (CBMs) have garnered increasing attention due to their ability to provide concept-based explanations for black-box deep learning models while achieving high final prediction accuracy using human-like concepts. However, the training of current CBMs is heavily dependent on the precision and richness of the annotated concepts in the dataset. These concept labels are typically provided by experts, which can be costly and require significant resources and effort. Additionally, concept saliency maps frequently misalign with input saliency maps, causing concept predictions to correspond to irrelevant input features - an issue related to annotation alignment. To address these limitations, we propose a new framework called SSCBM (Semi-supervised Concept Bottleneck Model). Our SSCBM is suitable for practical situations where annotated data is scarce. By leveraging joint training on both labeled and unlabeled data and aligning the unlabeled data at the concept level, we effectively solve these issues. We propose a strategy to generate pseudo labels and an alignment loss. Experiments demonstrate that our SSCBM is both effective and efficient. With only 10% labeled data, our model's concept and task accuracy on average across four datasets is only 2.44% and 3.93% lower, respectively, compared to the best baseline in the fully supervised learning setting.
-
Stable Vision Concept Transformers for Medical Diagnosis.
-
Editable Concept Bottleneck Models.
-
Faithful Interpretation for Graph Neural Networks.
Abstract
Currently, attention mechanisms have garnered increasing attention in Graph Neural Networks (GNNs), such as Graph Attention Networks (GATs) and Graph Transformers (GTs). It is not only due to the commendable boost in performance they offer but also its capacity to provide a more lucid rationale for model behaviors, which are often viewed as inscrutable. However, Attention-based GNNs have demonstrated instability in interpretability when subjected to various sources of perturbations during both training and testing phases, including factors like additional edges or nodes. In this paper, we propose a solution to this problem by introducing a novel notion called Faithful Graph Attention-based Interpretation (FGAI). In particular, FGAI has four crucial properties regarding stability and sensitivity to interpretation and final output distribution. Built upon this notion, we propose an efficient methodology for obtaining FGAI, which can be viewed as an ad hoc modification to the canonical Attention-based GNNs. To validate our proposed solution, we introduce two novel metrics tailored for graph interpretation assessment. Experimental results demonstrate that FGAI exhibits superior stability and preserves the interpretability of attention under various forms of perturbations and randomness, which makes FGAI a more faithful and reliable explanation tool.
News
All news →