Attention Analysis Limitations › Gradient-Based Alternatives

Because attention weights reflect mixing coefficients rather than causal importance, interpretability researchers have developed gradient-based methods that more directly measure how the output changes when inputs are perturbed. These methods approach the question differently: instead of asking "where did the model look?", they ask "which tokens, if changed, would most change the model's decision?"

The simplest gradient-based method computes how sensitive the output is to each input token embedding. Given a model output $y$ and input embedding $e_i$ for token $i$, the saliency score is:

$$s_i = \left\|\frac{\partial y}{\partial e_i}\right\|$$

where:

• $s_i$ is the saliency score for token $i$
• $y$ is the model's scalar output (e.g., the logit for the predicted class)
• $e_i$ is the embedding vector for token $i$
• $\|\cdot\|$ denotes the $\ell_2$ norm of the gradient vector

Gradient sensitivity provides the intuition: if slightly changing token $i$'s embedding substantially changes the output, then token $i$ has high causal importance. This is a direct measure of sensitivity, grounded in calculus rather than attention mechanics. If the gradient is large at a token, perturbing that token would change the output substantially, which is precisely what we mean when we say a token is "important" for the prediction.

Input gradient saliency has its own limitations. Gradients can be noisy and can saturate in regions of the input space where the model is highly confident. A gradient near zero might mean the token is unimportant, or it might mean the model is in a flat region where many perturbations produce the same prediction, not because the token lacks semantic content but because the model has already committed to a decision with high confidence. This saturation problem is particularly acute for deep networks where the loss surface has many flat regions.

Another issue with raw input gradients is instability. The gradient at a single point may not be representative of the token's overall importance, because the gradient measures local sensitivity rather than global influence. If you perturb the embedding slightly, the gradient at the new point may look quite different from the gradient at the original point.

Integrated Gradients (Sundararajan et al., 2017) addresses the saturation problem by integrating the gradient along a path from a baseline input (typically a zero embedding or a mask token) to the actual input:

