如何将 2D 注意力掩码传递给 HuggingFace BertModel？

Question

I would like to pass a directional attention mask to BertModel.forward , so that I can control which surrounding tokens each token can see during self-attention. This matrix would have to be 2D.

Here is an example with three input ids, where the first two tokens cannot attend to the last one. But the last one can attend to all tokens.

torch.tensor([
 [1, 1, 1]
 [1, 1, 1]
 [0, 0, 1]
])

Unfortunately, the documentation does not mention anything about supporting 2D attention masks (or rather 3D with batch dimension). It's possible to pass a 3D attention mask, but in my experiments performance of the model did not change much, regardless of how the mask looked like.

Is this possible, if so how?

Answer 1

If you can provide more details it will more clear. Any way this is my initila answer, to make things simple track in the implementation where the model uses the mask. For example in this line : As you notice if you tracked the expand function here then you will find that this is your case in this line At this point you can decide what do you need later.