Pruning/Compressing heads in attention blocks
Author: NamburiSrinath
Pull request: no
Labels: name: Feature request

### Feature request

I've a conceptual question

BERT-base has a dimension of 768 for query, key and value and 12 heads (Hidden dimension=768, number of heads=12). The same is conveyed if we see the BERT-base architecture

```
(self): BertSelfAttention(
              (query): Linear(in_features=768, out_features=768, bias=True)
              (key): Linear(in_features=768, out_features=768, bias=True)
              (value): Linear(in_features=768, out_features=768, bias=True)
              (dropout): Dropout(p=0.1, inplace=False)
)
```
Now, my question is:

Can I consider the first 64 neurons from the _out_features_ as the first-head, the next 64 neurons from the _out_features_ as the 2nd head and so on? (sec 3.2.2 from original paper; [Link](https://papers.nips.cc/paper/2017/file/3f5ee243547dee91fbd053c1c4a845aa-Paper.pdf))

Basically, I am wondering if the Linear module representing query matrix; which is 768x768 can be thought as (768x64), (768x64)...12 times? 

If so, is it possible to provide some starter code as I am unable to wrap around my head. Any help is appreciated (and I've some sample in the contribution section)

**P.S:** Here's the issue from StackOverflow ([link](https://datascience.stackexchange.com/questions/124233/understanding-multi-headed-attention-from-architecture-details))

### Motivation

Example applications include papers from academia such as - [Paper1](https://lena-voita.github.io/posts/acl19_heads.html), [Paper2](https://github.com/pmichel31415/are-16-heads-really-better-than-1)

I referred to some of the previous posts ([link](https://datascience.stackexchange.com/questions/88330/how-do-the-linear-layers-in-the-attention-mechanism-work)), but I would appreciate any validation on this thought-process as it's similar but not same.

### Your contribution

Here's a code which prunes a particular % in particular layer depending on _layer_index_ and _prune_percentage_

```
model = AutoModelForMaskedLM.from_pretrained(checkpoint)

linear_layers_list = []
for name, layer in model.named_modules():
    if name in model_layers_list:
        linear_layers_list.append(layer)
print(f"No of linear layers are: {len(linear_layers_list)}")

layer = linear_layers_list[layer_index]
if prune_type == 'ln_structured':
    # Ln structured with n=1 i.e L1 pruning
    prune.ln_structured(layer, name='weight', amount=prune_percentage, dim=0, n=n)
```
I can understand that I can basically pass the Linear module and prune x% of weights.

Now, I would like to prune/remove one head in a similar fashion. 

Thanks

Comments:
Hey! Would recommend you to have a look at the source code [here](https://github.com/ArthurZucker/transformers/blob/536e339b74b494615086b4286a3d14b98b42ac21/src/transformers/models/bert/modeling_bert.py#L399). The pruning logic is already implemented, and looking at the operations performed in the attention layers should help you clarify where the heads are, where the reshaping happens etc 😉 

Thanks @ArthurZucker, it was super helpful and I was successfully able to do it for BERT :)

But when I am trying to do it for Flan-T5, it threw the following error

```
AttributeError: 'T5ForConditionalGeneration' object has no attribute 
'_prune_heads' 
```

But when I did `help(model)`, FlanT5 base class does have `prune_heads` function.

Can you verify if it's a bug in HF? Because when the `prune_heads` works for BERT, I expected it to work for other class of models as well. 

Here's the code (ignore the print statements, it works when the model is _bert-base-uncased_):

```
from transformers import AutoModelForSeq2SeqLM, AutoTokenizer
model = AutoModelForSeq2SeqLM.from_pretrained("google/flan-t5-small")
tokenizer = AutoTokenizer.from_pretrained("google/flan-t5-small")
print(model)
print("-"*50)
print(help(model))
print("-"*50)
model.prune_heads({0: [1, 2]})
print(model)
```

In retrospection, I am not sure how it can be done. Because for _encoder-only_ models, having a prune dictionary `{0:[1, 2]}` makes sense, "_pick the first 2 heads from 1st layer_"

But in Encoder-decoder models, we will have 2 types of attentions:
1. Encoder self-attention
2. Decoder self-attention  
3. Decoder cross-attention

So, this numbering system of a dictionary of lists might be difficult! 

Correct if my understanding is wrong!

No prune heads is not supported for all models 😉 

Thanks for confirming @ArthurZucker, I would like to prune particular heads in Encoder-Decoder models  -- Example applications include papers from academia such as - [Paper1](https://lena-voita.github.io/posts/acl19_heads.html), [Paper2](https://github.com/pmichel31415/are-16-heads-really-better-than-1)

Can you suggest a way to do so? For example, a boiler-plate code which does this as I am unable to wrap around my head completely --

Here's my thought

```
model = AutoModelForSeq2SeqLM.from_pretrained(checkpoint)
linear_layers_list = []
model_layers_list = [f'encoder.block.{i}.layer.0.SelfAttention.q', 'encoder.block.{i}.layer.0.SelfAttention.k', 'encoder.block.{i}.layer.0.SelfAttention.v']

for name, layer in model.named_modules():
    if name in model_layers_list:
        linear_layers_list.append(layer)
layer = linear_layers_list[layer_index]

if prune_type == 'ln_structured':
    # Ln structured with n=1 i.e L1 pruning
    prune.ln_structured(layer, name='weight', amount=prune_percentage, dim=1, n=n)
```

In the above code, I prune the particular weight matrix (q, k, v) in _block 'i'_ at _output_dim_ (dim=1) with _prune_percentage_

Now, the changes that might be needed are:

1. Instead of _l1-structured_ with _prune_percentage_, prune the head (maybe have _head_index_ and make all the values in that 0)
2. If removing via _prune_heads_ logic as in Encoder-models, we might need to reshape. Or we can have a _weight_orig_ ([Pytorch link](https://pytorch.org/tutorials/intermediate/pruning_tutorial.html#remove-pruning-re-parametrization)), which has 0s corresponding to head positions that needs to be pruned and we can simply multiply internally.

Any help would be greatly appreciated! If you feel this adds value in HF, I can try contributing to a PR

Thanks a ton :)

Hi @ArthurZucker, 

It's been a while and I thought to follow-up. The below code tells that `prune_heads` is supported by T5 Class. If so, I am getting an error as mentioned in above thread.
 
```
AttributeError: 'T5ForConditionalGeneration' object has no attribute 
'_prune_heads' 
```
Code snippet:
```
from transformers import AutoModelForSeq2SeqLM, AutoTokenizer
model = AutoModelForSeq2SeqLM.from_pretrained("google/flan-t5-small")
print(help(model))
```
Output logs:
```
class T5ForConditionalGeneration(T5PreTrainedModel)
...
prune_heads(self, heads_to_prune: Dict[int, List[int]])
 |      Prunes heads of the base model.
 |      
 |      Arguments:
 |          heads_to_prune (`Dict[int, List[int]]`):
 |              Dictionary with keys being selected layer indices (`int`) and associated values being the list of heads
 |              to prune in said layer (list of `int`). For instance {1: [0, 2], 2: [2, 3]} will prune heads 0 and 2 on
 |              layer 1 and heads 2 and 3 on layer 2.
```

If it's not supported, then:

1. Please update the documentation reflecting the same
2. Help me if my understanding on removing heads is correct: Basically I would like to know if the heads internally are like this - [(768, 64), (768, 64) ... 12 times] -> (768, 768) i.e the heads are concatenated at dim=1 so, we need to create mask at dim=1 and remove the connections to prune a particular head. (Eg numbers from BERT-base but holds true for any Attention module)

Do correct if my understanding is wrong
Thanks a ton in advance!


Same issue appx a year ago which became stale - [https://github.com/huggingface/transformers/issues/19625](https://github.com/huggingface/transformers/issues/19625)

Issue raised in Huggingface forms - [https://discuss.huggingface.co/t/t5forconditionalgeneration-object-has-no-attribute-prune-heads/16003](https://discuss.huggingface.co/t/t5forconditionalgeneration-object-has-no-attribute-prune-heads/16003)

Hi @ArthurZucker,

I worked on this and would like to know if the following code is logically correct:

```
from transformers import AutoModelForSeq2SeqLM, AutoTokenizer
model = AutoModelForSeq2SeqLM.from_pretrained("google/flan-t5-small")
tokenizer = AutoTokenizer.from_pretrained("google/flan-t5-small") 
num_heads = 8

# A function which zeroes out particular columns depending on head_index!
def zero_out_columns(layer, head_index):
    in_features = layer.in_features
    out_features = layer.out_features
    start_col = head_index * (out_features // num_heads)
    end_col = (head_index + 1) * (out_features // num_heads)
    
    # Zero out the specified columns
    layer.weight.data[:, start_col:end_col] = 0

# Select the first Attention block in encoder
qkv_layers_list = ['encoder.block.0.layer.0.SelfAttention.q',  'encoder.block.0.layer.0.SelfAttention.k',
                        'encoder.block.0.layer.0.SelfAttention.v']

# Prune the 1st head
for name, layer in model.named_modules():
    if name in qkv_layers_list:
        zero_out_columns(layer, 0)
```

The above code is expected to _prune the first head from first attention block_ of FlanT5. If that's not correct, please help me to understand it in a better way!

**P.S:** No of heads (from T5 paper which I believe is the case for FlanT5 as well):
```
Small - 8
Base - 12
large - 16
XL - 32
XXL - 128 
```

If so, I would also appreciate adding the number of heads information to documentation (or maybe I am unable to figure out!)

Would be happy to open a PR if this is useful :)

Hey @NamburiSrinath, thanks for your interest in the `transformers` codebase and for your questions. I understand that you might be working on something urgent, but please refrain from pinging as many different persons, I'm here interacting with you, and I'll try to get you through this. 

1.  T5 seems to support head pruning, with the `T5Model` but not the `T5ForConditionalGeneration`. That is not necessarly expected and I think we should support it.
2. Regarding your understanding of how the head works, how they are supposed to be pruned etc: I can serve you an answer of course, but I'll try to give you some hints first: 
   - [Here](https://github.com/huggingface/transformers/blob/06146e312a5e2d588fcb10810c9c06410d1fa078/src/transformers/pytorch_utils.py#L53) you have an example of how we prune heads. As you can see, the new layer is a bit different from the old one. Based on this, I would say that zero-ing out the weights as you are doing might not be optimal. It's also not applicable for layers that have a bias. Other than that looks alright to me. 
  

Thanks for the response @ArthurZucker, sorry about pinging multiple people, will refrain doing that :)

1. If you can help me, I can open a PR as you also believe it needs to be supported
2. Zeroing out incase bias is there can be done simply by mapping the start and end indices! But as you said, I would like to know how it's not optimal. I am making changes in-place, and I think that's not ideal for a PR. I can understand that!

If my understanding is correct, the code you shared essentially creates a copy of the W, b and inserts those values to a new Linear layer module respecting the dimension of the prune arguments given i.e we won't select the values corresponding to the prune indices while placing it to the new module!

If my understanding is correct, I am still having trouble understanding on integrating this to `transformers/modeling_utils.py` 

Hi @ArthurZucker,

Gently pinging for a follow-up.

If the solution I provided is not correct -- Suggest me how the data structure looks like!
If the solution I provided is correct but not optimal -- That's great, it resolves my task. But I would like to hear your thoughts on how to make it optimal and do a PR given there're lots of previous issues which became stale :( 

Hey, #20106 gives you an idea of how it should be implemented. Feel free to take it over 😉 
A lot of other models integrate this in the attention layer in a similar way, I recommend you to check them out! 


This issue has been automatically marked as stale because it has not had recent activity. If you think this still needs to be addressed please comment on this thread.

Please note that issues that do not follow the [contributing guidelines](https://github.com/huggingface/transformers/blob/main/CONTRIBUTING.md) are likely to be ignored.