동적 청킹 확산 트랜스포머 (Dynamic Chunking Diffusion Transformer)

Dynamic Chunking Diffusion Transformer

확산 트랜스포머는 고정된 길이의 토큰 시퀀스로 이미지를 처리하며, 이 토큰은 정적인 패치화(patchify) 연산을 통해 생성됩니다. 이러한 설계는 효과적이지만, 이미지의 모든 영역에 동일한 연산량을 투입하여 이미지 내의 다양한 세부 묘사 수준을 무시하고, 또한 노이즈 제거 과정이 초기 타임 스텝에서는 거친 구조에서 시작하여 후기 타임 스텝에서는 미세한 디테일로 진행된다는 점을 간과합니다. 본 논문에서는 데이터 의존적인 방식으로 2차원 입력 데이터를 더 짧은 토큰 시퀀스로 적응적으로 압축하는 청킹(chunking) 메커니즘을 학습하는 인코더-라우터-디코더 구조를 DiT(Diffusion Transformer)의 핵심 구조에 추가한 Dynamic Chunking Diffusion Transformer (DC-DiT)를 제안합니다. 이 메커니즘은 균일한 배경 영역을 더 적은 토큰으로, 디테일이 풍부한 영역을 더 많은 토큰으로 압축하도록 학습하며, 명확한 시각적 분할이 명시적인 감독 없이도 나타납니다. 또한, DC-DiT는 노이즈가 심한 단계에서는 더 적은 토큰을 사용하고, 미세한 디테일이 나타나는 단계에서는 더 많은 토큰을 사용하여 압축 방식을 확산 타임 스텝에 따라 적응적으로 조절합니다. ImageNet $256{ imes}256$ 데이터셋에 대한 클래스 조건부 실험에서, DC-DiT는 동일한 파라미터 또는 FLOP(floating point operations) 수를 가진 DiT 모델을 기준으로 $4{ imes}$ 및 $16{ imes}$ 압축 비율에서 FID(Fréchet Inception Distance) 및 Inception Score를 지속적으로 향상시켜, 이 기술이 유망하며 픽셀 공간, 비디오 및 3D 생성 등 다양한 분야에 적용될 가능성이 있음을 보여줍니다. DC-DiT는 정확성 향상 외에도 실용적입니다. 사전 학습된 DiT 모델의 체크포인트로부터 최소한의 추가 학습 연산으로 업사이클링이 가능하며(최대 $8{ imes}$ 더 적은 학습 단계), 다른 동적 연산 방법과 결합하여 생성 과정의 FLOPs를 더욱 줄일 수 있습니다.

Diffusion Transformers process images as fixed-length sequences of tokens produced by a static $\textit{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.