Динамический чанковый диффузионный трансформер

Аннотация

Диффузионные трансформаторы обрабатывают изображения как последовательности токенов фиксированной длины, создаваемые статической операцией разбиения на патчи. Несмотря на эффективность, такая конструкция равномерно распределяет вычислительные ресурсы как на области с низкой, так и с высокой информативностью, игнорируя тот факт, что изображения содержат регионы с разной степенью детализации, а процесс денойзинга прогрессирует от грубой структуры на ранних шагах к мелким деталям на поздних. Мы представляем Dynamic Chunking Diffusion Transformer (DC-DiT), который расширяет основу DiT с помощью обучаемого каркаса «кодировщик-маршрутизатор-декодировщик», адаптивно сжимающего двумерный вход в более короткую последовательность токенов зависимым от данных образом с использованием механизма чанкинга, обучаемого сквозным образом вместе с диффузионной тренировкой. Механизм обучается сжимать однородные фоновые области в меньшее количество токенов, а богатые деталями регионы — в большее, при этом осмысленные визуальные сегментации возникают без явного контроля. Кроме того, он также учится адаптировать степень сжатия в зависимости от шага диффузии, используя меньше токенов на зашумленных стадиях и больше — по мере проявления мелких деталей. На данных ImageNet 256×256 с условиями по классам DC-DiT стабильно улучшает FID и Inception Score по сравнению с базовыми моделями DiT, согласованными как по параметрам, так и по FLOPs, при степенях сжатия 4× и 16×, что демонстрирует перспективность методики с потенциальным применением в генерации в пиксельном пространстве, для видео и 3D. Помимо точности, DC-DiT практичен: модель можно дообучить из предварительно обученных чекпоинтов DiT с минимальными вычислительными затратами (до 8× меньше шагов обучения), а также комбинировать с другими методами динамических вычислений для дальнейшего сокращения FLOPs при генерации.

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.