how can I interpret attention weights matrix? Are they reliable?

Question

I've fine-tuned two different models (Bert and Roberta) on a dataset for a binary classification task and I'm comparing the sentences where the models predict wrong. I decided to use attention weights as explainability method to understand which tokens are contributing the most to the model's output. I have a function that visualizes the attention matrix:

def show_attention_matrix(text, model):
    
    input_ids = tokenizer(text, return_tensors="pt")["input_ids"].to(device)
    attention_mask = tokenizer(text, return_tensors="pt")["attention_mask"].to(device)

    tokens = tokenizer.convert_ids_to_tokens(input_ids.view(-1))

    attentions = model(input_ids, attention_mask)["attentions"]
    last_layer_attention = attentions[-1]
    last_layer_attention = last_layer_attention.squeeze(0)
    mean_attentions = torch.mean(last_layer_attention, dim=0)

    mean_attentions = mean_attentions.cpu().detach().numpy()

    df = pd.DataFrame(mean_attentions)

    plt.figure(figsize=(20, 10))
    heatmap = sns.heatmap(df, annot=True, cmap="viridis", fmt=".3f", cbar=True, xticklabels=tokens, yticklabels=tokens)
    heatmap.xaxis.tick_top()
    plt.show()

The code above generates a plot like the following:

Since I'm getting the embeddings from the [CLS] token and pass it to a classification head, does it make sense to also look at the weights of the [CLS] and find the tokens with the highest scores?

Answer 1

Looking at your code, I'm not too prone to agree with you on taking the mean of the various heads (I see why it makes sense, though I think you will loose some interpretability)

Indeed, as you are saying, you will later use a classifier to do classification.
Now, let's consider every possible classification independently: given an output neuron, that neuron will tell you $P(y = class_i | x)$ and thus also the probability of not being that class ($1-P(...)$), therefore you can consider it as a binary classification.
At this point, what you can do, is to consider a single neuron, check the weights of that neuron, and use them to infer the weight of each head; at that point, you can do a weighted average of the attention heads

Consider that this is "sound" if you don't do the $W_o$ final transformation fo the multihead attention layer that they use in the original paper