This repository contains the official implementation of paper A Reparameterized Discrete Diffusion Model for Text Generation.
The codebase is implemented with FairSeq. To install the dependencies, run (recommended in a virtual environment) the following commands:
pip install -r requirements.txt
# install our package of discrete diffusion models
pip install -e discrete_diffusion
# install our fork of fairseq
cd fairseq
python3 setup.py build develop
cd ..
Note The environment is tested with Python 3.8.10, PyTorch 1.10.0/1.12.0, and CUDA 11.3. Also note our fork of fairseq modifies several files in the original codebase; using more recent versions of fairseq might lead to unexpected dependency conflicts.
We implement discrete diffusion models in a self-contained library discrete_diffusion
for general use. The library provides implementations of various typical discrete diffusion models, consisting of
(Vanilla/Reparameterized) multinomial diffusion
: diffusion processes that inject uniform
noise to the token sequence. The implementation of vanilla multinomial diffusion closely follows the codebase of the original paper;(Vanilla/Reparameterized) absorbing diffusion
: diffusion processes where tokens within the sequence could get absorbed to the masking
state, as described in the D3PM paper.By passing --decoding-strategy default
, the vanilla sampling scheme (specific to each discrete diffusion process) is used.
A more advanced decoding approach can be invoked by passing --decoding-strategy reparam-<conditioning-of-v>-<topk_mode>-<schedule>
. This approach is based on the proposed reparameterization in our paper and allows for more effective decoding procedures. The options specify the decoding algorithm via
<conditioning-of-v>
: uncond
or cond
(default uncond
): whether to generate the routing variable $v_t$ in a conditional or unconditional manner;<topk_mode>
: stochastic<float>
or deterministic
(default deterministic
): whether to use stochastic or deterministic top-$k$ selection. The float value in stochastic<float>
specifies the degree of randomness in the stochastic top-$k$ selection;<schedule>
: linear
or cosine
(default cosine
): the schedule for $k$ during our denoising procedure, which is used to control the number of top-$k$ tokens to be denoised for the next decoding step.See the implementation for more details about the options.
Please see the scripts below for details.
Note
- Note that all tasks considered in this work operate on the original data and do not adopt Knowledge Distillation (KD).
We follow the standard pre-processing in fairseq/examples to prepare the binarized data:
# fetch and preprocess the data to BPE codes
cd examples/translation/
bash prepare-iwslt14.sh
cd ../..
# binarize the data
TEXT=examples/translation/iwslt14.tokenized.de-en
fairseq-preprocess --joined-dictionary --source-lang de --target-lang en \
--trainpref $TEXT/train --validpref $TEXT/valid --testpref $TEXT/test \
--destdir data-bin/iwslt14.tokenized.de-en \
--workers 20
We use the data released in fairseq/examples to prepare the dataset:
wget http://dl.fbaipublicfiles.com/nat/original_dataset.zip
unzip original_dataset.zip
TEXT=wmt14_ende
fairseq-preprocess --joined-dictionary \
--source-lang en --target-lang de \
--trainpref $TEXT/train.en-de --validpref $TEXT/valid.en-de --testpref $TEXT/test.en-de \
--destdir data-bin/wmt14_ende --thresholdtgt 0 --thresholdsrc 0 \
--workers 20
For this dataset, we use the raw data wmt16.tar.gz as pre-processed in this repository.
tar xzvf wmt16.tar.gz
TEXT=wmt16/en-ro
# move train/ dev/ test/ bpe codes into the $TEXT folder
mv $TEXT/train/corpus.bpe.en $TEXT/train.bpe.en
mv $TEXT/train/corpus.bpe.ro $TEXT/train.bpe.ro
mv $TEXT/dev/dev.bpe.en $TEXT/dev.bpe.en
mv $TEXT/dev/dev.bpe.ro $TEXT/dev.bpe.ro
mv $TEXT/test/test.bpe.en $TEXT/test.bpe.en
mv $TEXT/test/test.bpe.ro $TEXT/test.bpe.ro
# binarize the data
fairseq-preprocess --joined-dictionary \
--source-lang en --target-lang ro \
--trainpref $TEXT/train.bpe --validpref $TEXT/dev.bpe --testpref $TEXT/test.bpe \
--destdir data-bin/wmt16_enro --thresholdtgt 0 --thresholdsrc 0 \
--workers 20
We first get into the fairseq
folder and then run the following commands to train the models.
######## training scripts for IWSLT'14 , WMT'14, and WMT'16
# first cd to fairseq
# we use 1 GPU for IWSLT'14, 4 GPUs for WMT'14 and 2 GPUs for WMT'16 datasets respectively.
CUDA_VISIBLE_DEVICES=0 bash experiments/mt_train.sh -m absorbing -d <iwslt/wmt14/wmt16> -s default -e True --store-ema --label-smoothing 0.1
CUDA_VISIBLE_DEVICES=1 bash experiments/mt_train.sh -m multinomial -d <iwslt/wmt14/wmt16> -s default -e True --not-diffusing-special-sym --store-ema --label-smoothing 0.0
CUDA_VISIBLE_DEVICES=2 bash experiments/mt_train.sh -m reparam-absorbing -d <iwslt/wmt14/wmt16> -s default -e True --q-sample-mode coupled --store-ema --label-smoothing 0.1 --reweighting-type linear
CUDA_VISIBLE_DEVICES=3 bash experiments/mt_train.sh -m reparam-multinomial -d <iwslt/wmt14/wmt16> -s default -e True --not-diffusing-special-sym --q-sample-mode coupled --store-ema --label-smoothing 0.1 --reweighting-type linear
Note
-s <str>
is used to specify the name of the experiment.- We could pass custom arguments that might be specific to training by appending them after
-e True
.
The evaluation pipeline is handled by experiments/mt_generate.sh
. The script will generate the translation results and evaluate the BLEU score.
########### IWLS'14, WMT'14, and WMT'16 datasets
# we recommend putting each checkpoint into a separate folder
# since the script will put the decoded results into a file under the same folder of each checkpoint.
CUDA_VISIBLE_DEVICES=0 bash experiments/mt_generate.sh -a false -c <checkpoint_path> -d <iwslt/wmt14/wmt16>
Arguments:
-a
: whether to average multiple checkpoints-c
: indicates the location of the checkpoint.
If -a false
(not to average checkpoints), pass the checkpoint path;
if -a true
, pass the directory that stores multiple checkpoints at different training steps for averaging.-d
: the dataset nameWe also provide the checkpoints of our trained models.
Dataset | Model | Checkpoint link |
---|---|---|
IWSLT'14 | Multinomial | link |
IWSLT'14 | Absorbing | link |
IWSLT'14 | Reparam-multinomial | link |
IWSLT'14 | Reparam-absorbing | link |
WMT'14 | Multinomial | link |
WMT'14 | Absorbing | link |
WMT'14 | Reparam-multinomial | link |
WMT'14 | Reparam-absorbing | link |
WMT'16 | Multinomial | link |
WMT'16 | Absorbing | link |
WMT'16 | Reparam-multinomial | link |
WMT'16 | Reparam-absorbing | link |
We follow the experimental setup in DiffuSeq for question generation and paraphrasing tasks .
The raw data of these two tasks can be fetched from the original DiffuSeq repository. We then binarize the data via the provided script.
# put the raw data in the directory ``diffuseq_data/QG``
# Preprocess the question generation dataset
bash diffusion_mt/scripts/preprocess_diffuseq_datasets.sh QG
# put the raw data in the directory ``diffuseq_data/QQP``
# Preprocess the paraphrasing dataset
bash diffusion_mt/scripts/preprocess_diffuseq_datasets.sh QQP
# QQP or QG datasets
# first cd to fairseq
CUDA_VISIBLE_DEVICES=0,1 bash experiments/diffuseq_train.sh -m absorbing -d <qqp/qg> -s default -e True --store-ema --label-smoothing 0.1
CUDA_VISIBLE_DEVICES=2,3 bash experiments/diffuseq_train.sh -m multinomial -d <qqp/qg> -s default -e True --not-diffusing-special-sym --store-ema --label-smoothing 0.0
CUDA_VISIBLE_DEVICES=0,1 bash experiments/diffuseq_train.sh -m reparam-multinomial -d <qqp/qg> -s default -e True --not-diffusing-special-sym --q-sample-mode coupled --store-ema --label-smoothing 0.1 --reweighting-type linear
CUDA_VISIBLE_DEVICES=2,3 bash experiments/diffuseq_train.sh -m reparam-absorbing -d <qqp/qg> -s default -e True --q-sample-mode coupled --store-ema --label-smoothing 0.1 --reweighting-type linear
We closely follow the generation & evaluation protocols as in DiffuSeq to ensure a head-to-head comparison. The whole pipeline is re-implemented in fairseq/diffusion_mt/scripts/decode_diffuseq.py
and fairseq/diffusion_mt/scripts/eval_diffuseq.py
respectively to be compatible with Fairseq. Run the following commands:
# we recommend putting each checkpoint into a separate folder
# since the script will put the decoded results into a file under the same folder of each checkpoint.
CUDA_VISIBLE_DEVICES=0 bash experiments/diffuseq_generate.sh -a false -b true -c <checkpoint_path> -d <qqp/qg>
Arguments:
-a
: whether to average multiple checkpoints-b
: whether to use multiple samples for MBR decoding-c
: indicates the location of the checkpoint. If -a false
(not to average checkpoints), pass the checkpoint path; if -a true
, pass the directory that stores multiple checkpoints at different training steps for averaging.-d
: the dataset nameWe also provide the checkpoints of our trained models.
Dataset | Model | Checkpoint link |
---|---|---|
QG | Multinomial | link |
QG | Absorbing | link |
QG | Reparam-multinomial | link |
QG | Reparam-absorbing | link |
QQP | Multinomial | link |
QQP | Absorbing | link |
QQP | Reparam-multinomial | link |
QQP | Reparam-absorbing | link |
@article{zheng2023rdm,
title={A Reparameterized Discrete Diffusion Model for Text Generation},
author={Zheng, Lin and Yuan, Jianbo and Yu, Lei and Kong, Lingpeng},
journal={arXiv preprint arXiv:2302.05737},
year={2023}
}