Генерация изображения диффузионной моделью устроена расточительно. Десятки шагов денойзинга, и на каждом шаге трансформер заново обрабатывает тысячи токенов, хотя между соседними шагами большая их часть почти не меняется. Вопрос «какие токены можно объединить, а какие обязаны остаться нетронутыми» вынесен в заголовок свежей статьи на arXiv (2610.22370). Ответ авторы ищут не в тексте промпта, а в динамике самого процесса диффузии. Если сжатие токенов сработает, подешевеет каждый шаг генерации, а вместе с ним и каждая сгенерированная картинка.

Объединение токенов (token merging) это training-free техника ускорения: похожие токены группируются, их представления сливаются в один общий токен, и дальше слои считают меньше элементов. Веса модели не трогают, ничего не дообучают, меняется только то, как выполняются вычисления на этапе генерации. Работает это в диффузионных трансформерах (DiT), на которых построены PixArt-α и FLUX, и в классических U-Net-моделях вроде Stable Diffusion.

Почему объединять токены сложно

Первый популярный метод, ToMeSD из 2023 года, группировал токены просто по схожести признаков, без учёта содержания. В объединение попадали важные детали, и изображение «разъезжалось»: локальные структуры растворялись в соседях. Следующий шаг сделал ставку на classifier-free guidance: если токен сильно управляется текстом промпта, его защищают как якорь. Сигнал достаётся бесплатно, из уже существующих прогонов модели на каждом шаге. Параллельно развивались методы вроде ToMA, которые берут важность токенов из карт внимания: они точнее простой схожести признаков, но требуют дополнительных вычислений на каждом шаге.

Логика рабочая, но с перекосом. CFG измеряет только связь с промптом, поэтому всё, что в промпте не упомянуто, будь то фон, тени или мелкие детали на периферии, остаётся без защиты. Авторы DaTMe называют это prompt-centric bias и показывают цену ошибки в цифрах.

На PixArt-α при компрессии 50% (r = 0.5) у ToMeSD метрика FID ухудшается с 42.1 у несжатой модели до 94.9, у CFG-метода до 52.3. Для тех, кто не следит за метриками: FID измеряет расстояние между распределением сгенерированных и реальных изображений, и рост более чем вдвое означает заметную деградацию. DaTMe в тех же условиях держит 43.4.

Tweedie-оценка как источник сигнала важности

Идея авторов: важность токена надо измерять по тому, как он ведёт себя в процессе генерации. В основе метода Tweedie-оценка: на каждом шаге денойзинга модель может восстановить приблизительное «чистое» изображение из текущего зашумлённого. Формула известна как формула Твиди и появилась в статистике задолго до диффузионных моделей, а в DaTMe она стала основой для оценки важности токенов.

Из оценки извлекаются два сигнала. Временной: если предсказание токена сильно меняется от шага к шагу, модель ещё дорабатывает эту область, и объединять её нельзя. Пространственный: если токен содержит высокочастотные детали, границы и текстуры, он несёт структуру. Безопасно объединять только те токены, которые одновременно стабильны во времени и однородны в пространстве. Токены, которые «застыли» и ничего не держат, можно смело сливать с соседями.

Напрашивался и более простой сигнал, наклон score-функции, но авторы показывают в приложении, почему он не годится. Наклон раскладывается на два слагаемых: разницу предсказаний модели между шагами и масштабный член, который целиком задаётся глобальным расписанием шума. На поздних шагах второй член доминирует, и наклон начинает измерять шумовое расписание вместо изменений картинки. Аблация подтверждает: замена Tweedie-разности на score slope ухудшает SSIM с 0.811 до 0.803 и FID с 42.4 до 43.0.

Как работает DaTMe

Метод получил имя DaTMe, от dynamics-aware token merging, объединение токенов с учётом динамики. На каждом шаге он строит карту важности и распределяет токены по ролям. Пересечение верхних долей временной и пространственной карт даёт набор якорей, которые сохраняются всегда. Пересечение нижних долей даёт безопасные источники, кандидатов на объединение. Между этими двумя наборами остаётся серая зона: её добирают до нужного бюджета компрессии равномерно по сетке, чтобы источники не собирались в одном углу кадра. Роли пересчитываются на каждом шаге заново: токен, который на ранних шагах был важным, позже может стать безопасным кандидатом на объединение.

Финальный шаг, локально ограниченное сопоставление: объединяются только соседние токены, поэтому «склеить» далёкие области кадра невозможно. Всё вместе это training-free надстройка над inference: на PixArt-α её применяют к средним слоям self-attention в DiT, на FLUX.1-dev ещё и к блокам joint-attention ближе к выходу. Бюджет задаётся вручную: объединяется 30, 50 или 70% токенов, и от этого зависит баланс между экономией и качеством.

Что показали эксперименты

Эксперименты шли на PixArt-α, где использовали 10 тысяч пар «изображение-промпт» из валидации COCO 2014, и на FLUX.1-dev, а в приложении ещё на SD 1.5 и SDXL. Всё на одной GPU RTX PRO 6000 Blackwell. Метрики стандартные: CLIP-T и HPSv2 за следование промпту, FID за реализм, SSIM, PSNR и LPIPS за верность исходной генерации, плюс латентность и TFLOPs.

Главный результат на PixArt-α: при компрессии 30% DaTMe почти неотличим от несжатой модели. FID 42.4 против 42.1 у базовой, CLIP-T 31.14 против 31.25, HPSv2 28.6 против 29.3. Конкуренты в тех же условиях теряют заметно больше:

