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)

The code above generates a plot like the following:enter image description here

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?


1 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


You must log in to answer this question.

Not the answer you're looking for? Browse other questions tagged .