Prompt Reinjection: Alleviating Prompt Forgetting in Multimodal Diffusion Transformers

이 논문은 멀티모달 확산 트랜스포머에서 발생하는 프롬프트 망각 현상을 분석하고, 초기 층의 프롬프트 표현을 후기 층에 재주입하는 훈련 없는 '프롬프트 재주입' 기법을 제안하여 텍스트 - 이미지 생성의 지시 따르기 능력과 전반적인 품질을 향상시켰습니다.

핵심

🎨 핵심 주제: "생각이 사라지는 AI"와 "기억을 되살리는 기술"

1. 문제 상황: AI 의 '기억 상실증' (Prompt Forgetting)

생각해 보세요. 당신이 그림을 그리라고 AI 에게 아주 구체적인 지시를 내렸다고 가정해 봅시다.

"파란색 우산과 노란색 개가 있는, 오른쪽에 있는 빨간 자전거."

AI 는 이 지시를 듣고 그림을 그리기 시작합니다. 하지만 이 최신 AI 모델들은 그림을 그리는 과정에서 **수십 번의 단계 (레이어)**를 거치며 이미지를 다듬습니다.

• 비유: 마치 긴 여정을 떠나는 것과 같습니다.
  • 출발할 때 (초기 단계) 는 "파란 우산, 노란 개, 빨간 자전거"라는 지시를 아주 선명하게 기억하고 있습니다.
  • 하지만 여정이 길어질수록 (레이어가 깊어질수록), AI 는 "아, 파란 우산이 필요했지... 아, 노란 개도 있었지..."라고 생각하다가, 중간 단계에서 지시 내용을 조금씩 잊어버리기 시작합니다.
  • 최종적으로 그림이 완성될 때는, "빨간 자전거"가 "초록색 자전거"가 되거나, "노란 개"가 "검은 고양이"로 변해버리는 일이 생깁니다.

이 현상을 논문에서는 **'프롬프트 망각 (Prompt Forgetting)'**이라고 부릅니다. AI 가 그림을 그리는 과정에서 원래의 지시 (텍스트) 가 가진 세부적인 의미들이 점점 희미해져서 사라져버리는 것입니다.

2. 원인: 왜 잊어버릴까요?

AI 의 뇌 구조를 살펴보면, '텍스트 (지시)'와 '이미지 (그림)'가 같은 공간에서 함께 처리됩니다. 하지만 AI 를 훈련시킬 때는 오직 '그림'이 잘 그려지도록만 점수를 매깁니다.

• 비유: **요리사 (AI) 와 메뉴판 (지시)**의 관계입니다.
  • 요리사는 메뉴판 ("소금 적게, 고기 많이") 을 보고 요리를 시작합니다.
  • 하지만 요리사가 완성된 요리를 평가할 때, "메뉴판을 얼마나 잘 따랐는가?"는 중요하지 않고, **"요리 맛이 좋은가?"**만 평가합니다.
  • 그래서 요리사는 맛만 좋으면 메뉴판의 세부 지시 ("소금 적게") 를 무시하고, 임의로 소금을 더 넣거나 빼도 됩니다.
  • 결과적으로, 요리가 끝날 때쯤이면 메뉴판의 원래 지시는 요리사의 기억에서 사라져버립니다.

3. 해결책: "기억 주입 (Prompt Reinjection)"

저자들은 이 문제를 해결하기 위해 **'프롬프트 리인젝션 (Prompt Reinjection, 프롬프트 재주입)'**이라는 방법을 제안했습니다. 이 방법은 AI 를 다시 훈련시킬 필요 없이, 그림을 그리는 **실시간 (추론 단계)**에 적용할 수 있습니다.

• 비유: "여정 중에도 지도를 다시 확인하기"
  • AI 가 그림을 그리는 여정 (레이어) 을 진행할 때, **가장 처음에 받았던 선명한 지시 (초기 텍스트 정보)**를 중간중간 다시 꺼내서 AI 에게 보여줍니다.
  • 마치 긴 여행 중에도 "아, 우리는 '파란 우산'을 그려야 했지!"라고 초기 지도를 다시 펼쳐서 현재 위치와 비교하며 방향을 수정하는 것과 같습니다.
  • 이렇게 하면 AI 가 중간에 지시를 잊어버리지 않고, 마지막까지 원래의 지시 ("노란 개", "빨간 자전거") 를 정확히 따를 수 있습니다.

4. 실제 효과: 어떻게 달라졌나요?

이 기술을 적용한 결과, AI 의 그림 실력이 놀라워졌습니다.

