[Paper] Diffusion Transformers에서 이상 토큰 제어

Source: arXiv - 2605.05206v1

Overview

논문 **“Taming Outlier Tokens in Diffusion Transformers”**는 Vision Transformer(ViT)를 사용하는 최신 확산 기반 이미지 생성기에서 숨겨진 결함을 밝혀냅니다. 인코더와 디노이징 트랜스포머 모두가 “이상치” 토큰 몇 개—크기가 비정상적으로 크면서 시각적 정보는 거의 담고 있지 않은 벡터—를 생성할 수 있음을 보여줍니다. 이러한 토큰은 주의를 지배하지만 유용한 시각 정보를 거의 제공하지 못합니다. 저자들은 Dual‑Stage Registers (DSR) 라는 경량 레지스터 기반 수정을 도입하여 이러한 아티팩트를 크게 감소시키고 ImageNet 및 대규모 텍스트‑투‑이미지 모델에서 생성 품질을 향상시켰습니다.

주요 기여

• Diffusion Transformers (DiTs)에서 이상치 토큰을 식별한다. 고노름 토큰이 사전 학습된 ViT 인코더에만 나타나는 것이 아니라, 특히 중간 레이어에서 확산 디노이징 중에도 내부적으로 발생함을 보여준다.
• 단순 마스킹이 실패함을 보여준다. 고노름 토큰을 단순히 0으로 만들어도 결과가 개선되지 않으며, 문제는 극단값이 아니라 의미적 손상임을 나타낸다.
• Dual‑Stage Registers (DSR)를 제안한다. 두 단계의 레지스터 기반 개입:
  1. 학습 시 레지스터는 모델 학습 중 이상치 토큰을 교체하거나 수정하는 방법을 학습한다.
  2. 재귀적 테스트 시 레지스터는 추론 중 실시간으로 이상치를 감지하고 교체하며, 디노이저용 특수 diffusion register도 포함한다.
• 광범위한 실증 검증. 표준 ImageNet 생성 및 대규모 텍스트‑투‑이미지 벤치마크 전반에 걸쳐 DSR은 시각적 아티팩트를 지속적으로 감소시키고 FID/IS 점수를 향상시킨다.
• 새로운 연구 방향을 연다. 이상치 토큰 제어가 견고한 확산 기반 생성 모델을 구축하기 위한 핵심이자 이전에 간과된 요소임을 강조한다.

Methodology

1. Diagnosing the problem

   • 저자들은 먼저 Representation Autoencoder‑DiT (RAE‑DiT)의 인코더‑디코더 파이프라인 전반에 걸쳐 토큰 노름을 분석합니다.
   • 어텐션 맵을 시각화한 결과, 몇몇 토큰이 어텐션 분포를 지배하면서도 모호하거나 잡음이 섞인 패치를 나타내는 것을 발견했습니다.

2. Baseline experiments

   • 간단한 마스킹(노름 임계값을 초과하는 토큰을 0으로 만들기)과 노름 클리핑을 적용했지만, 생성 품질에 거의 영향을 주지 않거나 오히려 부정적인 영향을 미쳤습니다.

3. Dual‑Stage Registers (DSR)

   • Training‑stage registers: 작은 학습 가능한 벡터(‘레지스터’)를 토큰 시퀀스에 추가합니다. 학습 중에 게이팅 네트워크가 이상치 토큰을 레지스터 항목으로 대체할 시점을 학습하여, 손상된 의미를 효과적으로 “복구”합니다.
   • Test‑time registers: 추론 시, 재귀적 탐지 모듈이 각 레이어에서 고노름 토큰을 스캔하고 가장 적절한 레지스터 항목으로 교체한 뒤, 수정된 시퀀스를 이후 레이어에 다시 입력합니다.
   • Diffusion registers: 노이즈 제거 트랜스포머 전용 레지스터 세트를 별도로 학습시켜, 확산 과정 자체에서 발생하는 이상치를 바로잡을 수 있게 합니다.

4. Evaluation

   • 파이프라인을 무조건적인 ImageNet 생성(256×256) 및 대규모 텍스트‑투‑이미지 모델(예: Stable Diffusion‑유사 아키텍처)에서 테스트했습니다.
   • 표준 지표(FID, IS, CLIP‑Score)와 정성적인 시각 검토를 통해 개선 효과를 평가했습니다.

결과 및 발견

