Decrease Longformer window size / computational cost
Author: iliaschalkidis
Pull request: no

Hi there,

I would like to use Longformer instead of BERT or ROBERTA for longer documents, e.g., 1024 subword units. My goal is to fit a batch of equal size in the same GPU card for all models. In my understanding, this cannot happen with the default configuration, which used windows of 512 subword-units for local attention. In other words, this is by default more computationally expensive than running BERT or ROBERTA. So I thought the solution would be to decrease the size of window proportionally to the increase of input sequence size. This will lead to equal or less computations.  I run the following experiments in a single RTX 2080Ti:

**Train ROBERTA with `batch_size=6` and `max_len=512` (SUCCESS)**

```python
from transformers import TFLongformerModel, LongformerConfig, TFRobertaModel
import tensorflow as tf
import numpy as np
import logging
logging.getLogger("tensorflow").setLevel(logging.ERROR)
logging.getLogger("transformers").setLevel(logging.ERROR)


class Classifier(tf.keras.Model):

    def __init__(self, bert_encoder, *args, **kwargs):
        super(Classifier, self).__init__(*args, **kwargs)
        self.classifier = tf.keras.layers.Dense(2)
        self.bert_encoder =  bert_encoder

    def call(self, inputs):
        bert_encodings = self.bert_encoder(inputs)
        return self.classifier(tf.squeeze(bert_encodings[0][:, 0:1, :], axis=1))

# Train ROBERTA for 512 TOKENS

roberta = TFRobertaModel.from_pretrained('roberta-base')
roberta_classifier = Classifier(bert_encoder=roberta)
dummy_inputs = np.zeros((6, 512), dtype=np.int32)
dummy_outputs =  np.zeros((6, 2), dtype=np.int32)

roberta_classifier.compile(optimizer='adam', loss='categorical_crossentropy')
roberta_classifier.fit(dummy_inputs, dummy_outputs, batch_size=8)
print('Roberta (512) trained successfully!')
```

**Train LONGFORMER with `batch_size=6` and `max_len=512` and `attention_window=512` (SUCCESS)**

```python
# Train LONG-FORMER for 512 TOKENS
config =  LongformerConfig.from_pretrained('allenai/longformer-base-4096')
config.attention_window = [512] * 12
longformer = TFLongformerModel(config)
longformer_classifier = Classifier(bert_encoder=longformer)
dummy_inputs = np.zeros((6, 512), dtype=np.int32)
dummy_outputs =  np.zeros((6, 2), dtype=np.int32)

longformer_classifier.compile(optimizer='adam', loss='categorical_crossentropy')
longformer_classifier.fit(dummy_inputs, dummy_outputs, batch_size=8)
print('Longformer (512) trained successfully!')
```

**Train LONGFORMER with `batch_size=6` and `max_len=1024` and `attention_window=128` (FAILED-OOM)**

```python
# Train LONG-FORMER for 1024 TOKENS
config =  LongformerConfig.from_pretrained('allenai/longformer-base-4096')
config.attention_window = [128] * 12
longformer = TFLongformerModel(config)
longformer_classifier = Classifier(bert_encoder=longformer)
dummy_inputs = np.zeros((6, 1024), dtype=np.int32)
dummy_outputs =  np.zeros((6, 2), dtype=np.int32)

longformer_classifier.compile(optimizer='adam', loss='categorical_crossentropy')
longformer_classifier.fit(dummy_inputs, dummy_outputs, batch_size=8)
print('Longformer (1024) trained successfully!')
```

The last script is failing with an OOM issue. The same happens for  `attention_window in [32,64]`

@patrickvonplaten  do I miss something? Thanks!

Comments:
Hey @iliaschalkidis,

thanks for your issue! The memory usage in Longformer does not decrease linearly when reducing the attention_window...but I'm a bit surprised that you are experiencing OOM in your set-up...Does the same happen for your in eager mode? I'll try to look into it a bit next week. One thing that would be of great help is if you find time to benchmark the memory usage of `TFLongformer`, for:

- eager mode
- compiled

for different settings of the window size

Hi @patrickvonplaten, how do you recommend to perform this benchmarking? Any suggestion (best practice/ tool)? I also had the impression that TF2 is in eager execution by default...

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/master/CONTRIBUTING.md) are likely to be ignored.