AB-Attention ускорил генерацию видео до 2,23× и пропускает 92% внимания

Команда InfRec вместе с исследователями из CUHK Shenzhen, Шаньдунского, Токийского и Фуданьского университетов, Беркли, HKUST и NUS выпустила AB-Attention. Это способ ускорить генерацию видео в диффузионных трансформерах (Wan2.1 и HunyuanVideo) без дообучения весов. Самая дорогая операция в таких моделях — внимание: каждый фрагмент видео сверяется со всеми остальными по пространству и времени, поэтому стоимость растёт квадратично с длиной последовательности. AB-Attention пропускает 89–92% этих вычислений. Видео при этом получается ближе к «честному» полному вниманию, чем у конкурентов. PSNR (метрика сходства с эталонным роликом в децибелах, чем выше, тем ближе) лучше на 1,6–4,4 дБ, а вся генерация на одной видеокарте NVIDIA H20 идёт в 1,84–2,23 раза быстрее.

Где ломаются прежние разреженные схемы

Разреженное внимание считает не все пары «запрос — ключ», а только часть. Ранние методы задавали шаблон заранее, например локальные 3D-окна. Они не знают, что важно для конкретного промпта, головы внимания или шага шумоподавления, поэтому вынуждены оставлять много лишнего. Динамические методы выбирают блоки ключей и значений по текущим активациям, но авторы насчитали у них три слабых места:

• Оценки важности грубые или дорогие. Усреднение блока в один вектор смешивает токены с разным поведением, а глобальная кластеризация (SVG2) требует k-means по всей последовательности и перестановки токенов.
• Отбор top-k отрезает хвост. Выброшенный блок пропадает и из числителя внимания, и из нормировки softmax. Лучшие оценки этого не исправят: даже идеальный отбор top-k при 10,8% ключей сохраняет лишь около 69% массы внимания.
• Единица отбора равна плитке вычислений. Блок целиком оставляют или отбрасывают, даже если нужны несколько его ключей.

Три приёма AB-Attention

Первый приём, IBCS (Intra-Block Clustering Selector), оценивает важность блоков ключей. Запросы делятся на блоки по 128 токенов. Внутри каждого блока они группируются на 8 кластеров, у каждого есть центроид и вес, равный доле запросов в кластере. Блок ключей получает оценку как взвешенное внимание центроидов к его усреднённым ключам. Кластеры используются только для оценки, сам расчёт идёт в исходном порядке токенов. Так блок, на который смотрит лишь один кластер, не теряется, как при усреднении всего блока. Разрыв в полноте отбора между усреднением и идеальным top-k сокращается вдвое.

Второй приём, SIP (Sampling with Inclusion-Probability Weighting), заменяет жёсткий отбор выборкой. Блок попадает в выборку с вероятностью, пропорциональной оценке. Самые важные блоки берутся всегда, у остальных шанс ненулевой. Каждый попавший блок умножается на обратную вероятность попадания: блок с вероятностью 25% считается как четыре. Это классическая оценка Хорвица — Томпсона, и она даёт несмещённые оценки и числителя, и нормировки softmax. Ошибка внимания получается на 41% ниже, чем у идеального top-k при том же числе блоков. Дополнительного прохода для поправки не нужно: расчёт остаётся обычным softmax-вниманием со сдвигом логитов.

Третий приём — тонкий отбор на «толстом» железе. Блоки ключей и значений всего по 4 токена, но ядро собирает выбранные блоки по индексам в полноразмерные плитки тензорных ядер. Токены не переставляются, лишних копий не создаётся. Уменьшение блока со 128 до 4 токенов поднимает долю массы внимания, удержанную идеальным отбором, с 59% до 69%, а число плиток при фиксированном бюджете не растёт. План выборки пересчитывается раз в четыре шага шумоподавления, само внимание — на каждом.

Как мерили и что вышло

Сравнивали три модели на 720p: две версии Wan2.1 (81 кадр) и HunyuanVideo (129 кадров), на 600 промптах бенчмарка Penguin Video. Опорой служило видео с полным вниманием при тех же промпте и зерне генерации. Помимо PSNR, считали SSIM и LPIPS, а ещё пять оценок самого ролика: эстетику, согласованность фона, качество изображения, плавность движения и постоянство объекта. Конкуренты: RadialAttention (статичная маска с убыванием энергии), SVG2 и DFSAttn, все на официальном коде и бэкенде FlashInfer с FlashAttention-3. Первый блок трансформера и первые 10 из 50 шагов у всех методов идут с полным вниманием.

• Разреженность: 89,16% на Wan2.1 и 92,35% на HunyuanVideo.
• Конкуренты при этом считают в 1,8–3,6 раза больше внимания.
• AB-Attention лучший по PSNR, SSIM и LPIPS на всех трёх моделях.
• Сквозное ускорение: 1,84–2,23× на одной H20 в BF16.

Что это даёт на практике

Метод не требует дообучения, поэтому его можно подключить к готовым весам. Для команд, которые запускают Wan2.1 или HunyuanVideo, это почти вдвое меньше времени на ролик при близком к эталону качестве. Выигрыш вырастет на длинных и высоких по разрешению видео, где доля внимания в задержке больше. Есть и ограничения: замеры сделаны на H20 и на трёх моделях, а ядро и план выборки рассчитаны на тензорные ядра NVIDIA. Код открыт на GitHub в репозитории InfRec-Cardinal/AB-Attention, так что результаты можно перепроверить на своём железе.