Transformador de Difusão com Segmentação Dinâmica

Resumo

Os Diffusion Transformers processam imagens como sequências de tokens de comprimento fixo produzidas por uma operação estática de divisão em *patches*. Embora eficiente, este projeto aplica computação uniforme tanto em regiões de baixa quanto de alta informação, ignorando que as imagens contêm regiões com detalhes variados e que o processo de remoção de ruído progride de estruturas grosseiras nos primeiros *timesteps* para detalhes finos nos *timesteps* finais. Apresentamos o *Dynamic Chunking Diffusion Transformer* (DC-DiT), que amplia a arquitetura DiT com um *scaffold* aprendido de codificador-roteador-decodificador que comprime adaptativamente a entrada 2D numa sequência mais curta de tokens de forma dependente dos dados, utilizando um mecanismo de agrupamento (*chunking*) aprendido de forma *end-to-end* com o treino de difusão. O mecanismo aprende a comprimir regiões de fundo uniformes em menos tokens e regiões ricas em detalhes em mais tokens, com segmentações visuais significativas emergindo sem supervisão explícita. Além disso, também aprende a adaptar a sua compressão ao longo dos *timesteps* de difusão, usando menos tokens em fases ruidosas e mais tokens à medida que os detalhes finos emergem. Na geração condicionada por classe do ImageNet 256×256, o DC-DiT melhora consistentemente o FID e o *Inception Score* em relação às *baselines* DiT com igual número de parâmetros e igual número de FLOPs, através de compressões de 4× e 16×, mostrando que esta é uma técnica promissora com potenciais aplicações futuras em geração no espaço de pixels, vídeo e 3D. Para além da precisão, o DC-DiT é prático: pode ser reaproveitado a partir de *checkpoints* DiT pré-treinados com computação pós-treino mínima (até 8× menos passos de treino) e combina-se com outros métodos de computação dinâmica para reduzir ainda mais os FLOPs de geração.

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.