Sol-Attn ускоряет генерацию видео нейросетями за счёт разреженного внимания на лету
Диффузионные трансформеры (diffusion transformers), основа современной генерации видео высокого качества, но при работе с длинными последовательностями токенов именно механизм внимания становится главным узким местом по скорости вывода. Один из способов ускорения без дообучения модели, динамическое разреженное внимание: вместо всех пар токенов считаются только выбранные блоки ключей и значений (key-value). Но у существующих методов такого рода, по словам авторов, две проблемы. Во-первых, маршрутизация блоков либо жёсткая, либо непредсказуемая: либо берут фиксированную долю блоков с лучшей прокси-оценкой (фиксированный бюджет вычислений), либо набирают блоки до достижения целевой суммарной вероятностной массы (бюджет плавающий, но может получаться несбалансированным), и в обоих случаях сам расчёт и хранение прокси-оценок создаёт заметные накладные расходы. Во-вторых, невыбранные блоки просто отбрасываются целиком, что при высокой степени разреженности заметно портит точность.
Авторы предлагают Sol-Attn (Sparsifying online attention), метод без дообучения, который объединяет динамическую маршрутизацию блоков, разреженные вычисления и коррекцию приближения в одном проходе online softmax (онлайн-вычисления softmax). Ядро метода, пороговая фильтрация блоков «на лету» с повторным использованием прокси-оценок: критические блоки отбираются сравнением их прокси-оценки с порогом прямо во время прохода online softmax. Это даёт динамический, но управляемый бюджет вычислений без необходимости заранее строить полную карту прокси-оценок для всех блоков. Кроме того, прокси-оценки невыбранных блоков не выбрасываются, а используются повторно, чтобы приближённо учесть их вклад в результат, вместо полного отбрасывания, как в прежних методах.
В экспериментах на задачах генерации изображений и видео Sol-Attn показал лучшее соотношение качества и скорости среди методов разреженного внимания без дообучения: ускорение работы «из конца в конец» в 2,1 раза для генерации видео и в 2,3 раза для редактирования видео при сохранении визуального качества.
Ключевые факты
- Sol-Attn, новый метод разреженного внимания для диффузионных трансформеров в задачах генерации видео, не требующий дообучения модели
- Решает две проблемы прежних методов: дорогую и негибкую маршрутизацию блоков внимания и потерю точности из-за полного отбрасывания невыбранных блоков
- Ключевая идея, пороговая фильтрация блоков внутри одного прохода online softmax с повторным использованием прокси-оценок невыбранных блоков
- На задачах генерации и редактирования видео метод даёт ускорение в 2,1 и 2,3 раза соответственно при сохранении визуального качества
- Первый автор работы, Хаопэн Ли (Haopeng Li)
Почему это важно
Генерация видео диффузионными трансформерами требует обработки очень длинных последовательностей токенов, и механизм внимания, самая тяжёлая по вычислениям часть вывода. Ускорить его без переобучения модели, практическая задача для всех, кто считает вывод дорогим. Sol-Attn предлагает способ, который одновременно снижает вычислительные накладные расходы на маршрутизацию блоков и не жертвует точностью так сильно, как прежние варианты разреженного внимания.
Кому это важно
Инженерам и исследователям, которые занимаются инференсом диффузионных моделей для видео, генерацией и редактированием, а также компаниям, для которых стоимость и скорость вывода видео-моделей, прямая статья расходов на инфраструктуру.
Как это применить
Метод не требует дообучения и может применяться к уже готовым диффузионным трансформерам как оптимизация вывода: пороговая фильтрация блоков внимания встраивается прямо в проход online softmax. По данным экспериментов авторов, это даёт ускорение в 2,1 раза для генерации видео и в 2,3 раза для его редактирования без заметной потери визуального качества.
Можно ли доверять
Это исследовательская работа, опубликованная как препринт на Hugging Face Papers; в источнике доступен только текст аннотации без деталей о выборке бенчмарков, конкретных базовых моделях и статистической значимости результатов. На момент сбора у публикации скромный отклик, 28 очков и 3 комментария, независимой проверки результатов пока нет.
Риски и подводные камни
Заявленные цифры ускорения получены авторами метода на собственных экспериментах, независимого воспроизведения результатов нет. Из доступного текста неясно, опубликован ли код, на каких именно моделях и датасетах проверялся метод и насколько устойчиво приближение прокси-оценок ведёт себя при экстремальной разреженности внимания.