FlexiDiT: Ваш трансформер диффузии способен легко генерировать высококачественные выборки с меньшими вычислительными затратами

FlexiDiT: Your Diffusion Transformer Can Easily Generate High-Quality Samples with Less Compute

February 27, 2025

Авторы: Sotiris Anagnostidis, Gregor Bachmann, Yeongmin Kim, Jonas Kohler, Markos Georgopoulos, Artsiom Sanakoyeu, Yuming Du, Albert Pumarola, Ali Thabet, Edgar Schönfeld

cs.AI

Аннотация

Несмотря на выдающуюся производительность, современные Diffusion Transformers сталкиваются с существенными требованиями к ресурсам во время вывода, что обусловлено фиксированным и значительным объемом вычислений, необходимых для каждого шага удаления шума. В данной работе мы пересматриваем традиционную статическую парадигму, которая выделяет фиксированный бюджет вычислений на каждую итерацию удаления шума, и предлагаем вместо этого динамическую стратегию. Наш простой и эффективный по выборкам фреймворк позволяет преобразовать предобученные модели DiT в гибкие — названные FlexiDiT — что дает им возможность обрабатывать входные данные с различными бюджетами вычислений. Мы демонстрируем, как единая гибкая модель может генерировать изображения без потери качества, при этом сокращая необходимые FLOPs более чем на 40% по сравнению с их статическими аналогами, как для генерации изображений с условием по классам, так и с условием по тексту. Наш метод является универсальным и независимым от модальностей входных данных и условий. Мы показываем, как наш подход может быть легко расширен для генерации видео, где модели FlexiDiT генерируют образцы с уменьшением вычислений до 75% без ущерба для производительности.

English

Despite their remarkable performance, modern Diffusion Transformers are hindered by substantial resource requirements during inference, stemming from the fixed and large amount of compute needed for each denoising step. In this work, we revisit the conventional static paradigm that allocates a fixed compute budget per denoising iteration and propose a dynamic strategy instead. Our simple and sample-efficient framework enables pre-trained DiT models to be converted into flexible ones -- dubbed FlexiDiT -- allowing them to process inputs at varying compute budgets. We demonstrate how a single flexible model can generate images without any drop in quality, while reducing the required FLOPs by more than 40\% compared to their static counterparts, for both class-conditioned and text-conditioned image generation. Our method is general and agnostic to input and conditioning modalities. We show how our approach can be readily extended for video generation, where FlexiDiT models generate samples with up to 75\% less compute without compromising performance.