Диффузионные языковые модели получили register-токены для рассуждений без полного контекста
Маскированные диффузионные языковые модели (dLLM) генерируют текст, итеративно «расшумляя» замаскированные токены при помощи двунаправленного внимания. Чтобы рассуждение продолжалось между кусками (чанками) генерации, обычно нужно хранить в контексте весь ранее сгенерированный текст. Альберт Ге с соавторами предложили альтернативу: несколько register-токенов, токенов с фиксированной позицией, чьи непрерывные скрытые состояния специально дообучают переносить прогресс рассуждения между чанками, так что в контексте достаточно хранить только эти токены, а не весь сгенерированный ранее текст.
Модель дообучают декодировать очередной чанк текста, затем очищать его из контекста, сохраняя при этом значения register-токенов, и продолжать декодирование дальше, опираясь только на исходный запрос и перенесённое состояние. В основных сравнениях на моделях LLaDA и Dream перенос через register-токены обошёл перенос через обычный текст на каждом из тестовых бенчмарков: прирост баллов относительно базового варианта с переносом текста доходил до 8,5 на математических задачах и до 19,5 на задачах по коду.
Особенно заметен эффект на генерации кода ограниченной длины, где корректная программа обычно занимает несколько чанков подряд. Авторы также показывают, что register-токены можно дополнительно улучшать обучением с подкреплением на задачах с длинным горизонтом рассуждения.
Ключевые факты
- Register-токены, фиксированные токены с непрерывными скрытыми состояниями, которые дообучают переносить прогресс рассуждения между чанками генерации в диффузионных языковых моделях (dLLM).
- Модель декодирует чанк текста, стирает его из контекста, но сохраняет register-токены, и продолжает рассуждать дальше только по запросу и перенесённому состоянию.
- На моделях LLaDA и Dream перенос через register-токены обошёл перенос через обычный текст на каждом бенчмарке: прирост до 8,5 балла на математике и до 19,5 балла на коде.
- Сильнее всего эффект проявляется в генерации кода ограниченной длины, где верная программа занимает несколько чанков подряд.
- Register-токены можно доучивать обучением с подкреплением на задачах с длинным горизонтом рассуждения.
Почему это важно
Диффузионные языковые модели, альтернатива авторегрессионным LLM, но перенос длинного рассуждения между чанками генерации у них упирался в память: чтобы не терять нить, приходилось держать в контексте весь ранее сгенерированный текст. Register-токены сжимают этот перенос до фиксированного набора скрытых состояний вместо растущего текста, что снимает эту зависимость от размера контекста.
Кому это важно
Разработчикам и исследователям диффузионных языковых моделей, которые работают над длинными цепочками рассуждений и генерацией кода: метод напрямую адресует ограничение архитектуры dLLM, а не общую задачу LLM.
Как это применить
Метод, это способ дообучения (post-training) уже существующей dLLM: модель учат декодировать чанк, очищать текст из контекста с сохранением значений register-токенов и продолжать рассуждение по перенесённому состоянию. Это исследовательская техника, а не готовый продукт или API, интегрировать её можно только на этапе обучения/дообучения модели.
Можно ли доверять
Материал, препринт на HuggingFace Papers, результаты сравнения приводят сами авторы на моделях LLaDA и Dream; независимого воспроизведения или рецензирования в тексте не упомянуто. Конкретные названия бенчмарков и абсолютные показатели точности в источнике не приведены, есть только относительный прирост над базовым вариантом с переносом текста.
Риски и подводные камни
В источнике нет деталей о том, сколько именно register-токенов используется и как устроена процедура обучения помимо общей схемы «декодировать → очистить → продолжить», а также нет дат и институциональной принадлежности авторов. Результаты пока получены только на двух семействах диффузионных моделей (LLaDA и Dream), поэтому насколько метод переносится на другие архитектуры, не показано.