Fast Explanations via Policy Gradient-Optimized Explainer

1 Introduction

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.

2 RelatedWork

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. CoRTX[Chuang 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.

3 ProposedMethod: Fast Explanation (FEX)

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.

3.1 Empirical Attribution

For an input comprising N features, there exists 2N distinct feature selection combinations. We represent each feature combination using a mask m ∈ {0, 1}N . For clarity, we define a mask entry of 0 to indicate that a feature is masked (removed), while a mask entry of 1 denotes that a feature is retained.

Consider an input vector x = (x1, ..., xN)⊤, with a binary classification function f : X → [0, 1]. For each masked version of the input, represented as m⊙ x, we obtain a corresponding prediction f(m⊙ x) ∈ [0, 1].

Intuitively, the native contribution of a feature xi to f(m ⊙ x) can be represented by the average contribution from all features present inm, i.e.,

where Km denotes the number of non-zero entries in m, representing the number of retained features in the masked input. This intuition for naive contribution comes from the philosophy that all features should be viewed as equally important if there isn't any prior knowledge.

Assume that there is a subset Mi = {m|m ∈ {0, 1}N , mi = 1}, which comprises all masks that retain feature xi. Then the sum of naive contributions of xi over Mi is a good indicator of the importance of feature xi.

Building on this, we define the empirical attribution as the total naive contribution from all masked inputs that retain the feature xi. Specifically:

Definition 1. (Empirical Attribution) Given an input vector x ∈ X, a set of masks M = {0, 1}N , and the prediction function f(x) : X → [0, 1], the empirical attribution of a feature xi is defined as:

where Mi = {m ∈ M|mi = 1} represents the set of masks retain feature xi, and Km represents the number of non-zero elements inm.

The cumulative empirical attribution can also be expressed in vector form when denoting ϕ(x) by:

Since m = (m1,m2, ...,mN)T , the vector form of the empirical attribution becomes:

where ϕ(x) ∈ RN+ .

3.2 Attribution as an Expectation

Calculating the empirical attribution is computationally prohibitive due to its exponential complexity of O(2N ) . A common approach to address this challenge is to approximate it via Monte Carlo simulation, as has been similarly demonstrated in the RISE [Petsiuk et al., 2018] framework. However, Monte Carlo methods remain computationally intensive, and the quality of the attribution is highly dependent on the number of simulation steps.

Note that the empirical attribution ϕ(x) ∈ RN+ in Eq 4 can be written as a form of probability expectation with an appropriate normalization factor A(x) = ∑m f (m⊙x) Km . Specifically, we have:

where the distribution p(m|x) is defined by

Although it is still intractable to directly calculate the expectations, it is possible to obtain the expectation as a closed form if p(m|x) follows some specific distribution families.

3.3 Tractable Bernoulli Surrogate

For the purpose of explanation, we chose a multivariate Bernoulli distribution for its natural alignment with the disentanglement of different features. It serves as a surrogate to p(m|x):

where λ ∈ [0, 1]N is the mean parameter of the Bernoulli distribution, parameterized by a neural network g(x).

Its expectation has a closed form. Specifically, for m sampled from q, we have:

This property is particularly advantageous because it allows the mean parameter λ = g(x) to represent the empirical attribution directly if q can approximate p, as defined in Eq 5.

However, p(m|x) is not directly computable, instead, we start by defining a score function c(m, x) as

Note that the expectation is primarily influenced by regions of the probability distribution with high density. Therefore, to approximate p with q, we need to optimize q such that the high density region matches p. Therefore, assume there are T masks m1, ...,mT sampled from q, we aim to maximize the following objective:

3.4 Policy Gradient

In the reinforcement learning literature, objectives with a structure similar to Eq. 10 can be effectively optimized using policy gradient methods. The policy gradient framework is characterized by four fundamental components: states (of the environment), actions (by the agent), the policy (for generating actions), and the return (of a series of actions).

Policy Gradients Adaptation

To adapt our framework to the policy gradient methodology, it is crucial to establish a clear correspondence between the key concepts in our approach and those traditionally utilized in the policy gradient literature. This section provides a detailed mapping of these conceptual alignments.

Lets rephrase our problem as follows: Given the score function c, and original model input x, we need to find a distribution q such that it maximizes Eq 10.

Input x as Static States: In policy gradient framework, a state st represents the current situation or configuration of the environment with which the agent interacts. In our context, since an input sample doesn't change over the masking actions, we consider these samples as static states. Formally, st = x, where t = 0, 1, ..., T .

Mask m as Actions In our framework, applying masks to static input samples can be viewed as actions applied towards the states. i.e., at =mt.

Bernoulli Surrogate q as the Policy The policy in reinforcement learning generates actions. Similarly, in our context, the mask distribution q can be viewed as the policy that generates the masks. Specifically, mt ∼ q.

Return Consequently, the weighted score function 1 T c(m, x) performs as the reward given an action m upon state x. Furthermore, if we define τ as a trajectory of a mask sequence m1, ...,mT , the return R can be computed by

Objective Function The objective function can be formally expressed as

where τ is a trajectory sampled from q.

Policy Gradient Formulation

Considering the above terminology connections between our framework and the policy gradient method, the gradient of the objective in Eq 12 can be expressed as:

The advantage function Aq (x,m) is the difference between Action-Value function (Q-function) and Value function (V-function). The Q-function is defined as the expected return with the first action (mask) being m

where the factor T is multiplied on both sides for conciseness. Similarly, the V-function is defined by the expected return if the first action m is sampled from q:

Therefore, the advantage function can be obtained by subtracting V fromQ . Specifically,

Note that V q (x) can also be approximated by a neural network v(x), which can be trained by minimizing the following loss function:

Proximal Policy Optimization: Policy gradient methods may suffer the issue of performance collapse when the policy changes too much during a single update. Therefore, we facilitate the clip trick used in PPO (Proximal Policy Optimization)[Schulman et al., 2017] that constrains the update within each step. Consequently, the gradient in Eq 13 can be written by:

where qℓ represents the policy from the last updating step ℓ. Additionally, an entropy regularization term H(q) is also added to balance the exploration and the exploitation during the reinforcement learning steps.

Combining the PPO objective, the entropy, and the MSE loss for v(x), the objective function can then be written by

3.5 Generalizability

Generalizability in our framework involves two key aspects: (1) generalization over the distribution of all samples and (2) generalization across different output classes. Both are crucial for creating robust explainers that go beyond individual input-prediction pairs.

Generalization Over Sample Distribution: Generalization across samples ensures the explainer g(x) consistently provides meaningful explanations over a dataset X. When trained on a diverse dataset, our framework allows effective adaptation to diverse inputs.

Generalization Over Class Distribution: In multi-class classification, the prediction function f outputs a probability vector (f1, . . . , fK)⊤ over K classes, requiring K explainers g1, . . . , gK. Intuitively, when fi dominates the predicted probabilities, the corresponding explainer gi ∈ (0, 1) should also have a dominant average score. To enforce this alignment, KL-divergence is used to match the average explainer scores with the predicted class probabilities:

Algorithm 1 PPO for Fast Explanations

• 1: Input: training samples set X, prediction function f = (f1, ..., fK)T , initial explainer network g = (g1, ..., gK)T , initial value network v = (v1, ..., vK)T , and hyperparameters λen, λv, λkl
• 2: for i = 1, 2, . . . do
• 3: Get a batch of input-output pairs Xi ⊂ {(x, y)|x ∈ X, y = arg maxk fk(x)}
• 4: for ℓ = 0, 1, 2, . . . do
• 5: Collect a set of trajectories Dℓ = {τj} by running policy q = Bern(gy(x)) for all (x, y) pairs.
• 6: Compute the Advantage Aq (mt, x) by Eq 18.
• 7: Obtain the PPO-Clip objective Lppo by Eq 20, value network MSE loss Lv by Eq 19, KL-divergence Lkl by Eq 24, and entropy H(q).
• 8: Update the policy by minimizing the objective L in Eq 25
• 9: end for
• 10: end for

Averaging the log values of the explainer scores ensures a more stable computation, as log probabilities are additive by nature. The softmax function is then applied to form a valid probability distribution. This approach guarantees consistency between the explainers and the classifier's output, facilitating robust and scalable explanations across classes. Consequently, the overall objective function for the policy gradient adaptation becomes:

3.6 Efficiency and Capabilities

Table 1 provides a detailed comparison of our Fast Explanation method (FEX) with several related explanation methods.

FEX distinguishes itself by requiring only O(1) forward passes of g(x) during inference, ensuring exceptional computational efficiency. In contrast, other model-agnostic baselines, such as RISE, IG, and GradSHAP, require O(K) queries to the prediction model. While methods like GradCAM and AttLRP achieve similar O(1) efficiency in terms of model queries, they are inherently model-specific and therefore cannot be applied in a black-box setting.

Beyond its computational efficiency and model-agnostic nature, FEX offers an additional advantage: it does not depend on pseudo-labels generated by proxy explainers. For example, although FastSHAP achieves O(1) efficiency, its reliance on pseudo-labels from SHAP introduces potential limitations, as its performance is constrained by the accuracy of the proxy explainer.

4 Experiments

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(x) is set to the same architecture as the predictor f, with appended MLP prediction heads, and the hyperparameters are set to λen = 10-5 , λv = 0.5 and λkl = 1.

4.1 Baselines

For the image classification task, we evaluate our proposed method against six baseline approaches, encompassing model-specific, model-agnostic, and amortized explanation techniques. The model-specific baselines include GradCAM, where we use the last hidden state as the target feature map, and AttLRP , where the default configurations from the original work are utilized. The model-agnostic baselines include GradSHAP, RISE and Integrated Gradients (IG). They require a number of queries (K) to the prediction model. In our experiments, Ks are set to 100 for all model-agnostic baselines. For the amortized methods, we include FastSHAP, where the explainer is implemented as a U-Net generating a 14×14 heatmap and is trained on 50,000 ImageNet samples (We are not able to train FastSHAP on the full ImageNet dataset because it's extremely slow).

For the text classification task, due to the discrete nature of text tokens, FastSHAP, IG and GradSHAP are not directly applicable. Hence, we only compare our method with RISE, GradCAM and AttLRP, with random attribuiton as a reference.

4.2 Metrics and Results

Figure 2 illustrates qualitative comparisons of various explanation methods for the image classification task. Our approach achieves comparable visual quality to model-specific methods such as AttLRP and GradCAM, while significantly surpassing model-agnostic baselines like IG and GradSHAP.

For quantitative evaluation, we follow the strategies outlined in [Chefer et al., 2021]. For image classification, attribution performance is assessed using the area under the curve (AUC) of prediction accuracies, computed by progressively masking features (from 0% masked to 100% masked) based on their attributed importance. Positive AUC is calculated by masking the most important features first, whereas Negative AUC begins with the least important features. These evaluations are conducted on a randomly selected subset of 5,000 images from the ImageNet validation set.

To further assess the quality of explanations, we use an annotated image segmentation dataset [Guillaumin et al., 2014] comprising 4,276 images across 445 categories. Segmentation labels serve as ground truth for attribution scores in the classification task. Performance is evaluated using proxy metrics: pixel accuracy, which measures the proportion of the most important patches falling within segmentation boxes; mean intersection over union (mIoU), quantifying overlap between segmentation boxes and features with above-average scores; and mean average precision (mAP), representing the area under the precision-recall curve with attribution scores as predictions.

Results for FEX and FastSHAP are averaged over three trained explainers, RISE, IG, GradSHAP, and Random are averaged over three runs, while only one run for GradCAM and AttLRP as they are deterministic. Table 2 demonstrate the superior performance of our proposed FEX method compared to other baselines.

For text classification, we adopt the ERASER benchmark [DeYoung et al., 2019] and evaluate sentiment predictions on the Movie Reviews dataset [Zaidan and Eisner, 2008]. Explanations are evaluated by plotting F1 score curves as text tokens are progressively inserted based on their attributed importance. As shown in Figure 3, our method also achieves better performance on the text classification task.

In terms of efficiency, fine-tuning ViT on 1.3M ImageNet samples for a single epoch takes approximately 5 hours on a single A100 GPU, and one epoch is generally sufficient to achieve high-quality attribution scores. Notably, inference with FEX is lightweight, requiring only a single forward pass of g(x). According to Table 3, our framework achieves the same level of inference efficiency as FastSHAP. It reduces inference time by over 97% and memory usage by 70% compared to traditional model-agnostic approaches (RISE, IG and GradSHAP).

5 Ablation Study

Effect of Trajectory Size The trajectory size sampled from the policy can impact its optimization. While longer trajectories can provide richer information for policy learning, computational constraints limit their feasibility during training. Striking a balance between trajectory size and computational efficiency is thus critical. As presented in Table 4, performance improves when the trajectory length increases from s = 1 to s = 5; however, it saturates when extending the trajectory further to s = 10. These results indicate that sampling excessively long trajectories is unnecessary.

Effect of Training Data Size Training data size is another crucial factor in achieving robust performance. The results in Table 4 highlight that the explainer trained on 1.3 million ImageNet samples (FEX-1.3M) significantly outperforms the one trained on only 50,000 samples (FEX-50k). This underscores the importance of using a sufficiently large dataset to enhance the explainer's generalization and reliability.

Effect of KL-Divergence Regularization The inclusion of KL-divergence regularization enhances the generalizability of the explainer across different classes. As shown in Figure 4, the absence of KL regularization results in a trained explainer that cannot effectively distinguish between classes. Additionally, the results in Table 4 indicate that introducing the KL-divergence regularization leads to improved performance.

Figure 4: The top two predictions for this image are "golden retriever" and "Siamese cat". When λkl = 0, the explainer cannot differentiate these two classes. While when the KL regularization is introduced, it gains the ability to generalize over different classes.

Impact of g(x) Selection g(x) can be implemented as any neural network that takes x as input and outputs λ ∈ [0, 1]N , making it particularly suitable for scenarios where the prediction model f is treated as a black box. To evaluate the effect of different g(x) choices, we compared UNet [Ronneberger et al., 2015] and Vision Transformer (ViT) architectures. The results in Table 4 indicate no significant differences in performance, suggesting that the specific model structure is less critical as long as its capacity (e.g., parameter size) is adequate.

6 Conclusion

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.

Limitations

Similar to amortized methods, our framework requires training on a large and diverse dataset to achieve better quality, which may pose challenges when data privacy or data acquisition is a concern. A potential mitigation strategy is to train the explainer jointly with the predictor. This approach not only facilitates explainability for any prediction model but also ensures alignment between the explainer's domain and the predictor's application domain.

Broader Impact

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.