A-SelecT: 디퓨전 트랜스포머 표현 학습을 위한 자동 타임스텝 선택

A-SelecT: Automatic Timestep Selection for Diffusion Transformer Representation Learning

디퓨전 모델은 생성형 인공지능 분야에 큰 영향을 미쳤으며, 최근에는 판별적 표현 학습 능력을 갖춘 모델로도 활발히 연구되고 있습니다. 디퓨전 트랜스포머(DiT)는 기존 U-Net 기반 디퓨전 모델의 유망한 대안으로 떠오르며, 생성적 사전 학습을 통해 다양한 판별적 작업에 활용될 가능성을 보여줍니다. 그러나 현재 DiT의 학습 효율성과 표현 능력은 부적절한 타임스텝 탐색과 DiT 특유의 특징 표현 활용 부족으로 인해 제한적인 상황입니다. 이러한 문제점을 해결하기 위해, 우리는 자동으로 선택된 타임스텝(A-SelecT)을 제안합니다. A-SelecT는 단일 실행 내에서 DiT의 가장 유용한 정보가 담긴 타임스텝을 동적으로 파악하여, 계산 집약적인 모든 타임스텝 탐색 및 최적 이하의 판별적 특징 선택의 필요성을 없앱니다. 분류 및 분할 벤치마크에 대한 광범위한 실험 결과, A-SelecT를 통해 향상된 DiT는 기존의 모든 디퓨전 기반 모델보다 효율적이고 효과적으로 성능을 뛰어넘는 것으로 나타났습니다.

Diffusion models have significantly reshaped the field of generative artificial intelligence and are now increasingly explored for their capacity in discriminative representation learning. Diffusion Transformer (DiT) has recently gained attention as a promising alternative to conventional U-Net-based diffusion models, demonstrating a promising avenue for downstream discriminative tasks via generative pre-training. However, its current training efficiency and representational capacity remain largely constrained due to the inadequate timestep searching and insufficient exploitation of DiT-specific feature representations. In light of this view, we introduce Automatically Selected Timestep (A-SelecT) that dynamically pinpoints DiT's most information-rich timestep from the selected transformer feature in a single run, eliminating the need for both computationally intensive exhaustive timestep searching and suboptimal discriminative feature selection. Extensive experiments on classification and segmentation benchmarks demonstrate that DiT, empowered by A-SelecT, surpasses all prior diffusion-based attempts efficiently and effectively.