Dynamischer Chunking-Diffusion-Transformer

Zusammenfassung

Diffusion Transformers verarbeiten Bilder als festlange Sequenzen von Tokens, die durch eine statische Patch-Operation erzeugt werden. Diese effektive Methode wendet jedoch gleichmäßige Rechenleistung auf sowohl informationsarme als auch informationsreiche Regionen an und ignoriert dabei, dass Bilder Bereiche mit unterschiedlichem Detaillierungsgrad enthalten und dass der Denoising-Prozess von grober Struktur in frühen Zeitschritten zu feinen Details in späten Zeitschritten fortschreitet. Wir stellen den Dynamic Chunking Diffusion Transformer (DC-DiT) vor, der das DiT-Grundgerüst mit einem gelernten Encoder-Router-Decoder-Scaffold erweitert. Dieses komprimiert den 2D-Eingang adaptiv in eine kürzere Token-Sequenz, und zwar datenabhängig mittels eines Chunking-Mechanismus, der end-to-end mit dem Diffusionstraining erlernt wird. Der Mechanismus lernt, uniforme Hintergrundregionen in weniger Tokens und detailreiche Regionen in mehr Tokens zu komprimieren, wobei aussagekräftige visuelle Segmentierungen ohne explizite Supervision entstehen. Darüber hinaus lernt er, seine Komprimierung über die Diffusion-Zeitschritte hinweg anzupassen, indem er in verrauschten Phasen weniger Tokens und beim Auftauchen feiner Details mehr Tokens verwendet. Bei klassenkonditionierter ImageNet 256×256 verbessert DC-DiT konsistent FID und Inception Score gegenüber parameter- und FLOP-angepassten DiT-Baselines bei 4-facher und 16-facher Kompression, was zeigt, dass dies eine vielversprechende Technik mit potenziell weiteren Anwendungen in der Pixelraum-, Video- und 3D-Generierung ist. Über die Genauigkeit hinaus ist DC-DiT praktisch: Er kann aus vortrainierten DiT-Checkpoints mit minimalem Rechenaufwand nach dem Training (bis zu 8-mal weniger Trainingsschritte) weiterentwickelt werden und lässt sich mit anderen Methoden zur dynamischen Berechnung kombinieren, um die Generierungs-FLOPs weiter zu reduzieren.

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.