Я им жалоту у портых и ненких мудств

Некоторое время назад вышла Gemini Diffusion — текстовая диффузионная модель, способная очень быстро генерить (достаточно бредовый) текст. Технология была, имхо, интересной, но из-за загруза на работе я так и не разобрался как оно работает. Сейчас у меня появилось свободное время, карточки и желание, так что я полез разбираться.

У нас есть BERT'ы. Это давно известная всем штука без особо сложной теории за ними — ну, трансформер, с немаскированным аттеншном, предобучается на MLM или в составе энкдека, тюнится на NLI+контрастиве, получается эмбеддер для инфопоиска. В генеративных задачах, как правило, его не используют — потому что зачем? Текст то мы читаем слева направо, он авторегрессивен по своей природе, так что отдайте это декодерам.

С другой стороны, I have a soft spot for broken and puny things. Мне очень нравится классические ML и CV из-за своей простоты, BM25 и WordLLaMA это классные экономные штуки для решения специализированных задач — в общем, классно, когда less is more. В этом смысле меня всегда задевало, что про генеративные возможности BERT'ов часто забывают. А какие у них вообще генеративные возможности?

Претрейнинг бертов, как правило, начинается с Masked Language Modeling (MLM) — подхода, когда мы маскируем небольшую часть токенов в предложении (обычно от 15%) и пытаемся их восстановить по окружающему контексту. Делается это за один пасс: мы подаем эмбеддинг каждого маскированного токена в единую для всех голову и получаем предикты.

У такого обучения есть плюсы и минусы. Из плюсов -- модель учится смотреть на все предложение сразу, так что эмбеддинг получается более информативным и эмбеддеры с полным аттеншном, как правило, лучше. Из минусов — такое обучение менее эффективно, чем обычный CLM для декодеров. Из-за того, что на каждый семпл данных получается всего 0.15 * n_tokens, а не n_tokens обновлений весов, до той же перплексии модель доходит за большее число эпох. Увы.

Но что если мы заскейлим число маскированных токенов? Например, сделаем два пасса: на первом пассе мы будем заменять какую то случайную (а не фиксированную) долю токенов в предложении на [MASK], получая последовательность все более и более закоррапченных текстов, а потом развернем последовательность и будем учить берт размаскировать эти токены?

Получится дискретная текстовая диффузия. Всё просто.

Как вывод из предыдущего текста — если у нас есть обученная на MLM модель, то мы можем очень дёшево и быстро её доучить на диффузию. Больше того, в этом райтапе (без впн не откроется 😡) показано, что даже не обязательно претрейнить на диффузию берты, которые учились на MLM: из-за того, что MLM это частный случай диффузии (со всего одним шагом денойзинга), мы можем сразу же учить модели в диалоговом режиме через SFT, зашумляя и расшумляя ответ ассистента, оставляя запрос юзера нетронутым.

Авторы того же райтапа сделали прикольную репу с трейнерами для диффузии и код для запуска, так что в качестве "магнитика на холодильник" я конвертнул три модели в диффузионные: RuModernBERT-base (150M параметров, повторение эксперимента авторов райтапа на русском), mmBERT-base (307M параметров, вдруг скейлинг поможет) и ettin-encoder-1b (1b параметров, нет русского, но посмотреть на экстремальный скейлинг было бы интересно), обучив их на равной смеси Tulu-sft (английский) и Grandmaster-Pro-Max (русский). Я ещё попробовал дообучить FRIDA, но, видимо, из-за контрастивного обучения способности к MLM она растеряла, так что модель генерила совсем уж рандомный текст.