rtd-gdes 0.1.4.1

ELECTRA-style RTD pretraining with Gradient-Disentangled Embedding Sharing (GDES)

GDES-DeBERTaV3

ELECTRA-style training with Gradient-Disentangled Embedding Sharing (GDES) in native PyTorch using HuggingFace Transformers.

Overview

GDES-DeBERTaV3 implements the replaced token detection (RTD) pretraining objective from ELECTRA with the gradient-disentangled embedding sharing technique introduced in DeBERTaV3. Rather than relying on a separate generator and discriminator network, GDES shares embeddings between the two while disentangling their gradient flows — enabling more parameter-efficient training without the instability of naive weight tying.

The training loop performs two forward passes per step:

1. Generator pass — predict masked tokens via MLM
2. Discriminator pass — classify each token as original or replaced, with embedding gradients frozen to disentangle the two objectives

Installation

git clone https://github.com/rsyue/rtd-gdes
cd rtd-gdes

pip install .

If you intend to contribute or modify the source, an editable install is recommended so changes are reflected immediately without reinstalling:

pip install -e .

Requirements

• Python ≥ 3.10
• PyTorch ≥ 2.0 (CUDA recommended, ROCm builds supported)
• Transformers
• Datasets
• Safetensors
• scikit-learn
• tqdm

Quick Start

python -m rtd_gdes.train \
  --model microsoft/deberta-v3-base \
  --lambda_disc 0.5 \
  --batch_size 8 \
  --epochs 5 \
  --learning_rate 2e-5 \
  --weight_decay 0.01 \
  --gamma 0.9 \
  --output_dir ./checkpoints \
  --bf16

Usage

CLI Arguments

Argument	Flag	Type	Default	Description
--model	-m	str	microsoft/deberta-v3-base	Pretrained model to train with RTD + GDES
--lambda_disc	-ld	float	0.5	Lambda coefficient scaling the discriminator loss
--batch_size	-bs	int	8	Batch size for training and evaluation
--epochs	-ep	int	5	Number of training epochs
--learning_rate	-lr	float	2e-5	Learning rate for AdamW
--weight_decay	-wd	float	0.01	Weight decay for AdamW
--gamma	-g	float	0.9	Gamma for exponential LR scheduler
--dataset		str	imdb	HuggingFace dataset name
--output_dir	-o	str	None	Directory to save the model and tokenizer. The model is saved under <output_dir>/<save_name>/. Defaults to ./<save_name>/
--fp16		flag	False	Enable FP16 mixed precision
--bf16		flag	False	Enable BF16 mixed precision
--compile	-c	flag	False	Run torch.compile with max-autotune mode

Training Details

The script trains on the IMDB unsupervised split by default, with a configurable 90/10 train/eval split. The dataset can be changed via --dataset. The combined loss is computed as:

$$\mathcal{L} = \mathcal{L}{\text{gen}} + \lambda \cdot \mathcal{L}{\text{disc}}$$

where $\mathcal{L}{\text{gen}}$ is the standard MLM cross-entropy loss and $\mathcal{L}{\text{disc}}$ is binary cross-entropy over token-level replaced/original predictions.

Evaluation reports discriminator loss, accuracy, and F1 score on the held-out set.

Saved Outputs

After training, the model and tokenizer are saved under <output_dir>/<save_name>/ where save_name is derived from the model id (e.g. deberta_v3_base_gdes). If --output_dir is not specified, the model is saved to ./<save_name>/ in the current working directory.

Development

Install with dev dependencies:

pip install -e ".[dev]"

Run the test suite:

pytest tests/ -v --cov=rtd_gdes

Project Structure

rtd-gdes/
├── src/
│   └── rtd_gdes/
│       ├── config.py          # TrainConfig dataclass — all hyperparameter defaults
│       ├── train.py           # Entry point and CLI
│       └── gdes/
│           ├── data.py        # Dataset loading and DataLoader construction
│           ├── model.py       # DebertaV3GDES — generator + discriminator
│           ├── trainer.py     # train_one_epoch and evaluate loops
│           └── utils.py       # Shared exceptions
├── tests/
│   └── test_gdes.py           # Model, trainer, and config unit + integration tests
└── pyproject.toml

Roadmap

• Distributed training (DDP / FSDP)
• Publish as PyPI package
• Support additional model architectures beyond DeBERTaV3

Citation

If you use this code, please cite the original papers:

@article{he2021debertav3,
  title={DeBERTaV3: Improving DeBERTa using ELECTRA-Style Pre-Training with Gradient-Disentangled Embedding Sharing},
  author={He, Pengcheng and Liu, Jianfeng and Gao, Jianfeng and Chen, Weizhu},
  journal={arXiv preprint arXiv:2111.09543},
  year={2021}
}

@article{clark2020electra,
  title={ELECTRA: Pre-training Text Encoders as Discriminators Rather Than Generators},
  author={Clark, Kevin and Luong, Minh-Thang and Le, Quoc V. and Manning, Christopher D.},
  journal={arXiv preprint arXiv:2003.10555},
  year={2020}
}

License

MIT