Fast Explanations via Policy Gradient-Optimized Explainer ¹¹1Preprint, under review

While deep classification models demonstrate superior performance across a range of tasks, their ”black-box” nature often hinders their acceptance and deployment in critical areas such as healthcare, finance, and autonomous systems Miotto et al. (2018); Ozbayoglu et al. (2020); Grigorescu et al. (2020). In these high-stakes contexts, it is essential not only to achieve high predictive accuracy but also to provide clear, understandable explanations of the models’ decisions to foster trust and ensure accountability.

Despite the progress in explainable AI (XAI) research Atakishiyev et al. (2021); Singh et al. (2020), achieving explainability in real-world, large-scale applications remains challenging. A significant barrier is the difficulty of providing efficient explanations that can scale without imposing prohibitive computational costs. Current attribution-based explanation methods often require substantial computation during inference, making them impractical for time-sensitive tasks and large-scale deployment Chuang et al. (2023a); Lundberg et al. (2020). Therefore, improving the efficiency of explanations is critical for enabling their broader adoption in real-world applications.

In practice, when working with black-box models or complex architectures, it is expected to use model-agnostic explanation methods Ribeiro et al. (2016); Lundberg and Lee (2017); Petsiuk et al. (2018); Fong and Vedaldi (2017). These methods have the advantage of being applicable to a wide range of models, but they often require numerous additional forward passes or gradient computations, making them inefficient and costly for real-world applications.

In scenarios where we have full access to the model’s architecture, such as CNNs or Transformers, model-specific explanation methods can be employed to provide rapid explanations Selvaraju et al. (2017); Chefer et al. (2021); Qiang et al. (2022). These methods are tailored to specific model structures, leveraging the unique behaviors of certain architectures to achieve efficient explanations. However, in real-world settings, models are often either black-box or not easily categorized into standard architectures, which limits the application of model-specific explanations.

To address the inefficiency of model-agnostic methods and the limitation of model-specific approaches, amortized explanation techniques have been proposed by training a deep neural network (DNN) to approximate an explanation distribution, thereby accelerating model-agnostic explanations to a single forward pass during inference Chuang et al. (2023a); Jethani et al. (2021); Chen et al. (2018). However, these methods rely heavily on approximating specific proxy explanation methods, such as SHAP, by treating their explanations as ground truth (or pseudo-label). This introduces limitations: the performance of amortized methods is inherently capped by the quality of the proxy explanations; and they also rely on the assumptions made by the proxy methods.

In this work, we propose a novel policy gradient based approach to learn a model-specific explainer that is not only capable of making fast explanation to any black-box models, but also have no reliance on pseudo-labels from existing proxy explanation methods.

Our main contributions are summarized as follows: 1) To the best of our knowledge, this is one of the first work that leverages reinforcement learning to directly learn an efficient explainer directly from data and the prediction model. 2) Unlike other amortized methods, our method doesn’t rely on the pseudo-labels provided by any proxy explanation method, such as SHAP. 3) A KL-divergence regularization is also introduced to enhance the generalizability of the learned explainer. 4) Comprehensive qualitative and quantitative experiments across multiple datasets demonstrate the superior quality and efficiency of our approach.

In this work, we focus on explanations in the format of feature attribution Linardatos et al. (2020), i.e., finding the importance score for individual input features that influence a prediction. Therefore, we review three categories of feature attribution methods that are closely related to our approach: model-agnostic approaches, model-specific approaches, and amortized approaches.

Model-agnostic approaches: Model-agnostic approaches are designed to be broadly applicable, making minimal or no assumptions about the to-be-explained prediction models. One common strategy involves using an explainable surrogates to approximate the local behavior of models, which is particularly useful for black-box models. For instance, LIME Ribeiro et al. (2016) fits a surrogate interpretable model (such as a linear model) to explain predictions locally by perturbing the input data and observing the changes in predictions. Similarly, SHAP Lundberg and Lee (2017) leverages Shapley values from game theory to ensure a unique surrogate solution with desirable properties such as local accuracy, missingness, and consistency. RISE Petsiuk et al. (2018) generates saliency maps by sampling randomized masks and evaluating their impact on the model’s output.

