BERT for Regression predicts constant
## Environment info
<!-- You can run the command `transformers-cli env` and copy-and-paste its output below.
     Don't forget to fill out the missing fields in that output! -->

- `transformers` version: 4.3.2
- Platform: Ubuntu?
- Python version: 3.8
- Tensorflow version (GPU?): 2.4.1 (Yes)
- Using GPU in script?: Yes
- Using distributed or parallel set-up in script?: No

### Who can help

- albert, bert, xlm: @LysandreJik
- tensorflow: @jplu

## Information

Model I am using (Bert, XLNet ...): BERT.

The problem arises when using:
* [ ] the official example scripts: (give details below)
* [X] my own modified scripts:

I am doing BERT for a regression `]0-1]` in a deep (sequential) genomic data. Similar to [DNABERT](https://github.com/jerryji1993/DNABERT) and [this medium post](https://towardsdatascience.com/bringing-bert-to-the-field-how-to-predict-gene-expression-from-corn-dna-9287af91fcf8)

My code is quite simple, basically I am using BERT from the lib, averaging the outputs and passing it to a head model.

```python
class Predictor(Model):

    def __init__(
        self, batch_size: int, sequence_size: int, hidden_layers: List[int], bert_params: dict
    ):
        super().__init__()
        self._embedder = Embedder(sequence_size, bert_params)
        self._head_model = _create_head_model(batch_size, hidden_layers, bert_params["hidden_size"])

    def call(self, inputs):
        embedding = self._embedder(inputs)
        return self._head_model(embedding)


class Embedder(Model):
    def __init__(self, sequence_size: int, bert_params: dict):
        super().__init__()
        self._bert = _create_bert_model(sequence_size, bert_params)
        self.avg_pooling = GlobalAveragePooling1D()

    def call(self, sequence):
        x = self._bert(sequence)
        return self.avg_pooling(x.last_hidden_state)


def _create_bert_model(sequence_size: int, bert_params: dict) -> TFBertModel:
    tokenizer = KmerTokenizer.load()
    sequence_length = sequence_size - tokenizer.k
    config = BertConfig(
        vocab_size=tokenizer.vocab_size,
        max_position_embeddings=sequence_length,
        **bert_params,
    )
    return TFBertModel(config)


def _create_head_model(batch_size: int, n_mixtures: int, hidden_layers: List[int], input_size: int):
    embedding = Input(input_size, batch_size)
    x = embedding
    for n_neurons in hidden_layers:
        x = Dense(n_neurons, activation=nn.gelu)(x)
    output = Dense(1, x)
    return Model(embedding, output)
```

The tasks I am working on is:
* [ ] an official GLUE/SQUaD task: (give the name)
* [X] my own task or dataset: (give details below)

```
X -> Tokenized sequence of integers ([1,5,10010, 2,200, 304,1001,535,341])
y -> Float, ]0,1]
```

My Y variable is distributed like this:
![image](https://user-images.githubusercontent.com/11489228/111401808-c30d7980-86a8-11eb-86f4-21ce677ecf7d.png)

## Problem

The problem I am having is that my model is predicting constant.

Scatter (y_true x y_pred)

![image](https://user-images.githubusercontent.com/11489228/111402202-8a21d480-86a9-11eb-96fc-4679c7e81217.png)

Predictions histogram:

![image](https://user-images.githubusercontent.com/11489228/111402564-2e0b8000-86aa-11eb-9b77-31239674e377.png)


Parameters:

```yaml
    batch_size: 16
    training_steps: 20
    sequence_size: 600
    bert:
        hidden_size: 32
        num_attention_heads: 8
        num_hidden_layers: 2
    hidden_layers: [32]
    early_stopping:
        patience: 15
    optimizer:
        type: "RectifiedAdam"
        lr: 0.0004
        epsilon: 0.000001
        beta_2: 0.98
        total_steps: 100
        weight_decay: 0.01
    loss: "mean_squared_error"
```

## What I have tried

* Smaller/Larger learning rate
* Smaller/Larger batch size
* Shallower/Deeper network
* Changing Y distribution (Std Scaling)
* Mixture Density Networks

I simply cannot get through this constant prediction

Hello!

I think this more a question to address on the forum https://discuss.huggingface.co/ as it doesn't looks like to be related to a bug in the library.