$$\text{IG}_i = (e_i - e'_i) \cdot \int_0^1 \frac{\partial F(e' + \alpha(e - e'))}{\partial e_i} \, d\alpha$$

where:

• $e_i$ is the actual embedding of token $i$
• $e'_i$ is the baseline embedding (often zeros)
• $F(\cdot)$ is the model's output function
• $\alpha$ is a scalar interpolation parameter ranging from 0 to 1
• The integral accumulates gradient information across the full interpolation path from baseline to actual input

Integrated Gradients satisfies an important axiom called "completeness": the sum of attributions across all inputs equals the difference in model output between the actual input and the baseline. Formally:

$$\sum_i \text{IG}_i = F(e) - F(e')$$

This gives the attributions a consistent interpretation and prevents the method from attributing importance that does not exist. If the model's output is 0.9 for the actual input and 0.1 for the baseline, the attributions must sum to 0.8, accounting for the full difference in prediction.

The practical advantage over simple input gradients is robustness to saturation. By integrating over the path, you capture gradient information from regions where the model is less confident, including the interpolated inputs where the model is transitioning from its baseline state to its actual prediction. This avoids the zero-gradient trap where a confident model appears to be insensitive to all inputs.

The choice of baseline matters significantly for the interpretation. Using a zero embedding baseline treats the "absence" of a token as the reference point, measuring how much each token contributes relative to having no input. Using a mask token baseline treats the masked-out token as the reference point, measuring how much information each specific token provides beyond what the model would predict with unknown inputs. These give different attributions and both are defensible, depending on what question you are trying to answer.

Abnar and Zuidema (2020) proposed methods that attempt to fix attention-based attribution rather than abandon it. Their observation: because transformers stack multiple attention layers, the attention at layer $L$ is not directly connected to the input. Information flows through all intermediate layers, and examining attention at a single layer misses how that information was aggregated across the full depth of the network.

Attention Rollout computes the effective attention from output to input by multiplying attention matrices across layers:

$$A_\text{roll}^{(L \to 0)} = A^{(1)} \cdot A^{(2)} \cdots A^{(L)}$$

where $A^{(l)}$ is the attention matrix at layer $l$. Because each layer also has a residual connection, they add an identity matrix to account for the direct pass-through:

$$\hat{A}^{(l)} = 0.5 \cdot A^{(l)} + 0.5 \cdot I$$

where $I$ is the identity matrix. The rollout then uses these adjusted matrices. The resulting $A_\text{roll}$ estimates how much each input token contributes to each output position across the full model depth. This is a significant improvement over single-layer attention because it accounts for how attention patterns compose across layers.

Attention Flow further improves on Attention Rollout by treating the attention graph as a flow network and computing maximum flow rather than simple matrix products, which better handles the normalization properties of attention. In a flow network, information can travel along multiple paths, and the maximum flow captures the bottleneck capacity of each path. This is more faithful to how information propagates through the transformer, where a single token might contribute to an output through several different intermediate tokens rather than through a single direct attention path.

Both methods have their own limitations. Attention Rollout assumes that attention matrices at different layers can be validly multiplied, which requires that the value vectors at each layer have similar properties. Attention Flow is more principled but computationally expensive and harder to scale to large models. Neither method fully accounts for the role of feed-forward layers and residual connections in information storage and retrieval.

A pragmatic hybrid approach is gradient-weighted attention. Rather than using raw attention weights, you weight each attention head's output by the gradient of the loss with respect to that head's output:

$$\tilde{\alpha}_{ij}^{(h)} = \alpha_{ij}^{(h)} \cdot \left|\frac{\partial y}{\partial \alpha_{ij}^{(h)}}\right|$$

where:

• $\tilde{\alpha}_{ij}^{(h)}$ is the gradient-weighted attention from position $i$ to $j$ in head $h$
• $\alpha_{ij}^{(h)}$ is the raw attention weight
• The gradient term measures how much the loss changes when you perturb that specific attention weight

This hybrid retains the spatial structure of attention (which tokens attend to which) while incorporating gradient information about which of those attention connections matter for the prediction. It tends to outperform raw attention visualization as an explanation method because it filters out attention connections that are high-weight but causally inert, keeping only the connections where both the attention weight is substantial and the gradient confirms that connection contributes to the prediction.

The intuition behind gradient-weighted attention is similar to Grad-CAM in computer vision, which weights spatial feature maps by their gradient with respect to the class score to produce localization maps. In both cases, the spatial structure (attention weights or feature maps) is refined by gradient information to produce a more causally grounded attribution.

SHapley Additive exPlanations (SHAP) values provide yet another approach to token attribution, grounded in cooperative game theory. The Shapley value of a token measures its marginal contribution to the model's prediction, averaged over all possible subsets of tokens. Given a model $F$ and a set $S$ of tokens, the Shapley value of token $i$ is:

$$\phi_i = \sum_{S \subseteq T \setminus \{i\}} \frac{|S|! \, (|T| - |S| - 1)!}{|T|!} \left[ F(S \cup \{i\}) - F(S) \right]$$

where $T$ is the full set of tokens. This computes a weighted average of the marginal contribution of token $i$ across all possible subsets of the other tokens. The weighting ensures that the Shapley values satisfy several fairness properties, including efficiency (the values sum to the model's output) and symmetry (tokens that contribute equally receive equal values).

SHAP values are computationally expensive to compute exactly for transformer models because the number of subsets is exponential in the number of tokens. In practice, approximations are used, including KernelSHAP (which approximates Shapley values using a weighted linear model) and TreeSHAP (which is exact for tree-based models but requires approximations for neural networks). These approximations introduce their own assumptions and potential inaccuracies.

Despite the computational expense, SHAP values have attractive properties for interpretability. They are model-agnostic (the same method applies regardless of architecture), they satisfy consistency and fairness axioms, and they produce attributions that sum to the model's output. When comparing across different models or architectures, SHAP values provide a common currency for attribution that attention weights cannot provide because attention is architecture-specific.