Another category of model-agnostic techniques leverages gradient information from white-box models to provide explanations. Instead of learning surrogates, these methods exploit locally smoothed gradients to approximate the model’s local behavior. The smoothing strategies vary among approaches. For instance, Integrated Gradients Sundararajan et al. (2017) computes explanations by averaging gradients of interpolated samples between a baseline input and the target input. AGI Pan et al. (2021) refines this concept by averaging gradients along multiple adversarial attack trajectories, while NeFLAG Li et al. (2023) utilizes gradients averaged over a hyperspherical neighborhood.

Although these methods are usually widely applicable, they are resource-intensive during inference due to the necessity of a large number of additional model queries.

Model-specific approaches: Model-specific approaches are typically tailored to specific model architectures, enabling efficient explanations by utilizing attention weights, convolutional feature maps or custom layers. GradCAM Selvaraju et al. (2017), for instance, uses the weighted average of the convolutional feature maps to generate attributions, effectively working on CNNs. Similarly, methods like AttLRP Chefer et al. (2021) and AttCAT Qiang et al. (2022) are designed specifically for transformer-based models, relying on attention weights from various attention heads and layers to compute final explanations. DeepLIFT Shrikumar et al. (2017) provides a framework for explaining deep learning models under the condition that propagation rules can be adapted.

Amortized approaches: Amortized explanation methods approximate the explanations from the resource-heavy model-agnostic methods (proxy methods) by a single forward pass. For example, FastSHAP Jethani et al. (2021), which amortizes the cost of fitting kernelSHAP by stochastically training a neural network to approximate it globally. CoRTXChuang et al. (2023b), on the other hand, learns the explanation-oriented representation in a self-supervised manner and reduces the dependence of training on pseudo-labels from proxy methods. Overall, this type of methods achieves efficiency via an additional global surrogate function on top of the surrogates in model-agnostic methods, which introduces additional uncertainty.

In this paper, we propose a novel reinforcement learning framework that learns a distribution-based explainer that achieves the universality of model-agnostic approaches and the efficiency of model-specific approaches without relying on any proxy methods.

We lead with a discussion of an intractable empirical attribution, which relies on an exhaustive search over all possible feature combinations. Then, we interpret this empirical attribution as an expectation of a probability distribution and approximate this distribution via a policy gradient approach. Figure 1 illustrates the proposed model.

We conduct experiments on both image and text classification tasks. For image classification, we use the ViT model Dosovitskiy et al. (2020) fine-tuned on the ImageNet dataset Deng et al. (2009) as the prediction model. The FEX explainer is finetuned on the full ImageNet dataset with 1.3M samples (FEX-1.3M) or a subset of 50,000 samples (FEX-50k) for one epoch . For text classification, we use the BERT model Devlin et al. (2018) fine-tuned on the SST2 dataset Socher et al. (2013) for sentiment analysis. The FEX explainer is finetuned on the Movies Reviews Zaidan and Eisner (2008) dataset for one epoch with batch size 256. Unless otherwise specified, in all experiments, the $g(\boldsymbol{x})$ is set to the same architecture as the predictor $f$, with appended MLP prediction heads, and the hyperparameters are set to $\lambda_{en}=10^{-5}$, $\lambda_{v}=0.5$ and $\lambda_{kl}=1$.

To address the challenge of balancing general applicability with inference speed in explainable AI (XAI), we proposed FEX framework that bridges the gap between the slow inference speed of model-agnostic methods and the limited applicability of model-specific methods. And unlike amortized approaches, which require existing model-agnostic methods as proxy explainers, our framework has no reliance on any proxy explainers. Experiments demonstrate that our method outperforms other baselines in both explanation quality and inference speed across various metrics.

Our framework enhances transparency and trust in AI, crucial for applications in sectors like healthcare. It aids debugging and bias identification, supporting ethical AI use and regulatory compliance. However, risks include potential oversimplification of explanations and exposure of proprietary model details. Addressing these challenges is key to maximizing positive impact.