Self-attention masking for T5 encoder?

ronen · February 27, 2022, 6:00pm

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:

github.com

huggingface/transformers/blob/master/src/transformers/models/bert/modeling_bert.py#L261

      
        
                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):

github.com

huggingface/transformers/blob/84eaa6acf582206dba33135727dc3bfff05a7e9c/src/transformers/models/t5/modeling_t5.py#L945

      
        
            # 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!