Simpletransformers: Electra loss computation

[Laksh1997]

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?

[Laksh1997]

Also, wondering if there is weight sharing between the discriminator and generator's embedding layers?

[Laksh1997]

Also, the NLL loss: BCE loss ratio is 1:50 (as the NLL loss is a 30-000 way classification task, so the loss is much higher)

[ThilinaRajapakse]

I'll look into it. It's possible that I missed this.

[ThilinaRajapakse]

The ELECTRA model being used in the language_modeling_model.py file is ElectraForPreTraining. This returns the full loss of the ELECTRA pretraining objective.

[aced125]

Hi @ThilinaRajapakse which line is this on? On line 256 and 297 in language_modeling_model.py it shows ElectraForLanguageModelingModel being initialised.

Even so, ElectraForPreTraining doesn't have the full loss I believe.

[aced125]

https://github.com/huggingface/transformers/blob/master/src/transformers/modeling_electra.py#L523

The loss in ElectraForPretraining is only the discriminator BCE loss.

[ThilinaRajapakse]

Yeah, my bad. I didn't look properly. It looks like there is an issue with the Hugging Face language modeling script when using ELECTRA as well.

I think we just need to return the combined loss from ElectraForLanguageModelingModel as suggested by @Laksh1997 .

[ThilinaRajapakse]

That might not work either. :thinking:

[aced125]

Everything is fine as it is in ElectraForLanguageModelingModel.py.

All that needs to be done is total_loss = g_loss + 50 * d_loss and optimize total_loss

[aced125]

Another thing to be done: embedding layers in the generator and discriminator need to be tied

[ThilinaRajapakse]

I think #403 should fix the loss issue. However, I am not sure how to tie the embedding layers between the two models.

[aced125]

Let me write a code snippet for that

[ThilinaRajapakse]

Thanks.

[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
        )

[ThilinaRajapakse]

Great! Would you like to commit this to the PR?

[aced125]

Sure

[WP-LKL]

Good job!

Any idea as to how this change will affect Electra-model performance?

[aced125]

For backward compatibility, I will add a flag in the init to choose whether to tie gen and disc embeddings or not. Otherwise old models will give different output.

@WP-LKL in the old pretraining method, only the generator was being trained. The discriminator was not being trained. So when we used Electra discriminator on downstream tasks, it would effectively be like starting from random initialization. I am surprised it was able to get good results on the NER task in @ThilinaRajapakse 's blog post. Might be that the data distribution is easy to classify.

[ThilinaRajapakse]

Shall we put the flag in the args dict instead? Something like tie_generator_disciminator_embeddings: True. (Default True because that's the approach in the paper).

Yeah, that task was also my "test" to see if the model is working properly. :sweat_smile:

[aced125]

Raised PR: https://github.com/ThilinaRajapakse/simpletransformers/pull/404

[aced125]

I would advise to set default to true, I'll edit my PR accordingly. I imagine the flag will come through **kwargs? @ThilinaRajapakse

In the paper, they note a small improvement with tying weights:

GLUE scores are 83.6 for no weight tying, 84.3 for tying token embeddings, and 84.4 for tying all weight

The reason why they don't tie all weights is that restricts them to having the same generator and discriminator sizes (but this is not ideal as they would prefer a larger discriminator).

[ThilinaRajapakse]

Yup, default to True is my suggestion as well. Yeah, let's send it through **kwargs so that nothing else needs to change.

[aced125]

@ThilinaRajapakse reflected above changes in the latest commit

[ThilinaRajapakse]

Everything should be good now in 0.29.1.

Thank you guys for your support, @aced125 @Laksh1997 t!

[ThilinaRajapakse]

@all-contributors add @aced125 for code

[allcontributors[bot]]

@ThilinaRajapakse

I've put up a pull request to add @aced125! :tada:

[ThilinaRajapakse]

@all-contributors add @Laksh1997 for code

[allcontributors[bot]]

@ThilinaRajapakse

I've put up a pull request to add @Laksh1997! :tada:

[stale[bot]]

This issue has been automatically marked as stale because it has not had recent activity. It will be closed if no further activity occurs. Thank you for your contributions.