Метод (компрессия 30%)	FID, меньше лучше	SSIM, больше лучше	Латентность, с
Без компрессии	42.1	н/д	1.94
ToMeSD	61.5	0.731	1.69
CFG-метод	44.6	0.714	1.75
DaTMe	42.4	0.811	1.80

Дальше разрыв только растёт: при компрессии 70% у ToMeSD FID взлетает до 123, у CFG до 70.7, у DaTMe до 47.1. При таком сжатии конкуренты фактически теряют изображение, а DaTMe всего лишь слегка его упрощает.

На FLUX.1-dev картина по скорости скромнее, потому что объединяются не все блоки: латентность падает с 14.6 до 13.8 секунды, TFLOPs с 2098 до 2033. Зато по верности исходной генерации DaTMe впереди: SSIM 0.774 против 0.741 у ToMA и 0.640 у ToMeSD, PSNR 21.6 против 20.6 и 17.5. По FID методы близки, разброс в пределах половины пункта.

На U-Net-архитектурах метод тоже проверяли. На SD 1.5 DaTMe оказывается даже быстрее несжатой модели: 0.92 секунды против 1.02 при компрессии 30%, при этом SSIM выше, чем у конкурентов. На SDXL латентность падает с 5.08 до 4.25 секунды. Экономия в пределах 10-16%, и это честная цена за то, что метод трогает только часть слоёв. По эстетической оценке картинок (предиктор Aesthetic V2.5) DaTMe тоже впереди на всех уровнях компрессии: 6.20 против 5.99 у ToMeSD и 6.15 у CFG при 30%.

Аблации и человеческие оценки

Аблации показывают вклад каждой части. Если взять пайплайн конкурента и заменить только сигнал важности на карту DaTMe, метрики растут: SSIM с 0.714 до 0.764, PSNR с 17.8 до 20.1. Оба сигнала нужны вместе: только временной даёт FID 42.5, только пространственный 42.6, вместе 42.4. Без якорей, когда набор защищённых токенов пуст, SSIM падает с 0.811 до 0.804, а FID растёт с 42.4 до 43.8.

Параметры подобраны как компромисс. Сетка из четырёх ячеек лучше всего сохраняет структуру, а более плотная сетка уже замедляет генерацию: при 64 ячейках метод обгоняет по времени несжатую модель, 2.44 секунды против 1.94. Масштаб сглаживания s = 4 выбран из того же принципа.

Отдельно авторы провели человеческое исследование: 20 случайных кейсов PixArt-α, 32 оценщика, компрессия 30%, несжатый референс не показывали. DaTMe получил лучшие оценки и за следование промпту, и за визуальное качество, и в общем предпочтении.

Ограничения

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

Что это значит на практике

Для тех, кто генерирует изображения на своих GPU, ценность DaTMe не столько в скорости, сколько в устойчивости. Скорость растёт умеренно, на 7-16%, зато при агрессивной компрессии метод не разваливается: там, где ToMeSD уходит в FID выше 120, DaTMe держится ниже 50. Это позволяет сжимать сильнее и выжимать больше экономии там, где конкуренты уже непригодны.

Заодно это хорошая иллюстрация общего тренда: инженерия эффективности переезжает в математику процесса генерации. Исследователи всё чаще смотрят не на вход и выход модели, а на её внутреннюю динамику. Tweedie-оценка, score-функция, расписание шума из скучных деталей реализации превращаются в рабочую поверхность для оптимизации. Похожий сюжет мы разбирали с адаптивным разреженным вниманием для видеогенерации: там экономия тоже строится на понимании, какие вычисления действительно нужны на конкретном шаге.

Код DaTMe открыт на GitHub, статья доступна на arXiv.

Частые вопросы

Это требует дообучения модели?

Нет. DaTMe это training-free метод: он вмешивается только в inference, на этапе денойзинга, и не меняет веса. Метод встраивается в существующие пайплайны PixArt-α, FLUX и Stable Diffusion, не требуя ни часов на обучение, ни изменения архитектуры. Единственное, что нужно настроить, это бюджет компрессии и параметры сетки.

Насколько быстрее генерируются изображения?

Скромно, если считать только латентность: 1.94 против 1.80 секунды на PixArt-α при 30% компрессии, 1.02 против 0.92 на SD 1.5, 5.08 против 4.25 на SDXL. Основная экономия в другом: DaTMe сохраняет качество при агрессивной компрессии, поэтому сжимать можно сильнее, чем позволяют конкуренты, и экономить больше вычислений.

Работает ли это для видео и других модальностей?

В статье только изображения: PixArt-α, FLUX.1-dev, SD 1.5 и SDXL. Но логика метода не привязана к картинкам: динамика денойзинга есть у любого диффузионного процесса, включая видео. Авторы этот вопрос не проверяли, так что это скорее направление, чем готовый ответ. У видео свои источники избыточности, например согласованность соседних кадров, так что перенос потребует отдельной проверки.

Итог

DaTMe отвечает на вопрос «какие токены объединять» через динамику самой диффузии, а не через привязку к промпту. Карта важности из Tweedie-оценки позволяет сжимать до 70% токенов, сохраняя структуру и фон, при почти нулевой потере качества на умеренных компрессиях. Метод training-free, код открыт, так что проверить его на своём пайплайне можно без переобучения. Если работаете с диффузионными пайплайнами, начните с компрессии 30%: там результат почти неотличим от несжатой генерации, а риск минимален.