Diffusion Transformers (DiT) 의 이상치 토큰 제어 연구

이미지 생성을 위한 Diffusion Transformers (DiT) 에서 이상치 토큰 (outlier tokens) 을 연구합니다. 기존 연구는 Vision Transformers (ViT) 가 높은 노름 (norm) 을 가진 소수의 토큰을 생성하여 불균형한 주의를 끌지만, 제한된 지역 정보를 운반한다는 것을 보여주었습니다. 그러나 이러한 현상이 생성 모델에서의 역할은 아직 충분히 탐구되지 않았습니다.

우리는 이 현상이 현대 Representation Autoencoder (RAE)-DiT 파이프라인의 인코더 (encoder) 와 디노이저 (denoiser) 모두에서 나타낸다고 보여줍니다: 사전 학습된 ViT 인코더는 이상치 표현을 생성할 수 있으며, DiT 자체도 특히 중간 층에서 내부 이상치 토큰을 발전시킬 수 있습니다. 또한 단순히 높은 노름 토큰을 마스킹 (masking) 하더라도 성능이 개선되지 않는다는 것은 문제가 일부 극단적인 값에 의해서만 유발되는 것이 아니라, 손상된 지역 패치 세맨틱 (corrupted local patch semantics) 과 더 밀접한 관련이 있음을 시사합니다.

이 문제를 해결하기 위해 우리는 훈련된 레지스터가 존재할 경우와 그렇지 않은 경우를 모두 처리하는 레지스터 기반 개입인 Dual-Stage Registers (DSR) 을 소개했습니다. 디노이저의 경우 확산 레지스터 (diffusion registers) 를 사용합니다. ImageNet 과 대규모 텍스트-이미지 생성을 통해 이러한 개입은 일관되게 이상치 아티팩트를 줄이고 생성 품질을 향상시킵니다. 우리의 결과는 이상치 토큰 제어를 더 강력한 DiT 를 구축하는 중요한 요소임을 강조합니다.