확산 트랜스포머의 문맥 토큰을 자연어로 읽어 생성 품질을 끌어올린다

무엇인가

멀티모달 확산 트랜스포머(MM-DiT)는 텍스트와 시각 표현을 생성 전 과정에서 함께 처리한다. 이 구조에서 텍스트 토큰은 고정된 조건 벡터로 남지 않고 멀티모달 어텐션을 통해 반복적으로 갱신되며, 생성 중인 이미지로부터 정보를 흡수해 '문맥 토큰(contextual token)'으로 진화한다. 문제는 시각 토큰과의 비대칭성이다. 시각 토큰은 이미지의 공간 표현을 이루고 디노이징 목적함수로 직접 감독되기 때문에 중간 특징을 이미지로 시각화하거나 개입할 수 있다. 반면 문맥 토큰은 공간 대응이 없고 전용 감독도 받지 않아, 무엇을 인코딩하는지, 그 동역학이 생성을 어떻게 조종하는지 알려지지 않았다. 이 논문은 그 문맥 공간을 자연어 질의로 직접 읽는 프레임워크를 제안한다.

어떻게 동작하나

핵심 도구는 Contextual Reader다. 디노이징 타임스텝 t에서 각 트랜스포머 블록 l은 문맥 토큰 시퀀스 c(l)(t) ∈ R^(M×d)를 만든다(M은 패딩 포함 텍스트 시퀀스 길이, d는 은닉 차원). 저자들은 L개 블록 전체의 문맥 토큰을 모으고, 각 표현이 어느 레이어·어느 시점에서 추출됐는지 보존하기 위해 학습 가능한 레이어 임베딩과 타임스텝 임베딩을 덧붙인 뒤 시퀀스 차원으로 이어 붙여 C(t) ∈ R^(LM×d)를 만든다. 이어 LLaVA식 언어 인터페이스를 차용해, Q개의 학습 가능한 쿼리 토큰을 가진 2층 Q-Former 병목 네트워크가 크로스 어텐션으로 C(t)를 압축하고(Q ≪ L·M), 선형 투영이 그 결과를 동결된 LLM의 입력 임베딩 공간으로 매핑해 Contextual Descriptor p ∈ R^(Q×d_LLM)를 만든다. LLM은 이 디스크립터와 토큰화된 질문을 받아 답을 자기회귀적으로 생성한다. 학습되는 것은 병목 네트워크뿐이고 MM-DiT와 LLM은 모두 동결된다.

무엇과 다른가

감독 데이터는 MM-DiT 생성 과정 자체에서 만든다. 단일 인물 피사체가 담긴 MS-COCO 프롬프트를 사용해 프롬프트당 5개 시드를 생성했고, 2,000개 학습 프롬프트에서 10,000개 생성물, 500개 프롬프트에서 2,500개 홀드아웃 생성물을 확보했다. 질문은 카메라 앵글·배경 같은 거친 장면 속성부터 포즈·외형·사물 상호작용 같은 세밀한 속성까지 15개이며, 정답은 기성 VLM이 생성 이미지에 답한 결과를 쓴다. 프롬프트와 초기 노이즈를 저장해 두었다가 임의 타임스텝의 문맥 표현을 복원하는 방식으로, 리더는 5,000회 반복·배치 256으로 학습됐고 검증 손실은 처음 2,000회에서 급격히 감소한 뒤 완만해졌다.

어떻게 쓰나

읽어낸 내용은 구체적이다. 조건부 설정에서 같은 프롬프트로 만든 두 시드는 프롬프트가 명시하지 않은 속성을 서로 다르게 해소하는데, 리더는 그 시드별 차이를 문맥 토큰에서 구분해 낸다. "A person holding a scarf"에서는 한 시드의 배경을 "푸른 덤불과 나무", 다른 시드를 "민무늬 흰 벽"으로 답했고, "A person looking into a mirror"에서는 한 시드에서 스마트폰을 들고 있음을 짚어냈으며 손이 화면 밖으로 나간 다른 시드에서는 "Unknown"이라고 답했다. 이런 구분은 해당 시각 예측에서 뚜렷이 보이기 전인 이른 디노이징 단계에서 이미 가능하다. 프롬프트를 아예 비운 무조건부 설정에서도 8% 시점에 인물이 스키를 타고 있음을, 20% 시점에 고글을 쓰고 있음을 읽어냈다. VLM-as-a-judge로 정량화하면 두 아키텍처 모두에서 가독성은 디노이징 내내 상승해 80% 지점에서 최고치를 찍고 마지막 100%에서 소폭 하락한다. 빈 프롬프트 쪽은 훨씬 낮게 출발해 빠르게 상승하며, FLUX.2에서는 조건부 리더를 최종적으로 앞지르고 SD3.5에서는 조건부 성능에 근접한다. 또한 프롬프트별 5개 생성물을 리더 점수로 순위 매겨 상위와 하위의 HPS 차이를 계산하면, 아키텍처와 조건 설정 전반에서 ΔHPS가 일관되게 양수다. 문맥 가독성이 높은 생성물이 인간 선호 점수도 높다는 뜻이다.

