Explainable AI: Context-Aware Layer-Wise Integrated Gradients for Explaining Transformer Models

TL;DR

The paper introduces CA-LIG, a hierarchical attribution framework that enhances the interpretability of Transformer models by integrating layer-wise relevance with context-aware, class-specific attention gradients, providing clearer and more faithful explanations.

Contribution

It proposes a novel unified hierarchical attribution method, CA-LIG, that captures relevance flow across layers and structural components, improving explainability of Transformer models.

Findings

CA-LIG produces more faithful attributions across diverse tasks.

It shows stronger sensitivity to contextual dependencies.

It yields clearer, more semantically coherent visualizations.

Abstract

Transformer models achieve state-of-the-art performance across domains and tasks, yet their deeply layered representations make their predictions difficult to interpret. Existing explainability methods rely on final-layer attributions, capture either local token-level attributions or global attention patterns without unification, and lack context-awareness of inter-token dependencies and structural components. They also fail to capture how relevance evolves across layers and how structural components shape decision-making. To address these limitations, we proposed the \textbf{Context-Aware Layer-wise Integrated Gradients (CA-LIG) Framework}, a unified hierarchical attribution framework that computes layer-wise Integrated Gradients within each Transformer block and fuses these token-level attributions with class-specific attention gradients. This integration yields signed, context-sensitive attribution maps that capture supportive and opposing evidence while tracing the hierarchical flow of relevance through the Transformer layers. We evaluate the CA-LIG Framework across diverse tasks, domains, and transformer model families, including sentiment analysis and long and multi-class document classification with BERT, hate speech detection in a low-resource language setting with XLM-R and AfroLM, and image classification with Masked Autoencoder vision Transformer model. Across all tasks and architectures, CA-LIG provides more faithful attributions, shows stronger sensitivity to contextual dependencies, and produces clearer, more semantically coherent visualizations than established explainability methods. These results indicate that CA-LIG provides a more comprehensive, context-aware, and reliable explanation of Transformer decision-making, advancing both the practical interpretability and conceptual understanding of deep neural models.