MC-Sparse ускоряет диффузионные трансформеры до 2,32 раза без дообучения
Разреженное внимание, один из основных способов снизить задержку диффузионных трансформеров в задачах с длинными последовательностями, таких как генерация видео и 3D-объектов высокого разрешения. Проблема в том, что при высокой степени разреженности существующие методы могут ухудшать качество и точность генерации.
Авторы с помощью контролируемых оракульных сравнений прослеживают эту деградацию до трёх источников: ограничения, которые накладывает группировка токенов; неточный выбор взаимодействий; вклады внимания, теряемые при отбрасывании токенов.
На основе этого анализа они предлагают Meta-Cached Sparse Attention (MC-Sparse), фреймворк без дообучения. Он выбирает отдельные токены ключей и значений (KV), а похожие запросы объединяет в группы, выровненные по тайлам (блокам), чтобы эффективно выполняться на GPU. Метод кэширует метаданные: группы запросов, индексы KV, выбранные по точным вероятностям внимания, и остатки между выходами плотного и разреженного внимания. Эти данные затем переиспользуются на последующих шагах шумоподавления.
По заявлению авторов, на видео- и 3D-моделях MC-Sparse точнее воспроизводит результат плотного внимания и даёт более высокое ускорение, чем существующие базовые методы разреженного внимания, без заметной деградации качества. Относительно плотного внимания метод даёт ускорение шумоподавления в 1,80 раза на Minimax-H3-Base и ускорение в 2,32 раза на генерации 3D-объектов, в обоих случаях с пренебрежимой потерей качества.
Ключевые факты
- MC-Sparse (Meta-Cached Sparse Attention), метод разреженного внимания для диффузионных трансформеров, не требующий дообучения.
- Авторы выделяют три источника потери качества у существующих методов: ограничения группировки токенов, неточный выбор взаимодействий и потерянные вклады внимания от отброшенных токенов.
- Метод выбирает отдельные KV-токены, группирует похожие запросы в блоки, выровненные по тайлам для GPU, и кэширует группы запросов, индексы KV и остатки между плотным и разреженным вниманием.
- Заявленное ускорение относительно плотного внимания: в 1,80 раза (шумоподавление) на Minimax-H3-Base и в 2,32 раза на генерации 3D-объектов, обе цифры, с пренебрежимой потерей качества.
Почему это важно
Диффузионные трансформеры в генерации видео и 3D-объектов работают с длинными последовательностями, и внимание там, главный источник задержки. Разреженное внимание ускоряет работу, но на высокой разреженности страдает качество. MC-Sparse пытается закрыть этот разрыв с плотным вниманием: авторы утверждают, что метод даёт большее ускорение и большую близость к результату плотного внимания, чем существующие разреженные базовые методы, не требуя при этом обучения.
Кому это важно
Исследователям и инженерам, которые занимаются инференсом диффузионных трансформеров для видео и 3D-генерации и ищут способы снизить задержку без переобучения модели. Также статья может быть интересна тем, кто занимается эффективными ядрами внимания для GPU: выравнивание групп запросов по тайлам заявлено именно ради эффективного исполнения.
Как это применить
Метод заявлен как не требующий обучения, то есть дообучать модель для него не нужно. Идея, которую можно взять на заметку, кэшировать между шагами шумоподавления метаданные разреженного внимания (группы запросов, выбранные индексы KV, остатки между плотным и разреженным выходами). Выпуск кода или моделей в тексте статьи не упоминается, поэтому о готовой реализации говорить рано.
Можно ли доверять
Все цифры и выводы, заявления самих авторов статьи. Значения метрик качества не приведены: есть только формулировки «без заметной деградации» и «пренебрежимая потеря качества». Существующие методы, с которыми идёт сравнение, в тексте не названы, а модель для генерации 3D-объектов не указана. Для проверки нужны полный текст статьи и независимое воспроизведение.
Риски и подводные камни
Ускорения 1,80 и 2,32 раза даны относительно плотного внимания и относятся к разным задачам, поэтому их нельзя складывать или усреднять; сквозное время генерации в тексте не указано. Не названы оборудование, длина последовательности, разрешение и уровень разреженности, при которых получены цифры, поэтому перенос результата на другие модели и конфигурации не гарантирован. Для цифры 2,32 раза в тексте сказано лишь «ускорение на генерации 3D-объектов», без слова «шумоподавление».
Компания Meta Platforms признана экстремистской организацией, её деятельность на территории РФ запрещена.