Transformateur de Diffusion par Segmentation Dynamique

Résumé

Les Transformers de Diffusion traitent les images comme des séquences de tokens de longueur fixe produites par une opération statique de mise en patch. Bien qu'efficace, cette conception applique un calcul uniforme sur les régions à faible et haute information, ignorant que les images contiennent des régions de détails variables et que le processus de débruitage évolue d'une structure grossière aux premiers pas de temps vers des détails fins aux derniers pas de temps. Nous présentons le *Dynamic Chunking Diffusion Transformer* (DC-DiT), qui augmente l'architecture DiT de base avec un échafaudage encodeur-routeur-décodeur appris, comprimant adaptativement l'entrée 2D en une séquence de tokens plus courte de manière dépendante des données, via un mécanisme de regroupement (*chunking*) appris de bout en bout avec l'entraînement par diffusion. Le mécanisme apprend à compresser les régions d'arrière-plan uniformes en moins de tokens et les régions riches en détails en plus de tokens, des segmentations visuelles significatives émergeant sans supervision explicite. De plus, il apprend aussi à adapter sa compression au cours des pas de temps de diffusion, utilisant moins de tokens aux étapes bruyantes et plus de tokens à mesure que les détails fins apparaissent. Sur ImageNet 256×256 conditionné par classe, DC-DiT améliore constamment le FID et l'Inception Score par rapport aux lignes de base DiT, à nombre de paramètres égal et à nombre de FLOPs égal, pour des compressions de 4× et 16×, montrant qu'il s'agit d'une technique prometteuse avec des applications potentielles étendues à la génération dans l'espace pixel, vidéo et 3D. Au-delà de la précision, DC-DiT est pratique : il peut être recyclé à partir de checkpoints DiT pré-entraînés avec un calcul post-entraînement minimal (jusqu'à 8× moins d'étapes d'entraînement) et se combine avec d'autres méthodes de calcul dynamique pour réduire davantage les FLOPs de génération.

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.