전제와 한계

이 관찰을 학습 신호로 바꾼 것이 Contextual Alignment(CoAl)다. 동결된 시맨틱 인코더가 깨끗한 이미지를 전역 teacher 임베딩으로 매핑하고(대조 학습으로 캡션 임베딩과 정렬되도록 훈련된 pooled SigLIP 이미지 임베딩을 사용), Aligner Network라 부르는 소형 Q-Former가 단일 레이어의 문맥 토큰을 예측 임베딩으로 매핑한 뒤 음의 코사인 유사도 손실로 teacher와 정렬한다. Aligner는 학습에만 쓰이므로 MM-DiT의 추론 비용은 변하지 않는다. Aligner는 단일 레이어에 학습 쿼리 하나만 두는 극단적으로 작은 구조여서, 정렬 목표를 달성하려면 Aligner가 아니라 MM-DiT 자체가 적응해야 한다. 또한 이른 레이어 하나만 정렬해 나머지 레이어는 자유롭게 두는데, 이 선택과 teacher 임베딩 선택은 별도 실험으로 확인됐다고 밝힌다. 여기에 Null-Condition Weighting을 더한다. CFG를 위해 프롬프트를 드롭하는 반복에서 문맥 토큰은 빈 프롬프트에서 초기화되므로, 프롬프트 정보를 이용해 정렬 손실을 줄이는 지름길이 막히고 이미지 토큰에서 문맥 토큰으로 정보가 흘러가도록 강제된다. 저자들은 이 가중치를 전체 학습의 후반부나 파인튜닝 전 구간에 적용한다.

실험은 두 레짐으로 나뉜다. 하나는 MS-COCO 2014에서 듀얼 스트림 MM-DiT를 랜덤 초기화로 처음부터 학습하는 것이고, 다른 하나는 사전학습된 SD3를 장문 프롬프트와 전문 사진을 짝지은 Fine-T2I의 큐레이션 서브셋으로 파인튜닝하는 것이다. 베이스라인은 REPA와 그 후속인 HASTE, SRA다. 전체 학습에서는 FID, CLIP, HPS, DINOv2 Precision/Recall을 보고하며, CoAl은 세 정렬 베이스라인 모두를 모든 지표에서 개선해 전역 시맨틱 정렬이 기존 시각 정렬과 상보적임을 보인다. 파인튜닝에서는 CLIP과 HPS를 제외하는데, 정답 사진 자체가 베이스 모델보다 이 지표들에서 낮게 나와 비교가 성립하지 않기 때문이다. 대신 FID, KID, DINOv2 Precision/Recall로 목표 분포 적합도와 커버리지를 본다. 이 설정에서 REPA·HASTE·SRA는 vanilla 파인튜닝 대비 유의미한 이득이 없고 셋 다 FID가 오히려 나빠지지만, CoAl은 최고 FID와 Recall을 기록하고 Precision에서도 경쟁력 있는 수준을 유지한다. 다만 본문 발췌에는 표의 구체적 수치가 제시되지 않아 개선폭의 크기는 확인할 수 없다.

개발자 관점에서 이 논문이 주는 실용적 함의는 두 가지다. 첫째, FLUX나 SD3.5, Qwen-Image처럼 멀티모달 어텐션을 쓰는 생성 모델을 다룰 때 문맥 토큰은 디버깅 창이 된다. 프롬프트가 비어 있어도, 그리고 시각 예측에 드러나기 전에도 모델이 무엇을 그릴지 이미 담고 있으므로, 실패한 생성의 원인을 사후 이미지가 아니라 중간 표현에서 추적할 수 있다. 둘째, CoAl은 추론 비용을 전혀 늘리지 않는 학습 시 정규화 항이다. 특히 REPA 계열 정렬이 효과를 잃는 파인튜닝 국면에서 여전히 작동한다는 점이 실무적으로 중요하다. 다만 리더 자체는 단일 인물 피사체와 15개 질문으로 학습됐으므로, 다른 도메인에 그대로 적용하려면 별도 검증이 필요하다.

저자들이 밝힌 한계도 분명하다. 이 내부 표현이 어떻게 형성되고 진화하며 최종 이미지를 어떻게 빚어내는지는 여전히 거의 탐구되지 않았다고 결론에서 인정한다. 의미 정보가 지각 가능해지기 훨씬 전에 나타난다는 사실 자체가 생성 과정에 대한 새로운 질문을 던진다는 것이다. 또한 리더의 감독 데이터는 단일 인물 중심 이미지와 15개 질문으로 제한되며, 파인튜닝 평가에서 CLIP과 HPS를 배제한 것은 목표 데이터셋 특성과 일반 선호 지표가 어긋나기 때문이라는 전제를 깔고 있다. CoAl의 teacher 임베딩과 정렬 레이어 선택은 경험적 확인에 의존한다는 점도 남는다.