This seems to be because the attention weights are averaged across all of the heads:
  

      github.com/pytorch/pytorch
  

  
    
  
	  
  

  
    
      nn.MultiHeadAttention should be able to return attention weights for each head. 9
    

    
      
        opened 
        
          
        
        Mar 10, 2020
      
      


      
        
          
          ironcadiz
        
      
    

    
        enhancement
        module: nn
        oncall: transformer/mha
        triaged
    
  


  
    ## 🚀 Feature

## Motivation
Currently when using the `nn.MultiHeadAttention` …layer, the `attn_output_weights` consists of an average of the attention weights of each head, therefore the original weights are inaccessible. That makes analysis like the one made in this [paper](https://arxiv.org/abs/1906.04341v1) very difficult.

## Pitch
When the `nn.MultiHeadAttention` forward is called with `need_weights=True` (and maybe a second parameter like `nead_attn_heads=True`),  `attn_output_weights` should be a tensor of size `[N,num_heads,L,S]`,with the weights of each head, instead of the average of size `[N,L,S]` (following the notation in the [docs](https://pytorch.org/docs/stable/nn.html#multiheadattention))

## Alternatives

## Additional context
A small discussion about this subject with a potential solution was made [here](https://discuss.pytorch.org/t/getting-nn-multiheadattention-attention-weights-for-each-head/72195)

If you guys agree, I'll gladly make a PR.