ToCa: Accelerating Diffusion Transformers with Token-wise Feature Caching, ICLR 2025

Ускорение диффузионных трансформеров в 2 раза! Есть код!

Для генерации одной картинки нужно десятки раз прогнать тяжеловесную нейросеть. Естественно, этот процесс хочется ускорить. Можно, например, уменьшить число шагов генерации, как сделали в https://t.me/layercv/122 с помощью дистилляции, а можно ускорить каждый шаг генерации.

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

Чтобы не вычислять новые значения каждого токена в каждом блоке трансформера и на каждом шаге диффузии, авторы переиспользуют часть посчитанных значений.
Токены для кэширования выбирают по нескольким критериям:

1. Насколько токен влияет на другие токены. Чем сильнее влияет, тем важнее вычислить точное его значение, а не использовать заранее закэшированное.
2. Насколько токен зависим от входного сигнала. Чем сильнее он меняется от изменения входного сигнала, тем он важнее.
3. Когда токен был закеширован. Если давно (и с точки зрения слоёв трансформера, и с точки зрения шагов диффузии), то его значение лучше вычислить заново.
4. Лучший ли он кандидат среди в своём окружении. Среди соседних токенов выбирают тот, который по предыдущим трём критериям лучше других. Это нужно, чтобы закэшированные токены были равномерно распределены по всему изображению, иначе качество генерации ухудшится.

Влияние на другие токены смотрят в Self-Attention блоках, а зависимость от входного сигнала в Cross-Attention блоках. В экспериментах с PixArt-alpha это возможно, а с более современными моделями (Stable Diffusion 3, Flux.1) нужно будет придумывать что-то другое: там обуславливание делается уже через MM-DiT блоки без Cross Attention.
В новой версии статьи авторы добавили и результаты с Flux.1 (dev и schnell версиями), но без пояснений, видимо просто без второго критерия.

Качество и скорость замеряли на text-to-image генерации моделью PixArt-alpha, text-to-video OpenSora и генерацию изображения по названию категории DiT-XL/2.

На PixArt-alpha, OpenSora, DiT-XL/2 ускорение в ~2 раза с небольшим ухудшением качества генерации, на Flux.1 - в 1.5 раза.