동적 청킹 확산 트랜스포머

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

왜 중요한가

핵심 기여

데이터 의존적 동적 청킹 메커니즘

이미지의 공간적 복잡도에 따라 토큰 수를 적응적으로 조절하는 encoder-router-decoder 구조를 도입하여 불필요한 연산을 제거함.

확산 타임스텝별 가변 압축 스케줄

노이즈가 많은 초기 단계에서는 높은 압축률을, 세부 묘사가 생성되는 후기 단계에서는 낮은 압축률을 자동으로 학습하여 연산 효율을 높임.

사전 학습된 DiT 모델의 업사이클링

처음부터 학습하는 대신 기존 DiT 체크포인트를 활용해 8배 적은 연산량으로 성능이 개선된 모델을 구축하는 방법론을 입증함.

핵심 아이디어 이해하기

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 맵을 통해 원래 격자 위치에 가장 가까운 경계 토큰의 표현을 할당한다.

실무 활용

고해상도 이미지 및 비디오 생성 시 연산 비용을 획기적으로 줄이면서도 품질을 유지하거나 개선할 수 있는 실용적인 아키텍처이다.

고해상도 이미지 생성 모델의 추론 속도 최적화
제한된 컴퓨팅 자원에서의 대규모 확산 모델 학습
비디오 생성 모델의 프레임 간 중복 정보 압축

코드 공개 여부: 비공개

키워드

DiT(확산 트랜스포머)Dynamic Chunking(동적 청킹)Adaptive Computation(적응형 연산)Image Generation(이미지 생성)Efficiency(효율성)

용어 해설

패치화(Patchify): — 이미지를 고정된 크기의 격자(패치)로 나누어 트랜스포머가 처리할 수 있는 토큰 시퀀스로 변환하는 과정이다. 기존 DiT에서는 모든 영역을 동일한 크기의 패치로 나누어 정보 밀도와 상관없이 균일한 연산을 수행하게 된다.
프레셰 인셉션 거리(FID): — 생성된 이미지의 품질과 다양성을 평가하는 표준 지표로, 실제 이미지 분포와 생성 이미지 분포 사이의 거리를 계산한다. 값이 낮을수록 실제 이미지와 통계적으로 유사하며 품질이 높음을 의미한다.
업사이클링(Upcycling): — 사전 학습된 모델의 가중치를 재사용하여 새로운 구조나 기능을 최소한의 추가 학습으로 이식하는 기법이다. 처음부터 학습하는 것보다 비용이 훨씬 저렴하며 기존 모델의 지식을 보존할 수 있다.
스트레이트 스루 에스티메이터(Straight-Through Estimator): — 이산적인 결정(토큰 선택 등)을 포함하는 연산에서 미분 불가능한 구간을 건너뛰어 그래디언트를 전달하는 기법이다. 이를 통해 불연속적인 라우팅 결정 과정에서도 역전파를 통한 학습이 가능해진다.
지식 증류(Distillation): — 복잡한 교사(Teacher) 모델의 출력을 가벼운 학생(Student) 모델이 모방하도록 학습시키는 기법이다. 본 논문에서는 사전 학습된 DiT의 활성화 값을 모방하여 새로운 모듈의 수렴 속도를 높이는 데 사용된다.

동적 청킹 확산 트랜스포머

왜 중요한가

핵심 기여

핵심 아이디어 이해하기

방법론

주요 결과

기술 상세

실무 활용

키워드

용어 해설

관련 토론

댓글

관련 피드

왜 중요한가

핵심 기여

핵심 아이디어 이해하기

방법론

주요 결과

기술 상세

실무 활용

키워드

관련 토론

댓글

관련 피드