cosmos

В репозитории опубликован код COSMOS — метод диффузионной генерации текста в сжатом и гладком латентном пространстве, представленный на конференции NeurIPS в 2025 году. Авторы переносят идею latent diffusion из генерации изображений в текст, где прямое моделирование токенов остаётся дорогим из-за высокой размерности и длины последовательностей. COSMOS сначала кодирует текст замороженным BERT, затем сжимает полученные представления через лёгкий compressor и обучает diffusion model уже в этом компактном пространстве. Важная часть работы не только в самом сжатии, а в том, как это пространство делается пригодным для диффузии. Автоэнкодер обучается одновременно на восстановление токенов, приближение к исходным BERT-представлениям, устойчивость к шуму, маскированию активаций и частичному занулению признаков внутри latent vectors. Такая схема делает пространство более гладким и уменьшает разрыв между тем, какие представления модель видит при обучении, и тем, какие появляются во время генерации. В экспериментах авторы показывают, что текст можно сжимать до 8 раз без заметной потери качества относительно token-level diffusion, а при умеренном сжатии COSMOS превосходит ряд diffusion-based и autoregressive baselines. Метод проверяется на ROCStories, XSum, ParaDetox и SQuAD2.0, а также на OpenWebText для более длинной генерации. На ROCStories версия COSMOS с большим числом latent vectors достигает MAUVE 0.940 при 0.953 у исходных текстов и заметно обходит GPT-2 по этой метрике. Для длинных последовательностей подход даёт особенно сильный выигрыш по скорости, так как диффузия работает с фиксированным числом шагов и не генерирует текст строго токен за токеном. Работа будет полезна исследователям генерации текста, text diffusion, latent diffusion и всем, кто изучает альтернативы авторегрессионным языковым моделям для быстрой и управляемой генерации.

статья | код