reparam discrete diffusion
1.0.0
Этот репозиторий содержит официальную реализацию бумаги репараметрированной дискретной диффузионной модели для генерации текста.
Кодовая база реализована с Fairseq. Чтобы установить зависимости, запустите (рекомендуется в виртуальной среде) следующие команды:
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 ..Обратите внимание, что среда проверяется с помощью Python 3.8.10, Pytorch 1.10.0/1.12.0 и CUDA 11.3. Также обратите внимание, что наша вилка Fairseq изменяет несколько файлов в исходной кодовой базе; Использование более поздних версий Fairseq может привести к неожиданным конфликтам зависимости.
Мы реализуем дискретные диффузионные модели в автономной библиотеке discrete_diffusion для общего использования. Библиотека предоставляет реализации различных типичных дискретных диффузионных моделей, состоящих из
(Vanilla/Reparameterized) multinomial diffusion : диффузионные процессы, которые вводят uniform шум в последовательность токена. Реализация ванильной многономиальной диффузии внимательно следует за кодовой базой исходной статьи;(Vanilla/Reparameterized) absorbing diffusion : процессы диффузии, где токены в последовательности могут поглощаться в состояние masking , как описано в документе D3PM. Эти диффузионные модели имеют один и тот же набор интерфейсов, позволяющих использовать внешнее использование. В частности, они определяются как подклассы класса DiscreteDiffusion , принимая следующую форму:
class DiscreteDiffusion ( nn . Module ):
"""
The parent class for discrete denoising diffusion probabilistic models.
It supports the following methods:
- q_sample()
Sample x_t ~ q(x_t | x_0) to construct noisy Transformer inputs.
- compute_losses()
Compute the loss L_t = KL(q||p) at t-th time step.
- sample_step()
Sample x_t ~ p(x_{t-1} | x_t, x_0) at t-th time step.
"""
def __init__ ( self , num_timesteps ):
super (). __init__ ()
self . num_timesteps = num_timesteps
def q_sample ( self , x_0 , t , ** kwargs ):
"""
Sample from q(x_t | x_0), which is used as the model inputs.
Args:
x_0: token ids with shape [B, N]
t: current time step, tensor with shape [B]
Returns:
return a dict of relevant outputs including x_t.
"""
def compute_losses ( self , inputs , ** kwargs ):
"""
Compute the loss objective KL(q||p) to train our generative process.
Args:
inputs: a dict that contains input types specific to different diffusion processes, containing
- x_t: token ids with shape [B, N]
- t: scalar timesteps, with shape [B]
Returns:
possibly return a dict of relevant outputs, including the loss used for training.
"""
def sample_step ( self , decoder_out , denoising_fn , ** kwargs ):
"""
Given a time step t, start from x_t and sample x_{t-k} from q(x_{t-k} | x_t).
Args:
decoder_out: a namedtuple that contains decoding info, including
- x_t: token ids with shape [B, N]
- t: scalar timesteps
- max_steps: the maximum number of decoding steps
- ...
denoising_fn: a function that takes in x_t and t and returns model logits
kwargs: other arguments that are used to control decoding.
Returns:
return a new decoder_out namedtuple.
""" Модель DiscreteDiffusion может быть создана путем настройки следующего:
--num-diffusion-timesteps <int> указывает все количество временных шагов диффузии (по умолчанию: 50)--diffusion-type <str> Указывает тип диффузионной модели (выбор: {absorbing, multinomial, reparam-absorbing, reparam-multinomial} )--noise-scheduler-type <str> указывает расписание шума только в многономической диффузии ванили/репарама (типичные варианты: {linear, cosine} ; по умолчанию: cosine )q_sample()--q-sample-mode <str> Определяет стратегию выборки (выбор: {default, coupled, multi-step, multi-sample} ; по умолчанию: default ). Мы предоставляем различные варианты для отбора проб из default : один образец рисуется как multi-step : образец два временных шага IID multi-sample : пример два образца IID coupled : также известный как кондиционированное обучение, которое подробно описано в Приложении F статьи. Это начинается с выборки двух временных шагов IID coupled выборки привносит значительные улучшения как для ванильной многономической/поглощающей диффузии, но усиление не является неизменно существенным в репараметизированных вариантах.--not-diffusing-special-sym указывает, включать ли специальные символы в процессе диффузии (по умолчанию: false)compute_losses() , включая--reweighting-type <str> тип {linear, reciprocal, none} linear--label-smoothing <float> Указывает скорость сглаживания метки (по умолчанию: 0,1)sample_step() , включая--argmax-decoding указывает, использовать ли декодирование Argmax для разэлектрического выхода трансформатора --temperature <float> указывает температуру --decoding-strategy <str> Указывает использование ванильного ( default ) / репараметизованного ( reparam-<options> ; см. Детиляцию) Стратегия декодирования (выбор: {default, reparam-<options>} ; по умолчанию: default )--load-ema-weights указывает, стоит ли загружать веса модели EMA для генерации (по умолчанию: false)--iter-decode-max-iter <int> указывает максимальное количество временных точек для декодирования (по умолчанию: 10)--iter-decode-with-beam <int> Указывает размер луча для декодирования нескольких последовательностей с различной длиной параллельно (по умолчанию: 1)--iter-decode-force-max-iter указывает, что итеративное декодирование должно запускать указанное количество итераций и не выходить. Рекомендуется установить этот флаг на True.Смотрите здесь для более полного списка аргументов.
При прохождении --decoding-strategy default схема отбора проб ванили (специфическая для каждого процесса дискретной диффузии).
Более продвинутый подход декодирования может быть вызван путем прохождения --decoding-strategy reparam-<conditioning-of-v>-<topk_mode>-<schedule> . Этот подход основан на предлагаемой репараметризации в нашей статье и позволяет обеспечить более эффективные процедуры декодирования. Параметры указывают алгоритм декодирования через
<conditioning-of-v> : uncond или cond (по умолчанию uncond ): создать переменную маршрутизации <topk_mode> : stochastic<float> или deterministic ( deterministic по умолчанию): использовать стохастический или детерминированный выбор Top- $ K $. Значение поплавки в stochastic<float> указывает степень случайности в стохастическом выборе Top-$ K $;<schedule> : linear или cosine ( cosine по умолчанию): расписание для Смотрите реализацию для более подробной информации о параметрах.
Пожалуйста, смотрите сценарии ниже для получения подробной информации.
Примечание
- Обратите внимание, что все задачи, рассмотренные в этой работе, работают на исходных данных, и не принимают дистилляцию знаний (KD).
Мы следуем стандартной предварительной обработке в Fairseq/примерах, чтобы подготовить бинаризованные данные:
# 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Мы используем данные, выпущенные в Fairseq/Примеры для подготовки набора данных:
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Для этого набора данных мы используем необработанные данные wmt16.tar.gz, как предварительно обработано в этом репозитории.
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 Сначала мы попадаем в папку fairseq , а затем запускаем следующие команды для обучения моделей.
# ####### 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/wmt 16> -s default -e True --store-ema --label-smoothing 0.1
CUDA_VISIBLE_DEVICES=1 bash experiments/mt_train.sh -m multinomial -d < iwslt/wmt14/wmt 16> -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/wmt 16> -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/wmt 16> -s default -e True --not-diffusing-special-sym --q-sample-mode coupled --store-ema --label-smoothing 0.1 --reweighting-type linearПримечание
-s <str>используется для указания названия эксперимента.- Мы могли бы передать пользовательские аргументы, которые могут быть специфичными для обучения, добавив их после
-e True.
Трубопровод оценки обрабатывается experiments/mt_generate.sh . Сценарий будет генерировать результаты перевода и оценивать оценку BLEU.
# ########## 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/wmt 16> Аргументы:
-a : Среднее ли среднее несколько контрольно -пропускных пунктов-c : Указывает местоположение контрольной точки. Если -a false (не для средних контрольных точек), пройдите путь контрольной точки; Если -a true , передайте каталог , который хранит несколько контрольно -пропускных пунктов на разных этапах обучения для усреднения.-d : имя набора данныхМы также предоставляем контрольно -пропускные пункты наших обученных моделей.
| Набор данных | Модель | Ссылка на контрольно -пропускной пункт |
|---|---|---|
| Iwslt'14 | Многономиальный | связь |
| Iwslt'14 | Поглощение | связь |
| Iwslt'14 | Репарам-мультиномиальный | связь |
| Iwslt'14 | Репарам-поглощение | связь |
| WMT'14 | Многономиальный | связь |
| WMT'14 | Поглощение | связь |
| WMT'14 | Репарам-мультиномиальный | связь |
| WMT'14 | Репарам-поглощение | связь |
| WMT'16 | Многономиальный | связь |
| WMT'16 | Поглощение | связь |
| WMT'16 | Репарам-мультиномиальный | связь |
| WMT'16 | Репарам-поглощение | связь |
Мы следуем экспериментальной установке в DiffuseQ для задач генерации вопросов и перефразирования .
Необработанные данные этих двух задач могут быть извлечены из исходного репозитория DiffuseQ. Затем мы бинаризируем данные с помощью предоставленного сценария.
# 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 Мы внимательно следим за протоколами генерации и оценки, как в DiffuseQ, чтобы обеспечить сравнение навсегда. Весь трубопровод переосмысливается в fairseq/diffusion_mt/scripts/decode_diffuseq.py и fairseq/diffusion_mt/scripts/eval_diffuseq.py соответственно, будет совместима с Fairseq. Запустите следующие команды:
# 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 > Аргументы:
-a : Среднее ли среднее несколько контрольно -пропускных пунктов-b : использовать несколько образцов для декодирования MBR-c : Указывает местоположение контрольной точки. Если -a false (не для средних контрольных точек), пройдите путь контрольной точки; Если -a true , передайте каталог , который хранит несколько контрольно -пропускных пунктов на разных этапах обучения для усреднения.-d : имя набора данныхМы также предоставляем контрольно -пропускные пункты наших обученных моделей.
| Набор данных | Модель | Ссылка на контрольно -пропускной пункт |
|---|---|---|
| QG | Многономиальный | связь |
| QG | Поглощение | связь |
| QG | Репарам-мультиномиальный | связь |
| QG | Репарам-поглощение | связь |
| QQP | Многономиальный | связь |
| QQP | Поглощение | связь |
| QQP | Репарам-мультиномиальный | связь |
| QQP | Репарам-поглощение | связь |
@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 }
}
