TiDAR: гибридная архитектура NVIDIA сочетает диффузию и автогрессию для ускорения инференса LLM

Почему TiDAR важен

TiDAR от NVIDIA объединяет диффузионное черновое предсказание токенов и автогрессивную верификацию в одной пропускной операции сети. Идея в том, чтобы использовать так называемые свободные слоты токенов на GPU, чтобы параллельно сдиффузировать много кандидатов, а затем в том же проходе проверить и принять часть из них автогрессивно. Это позволяет значительно увеличить пропускную способность при сохранении качества автогрессивной генерации.

Свободные слоты токенов и проблема качества

Классический автогрессивный декодер выдаёт по одному токену за шаг и часто ограничен памятью, поскольку задержка формируется загрузкой весов и KV кэша. Диффузионные модели добавляют замаскированные позиции и восстанавливают несколько токенов параллельно, но при этом токены внутри одного шага часто отбираются независимо, что ухудшает связность последовательности и точность фактов. Поэтому теоретическое преимущество по скорости часто нивелируется потерей качества.

TiDAR решает эту проблему, сохраняя эффективность диффузии и возвращая автогрессивную согласованность через этап верификации, при этом используя единый бэкенд трансформера.

Архитектура и маски внимания

На каждом шаге TiDAR делит последовательность на три области: префикс принятых токенов, область чернового предсказания предыдущего шага и маскируемая область для кандидатов следующего шага. Маска внимания устроена так, что префикс использует причинное внимание, а области чернового предсказания и маски используют двунаправленное внимание внутри блока. Это вариант Block Diffusion, где двунаправленна лишь блоковая зона декодирования.

При обучении последовательность удваивается: в причинной части находится исходный ввод со сдвинутыми метками для предсказания следующего токена, а в диффузионной части — искажённая копия. TiDAR применяет полную стратегию маскирования, заменяя все токены в диффузионной части специальным токеном маски, что даёт плотную диффузионную потерю и упрощает балансировку потерь одной весовой константой, обычно равной 1.

Однопроходная самоспекулятивная генерация

Генерация организована как самоспекулятивный процесс в одном вызове сети на шаг. На первом шаге модель причастно кодирует префикс и делает один шаг диффузии по маскируемым позициям, генерируя блок черновиков. На последующих шагах один проход одновременно проверяет черновики автогрессивными логитами с правилом отбраковки и предварительно черновит следующий блок с условием на все возможные результаты принятия текущего шага. Принятые токены добавляются в префикс и остаются в KV кэше, отклонённые токены и их записи удаляются. Поскольку чернование и верификация идут в одном бэкбоне с общей маской внимания, диффузионное вычисление использует свободные слоты токенов в том же проходе.

Модель поддерживает режимы с разной степенью доверия к диффузионным и автогрессивным предсказаниям. Для 8B варианта доверие диффузии улучшало результаты на некоторых математических задачах, при этом отбраковывание сохраняло автогрессивное качество.

Тренировка и реализация

TiDAR создавали непрерывным дообучением от Qwen2.5 1.5B и Qwen3 4B и 8B. Вариант 1.5B обучали на 50B токенов с размерами блоков 4, 8 и 16, вариант 8B — на 150B токенов с блоком 16. Использовались максимальная длина 4096, косинусный план скорости обучения, распределённый Adam в BF16 и модифицированный Megatron LM на NVIDIA H100. Оценка проводилась в lm_eval_harness по задачам кодирования, математики, рассуждений и знаний.

Результаты по качеству и пропускной способности

TiDAR 1.5B сопоставим с автогрессивным аналогом на задачах кодирования и математики, генерируя в среднем около 7.45 токенов за проход. TiDAR 8B показывает минимальную потерю качества по сравнению с Qwen3 8B и достигает около 8.25 токенов за проход. На одном H100 GPU при batch size 1 TiDAR 1.5B даёт примерно 4.71x ускорение декодирования по сравнению с Qwen2.5 1.5B, а TiDAR 8B — около 5.91x по сравнению с Qwen3 8B.

В сравнении с диффузионными моделями Dream и Llada TiDAR выигрывает и по эффективности, и по точности при условии, что диффузионные модели декодируют по одному токену за проход ради качества. В сравнении с методами спекулятивного декодирования и Block Diffusion TiDAR выгодно смотрится благодаря единому бэкбону и параллельному чернованию с верификацией.

Итоговые замечания

TiDAR показывает, что диффузионное черновое предсказание и автогрессивная проверка могут сочетаться в одной модели, эффективно используя свободные слоты токенов на GPU и повышая плотность вычислений при декодировании, при этом сохраняя точность на ключевых бенчмарках и поддерживая точную оценку правдоподобия через чисто причинную маску при валидации.