ТЕХНИЧЕСКОЕ ЗАДАНИЕ: Реализовать diffusion LLM (PLANNER)

1. Цель задачи

Разработать и обучить упрощённую версию diffusion language model с планировщиком (PLANNER), который позволяет генерировать несколько токенов за один шаг денойзинга. В отличие от классических авторегрессивных LLM, diffusion-подход итеративно превращает случайный шум в текст, а PLANNER предсказывает сразу несколько позиций (слов) в предложении, ускоряя генерацию коротких ответов.

Ключевой результат Рабочая модель, которая на коротких ответах (длиной до 128 токенов) генерирует текст в 2 раза быстрее (по latency) по сравнению с baseline-авторегрессивной моделью той же архитектуры, с сопоставимым качеством (perplexity не выше baseline на 15%).

2. Исходные данные

Если нет возможности обучить на большом датасете — симулируем:

1. Используем TinyStories (HuggingFace: roneneldan/TinyStories) — 2–3 эпохи обучения на подвыборке 50K примеров.
2. Baseline — это та же модель с заменой diffusion head на causal LM head, обученная на том же датасете (можно с тем же числом шагов).
3. Для замера ускорения используем один GPU (T4 или аналоги) и фиксированный batch_size=1, измеряем среднее время генерации 10 коротких запросов (длина ответа 32, 64, 128 токенов).

3. Технологический стек

4. Этапы выполнения

Этап 1: Исследование и архитектура (2 часа)

Действия

1. Изучить основные работы по дискретному diffusion для текста (D3PM, MDLM, DiffusionBERT). Выбрать простой подход: multinomial diffusion (каждый токен — one-hot вектор, шум — uniform по словарю).

2. Определить формат PLANNER: вместо предсказания всего предложения за один шаг — модель предсказывает маску для нескольких позиций (например, 25% токенов за шаг). На каждом шаге денойзинга планировщик выбирает случайное подмножество позиций и обновляет их, остальные остаются зашумлёнными.

3. Зафиксировать гиперпараметры:

   Параметр	Значение
   Число шагов денойзинга (T)	32
   Размер словаря	30522 (как у BERT)
   Макс. длина ответа	128
   Планировщик: доля обновляемых позиций за шаг	0.25
   Оптимизатор	AdamW (lr=5e-5)
   Batch size	16 (на GPU 16GB)

Ожидаемый результат этапа Документ с обоснованием выбора архитектуры, диаграмма процесса генерации (шум → итеративное денойзинг с PLANNER). Код пока не пишется.

Этап 2: Реализация процесса diffusion (3 часа)

Действия

1. Реализовать forward_noising — превращение чистого текста в последовательность шумовых векторов за T шагов (диффузионный процесс):

   • На входе: one-hot токены (batch, seq_len, vocab_size)
   • На шаге t: смешиваем исходный вектор с uniform-шумом с интенсивностью β(t) (cosine schedule)
   • Выход: зашумлённое распределение на каждом шаге t.

2. Реализовать reverse_sampling — обратный процесс:

   • Начинаем с чистого шума (uniform) на шаге T.
   • Для t от T-1 до 0:
     • Модель получает текущее зашумлённое состояние.
     • PLANNER выбирает маску (какие позиции обновлять).
     • Обновляем только выбранные позиции согласно предсказанию модели.
   • Возвращаем финальный текст (argmax по последнему шагу).

3. Написать тестовая проверка на одном предложении: зашумление → денойзинг → проверка, что восстановленный текст близок к исходному.

Ожидаемый результат этапа Модули diffusion_utils.py с функциями add_noise и sample_step. Проверочный скрипт, который показывает, что при T=32 и правильных предсказаниях текст восстанавливается.

Этап 3: Реализация PLANNER (4 часа)

Действия

1. Добавить к BERT дополнительную голову — планировщик:

   • Это MLP + линейный слой, который из скрытых состояний BERT предсказывает маску выбора (binary mask, sigmoid) + новые токены (softmax по словарю).
   • На вход: скрытые состояния после последнего слоя BERT (batch, seq_len, hidden_size).
   • На выход: logits_mask (batch, seq_len, 1) и logits_token (batch, seq_len, vocab_size).

2. Реализовать loss-функцию:

   • Для каждой позиции, если она выбрана маской (ground truth маска — те позиции, которые мы знаем, что нужно обновить), применяется cross-entropy loss к предсказанному токену.
   • Если не выбрана — loss только за то, что маска предсказала не обновлять (BCE).
   • Суммарный loss: L_total = L_ce + λ * L_mask, где λ=0.1.

