FlexiDiT: Seu Transformador de Difusão Pode Gerar Amostras de Alta Qualidade com Menos Recursos Computacionais

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

February 27, 2025

Autores: Sotiris Anagnostidis, Gregor Bachmann, Yeongmin Kim, Jonas Kohler, Markos Georgopoulos, Artsiom Sanakoyeu, Yuming Du, Albert Pumarola, Ali Thabet, Edgar Schönfeld

cs.AI

Resumo

Apesar do desempenho notável, os modernos Transformadores de Difusão são prejudicados por requisitos substanciais de recursos durante a inferência, decorrentes da quantidade fixa e grande de computação necessária para cada etapa de remoção de ruído. Neste trabalho, revisitamos o paradigma estático convencional que aloca um orçamento fixo de computação por iteração de remoção de ruído e propomos uma estratégia dinâmica em vez disso. Nosso framework simples e eficiente em amostras permite que modelos DiT pré-treinados sejam convertidos em modelos flexíveis - chamados FlexiDiT - permitindo que processem entradas em orçamentos de computação variáveis. Demonstramos como um único modelo flexível pode gerar imagens sem queda na qualidade, enquanto reduz os FLOPs necessários em mais de 40\% em comparação com seus equivalentes estáticos, tanto para geração de imagens condicionadas por classe quanto por texto. Nosso método é geral e agnóstico em relação às modalidades de entrada e condicionamento. Mostramos como nossa abordagem pode ser prontamente estendida para geração de vídeo, onde modelos FlexiDiT geram amostras com até 75% menos computação sem comprometer o desempenho.

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.