MD4: Маскированная диффузия для дискретных данных

Маскированная (или абсорбирующая) диффузия - перспективный подход в генеративном моделировании дискретных данных, предлагающий альтернативу авторегрессионным моделям.

MD4 (Masked Discrete Diffusion for Discrete Data) (https://arxiv.org/pdf/2406.04329) - метод, разработанный в Google DeepMind предлагает упрощенный и обобщенный подход к маскированной диффузии. Структура метода позволяет обучать обобщенные модели маскированной диффузии с гибкими схемами маскировки, зависящими от состояния данных.

В основе MD4 лежит «маскирующий» процесс, превращающий исходные данные в состояние «маски» в случайный момент времени. Обращение этого процесса позволяет синтезировать новые данные, сохраняющие распределение обучающей выборки.

Математически прямой процесс описывается как марковская последовательность дискретных случайных величин, индексируемых временным параметром от 0 до 1.

MD4 продемонстрировал превосходство над диффузионными языковыми моделями по показателю перплексии на наборе данных OpenWebText и значительно обошел существующие дискретные диффузионные модели по качеству пиксельного моделирования изображений, достигая 2,75 бит на измерение для CIFAR-10 и 3,40 бит на измерение для ImageNet 64 ? 64.

Эти результаты выше, чем показатели авторегрессионных моделей сопоставимого размера (GPT-2, PixelRNN, Gated PixelCNN, PixelCNN++, PixelSNAIL, Image Transformer, Sparse Transformer).

Несмотря на все преимущества метода, MD4 склонен к переобучению, что снижает его эффективность для задач с нулевой выборкой по сравнению с более простыми моделями.

Прикладная реализация MD4 опубликована в репозитории Google Deepmind (https://github.com/google-deepmind/md4), в котором представлена возможность повторить экспериментальное обучение на тексте или изображениях.

Batch size зависит от вычислительных ресурсов. Для обучения модели MD4-S с длиной последовательности 1024, 8 GPU A100 могут поддерживать максимальный batch size=128. При запуске на TPU, 8 чипов v5litepod, batch size=32.

Локальная установка и пример обучения на тексте и изображениях:

# Create & activate env  
 python -m venv md4_venv 
 
 source md4_venv/bin/activate 
 
 # Install required packages 
 
 pip install -r requirements_gpu.txt 
 
 # Include a path dir in the Python path 
 
 export PYTHONPATH="$PYTHONPATH:~/path/to/md4" 
 
 # Prepare openwebtext for training 
 
 mkdir data_dir 
 
 python prepare_openwebtext_data.py 
 
 # Train a MD4-S model over text data 
 
 python md4/main.py —config=md4/configs/md4/openwebtext.py —sharded=false —workdir=./expt 
 
 # Train a MD4-S model over image data via cifar10 
 
 python md4/main.py —config=md4/configs/md4/cifar10.py —sharded=false —workdir=./expt

Лицензирование: Apache 2.0 License.

Arxiv (https://arxiv.org/pdf/2406.04329)

GitHub (https://github.com/google-deepmind/md4)

Источник: github.com

Комментарии: