Incoherence in the pooled_output definition in ElectraBackbone and BertBackbone

[mbrhd]

Describe the bug

There is an incoherence in the definition of the pooled_output in ElectraBackbone and BertBackbone vs AlbertBackbone and FNet.

• In ElectraBackbone and BertBackbone the pooled_output is defined as the pooling of the cls token before the dense layer.
• In AlbertBackbone and FNet the pooled_output is defined as the output of dense layer which takes the cls token from the sequence output.

Expected behavior

The pooled_output should have one definition or follow the original implementation.

Additional context

The original implementation of Bert, Fnet, Albert

Would you like to help us fix it?
I would like to work on this issue.

[mattdangerw]

Thanks for flagging! Agreed we should just do the same thing here across all these models.

It seems to me like the two approaches are equivalent, but the albert one is more efficient.

# This only computes the dense for the output we need.
pooled_output = keras.layers.Dense(
    hidden_dim,
    ...
)(x[:, cls_token_index, :])

# This pushes all token hidden states through the dense, then discards all but one.
# It is slower!
x = keras.layers.Dense(
    hidden_dim,
    ...
)(x)
pooled_output = x[:, cls_token_index, :]

If that makes sense to you, let's go with the albert approach everywhere. Feel free to open a PR!