Replaced token detection

Overview

Replaced token detection (RTD) pre-trains an encoder in two parts. A small generator fills in masked positions with plausible tokens, and a discriminator labels every position as original or replaced 2111.09543 §3.1 eq. 3 eq. 4. Because the discriminator scores every token, it gets a training signal from the whole input, while masked language modeling (MLM) learns only from about 15% of tokens 2111.09543 §2.3.1.

The invariant all variants share is the pairing of a generator and a discriminator, with the two losses combined using a weight λ. The paper uses λ = 50, as in ELECTRA 2111.09543 §2.3.2 §3.1. The generator is half the discriminator's depth 2111.09543 §3.1.

The variants differ in how the two models share the token embedding table, and the sharing choice is where the paper's measurements matter.

Variants

Shared embeddings (ES)

Generator and discriminator use one embedding table. In the paper's 125k-step ablation, the discriminator's average embedding cosine similarity is 0.02, the same as the generator's 2111.09543 tab. 1. The authors trace this to the two losses pulling the shared vectors apart: MLM wants similar tokens close together, while RTD wants them separable 2111.09543 §3.2. ES reaches 88.8/88.4 on MNLI-m/mm and 86.3/83.5 on SQuAD v2.0 F1/EM in that setting 2111.09543 tab. 2. ELECTRA is the reference for this arrangement; no card here covers it, so its content is not used.

No sharing (NES)

The discriminator gets its own embedding table. The generator's cosine similarity rises to 0.45, but the discriminator's stays at 0.02 [arxiv:2111.09543 tab. 1]. The paper reports 88.3/87.9 on MNLI and 85.3/82.7 on SQuAD, below ES on both 2111.09543 tab. 2. Under the paper's measurements, dropping sharing removes the conflict but also loses what sharing gave.

Gradient-disentangled embedding sharing (GDES)

The discriminator reads the generator's embedding with its gradient stopped and adds a residual matrix that starts at zero and is the only part the discriminator trains 2111.09543 §3.3. The discriminator embedding is

E_D = \mathrm{sg}(E_G) + E_\Delta, \qquad E_\Delta \leftarrow 0 \text{ at init} \qquad [\S 3.3]

In the same ablation, GDES gives cosine similarity 0.45 for the generator and 0.29 for the discriminator, and reaches 89.3/89.0 on MNLI and 87.2/84.5 on SQuAD 2111.09543 tab. 1 tab. 2. On an ELECTRA-base reimplementation, the paper reports the same ordering: GDES 88.3/87.8 on MNLI against 87.9/87.4 for ES and 86.3/85.6 for NES 2111.09543 tab. 9 §6.3. Card: DeBERTaV3 (He, Gao and Chen, 2021).

When to use which

Implementation sketch

Language-agnostic pseudocode of one training iteration. The GDES variant is shown; for ES, use the generator's embedding directly in the discriminator, and for NES, give the discriminator its own table. Figure 1 of the paper draws the three gradient paths 2111.09543 fig. 1.

# One RTD iteration with GDES  [arxiv:2111.09543 §3.3] [eq. 2] [eq. 3] [eq. 4]
# Inputs: masked batch, original ids, generator G, discriminator D
#         E_G (generator embedding), E_delta (zero-initialised residual)

def rtd_iteration(batch, lam):
    # Generator forward pass and MLM loss  [arxiv:2111.09543 §3.1] [eq. 2]
    gen_logits = G(batch.masked_ids)
    mlm_loss = masked_cross_entropy(gen_logits, batch.original_ids)
    # Update the generator and E_G from the MLM loss only  [arxiv:2111.09543 §3.3]
    update(G, mlm_loss)
    # Sample a replacement at each masked position, not the argmax  [eq. 3]
    sampled = sample_from_softmax(gen_logits[batch.mask])
    corrupted = replace_at(batch.original_ids, batch.mask, sampled)
    # Discriminator embedding with stopped gradient  [arxiv:2111.09543 §3.3]
    E_D = stop_gradient(E_G) + E_delta
    # Binary label per token: replaced or original  [eq. 4]
    labels = corrupted != batch.original_ids
    rtd_loss = binary_cross_entropy(D(corrupted, E_D), labels)
    # Weight lambda = 50 in the paper; gradients reach E_delta only  [§2.3.2]
    update(D, lam * rtd_loss)

The labels mark a position as replaced only where the sampled token differs from the original, which matches the indicator in equation 4 [code: DeBERTa/apps/models/replaced_token_detection_model.py#L52-L60]. The released code's schedule runs the discriminator as a separate trainer after each generator update rtd_task.py, and its default λ is 10 (see Contradictions).

Contradictions

• Loss weight λ. The paper uses 50 2111.09543 §2.3.2 §3.1. The released code defaults to 10, and its RTD scripts do not override it [code: DeBERTa/apps/tasks/rtd_task.py#L392]. Set λ explicitly.
• Pre-training regime. The Section 3 ablations use Wikipedia and BookCorpus with batch 2,048 for 125k steps. The Section 4 full runs use 160 GB, batch 8,192 and 500k steps 2111.09543 §3.1 §4. The ablation numbers and the Section 4 numbers come from different regimes and are not comparable.
• Base learning rate. The text gives 5e-4 for base and small models; the hyperparameter table gives 6e-4 2111.09543 §4 tab. 10.
• Released scripts. The released rtd.sh uses batch 256 and a learning rate of 1e-4 for base and large, which is not the paper's Section 4 setting rtd.sh.

Papers

• DeBERTaV3 (He, Gao and Chen, 2021): replaces MLM with RTD and introduces gradient-disentangled embedding sharing (GDES).