• 숫자 맞추기: "4 마리의 개"를 그려달라고 하면, 예전에는 3 마리나 5 마리가 나왔다면, 이제는 정확히 4 마리를 그립니다.
• 위치 관계: "오른쪽에 있는 자전거"라고 하면, 왼쪽에 그리는 실수를 하지 않습니다.
• 색상과 질감: "금색 시계"와 "초록색 잔디"를 정확히 구분해서 그립니다.

논문에서 테스트한 결과, 기존에 가장 잘하는 것으로 알려진 최신 모델들 (SD3.5, FLUX 등) 도 이 기술을 적용하면 **지시 따르기 능력 (Instruction Following)**이 크게 향상되었습니다. 특히 공간 관계나 숫자 세기 같은 복잡한 지시를 잘 따르게 되었습니다.

5. 요약: 이 기술의 핵심 가치

1. 문제: 최신 AI 는 그림을 그리는 과정에서 원래의 지시 내용을 점점 잊어버립니다.
2. 원인: AI 가 그림의 '미학'만 중시하다 보니, '지시'의 세부 사항은 희생됩니다.
3. 해결: 그림을 그리는 중간중간 처음의 지시를 다시 주입하여 기억을 되살립니다.
4. 장점: AI 를 다시 훈련시킬 필요 없이, 단순히 지시만 다시 넣어주는 것만으로 훨씬 더 똑똑하고 정확한 그림을 그릴 수 있습니다.

한 줄 요약:

"AI 가 그림을 그리는 동안 잊어버리지 않도록, 여행 내내 초기 지시 (메뉴판) 를 계속 보여줘서 원하는 그림을 정확히 그려내게 만든 기술입니다."

기술 요약

논문 요약: Prompt Reinjection: Alleviating Prompt Forgetting in Multimodal Diffusion Transformers

이 논문은 텍스트-이미지 생성을 위한 멀티모달 확산 트랜스포머 (Multimodal Diffusion Transformers, MMDiTs) 에서 발생하는 '프롬프트 망각 (Prompt Forgetting)' 현상을 규명하고, 이를 해결하기 위한 학습이 필요 없는 (Training-free) 추론 시간 기법인 '프롬프트 재주입 (Prompt Reinjection)' 을 제안합니다.

---

1. 문제 정의: 프롬프트 망각 (Prompt Forgetting)

최근 등장한 SD3, SD3.5, FLUX, Qwen-Image 와 같은 MMDiTs 는 텍스트와 이미지 잠재 벡터 (Latent) 를 단일 트랜스포머 스택 내에서 결합하여 처리하며, 양방향 상호작용을 통해 복잡한 프롬프트를 이해하려 합니다. 그러나 저자들은 다음과 같은 근본적인 문제를 발견했습니다.

• 감독 불균형 (Supervisory Imbalance): 확산 모델의 손실 함수는 이미지 잠재 공간 (Visual Latent Space) 에만 정의됩니다. 따라서 이미지 토큰은 직접적인 감독을 받지만, 텍스트 토큰은 이미지 재구성을 통한 간접적인 피드백만 받습니다.
• 심층에서의 의미 소실: 이러한 불균형으로 인해, 트랜스포머 레이어가 깊어질수록 텍스트 표현 (Text Representations) 이 점진적으로 왜곡되거나 변형됩니다.
• 현상: 초기 레이어에서 명확했던 프롬프트의 세부 정보 (색상, 수량, 공간 관계 등) 가 깊은 레이어로 갈수록 복원 불가능해지거나 '망각'되는 현상이 발생합니다. 이를 프롬프트 망각이라고 명명했습니다.

2. 방법론: 프롬프트 재주입 (Prompt Reinjection)

저자는 이 문제를 해결하기 위해 모델의 가중치를 수정하지 않고 추론 (Inference) 시에 적용할 수 있는 프롬프트 재주입 기법을 제안합니다. 이 방법은 얕은 레이어의 고품질 텍스트 특징을 깊은 레이어로 다시 주입하여 망각을 방지합니다.

핵심 메커니즘

1. 얕은 레이어 특징 추출 (Origin Layer): 텍스트 인코더 출력 직후의 초기 MMDiT 블록 (얕은 레이어) 에서 텍스트 토큰 특징을 추출합니다. 이 시점의 특징이 프롬프트의 의미적 충실도 (Semantic Fidelity) 가 가장 높습니다.
2. 깊은 레이어 주입 (Target Layers): 추출된 특징을 깊은 레이어의 텍스트 특징에 잔차 (Residual) 형태로 더합니다.
3. 정렬 및 보정 (Alignment): 서로 다른 깊이의 특징은 분포 (Distribution) 와 기하학적 구조 (Geometry) 가 달라 직접 더하면 불안정할 수 있습니다. 이를 해결하기 위해 두 단계의 정렬을 수행합니다.
   • 분포 고정 (Distribution Anchoring): 레이어 정규화 (Layer Normalization) 를 통해 특징의 평균과 분산을 표준화한 후, 원래 레이어의 통계적 특성으로 다시 매핑합니다.
   • 기하학적 정렬 (Geometry Alignment): 직교 프로크루스테스 변환 (Orthogonal Procrustes Transform) 을 사용하여 원본 특징 공간과 목표 레이어 공간 간의 좌표계 회전 (Rotation) 차이를 보정합니다. 이는 COCO-5K 와 같은 데이터셋을 사용하여 한 번만 계산하면 됩니다.

