Grad-ELLM: Gradient-based Explanations for Decoder-only LLMs

TL;DR

Grad-ELLM introduces a gradient-based attribution method tailored for decoder-only LLMs, producing per-generation-step heatmaps by combining channel-level gradients and softened attention. It also proposes π-Soft-NC/NS, a normalization-based refinement of Soft-NC/NS to enable fair comparisons across attribution methods by fixing the expected information kept during soft perturbations. Through experiments on sentiment classification, QA, and open-generation tasks across multiple models, Grad-ELLM demonstrates superior faithfulness relative to baselines, supported by both quantitative metrics and qualitative analyses. The work advances interpretable AI for autoregressive transformers and provides practical metrics and tools for evaluating explanation quality in LLMs.

Abstract

Large Language Models (LLMs) have demonstrated remarkable capabilities across diverse tasks, yet their black-box nature raises concerns about transparency and faithfulness. Input attribution methods aim to highlight each input token's contributions to the model's output, but existing approaches are typically model-agnostic, and do not focus on transformer-specific architectures, leading to limited faithfulness. To address this, we propose Grad-ELLM, a gradient-based attribution method for decoder-only transformer-based LLMs. By aggregating channel importance from gradients of the output logit with respect to attention layers and spatial importance from attention maps, Grad-ELLM generates heatmaps at each generation step without requiring architectural modifications. Additionally, we introduce two faithfulneses metrics $π$-Soft-NC and $π$-Soft-NS, which are modifications of Soft-NC/NS that provide fairer comparisons by controlling the amount of information kept when perturbing the text. We evaluate Grad-ELLM on sentiment classification, question answering, and open-generation tasks using different models. Experiment results show that Grad-ELLM consistently achieves superior faithfulness than other attribution methods.

Figures (16)

• Figure 1: Framework of the proposed Grad-ELLM. An attribution map is generated by aggregating the values $\mathbf{v}$ as feature map in the attention layer with token weight $\lambda_i$ and channel weight $w_c$. This example shows the attribution map when generating the word "Negative".
• Figure 2: Proposed $\pi$-Soft-NS/NC vs. Retaining Probability on IMDb for different XAI methods with Llama.
• Figure 3: Explanations for negative sample from IMDb. The importance of words are visualized by the degree of green color. The retaining probability is $\pi=0.1$.
• Figure 4: Explanations for positive sample from SST2. The importance of words are visualized by the degree of green color. The retaining probability is $\pi=0.1$.
• Figure 5: $\pi$-Soft-NS/NC = 0.98 / 0.16