Training the ELECTRA Model: Pretraining Workflow and Loss Functions

ELECTRA distinguishes itself through a novel two-model pretraining architecture, comprising a generator and a discriminator. The generator is trained using the Masked Language Modeling (MLM) task, while the discriminator is trained on the Replaced Token Detection (RTD) task. This document elaborates on the training process and the underlying mathematical formulations.

1. Generator Training: Masked Language Modeling (MLM)

The first stage involves training a generator model. Given an input token sequence:

x = [x1, x2, x3, ..., xn]

A random subset of token positions, denoted by M, is selected for masking:

M ⊂ {1, 2, ..., n}

The input sequence is then transformed into x_masked, where tokens at positions in M are replaced with a special [MASK] token:

x_masked[i] = [MASK] if i ∈ M else x[i]

This x_masked sequence is fed into the generator. The generator's objective is to predict the original tokens that were masked. The loss function for the generator is the standard MLM loss:

$$ L_{gen}(\theta_G) = - \sum_{i \in M} \log P_G(x_i | x_{masked}; \theta_G) $$

Where:

• $ \theta_G $: Represents the trainable parameters of the generator model.
• $ P_G(x_i | x_{masked}; \theta_G) $: Denotes the probability assigned by the generator to the original token $x_i$ at position $i$, given the masked input sequence $x_{masked}$ and the generator's parameters $\theta_G$.

Example: If $x = \text{"The cat sat on the mat"}$ and masking is applied to "cat" and "mat", then: $x_{masked} = \text{"The [MASK] sat on the [MASK]"}$ The generator learns to predict "cat" and "mat" at their respective masked positions.

2. Corrupting the Input for Discriminator Training

Following the generator's prediction, a new, "corrupted" sequence $x_{replaced}$ is constructed. For positions not masked (i.e., $i \notin M$), the original token $x_i$ is retained. For the positions that were masked (i.e., $i \in M$), the generator's predicted token is used.

x_replaced[i] = x_i if i ∉ M else G(x_masked)[i]

Where $G(x_{masked})[i]$ is the token predicted by the generator for the $i$-th masked position.

Example (Continuing from above): If the generator predicts "dog" for the first [MASK] and "rug" for the second [MASK], then: $x_{replaced} = \text{"The dog sat on the rug"}$

This $x_{replaced}$ sequence serves as the input for the discriminator.

3. Discriminator Training: Replaced Token Detection (RTD)

The discriminator's task is Replaced Token Detection (RTD). It receives the x_replaced sequence and must classify each token as either original (real) or replaced (fake). The discriminator is trained using a binary classification loss across all token positions in the x_replaced sequence.

The loss function for the discriminator is:

$$ L_{disc}(\theta_D) = - \sum_{i=1}^{n} [ t_i \log P_D(\text{real} | x_{replaced}[i]) + (1 - t_i) \log P_D(\text{fake} | x_{replaced}[i]) ] $$

Where:

• $ \theta_D $: Represents the trainable parameters of the discriminator model.
• $ t_i $: A binary indicator variable. $t_i = 1$ if the token at position $i$ in $x_{replaced}$ is an original token from the input $x$. $t_i = 0$ if the token at position $i$ in $x_{replaced}$ is a replaced token generated by the generator.
• $ P_D(\text{real} | x_{replaced}[i]) $: The discriminator's predicted probability that the token at position $i$ is original.
• $ P_D(\text{fake} | x_{replaced}[i]) $: The discriminator's predicted probability that the token at position $i$ is a replaced token.

Example (Continuing from above): For $x_{replaced} = \text{"The dog sat on the rug"}$, assuming "dog" and "rug" were originally replaced:

• "The" (pos 1, 3, 5) is original: $t_1=1, t_3=1, t_5=1$. Discriminator should predict "real".
• "dog" (pos 2) was replaced: $t_2=0$. Discriminator should predict "fake".
• "rug" (pos 6) was replaced: $t_6=0$. Discriminator should predict "fake".

The discriminator is trained to correctly identify these original and replaced tokens.

4. Combined ELECTRA Loss Function

The ELECTRA model is trained by jointly optimizing the losses of both the generator and the discriminator. The total loss is a weighted sum of the discriminator's loss and the generator's loss:

$$ L_{ELECTRA}(\theta_G, \theta_D) = L_{disc}(\theta_D) + \lambda \cdot L_{gen}(\theta_G) $$

Where:

• $ \lambda $: A scaling factor, typically set to a small value (e.g., 1.0 or 0.1). This parameter controls the influence of the generator's loss on the overall training objective.
• $ \theta_G, \theta_D $: The parameters of the generator and discriminator models, respectively.

The model is trained on a large text corpus $C$, iterating through these steps.

Summary of Pretraining Workflow

1. Mask Input: Randomly select token positions in an input sequence $x$ and replace them with [MASK] to create $x_{masked}$.
2. Generator Prediction: Feed $x_{masked}$ into the generator to obtain predictions for the masked tokens.
3. Corrupt Sequence: Construct $x_{replaced}$ by replacing the masked tokens in $x_{masked}$ with the generator's predictions.
4. Discriminator Classification: Pass $x_{replaced}$ to the discriminator, which classifies each token as either original or replaced.
5. Loss Calculation & Update: Compute the generator loss ($L_{gen}$) and discriminator loss ($L_{disc}$). Update both models' parameters ($\theta_G, \theta_D$) based on the combined ELECTRA loss function.

SEO Keywords

• ELECTRA pretraining
• Replaced Token Detection (RTD)
• Masked Language Modeling (MLM)
• Generator-Discriminator model
• NLP pretraining
• ELECTRA loss function
• Transformer pretraining
• Efficient NLP models
• Language model pretraining

Interview Questions

• Motivation for Two-Model Approach: Can you explain the primary motivation behind ELECTRA’s two-model (generator-discriminator) approach compared to traditional MLM-only pretraining methods like BERT?
• Generator's Role: Describe the role of the generator in the ELECTRA pretraining process and the specific task it is trained on.
• RTD vs. MLM: What is Replaced Token Detection (RTD), and how does it differ from the Masked Language Modeling task used by the generator?
• Constructing x_replaced: Explain how the input sequence $x_{replaced}$ is constructed and why this step is crucial for training the discriminator.
• Discriminator Loss Formulation: Walk through the mathematical formulation of the discriminator’s loss function, $L_{disc}(\theta_D)$, and explain what each component represents.
• Discriminator Training Scope: Why is the discriminator in ELECTRA trained on all tokens in the $x_{replaced}$ sequence, rather than just the masked positions? What is the advantage of this approach?
• Role of $\lambda$: What is the purpose of the scaling factor $ \lambda $ in the combined ELECTRA loss function, $L_{ELECTRA}(\theta_G, \theta_D)$? How might different values of $ \lambda $ affect the training dynamics?
• Parameter Updates: During ELECTRA pretraining, are the generator and discriminator parameters updated simultaneously or in an alternating fashion? What are the implications of this choice?
• Leveraging Generator Output: How does the “corrupting the input” step in ELECTRA leverage the generator’s output for the discriminator’s training?
• Computational Efficiency: From a computational efficiency perspective, how does ELECTRA’s pretraining approach compare to models that rely solely on MLM, particularly concerning the amount of computation per input token?