스프린트: 효율적인 확산 변환기를 위한 희소-밀집 잔차 융합

초록

확산 트랜스포머(DiT)는 최고 수준의 생성 성능을 제공하지만, 시퀀스 길이에 따른 이차적인 학습 비용으로 대규모 사전 학습이 극도로 비싸다는 문제가 있습니다. 토큰 드롭핑은 학습 비용을 줄일 수 있으나, 단순한 전략은 표현력을 저하시키며, 기존 방법들은 매개변수가 많거나 높은 드롭 비율에서 실패합니다. 본 연구에서는 공격적인 토큰 드롭핑(최대 75%)을 허용하면서도 품질을 유지하는 간단한 방법인 SPRINT(Sparse-Dense Residual Fusion for Efficient Diffusion Transformers)를 제시합니다. SPRINT는 얕은 층과 깊은 층의 상호 보완적 역할을 활용합니다. 초기 층은 모든 토큰을 처리하여 지역적 세부 사항을 포착하고, 더 깊은 층은 희소한 토큰 집합만으로 연산을 수행하여 계산량을 절감하며, 이들의 출력은 잔차 연결을 통해 융합됩니다. 학습은 효율성을 위한 긴 마스크 사전 학습 단계와 학습-추론 간격을 줄이기 위한 짧은 전체 토큰 미세 조정 단계로 구성된 2단계 일정을 따릅니다. ImageNet-1K 256x256에서 SPRINT는 유사한 FID/FDD 성능을 유지하면서 9.8배의 학습 비용 절감을 달성했으며, 추론 시에는 Path-Drop Guidance(PDG)를 통해 FLOPs를 거의 절반으로 줄이면서 품질을 개선했습니다. 이러한 결과는 SPRINT가 효율적인 DiT 학습을 위한 간단하고 효과적이며 일반적인 해법임을 입증합니다.

English

Diffusion Transformers (DiTs) deliver state-of-the-art generative performance but their quadratic training cost with sequence length makes large-scale pretraining prohibitively expensive. Token dropping can reduce training cost, yet na\"ive strategies degrade representations, and existing methods are either parameter-heavy or fail at high drop ratios. We present SPRINT, Sparse--Dense Residual Fusion for Efficient Diffusion Transformers, a simple method that enables aggressive token dropping (up to 75%) while preserving quality. SPRINT leverages the complementary roles of shallow and deep layers: early layers process all tokens to capture local detail, deeper layers operate on a sparse subset to cut computation, and their outputs are fused through residual connections. Training follows a two-stage schedule: long masked pre-training for efficiency followed by short full-token fine-tuning to close the train--inference gap. On ImageNet-1K 256x256, SPRINT achieves 9.8x training savings with comparable FID/FDD, and at inference, its Path-Drop Guidance (PDG) nearly halves FLOPs while improving quality. These results establish SPRINT as a simple, effective, and general solution for efficient DiT training.

스프린트: 효율적인 확산 변환기를 위한 희소-밀집 잔차 융합

Sprint: Sparse-Dense Residual Fusion for Efficient Diffusion Transformers

초록

Support