⚡А вот и наша статья подоспела - Cache Me if You Can: Accelerating Diffusion Models through Block Caching

Как я уже упоминал, наша команда в Meta GenAI занимается ускорением диффузии. В этой статье мы ускоряем генерацию изображений до x1.8 раза без потери качества с помощью умного кеширования во время инференса.

Как?
Для генерации одной картинки обычно требуется сделать много прогонов через модель Unet, например 50 DDIM шагов. Мы тут подметили, что активации Spatial-Attention блоков довольно гладко меняются от шага к шагу, и паттерн их изменения не зависит от промпта. Естественно, мы подумали, почему бы не кешировать фичи тех блоков, которые меняются наиболее медленно и пересчитывать их только раз в несколько шагов. Понт в том, что львиная доля всех вычислений происходит именно в attention блоках, поэтому пропуская их вычисления хотя бы на некоторых шагах, мы сильно ускоряем генерацию.

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

Чтобы убрать мелкие артифакты после кеширования, мы дополнительно обучаем time-dependent scale и shift параметры для каждого выходнрго канала кешируемых блоков. Это помогает сгладить разницу распределений между "честно посчитанными" фичами и закеширвоанными.

В итоге, получили x1.5-1.8 ускорение, причем FID скор даже улучшился после применения кеширования. А результаты Human Eval показали, что при фиксированном времени генерации модель с кешированием выдает более качественные картинки чем бейзлайн.

На каких архитектурах тестировали:
- LDM 512x512 (та же архитектура как у SD 1.5, но натренированная нами внутри GenAI)
- Наша Emu 768x768 с 2.7B параметров.

То есть метод гибок и может подстраиваться под модели разных размеров, причем расписание кеширования разных блоков строится автоматически. Взлетит и на SDXL тоже.

Статья на Arxiv

@ai_newz