Transformador de Difusión con Fragmentación Dinámica

Resumen

Los Transformadores de Difusión procesan las imágenes como secuencias de longitud fija de tokens producidas por una operación estática de división en parches. Aunque es efectivo, este diseño emplea un cómputo uniforme tanto en regiones de baja como de alta información, ignorando que las imágenes contienen regiones de detalle variable y que el proceso de eliminación de ruido progresa desde una estructura gruesa en las primeras etapas temporales hasta detalles finos en las etapas finales. Presentamos el Transformador de Difusión con Segmentación Dinámica (DC-DiT), que aumenta la columna vertebral del DiT con un andamiaje aprendido de codificador-enrutador-decodificador que comprime adaptativamente la entrada 2D en una secuencia de tokens más corta de manera dependiente de los datos, utilizando un mecanismo de segmentación aprendido de extremo a extremo con el entrenamiento de difusión. El mecanismo aprende a comprimir regiones de fondo uniformes en menos tokens y regiones ricas en detalles en más tokens, surgiendo segmentaciones visuales significativas sin supervisión explícita. Además, también aprende a adaptar su compresión a lo largo de las etapas temporales de difusión, utilizando menos tokens en las etapas ruidosas y más tokens a medida que emergen los detalles finos. En ImageNet 256×256 condicionado por clase, DC-DiT mejora consistentemente el FID y el Inception Score respecto a las líneas base de DiT con igual número de parámetros y con igual número de FLOPS, a través de compresiones de 4× y 16×, lo que demuestra que esta es una técnica prometedora con posibles aplicaciones adicionales en generación en espacio de píxeles, vídeo y 3D. Más allá de la precisión, DC-DiT es práctico: puede ser reutilizado a partir de checkpoints de DiT preentrenados con un cómputo posterior al entrenamiento mínimo (hasta 8× menos pasos de entrenamiento) y se combina con otros métodos de cómputo dinámico para reducir aún más los FLOPS de generación.

English

Diffusion Transformers process images as fixed-length sequences of tokens produced by a static patchify operation. While effective, this design spends uniform compute on low- and high-information regions alike, ignoring that images contain regions of varying detail and that the denoising process progresses from coarse structure at early timesteps to fine detail at late timesteps. We introduce the Dynamic Chunking Diffusion Transformer (DC-DiT), which augments the DiT backbone with a learned encoder-router-decoder scaffold that adaptively compresses the 2D input into a shorter token sequence in a data-dependent manner using a chunking mechanism learned end-to-end with diffusion training. The mechanism learns to compress uniform background regions into fewer tokens and detail-rich regions into more tokens, with meaningful visual segmentations emerging without explicit supervision. Furthermore, it also learns to adapt its compression across diffusion timesteps, using fewer tokens at noisy stages and more tokens as fine details emerge. On class-conditional ImageNet 256{times}256, DC-DiT consistently improves FID and Inception Score over both parameter-matched and FLOP-matched DiT baselines across 4{times} and 16{times} compression, showing this is a promising technique with potential further applications to pixel-space, video and 3D generation. Beyond accuracy, DC-DiT is practical: it can be upcycled from pretrained DiT checkpoints with minimal post-training compute (up to 8{times} fewer training steps) and composes with other dynamic computation methods to further reduce generation FLOPs.