Text Template Tokens Are Implicit Semantic Registers in Diffusion Transformers

Text-to-image diffusion transformers (DiTs) jointly process text and image tokens, yet their internal computation during denoising remains poorly understood. We introduce a causal interpretability framework for modern large-scale DiTs that combines attention decomposition with targeted interventions across token spans, heads, and layers. Using it to separate prompt-content tokens from structural template tokens, we find that the structural tokens carry little prompt-specific information at the encoder output. Yet surprisingly, they emerge as dominant image-to-text attention sinks and causally maintain object identity inside the DiT, acting as implicit semantic registers. We show that they acquire this identity indirectly, with prompt semantics first injected into the image latents and then read back into the template tokens rather than transferred directly from the prompt tokens. Inspired by the above findings, we design a training-free pruning rule for DiTs. Heads that attend most strongly to prompt tokens are dispensable, and pruning them removes 20% of attention FLOPs with only a 1.4-point drop on GenEval. We further reveal how generative computation in DiTs is organized across heads and depth, separating semantic routing from visual synthesis and progressing from identity formation to propagation and refinement. Our work not only reveals that the tokens encoding semantics at input need not be those that maintain it during generation, but also provides a causal view of internal mechanisms in DiTs.

한 줄 요약

대규모 텍스트-이미지 디퓨전 트랜스포머(DiT)의 내부 메커니즘을 인과적 해석 프레임워크로 분석하고, 템플릿 토큰이 암묵적 의미 레지스터 역할을 한다는 사실을 밝힘.

핵심 기여도

인과 해석 프레임워크를 제안: 토큰 수준의 어텐션 분해, 스팬 조작, 헤드 이식, 레이어 마스킹을 결합.
템플릿 토큰이 이미지-텍스트 어텐션 경로에서 주요 흐름을 형성하며, 의미 레지스터 역할을 함.
훈련 없이도 적용 가능한 헤드 제거 규칙 제안: 20%의 어텐션 FLOPs 절감, GenEval 정확도 1.4점 하락.
디퓨전 트랜스포머의 생성 메커니즘을 헤드와 레이어별로 분리 분석: 의미 루팅과 시각 합성 분리, 객체 정체성 형성-전파-정제 단계 구분.

핵심 아이디어

기존 연구는 텍스트-이미지 디퓨전 트랜스포머(DiT)에서 텍스트 토큰이 의미 정보를 유지한다고 가정했으나, 본 연구는 템플릿 토큰이 생성 과정에서 더 중요한 역할을 한다는 점을 밝힘. 템플릿 토큰은 텍스트 인코더 출력에서 의미 정보를 거의 가지지 않지만, 디퓨전 과정에서 이미지 레이턴트를 통해 간접적으로 객체 정체성을 유지함. 이는 텍스트 토큰이 아닌 템플릿 토큰이 의미 레지스터 역할을 한다는 새로운 통찰을 제공함.

또한, 템플릿 토큰이 이미지-텍스트 어텐션 경로에서 주요 흐름을 형성하며, 이는 텍스트 조건이 이미지 생성에 어떻게 전달되는지를 이해하는 데 핵심적임. 이 발견을 바탕으로 훈련 없이도 적용 가능한 헤드 제거 규칙을 제안함. 텍스트 토큰에 강하게 어텐션을 주는 헤드는 생성에 거의 기여하지 않아, 이를 제거해도 정확도 하락이 적음.

기술적 접근법

**Causal Interpretability Framework**: 토큰 수준 어텐션 분해, 스팬 조작, 헤드 이식, 레이어 마스킹을 결합.
**Span Intervention**: 프롬프트-콘텐츠 토큰과 템플릿 토큰을 분리하여 조건 정보의 흐름을 분석.
**Head Transplantation**: 어텐션 헤드 간 역할을 비교하여 기능적 회로 분석.
**Layer-wise Causal Masking**: 레이어별로 조건 정보가 저장되는 방식을 추적.
**Training-Free Pruning Rule**: 텍스트 토큰에 강하게 어텐션을 주는 헤드를 제거. 20%의 어텐션 FLOPs 절감, GenEval 정확도 1.4점 하락.

주요 결과

**GenEval 데이터셋**에서 헤드 제거로 20%의 어텐션 FLOPs 절감, 정확도는 1.4점 하락.
템플릿 토큰이 이미지-텍스트 어텐션 경로에서 주요 흐름을 형성하며, 객체 정체성을 유지.
객체 정체성 형성은 초기 레이어에서, 전파는 중간 레이어에서, 정제는 후기 레이어에서 이루어짐.
의미 루팅과 시각 합성은 헤드 수준에서 분리됨.

의의 및 한계

본 연구는 디퓨전 트랜스포머의 내부 메커니즘을 인과적 관점에서 분석함으로써, 텍스트 조건이 이미지 생성에 어떻게 전달되는지를 명확히 밝힘. 특히, 템플릿 토큰이 의미 레지스터 역할을 한다는 점은 기존 인식을 재정의하며, 훈련 없이도 적용 가능한 헤드 제거 규칙은 생성 모델의 효율성 향상에 기여함.

한편, 본 연구는 특정 디퓨전 트랜스포머 아키텍처에 기반한 분석이며, 다른 아키텍처나 데이터셋에서의 일반화 가능성은 추가 연구가 필요함. 또한, 템플릿 토큰이 의미를 유지하는 방식에 대한 심층적 이해는 아직 부족함.

실용적 활용

본 연구는 대규모 텍스트-이미지 생성 모델의 내부 메커니즘을 이해하고, 헤드 제거를 통한 효율적 추론을 가능하게 함. 이는 생성 모델의 실시간 적용, 모바일 환경에서의 사용, 또는 대규모 배치 작업에서의 성능 향상에 활용될 수 있음. 특히, 템플릿 토큰의 역할을 이해하면 텍스트-이미지 조건링 방식을 더 정교하게 설계할 수 있음.