DiT의 ‘쓸모없어 보이는’ 템플릿 토큰, 의미를 붙잡고 있었다

도입

텍스트로 이미지를 생성할 때 우리는 보통 프롬프트의 내용 단어에 주목한다. “고양이”, “빨간 자동차”, “사이버펑크 도시” 같은 단어가 이미지의 의미를 결정한다고 보기 때문이다. 반면 텍스트 템플릿이나 구조적 토큰은 입력 형식을 맞추는 보조 요소처럼 보인다.

하지만 이 논문은 Diffusion Transformer, 즉 DiT 내부를 들여다보며 다른 그림을 제시한다. 입력에서 의미를 담는 토큰과 생성 중 의미를 유지하는 토큰은 같지 않을 수 있으며, 겉으로는 정보가 적어 보이는 템플릿 토큰이 실제로는 의미를 저장하는 레지스터처럼 작동한다는 것이다.

핵심 내용

• 템플릿 토큰은 처음에는 프롬프트 정보를 거의 담지 않는다. 연구진은 프롬프트 내용 토큰과 구조적 템플릿 토큰을 분리해 분석했다. 텍스트 인코더 출력 단계에서 템플릿 토큰은 특정 프롬프트에 관한 정보가 거의 없는 것으로 나타난다.

• DiT 내부에서는 어텐션의 주요 수신점이 된다. 저자들은 어텐션 분해와 토큰 구간, 헤드, 레이어별 표적 개입을 결합한 인과적 해석 프레임워크를 사용했다. 그 결과 템플릿 토큰이 DiT 내부에서 이미지에서 텍스트로 향하는 어텐션의 강한 sink로 부상한다는 점을 보였다.

• 객체 정체성 유지에 실제로 기여한다. 이는 단순한 관찰이 아니다. 개입 실험을 통해 템플릿 토큰이 생성 과정에서 객체의 정체성을 유지하는 데 인과적 역할을 한다는 점을 확인했다.

• 의미 전달 경로는 간접적이다. 흥미로운 점은 프롬프트 토큰의 의미가 템플릿 토큰으로 직접 복사되는 것이 아니라는 점이다. 논문에 따르면 의미는 먼저 이미지 latent에 주입되고, 이후 이미지 latent에서 템플릿 토큰으로 다시 읽혀 들어간다.

• 해석 결과가 프루닝으로 이어진다. 이 메커니즘을 바탕으로 저자들은 학습이 필요 없는 DiT 프루닝 규칙을 제안했다. 프롬프트 토큰에 가장 강하게 주목하는 어텐션 헤드는 오히려 제거 가능성이 높으며, 이를 제거하면 어텐션 FLOPs를 20% 줄이면서 GenEval 점수 하락은 1.4점에 그친다.

• 생성 계산에는 역할 분담이 있다. 논문은 헤드와 깊이에 따라 의미 라우팅과 시각 합성이 분리되고, 생성 과정이 정체성 형성에서 전파와 정교화로 진행된다고 설명한다.

의미와 영향

이 연구는 DiT의 내부 의미 흐름을 보다 인과적으로 이해하려는 시도라는 점에서 중요하다. 어텐션 맵을 단순히 시각화하는 데 그치지 않고, 특정 토큰과 헤드, 레이어에 개입해 실제 생성 결과에 어떤 영향을 주는지 분석했다.

모델 이해 측면에서는 의미 표현이 내부에서 이동할 수 있음을 보여준다. 프롬프트 토큰이 의미를 시작할 수는 있지만, 생성 중 그 의미를 계속 붙잡는 역할은 다른 토큰이 맡을 수 있다. 실용 측면에서는 이러한 해석이 계산량 절감으로 이어질 수 있음을 보여준다. 어떤 어텐션 헤드가 기능적으로 덜 중요한지 알면 재학습 없이도 효율화를 시도할 수 있기 때문이다.

출처: Hugging Face Daily Papers