수식적으로, 목표 레이어 $l$의 텍스트 특징 $T^{(l)}$은 다음과 같이 업데이트됩니다:
$$T^{(l)}_{new} = \text{Restoration}(\hat{T}^{(l)} + w \cdot \hat{T}^{(ori)} R)$$
여기서 $w$는 주입 강도, $R$은 정렬 행렬, $\hat{T}$는 정규화된 특징입니다.

3. 주요 기여 및 분석

• 망각 현상의 정량적 증명:
  • CKNNA (Conditional K-Nearest Neighbor Alignment): 입력 단계의 국소적 의미 구조가 깊은 레이어로 갈수록 유지되지 않음을 확인했습니다.
  • 레이어별 프로빙 (Layer-wise Probing): 각 레이어의 텍스트 특징에서 언어적 속성 (명사, 형용사, 공간 관계, 숫자 등) 을 분류하는 경량 분류기를 훈련시켰습니다. 그 결과, 깊이가 깊어질수록 분류 정확도가 단조 감소하는 것을 확인했습니다. 특히 공간 관계 (Spatial Relation) 와 숫자 (Counting) 정보가 가장 빠르게 망각되었습니다.
• 학습 불필요 (Training-free): 모델 재학습이나 미세 조정 (Fine-tuning) 없이 추론 시에만 적용 가능하여 적용 비용이 낮습니다.

4. 실험 결과

저자는 SD3-medium, SD3.5-large, FLUX.1-dev, Qwen-Image 등 4 가지 주요 MMDiT 모델에 대해 실험을 수행했습니다.

• 지시 따르기 (Instruction Following) 향상:
  • GenEval: 전체 점수가 SD3.5 는 6.48%, FLUX 는 5.64% 향상되었습니다. 특히 공간 관계 (Position) 와 숫자 세기 (Counting) 태스크에서 가장 큰 개선을 보였습니다.
  • DPG-Bench 및 T2I-CompBench++: 객체 속성 바인딩, 복잡한 구성, 공간 관계 등 다양한 벤치마크에서 일관된 성능 향상을 보였습니다.
• 이미지 품질 유지:
  • HPSv2, ImageReward, PickScore 와 같은 인간 선호도 지표와 CLIP 점수에서도 성능이 유지되거나 소폭 향상되었습니다. 이는 프롬프트 준수 능력이 향상되면서도 이미지 생성의 미적 품질이나 전반적인 일관성이 해치지 않음을 의미합니다.
• 정성적 분석:
  • "네 개의 개", "검은색 연과 초록색 곰", "공기 위쪽의 키" 등 복잡한 프롬프트에서 기존 모델이 속성이나 위치를 잘못 생성하는 반면, 제안된 방법은 프롬프트의 모든 제약을 정확하게 따르는 이미지를 생성했습니다.

5. 의의 및 결론

이 논문은 MMDiTs 아키텍처의 내부 동작 원리에 대한 중요한 통찰을 제공했습니다. 단순히 텍스트를 외부 조건으로 주입하는 것이 아니라, 텍스트와 이미지가 결합되어 처리되는 과정에서 발생하는 심층적 의미 소실을 규명하고, 이를 얕은 레이어의 특징을 재주입함으로써 해결할 수 있음을 증명했습니다.

• 실용성: 추가적인 학습 비용 없이 기존 모델의 성능을 즉시 향상시킬 수 있는 유효한 방법론을 제시했습니다.
• 확장성: 다양한 MMDiT 모델 (SD3, FLUX, Qwen 등) 에 적용 가능하며, 복잡한 다중 객체 구성 및 정밀한 공간 제어와 같은 고급 생성 작업에 필수적인 기술로 평가됩니다.

결론적으로, Prompt Reinjection은 텍스트-이미지 생성 모델이 프롬프트의 세부 사항을 깊이 있는 처리 단계까지 유지하도록 하여, 더 정확하고 복잡한 지시를 따르는 고품질 이미지를 생성할 수 있게 하는 획기적인 접근법입니다.