KAIST научил диффузионные VLM не смотреть на картинку: 35% вычислений при 99,6% качества

Группа из KAIST предложила способ ускорить диффузионные мультимодальные модели, то есть нейросети, которые понимают картинки и текст и генерируют ответ не слово за словом, а сразу пачками, постепенно «раскрывая» замаскированные токены. Метод называется LOWN (Look Only When Needed, «смотри, только когда нужно»). На десяти бенчмарках он сохраняет 99,6% баллов модели LLaDA-V при 35,3% её вычислительных затрат на декодирование. Дообучать модель и подбирать пороги не требуется.

Почему модель зря пересчитывает одну и ту же картинку

У диффузионных VLM (DVLM) есть конструктивная особенность. Внимание в них двунаправленное, поэтому привычный для авторегрессионных моделей точный кэш ключей и значений напрямую не работает. Обычный инференс на каждом шаге прогоняет всю последовательность целиком: картинку, промпт и ответ. Визуальные токены занимают в ней 89–97%, а проход без них стоит всего 3–10% от полного. Получается, что большая часть вычислений уходит на повторную обработку неизменного изображения.

Авторы проверили, как часто картинка действительно нужна. Они сравнили предсказания LLaDA-V с изображением и без него на одной траектории декодирования, на всех десяти бенчмарках:

• Для 83,4% позиций, которые модель собиралась раскрыть, самый вероятный токен совпал в обоих режимах. • 85,0% итоговых токенов уже появлялись как лучшее предсказание на своей позиции до момента раскрытия. • Картинка меняла раскрываемый токен лишь на 22,8% шагов в среднем. Разброс большой: 7,9% на MMBench и 43,8% на DocVQA, где нужно читать текст со сканов.

На восьми бенчмарках из десяти расхождения сосредоточены в начале генерации, но закономерность не универсальна, например для DeCapBench. Поэтому фиксированное расписание «смотрим на картинку каждый N-й шаг» работает плохо: потребность в изображении зависит и от задачи, и от конкретного шага.

Как LOWN решает, нужна ли картинка

После каждого полного прохода (с изображением) метод сохраняет его предсказания и уверенность как «якорь». На следующих шагах запускается дешёвый проход без визуальных токенов, причём позиционные индексы сохраняются. Дальше схема такая:

• Замаскированные позиции сортируются по уверенности якоря. • Токен считается подтверждённым, если предсказание без картинки совпало с якорем. • Берётся самый длинный префикс подряд идущих подтверждённых токенов. Первое же расхождение его обрывает. • Если префикс не короче обычного числа токенов за шаг, полный проход пропускается. Принимается весь префикс, то есть за один шаг можно раскрыть больше токенов, чем у базовой модели. • Если префикс короче, запускается полный проход. Он обновляет якорь, и первыми раскрываются позиции, где предсказания с картинкой и без неё разошлись.

Шагов не больше, чем у базовой модели, а длина принимаемого префикса определяется самим входом, без заданного вручную порога. Подтверждение иногда ошибается: свежий полный проход отменяет 2,2–4,7% токенов, раскрытых на шагах без картинки. Чтобы ошибка не копилась, якорь целиком заменяется при каждом полном проходе.

Что получилось в экспериментах

Проверка шла на десяти бенчмарках: тесты с выбором ответа, понимание документов и графиков, математические задачи, подробные описания изображений. Среди ускоряющих методов с бюджетом не меньше, чем у LOWN, он лидирует на восьми из десяти. В сравнении участвовали обрезка визуальных токенов (RedVTP, ZOO-Prune), расписание раскрытия KLASS и кэш Fast-dLLM.

Важная деталь: LOWN не заменяет другие методы, а складывается с ними. Он меняет, как часто нужен полный проход, а обрезка и кэш удешевляют сам проход. Вместе с обрезкой или кэшем он снижает затраты партнёра в 1,8–2,5 раза. В паре с кэшем декодирование в 5,6–7,0 раза быстрее, чем у исходной модели. То же правило переносится на LaViDa и VidLaDA, то есть на другие DVLM и видео.

Что это даёт инженерам

Если вы гоняете диффузионные VLM на документах или чартах, главная статья расходов — повторная обработка изображения. LOWN уменьшает её без переобучения и подбора гиперпараметров под задачу, поэтому его можно подключить на этапе инференса. Ограничение стоит учитывать: выигрыш зависит от бенчмарка. Чем больше визуальных токенов в запросе и чем реже ответу нужна картинка, тем больше экономия. На задачах вроде DocVQA, где изображение нужно почти на половине шагов, она скромнее. Код и страница проекта опубликованы авторами.