tencdm

В репозитории содержится код для воспроизведения экспериментов по генерации текста методом TEncDM (Text Encoding Diffusion Model) — диффузионной модели, работающей в пространстве выходных представлений предобученной языковой модели, а не в embedding-пространстве, как в большинстве предыдущих работ. Авторы демонстрируют, что использование таких представлений, содержащих контекстную информацию, существенно упрощает задачу денойзинга и повышает качество генерации. Ключевая особенность TEncDM — декодер, специально обученный восстанавливать текст из зашумленных латентных представлений, что позволяет компенсировать ошибки на этапах диффузии. Также авторы подробно исследуют влияние self-conditioning и scheduler’ов шума на качество модели. Предложен новый scheduler (tan-d), равномерно распределяющий сложность по всем шагам денойзинга. В экспериментах показано, что при использовании таких компонентов модель превосходит существующие SOTA подходы (DiffuSeq, AR-Diffusion и др.) на задачах перефразирования, суммаризации и упрощения текста (QQP, XSum, Wiki-Auto). Репозиторий предоставляет полный пайплайн: тренировка диффузионной модели в пространстве энкодингов, обучение декодера с corrupt-стратегией, настройка self-conditioning и различных схем шумов. Код открытый, реализован на PyTorch и включает запуск на множестве датасетов (ROCStories, Wikipedia и др.), поддерживая генерацию в условиях как с условием (conditional), так и без него. Работа может быть полезна исследователям в области генерации текста, особенно тем, кто занимается развитием диффузионных моделей, а также разработчикам, ищущим более интерпретируемые и мощные альтернативы автокорреляционным языковым моделям.

статья | код