ronen
February 27, 2022, 6:00pm
#1
Hi,
I’m interested in applying self-attention masking in the T5 encoder, meaning a tensor of shape (batch, max_input_len, max_input_len), which is a binary matrix for each tensor in the batch, specifying which tokens (i,j) in the input can attend to each other.
This idea is explored for a different Transformer architecture in the paper “Improving Compositional Generalization in Classification Tasks via Structure Annotations”. It’s implemented in TF, the figure below is from that paper:
In contrast, I’d like to use the Transformers PyTorch T5 implementation if possible. The option seems to be available for BERT:
self.is_decoder = config.is_decoder
def transpose_for_scores(self, x):
new_x_shape = x.size()[:-1] + (self.num_attention_heads, self.attention_head_size)
x = x.view(new_x_shape)
return x.permute(0, 2, 1, 3)
def forward(
self,
hidden_states,
attention_mask=None,
head_mask=None,
encoder_hidden_states=None,
encoder_attention_mask=None,
past_key_value=None,
output_attentions=False,
):
mixed_query_layer = self.query(hidden_states)
# If this is instantiated as a cross-attention module, the keys
# and values come from an encoder; the attention mask needs to be
But for T5 I haven’t been able to find an equivalent- the encoder_attention_mask
there is of dimension (batch_size, encoder_seq_len):
# required mask seq length can be calculated via length of past
mask_seq_length = past_key_values[0][0].shape[2] + seq_length if past_key_values is not None else seq_length
if use_cache is True:
assert self.is_decoder, f"`use_cache` can only be set to `True` if {self} is used as a decoder"
if attention_mask is None:
attention_mask = torch.ones(batch_size, mask_seq_length).to(inputs_embeds.device)
if self.is_decoder and encoder_attention_mask is None and encoder_hidden_states is not None:
encoder_seq_length = encoder_hidden_states.shape[1]
encoder_attention_mask = torch.ones(
batch_size, encoder_seq_length, device=inputs_embeds.device, dtype=torch.long
)
# initialize past_key_values with `None` if past does not exist
if past_key_values is None:
past_key_values = [None] * len(self.block)
# We can provide a self-attention mask of dimensions [batch_size, from_seq_length, to_seq_length]
# ourselves in which case we just need to make it broadcastable to all heads.
extended_attention_mask = self.get_extended_attention_mask(attention_mask, input_shape, inputs_embeds.device)
Anyone have experience with this?
Thanks!