인컨텍스트 이미지 생성용 토큰 프루닝 기법 ToPi

초록

본 논문은 Diffusion Transformer(DiT)에서 참조 이미지와 목표 라티스 토큰을 동시에 처리하는 인컨텍스트 생성 작업의 연산량 폭증 문제를 해결하기 위해, 사전 캘리브레이션 기반의 중요도 분석과 시간에 따라 동적으로 토큰을 선택·제거하는 훈련‑프리 토큰 프루닝 프레임워크 ToPi를 제안한다. 핵심 레이어를 선별하고, 각 레퍼런스 토큰의 기여도를 정량화한 영향도 지표를 이용해 불필요한 토큰을 제거함으로써 30 % 이상의 추론 속도 향상을 달성하면서도 시각적 일관성과 구조적 충실도를 유지한다.

상세 분석

ToPi는 인컨텍스트 이미지‑투‑이미지·비디오 생성 시, 레퍼런스 토큰과 타깃 라티스 토큰이 서로 다른 역할을 수행한다는 ‘역할 비대칭성’에 주목한다. 저자들은 이 비대칭성을 공간·시간·기능 세 차원에서 정량적으로 분석하였다. 공간적으로는 전체 DiT 레이어 중 소수의 ‘핵심 레이어’가 레퍼런스 토큰과의 상호작용을 집중적으로 담당한다는 사실을 발견했으며, 이는 레이어별 컨텍스트 민감도 점수(S(ℓ))를 통해 사전 캘리브레이션 단계에서 자동으로 식별된다. 시간적으로는 디노이징 초기 고노이즈 단계에서 레퍼런스 토큰의 중요도가 높고, 진행될수록 자체 토큰 간의 self‑attention이 주도적으로 변하면서 레퍼런스 의존도가 감소한다. 이를 ‘전체 주의 점수’의 감소 곡선으로 시각화하였다. 기능적으로는 작업 난이도에 따라 레퍼런스 토큰의 필요량이 달라진다. 예를 들어, 스타일 변환처럼 큰 자유도를 요구하는 경우는 빠르게 토큰을 축소해도 무방하지만, 리사이즈·디테일 보존과 같은 정밀 작업에서는 핵심 토큰을 오래 유지해야 한다.

이러한 관찰을 바탕으로 ToPi는 세 가지 메커니즘을 결합한다. 첫째, 오프라인 캘리브레이션을 통해 Top‑M 레이어를 선정하고, 이 레이어에서만 컨텍스트 민감도와 토큰‑레벨 주의 가중치를 수집한다. 둘째, ‘Attention‑Weighted Importance Scoring’ 단계에서 각 레퍼런스 토큰 j에 대해 목표‑인식 주의 점수 S_j를 계산하고, 이를 시간‑간격 ΔT마다 업데이트한다. 셋째, 매 디노이징 스텝마다 현재 마스크에 따라 토큰을 프루닝하고, 남은 토큰을 재정렬해 DiT에 입력한다. 이 과정은 훈련 없이도 실행 가능하도록 설계되었으며, 토큰 선택·제거가 디노이징 흐름에 실시간으로 반영된다.

실험에서는 Flux.1‑K와 Qwen‑Image‑Edit 두 대형 모델을 기반으로 AnyEdit 벤치마크에서 30 % 이상의 속도 향상을 기록했으며, FID·LPIPS·CLIP‑Score 등 정량 지표와 시각적 평가 모두에서 기존 전체 토큰 사용 대비 품질 저하가 거의 없음을 입증했다. 특히, 복잡한 레퍼런스가 다수 포함된 고해상도 이미지·비디오 생성 시, 토큰 프루닝 비율을 40 %까지 늘려도 구조적 일관성과 색감 재현성이 유지되는 점이 눈에 띈다.

ToPi의 주요 기여는 (1) 인컨텍스트 생성에서 토큰 비대칭성을 체계적으로 분석하고, (2) 레이어‑선별·시간‑가변 중요도 스코어링을 결합한 훈련‑프리 프루닝 프레임워크를 제시했으며, (3) 다양한 이미지·비디오 편집 작업에서 연산 효율과 품질을 동시에 개선한 점이다. 향후 연구에서는 프루닝 전략을 다른 확산 기반 멀티모달 모델에 확장하거나, 프루닝 과정에서 발생할 수 있는 미세한 디테일 손실을 보정하기 위한 복원 모듈을 도입하는 방향이 기대된다.