Electra loss computation

Hello - I looked at simpletransformers/language_modeling/language_modeling_model.py and it seems the loss computation for Electra does not take into account the discriminator loss.

On line 524 we have loss = outputs[0] whereas in line 470 of /simpletransformers/custom_models/models.py we are returning g_loss, d_loss, g_scores, d_scores, d_labels.

This seems as though only the generator loss is being optimized.

In the paper (https://arxiv.org/pdf/2003.10555.pdf) the authors combine the NLL loss from the generator and BCE loss from the discriminator (top of page 4).

Am I missing something?

[aced125]

Sure

[Laksh1997]

Hi - so huggingface already ties input and output embeddings of a language model.

All we need to do is tie the input embeddings of the generator and discriminator (the discriminator has no output embeddings).

The code is simple and straightforward:

class ElectraForLanguageModelingModel(PreTrainedModel):
    def __init__(self, generator_config, discriminator_config):
        super().__init__(discriminator_config)
        self.generator_config = generator_config
        self.discriminator_config = discriminator_config
        self.generator_model = ElectraForMaskedLM(generator_config)
        self.discriminator_model = ElectraForPreTraining(discriminator_config)
        self.vocab_size = generator_config.vocab_size
        self.tie_generator_and_discriminator_embeddings()

    def tie_generator_and_discriminator_embeddings(self):
        gen_embeddings = self.generator_model.electra.embeddings
        disc_embeddings = self.discriminator_model.electra.embeddings

        # tie word, position and token_type embeddings
        gen_embeddings.word_embeddings.weight = disc_embeddings.word_embeddings.weight
        gen_embeddings.position_embeddings.weight = (
            disc_embeddings.position_embeddings.weight
        )
        gen_embeddings.token_type_embeddings.weight = (
            disc_embeddings.token_type_embeddings.weight
        )