동적 청킹 확산 트랜스포머

기존 DiT는 이미지의 모든 영역에 동일한 연산을 수행하지만, DC-DiT는 복잡한 영역에 더 많은 토큰을 할당하고 단순한 배경은 압축하여 효율성을 극대화한다. 또한 노이즈가 많은 초기 단계에서는 거칠게, 세부 묘사가 필요한 후기 단계에서는 정밀하게 연산량을 조절하는 시간적 적응성을 제공한다.

Transformer 기반의 DiT는 이미지를 고정된 크기의 패치로 나누어 처리하는데, 이는 하늘과 같은 단순한 배경과 복잡한 물체 영역에 동일한 연산량을 쏟는 비효율을 초래한다. Attention 메커니즘은 모든 토큰 간의 관계를 계산하므로 토큰 수가 많아질수록 연산량이 제곱으로 늘어나는 한계가 있다. DC-DiT는 인접한 토큰 간의 유사도를 계산하여, 비슷한 정보를 가진 영역은 하나의 토큰으로 뭉치고 중요한 경계나 세부 묘사가 있는 곳은 개별 토큰을 유지한다. 이는 Embedding 공간에서 유사한 벡터들을 하나로 대표하여 시퀀스 길이를 줄이는 원리이다. 이 과정은 확산 공정의 진행 단계와도 연동된다. 전체적인 윤곽만 잡는 초기에는 적은 토큰으로 빠르게 처리하고, 디테일을 채우는 후기에는 더 많은 토큰을 사용하여 품질을 확보함으로써 전체적인 생성 효율과 품질을 동시에 잡는다.

DC-DiT는 기존 DiT 백본 주위에 Encoder-Router-Decoder 스캐폴드를 추가한 구조이다. Isotropic Encoder가 로컬 컨텍스트를 섞어 라우팅에 적합한 표현을 생성하면, Router가 각 토큰의 유지 여부를 결정한다. 이후 DiT 블록이 압축된 시퀀스를 처리하고, De-chunking 레이어와 Decoder가 원래 해상도로 복원한다. 공간적 청킹을 위해 3x3 depthwise convolution을 활용한다. [각 토큰의 Key 벡터를 입력으로] -> [주변 이웃의 Key 벡터 평균을 계산하여] -> [Query 벡터와의 내적을 수행해 유사도 점수 s_i를 얻고] -> [이 값이 높을수록 주변과 유사하므로 병합 대상임을 의미한다]. 최종적으로 p_i = (1 - s_i)/2 수식을 통해 경계 확률을 도출한다. 학습 시에는 타겟 압축률 N을 유도하기 위해 정규화 항 L_ratio를 추가한다. [라우팅 모듈의 출력 마스크와 확률값을 입력으로] -> [평균 다운샘플링 비율을 계산하여] -> [목표치 N과의 차이를 손실로 반환하고] -> [모델이 특정 압축률을 유지하면서도 중요한 정보를 보존하도록 유도한다].

ImageNet 256x256 벤치마크에서 DC-DiT-B 모델은 4배 압축 시 FID 13.51을 기록하여 파라미터가 유사한 DiT(19.45) 및 연산량이 유사한 DiT(15.78)를 모두 능가했다. XL 규모에서도 16배 압축 시 FID 13.60을 달성하며 고정 패치 방식(FID 16.35) 대비 월등한 효율성을 입증했다. 업사이클링 실험 결과, 사전 학습된 DiT-XL 가중치를 활용해 단 12.5%의 학습 예산만으로도 처음부터 학습한 모델과 DiT 베이스라인을 모두 뛰어넘는 성능을 보였다. 특히 Activation Distillation 기법을 적용했을 때 수렴 속도가 비약적으로 향상됨이 확인됐다.

Isotropic Encoder와 Decoder는 메인 트랜스포머 차원의 1/4 수준인 중간 차원에서 작동하는 2개의 컨볼루션 잔차 블록으로 구성되어 연산 오버헤드를 최소화한다. Router는 p > 0.5 임계값을 기준으로 토큰 유지 여부를 결정하며, 배치 학습을 위해 최대 토큰 수까지 패딩을 적용한다. De-chunking 레이어는 가우시안 커널 기반의 공간적 스무딩을 적용한다. [압축된 토큰의 위치와 확률값을 입력으로] -> [거리 기반 가중치 W_ij를 계산해 주변 정보를 섞어] -> [연속적인 표현을 생성하고] -> [이산적인 결정으로 인한 최적화의 어려움을 해결한다]. 이후 Plug-back 맵을 통해 원래 격자 위치에 가장 가까운 경계 토큰의 표현을 할당한다.