3. Имплементировать предиктор маски во время обучения: сначала модель предсказывает маску, затем мы применяем её к ground truth токенам (через teacher forcing) — на каждом шаге мы знаем, какие позиции должны быть обновлены (например, самые зашумлённые).

Ожидаемый результат этапа Класс DiffusionBERT с forward методом, который принимает (input_ids, noise_level_t) и возвращает (logits_mask, logits_token). Проход обучения для одного батча работает без ошибок.

Этап 4: Обучение и оценка (5 часов)

Действия

1. Подготовить датасет:

   • Загрузить TinyStories, отфильтровать примеры длиной ≤128 токенов.
   • Токенизировать BERT-токенизатором (add_special_tokens=True).
   • Создать DataLoader с padding до max_length и attention mask.

2. Реализовать цикл обучения (1000 шагов, валидация каждые 100 шагов):

   • Для каждого батча:
     • Выбрать случайный шаг t (0..T-1).
     • Применить шум к batch (get noisy_embedding, mask_gt — какие токены исходные, какие шумовые).
     • Прогнать модель, получить потери.
     • Backprop.
   • На валидации: измерять perplexity на зашумлённых данных (по предсказанию токенов).

3. После обучения зафиксировать модель и сохранить чекпоинт.

Ожидаемый результат этапа Обученная модель diffusion_bert_checkpoint.pt. График loss и perplexity на валидации.

Этап 5: Инференс и замер ускорения (3 часа)

Действия

1. Реализовать функцию generate_diffusion(prompt, max_length=128):

   • prompt токенизируется и паддится до max_length.
   • Инициализируем шум (uniform) по всем позициям.
   • Запускаем обратный процесс с PLANNER (T=32 шага).
   • Замеряем время генерации.

2. Реализовать baseline — авторегрессивная генерация с той же BERT-base (causal LM head):

   • Использовать transformers.AutoModelForCausalLM (можно дообучить на том же датасете).
   • Замеряем время для такой же длины ответа.

3. Сравнить latency для трёх длин: 32, 64, 128 токенов (по 20 повторений на каждую).

   • Вычислить среднее время и ускорение (time_baseline / time_diffusion).

4. Оценить качество сгенерированных ответов:

   • Perplexity по модели (встроенный loss).
   • BLEU-1, BLEU-2 на 100 примерах (если есть референсные ответы).
   • Ручная оценка осмысленности 10 ответов.

Ожидаемый результат этапа Таблица ускорения, графики, сгенерированные примеры. Ускорение 2x на short ответах — подтверждено.

5. Критерии приемки (Definition of Done)

• Реализован полный пайплайн: noising, sampling, обучение, инференс.
• Обученная diffusion-модель способна генерировать текст длиной до 128 токенов.
• Замеры latency показывают ускорение ≥2× для ответов 32-64 токена (сравнение с baseline).
• Perplexity на валидации не превышает baseline более чем на 15%.
• Код воспроизводим: README с инструкциями, requirements.txt, все скрипты в папке.
• Сгенерированные тексты визуально осмысленны (не полностью случайны).
• Есть отчёт в виде Jupyter Notebook с графиками и выводами.

6. Ожидаемый результат

Файлы и артефакты

Дополнительные результаты (опционально):

• Демо в Streamlit для интерактивной генерации.
• Сравнение с другими diffusion-подходами (MDLM, D3PM).
• Анализ влияния числа шагов T на скорость/качество.

7. Возможные сложности и их решение

8. Бюджет времени (оценка)

Примечание: Время указано для опытного ML-инженера. Первый раз может потребоваться +50% времени на отладку.

9. Связанные вопросы из базы знаний

10. Чек-лист самопроверки

• Я реализовал оба процесса (noising и sampling) и проверил на одном примере, что текст восстанавливается (хотя бы частично).
• Я обучил модель хотя бы 500 шагов и убедился, что loss уменьшается.
• Я измерил latency baseline и diffusion на одинаковых аппаратных условиях (GPU, batch_size=1).
• Я сравнил качество через perplexity на валидации — разница ≤15%.
• Я написал README с инструкцией по запуску и требованием к среде.
• Я зафиксировал seed (например, 42) для воспроизводимости.