Diffusion Transformer의 이상치 토큰 길들이기

Diffusion Transformer(DiT) 기반 이미지 생성 모델에서 일부 토큰이 지나치게 큰 노름(norm)을 갖고 어텐션을 과도하게 끌어모으는 이상치 토큰 현상이 나타나는 것으로 확인됐다. 이 토큰은 지역 패치 정보를 충분히 담지 못해 생성 과정에서 아티팩트를 유발할 수 있다.

연구진은 최신 Representation Autoencoder(RAE)-DiT 파이프라인의 ViT 인코더와 DiT 디노이저를 분석했다. 사전 학습된 ViT 인코더뿐 아니라 DiT 내부, 특히 중간 레이어에서도 이상치 토큰이 형성되는 현상을 관찰했다.

단순히 노름이 큰 토큰을 마스킹하는 방식은 성능을 개선하지 못했다. 문제의 핵심이 극단적인 값 자체보다 손상된 지역 패치 의미 정보에 더 가깝다는 분석이다.

이를 해결하기 위해 Dual-Stage Registers(DSR) 개입 기법을 제안했다. DSR은 다음과 같은 레지스터를 각 구성 요소에 적용한다.

• 사용 가능한 경우 학습된 레지스터
• 테스트 시 재귀적으로 생성하는 레지스터
• DiT 디노이저에 적용하는 diffusion 레지스터

ImageNet과 대규모 텍스트-이미지 생성 실험에서 DSR 개입은 이상치 아티팩트를 일관되게 줄이고 생성 품질을 개선했다. 연구진은 이상치 토큰 제어가 더 강력한 DiT를 구축하기 위한 중요한 요소라고 설명했다.