• 시각적 품질: DSR로 생성된 샘플은 이전에 이상 토큰에 기인한 것으로 추정되던 “blobby” 또는 “checkerboard” 아티팩트가 더 적게 나타납니다.
• 주의 분포: DSR 적용 후 어텐션 맵은 패치 전반에 걸쳐 보다 균형 잡히고 부드러운 분포를 보이며, 레지스터가 이상 토큰의 지배력을 성공적으로 희석함을 확인합니다.
• 효율성: 레지스터 모듈은 추론 시간에 < 2 % 정도의 오버헤드만 추가하여 실제 환경 배포에 실용적입니다.

실용적인 시사점

• 프로덕션 급 생성기를 위한 더 깨끗한 출력. AI 기반 이미지 생성 도구(예: 디자인 어시스턴트, 콘텐츠 생성 플랫폼)를 구축하는 기업은 전체 모델을 재학습하지 않고도 DSR을 통합하여 결함이 있는 아티팩트를 감소시킬 수 있습니다.
• 향상된 다운스트림 작업. 품질이 높은 잠재 표현은 diffusion 모델에 의존하는 이미지 편집, 인페인팅, 스타일 전송 등 다운스트림 파이프라인에서 더 높은 충실도로 이어집니다.
• 저비용 업그레이드 경로. DSR이 플러그인 형태로 작동하므로(레지스터를 기존 체크포인트 위에 학습 가능) 개발자는 최소한의 컴퓨팅 비용으로 레거시 diffusion 모델을 개조할 수 있습니다.
• 보다 안정적인 파인튜닝. 대형 diffusion 모델을 새로운 도메인(예: 의료 영상)으로 적용할 때 DSR은 종종 학습 불안정을 초래하는 이상 토큰의 발생을 완화할 수 있습니다.

Limitations & Future Work

• Scope of token types. The study focuses on visual tokens; extending the analysis to multimodal diffusion models (e.g., text‑image or video) remains open.
• Register capacity. A fixed small set of registers may eventually saturate for extremely large or highly diverse datasets; adaptive or hierarchical registers could be explored.
• Theoretical understanding. While empirical results are strong, a deeper theoretical explanation of why outlier tokens arise in diffusion dynamics is still lacking.
• Real‑time constraints. Although overhead is modest, ultra‑low‑latency applications (e.g., mobile inference) may need further optimization of the recursive detection step.

Bottom line: By shining a light on a subtle but pervasive issue—outlier tokens—in diffusion transformers, this work equips developers with a practical tool (DSR) to make generative models more reliable and visually appealing, paving the way for higher‑quality AI‑driven content creation.

제한 사항 및 향후 연구

• 토큰 유형의 범위. 본 연구는 시각 토큰에 초점을 맞추고 있으며, 멀티모달 확산 모델(예: 텍스트‑이미지 또는 비디오)로 분석을 확장하는 것은 아직 미해결 과제이다.
• 레지스터 용량. 고정된 소수의 레지스터는 매우 크거나 다양성이 높은 데이터셋에 대해 결국 포화될 수 있으므로, 적응형 또는 계층형 레지스터에 대한 탐구가 필요할 수 있다.
• 이론적 이해. 실험 결과는 강력하지만, 확산 역학에서 이상치 토큰이 발생하는 이유에 대한 보다 깊은 이론적 설명은 아직 부족하다.
• 실시간 제약. 오버헤드는 비교적 작지만, 초저지연 애플리케이션(예: 모바일 추론)에서는 재귀적 탐지 단계의 추가 최적화가 요구될 수 있다.

핵심 요약: 확산 트랜스포머에서 미묘하지만 널리 퍼진 문제인 이상치 토큰에 조명을 비춤으로써, 본 연구는 개발자들에게 실용적인 도구(DSR)를 제공하여 생성 모델을 보다 신뢰성 있고 시각적으로 매력적으로 만들 수 있게 한다. 이는 고품질 AI 기반 콘텐츠 제작을 위한 길을 열어준다.

저자

• Xiaoyu Wu
• Yifei Wang
• Tsu-Jui Fu
• Liang-Chieh Chen
• Zhe Gan
• Chen Wei

논문 정보

• arXiv ID: 2605.05206v1
• Categories: cs.CV, cs.AI, cs.LG
• Published: 2026년 5월 6일
• PDF: PDF 다운로드