FlexiDiT: Tu Transformer de Difusión Puede Generar Fácilmente Muestras de Alta Calidad con Menos Cómputo

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

Resumen

A pesar de su notable rendimiento, los Transformadores de Difusión modernos se ven limitados por los considerables requisitos de recursos durante la inferencia, derivados de la cantidad fija y grande de cómputo necesaria para cada paso de eliminación de ruido. En este trabajo, revisamos el paradigma estático convencional que asigna un presupuesto de cómputo fijo por iteración de eliminación de ruido y proponemos, en su lugar, una estrategia dinámica. Nuestro marco simple y eficiente en muestras permite convertir modelos DiT preentrenados en modelos flexibles —denominados FlexiDiT—, lo que les permite procesar entradas con distintos presupuestos de cómputo. Demostramos cómo un único modelo flexible puede generar imágenes sin pérdida de calidad, reduciendo los FLOPs requeridos en más de un 40\% en comparación con sus contrapartes estáticas, tanto para la generación de imágenes condicionadas por clase como por texto. Nuestro método es general y agnóstico a las modalidades de entrada y condicionamiento. Mostramos cómo nuestro enfoque puede extenderse fácilmente para la generación de videos, donde los modelos FlexiDiT generan muestras con hasta un 75\% menos de cómputo sin comprometer el rendimiento.

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.