[Paper] Prompt Reinjection: 멀티모달 디퓨전 트랜스포머에서 Prompt Forgetting 완화

Source: arXiv - 2602.06886v1

개요

논문 “Prompt Reinjection: Alleviating Prompt Forgetting in Multimodal Diffusion Transformers” 은 오늘날 최첨단 텍스트‑투‑이미지 모델(e.g., SD3, SD3.5, FLUX.1)에서 발견되는 미묘하지만 중요한 결함을 밝혀낸다. 확산 과정이 진행됨에 따라 모델은 원래 텍스트 프롬프트를 점차 “잊어버리게” 되며, 이로 인해 생성된 이미지가 사용자의 의도와 멀어질 수 있다. 저자들은 훈련이 필요 없는 “프롬프트 재주입(prompt reinjection)” 기법을 제안하는데, 이는 더 깊은 층에서 프롬프트의 영향을 복원하여 지시와 생성된 이미지 사이의 정렬을 눈에 띄게 개선한다.

핵심 기여

• 프롬프트 망각 식별: 경험적 분석에 따르면 멀티모달 디퓨전 트랜스포머(MMDiTs)의 텍스트 브랜치에서 프롬프트 표현의 의미 강도가 감소한다.
• 프롬프트 재주입 메커니즘: 초기 레이어 프롬프트 임베딩을 후속 레이어에 복사하는 간단한 추론‑전용 방법으로, 모델에게 원래 지시를 “상기”시킨다.
• 광범위한 경험적 검증: 세 가지 벤치마크 스위트(GenEval, DPG, T2I‑CompBench++)에 대한 실험에서 세 주요 MMDiT 패밀리 전반에 걸쳐 지시 수행, 미적 선호, 전반적인 생성 품질이 일관되게 향상됨을 보여준다.
• 학습 없이 배포: 이 기술은 추가 파인‑튜닝이나 추가 파라미터가 필요 없으며, 기존 파이프라인에 즉시 적용 가능하다.

Methodology

1. Probing Prompt Representations:

   • 저자들은 세 가지 인기 있는 MMDiTs에 대해 각 diffusion 단계에서 텍스트 브랜치의 숨겨진 상태를 추출합니다.
   • 언어 탐지기(예: 품사, 감성, 의미 유사도 분류기)를 사용하여 깊이가 증가함에 따라 원본 프롬프트 의미가 얼마나 유지되는지를 정량화합니다.

2. Prompt Reinjection Design:

   • “소스” 레이어를 선택합니다(보통 프롬프트가 아직 강하게 남아 있는 초기 레이어).
   • 이후 각 레이어에서 소스 프롬프트 임베딩을 현재 텍스트 토큰 임베딩에 연결하거나 더합니다.
   • 이 연산은 추론 단계에서만 수행되며, 학습된 가중치는 전혀 변경되지 않습니다.

3. Evaluation Protocol:

   • GenEval(일반 텍스트‑투‑이미지 생성), DPG(지시 프롬프트 생성), 그리고 T2I‑CompBench++(비교 벤치마크)를 사용해 지시 준수도, 미학 점수(예: CLIP‑기반 선호도), 그리고 전통적인 이미지 품질 지표(FID, IS)를 평가합니다.

결과 및 발견

• Instruction Following: 프롬프트 재주입은 미묘한 프롬프트(예: “빈티지 우주복 헬멧을 쓴 고양이”)를 정확히 반영하는 이미지 비율을 8‑12 % 상승시킵니다.
• Aesthetic & Preference Gains: 인간 평가자들은 일관되게 재주입 결과물을 선호하며, 이 기법이 관련성과 시각적 매력을 모두 향상시킴을 나타냅니다.
• Cross‑Model Consistency: 세 모델 모두 MMDiT가 혜택을 받아, 프롬프트 망각이 특정 아키텍처의 버그가 아니라 일반적인 현상임을 확인합니다.

Practical Implications

• Instant Upgrade for Existing Services: Companies running SD3/FLUX‑based APIs can integrate prompt reinjection with a single line of code, delivering sharper, more faithful images without retraining.
• Better User Experience in Creative Apps: Designers and marketers who rely on precise textual cues (e.g., “minimalist logo with teal accents”) will see fewer off‑target results, reducing iteration cycles.
• Improved Safety & Alignment: By keeping the model anchored to the original prompt, the risk of unintended or harmful content drift is lowered—important for moderation pipelines.
• Foundation for Future Research: The reinjection idea could inspire similar “memory‑preserving” tricks in other multimodal transformers (e.g., video generation, audio‑text synthesis).

Limitations & Future Work

• Layer Selection Heuristics: 현재 접근 방식은 고정된 초기 레이어를 선택합니다; 프롬프트 복잡도에 기반한 적응형 선택이 추가적인 향상을 가져올 수 있습니다.
• Potential Over‑reinforcement: 매우 짧거나 모호한 프롬프트의 경우, 동일한 임베딩을 반복적으로 주입하면 잡음이 증폭될 수 있습니다; 강화 강도의 균형은 아직 해결해야 할 과제입니다.
• Evaluation Scope: 벤치마크는 다양한 프롬프트를 다루지만, 실제 사용 환경(예: 다문장 지시, 인터랙티브 편집)에서는 더 깊은 연구가 필요합니다.
• Extending Beyond Diffusion Transformers: 저자들은 프롬프트 재주입을 자동회귀 멀티모달 모델 및 멀티모달 검색 시스템에서도 탐색할 것을 제안합니다.

Bottom line: Prompt reinjection은 현대 텍스트‑투‑이미지 디퓨전 모델의 숨겨진 약점을 조명하고, 재학습 비용 없이도 충실도, 미학, 안전성을 향상시키는 플러그‑인‑플레이 해결책을 제공합니다. 차세대 AI 기반 창작 도구를 구축하는 개발자에게는 즉시 적용해 볼 가치가 있는 저비용 업그레이드입니다.

저자

• Yuxuan Yao
• Yuxuan Chen
• Hui Li
• Kaihui Cheng
• Qipeng Guo
• Yuwei Sun
• Zilong Dong
• Jingdong Wang
• Siyu Zhu

논문 정보

• arXiv ID: 2602.06886v1
• 카테고리: cs.CV
• 발행일: 2026년 2월 6일
• PDF: Download PDF