기술 요약: TextLDM

문제 제기

다중 모달 모델링 분야는 시각적 콘텐츠와 텍스트 콘텐츠를 모두 매끄럽게 생성할 수 있는 통합 프레임워크를 추구합니다. VAE 잠재 공간에서 흐름 매칭 (flow matching) 으로 훈련된 Diffusion Transformers(DiT) 가 이미지 및 비디오 생성을 성공적으로 통합한 반면, 언어 모델링은 여전히 자기회귀 (AR) 패러다임이 지배하고 있습니다. 이해와 생성을 통합하려는 기존 시도들은 대부분 이산 토큰 기반 패러다임이나 이산 확산에 의존하여 방법론적 간극을 남겼습니다. 연속 잠재 확산을 텍스트에 적용하는 데 있어 핵심적인 과제는 확산 백본 자체가 아니라 연속 텍스트 표현의 품질입니다. 토큰 재구성을 목적으로만 훈련된 VAE 는 높은 충실도를 달성할 수 있지만, 조건부 탈노이즈에 적합하지 않은 잠재 공간을 생성하여 최적화되지 않은 생성 품질을 초래합니다.

방법론: TextLDM

저자들은 최소한의 아키텍처 수정으로 시각적 잠재 확산 레시피를 언어 모델링에 이전하는 TextLDM 프레임워크를 제안합니다. 이 시스템은 두 단계 프로세스로 작동합니다:

1. TextVAE: 연속 잠재 표현

토큰 기반 변분 오토인코더 (TextVAE) 는 이산 토큰을 연속 잠재 벡터로 매핑합니다.

• 아키텍처: 이전 접근법들이 토큰 시퀀스를 더 짧은 잠재 시퀀스로 압축하는 것과 달리, TextVAE 는 각 토큰이 하나의 연속 잠재 벡터에 대응하는 1 대 1 매핑을 유지합니다. 인코더 ($E_\phi$) 와 디코더 ($D_\psi$) 는 모두 Transformer 입니다.
• 표현 정렬 (REPA): 잠재 효과성의 병목 현상을 해결하기 위해 저자들은 REPA 를 도입합니다. VAE 인코더는 중간 표현에 대한 코사인 유사도 손실을 통해 프리트레인된 언어 모델 (Qwen3-1.7B) 과 정렬됩니다. 이는 재구성 충실도를 훼손하지 않으면서 확산 기반 생성에 더 적합한 잠재 기하학을 형성합니다.
• 훈련 목적 함수: VAE 는 교차 엔트로피 재구성 손실, KL 발산 정규화, 그리고 REPA 정렬 손실로 구성된 복합 손실 함수로 훈련됩니다.

2. TextDiT: 잠재 공간 내 흐름 매칭

시각적 대응물과 아키텍처적으로 동일한 표준 Diffusion Transformer(TextDiT) 가 학습된 잠재 공간 내에서 흐름 매칭을 수행합니다.

• 조건부 공식화: 이 모델은 언어 생성을 조건부 과정으로 간주하여, 깨끗한 컨텍스트 잠재 ($z_c$) 와 노이즈가 있는 대상 잠재 ($z_t$) 를 입력으로 받아 속도장을 예측합니다.
• 흐름 매칭: 모델은 가우시안 노이즈와 대상 잠재 사이를 보간하는 조건부 흐름 매칭 (CFM) 손실을 사용하여 최적화됩니다.
• 스케줄링 및 가이드: 시스템은 Stable Diffusion 3 에서 효과성이 입증된 logit-normal 타임스텝 스케줄을 채택하고, 생성 품질을 향상시키기 위해 Classifier-Free Guidance(CFG) 를 활용합니다.
• 추론: 생성은 토큰 수준의 자기회귀 디코딩을 피하면서 Euler ODE 솔버를 통해 병렬로 수행됩니다. 이는 시퀀스 길이와 무관한 일정한 추론 효율성 (일정한 함수 평가 횟수) 을 가능하게 합니다.

주요 기여

1. 통합 아키텍처: TextLDM 은 시각적 확산 레시피 (VAE + DiT + 흐름 매칭 + CFG) 를 언어 모델링에 성공적으로 이전하여 통합 다중 모달 아키텍처를 향한 한 걸음을 내딛었습니다. 전체 시스템은 확산 프로세스를 위한 프리트레인된 인코더나 디코더에 의존하지 않고 처음부터 훈련됩니다.
2. 표현 병목 현상의 식별: 저자들은 재구성 충실도만으로는 고품질 텍스트 확산에 불충분함을 식별했습니다. 그들은 표현 정렬 (REPA) 이 하류 생성 품질을 지원하도록 잠재 공간을 형성하는 데 결정적임을 입증했습니다.
3. 최첨단 성능: TextLDM 은 확산 언어 모델 간에 최첨단 결과를 달성하며, 동일한 설정 하에서 자기회귀 베이스라인 (GPT-2) 과 대등한 성능을 보입니다.
4. 효율성: 이 프레임워크는 자기회귀 모델에서 함수 평가 횟수의 선형 성장에 비해 더 일정한 시간 생성 프로파일을 제공합니다.

실험 결과

모델들은 OpenWebText2에서 처음부터 훈련되었으며, One Billion Words, TinyStories, Wikipedia, WikiSource의 네 가지 벤치마크에 걸친 텍스트 연속성 평가에서 검증되었습니다.

• 확산 베이스라인 대비 성능: TextLDM 은 모든 지표 (ROUGE, BERTScore, MAUVE) 에서 이전의 연속 (예: SSD-LM) 및 이산 (예: Block Diffusion) 확산 언어 모델을 크게 능가합니다.
• 자기회귀 베이스라인 대비 성능:
  • TinyStories와 One Billion Words에서 TextLDM 은 유사하거나 더 큰 크기의 GPT-2 모델과 ROUGE 지표에서 대등하거나 더 나은 성능을 보입니다.
  • 분포 외 벤치마크 (Wikipedia, WikiSource) 에서 768M TextLDM 변형은 모든 GPT-2 베이스라인을 능가합니다.
  • 일부 벤치마크에서 BERTScore 는 여전히 약간의 격차가 있으며, 여기서 AR 모델이 우위를 유지합니다.
• 확장 행동: TextLDM 은 일관된 확장 행동을 보여주며, DiT 크기가 114M 에서 768M 파라미터로 증가함에 따라 성능이 크게 향상됩니다.
• 훈련 효율성: 훈련 역학은 TextLDM 이 대부분의 벤치마크에서 유사한 수의 훈련 단계 내에서 GPT-2-medium 성능과 대등하거나 더 나은 성능을 보임을 나타내며, 이는 훨씬 더 많은 컴퓨팅 자원을 필요로 했던 이전 연속 확산 모델에 비해 상당한 개선입니다.
• 애블레이션 통찰:
  • REPA: REPA 를 제거하면 생성 품질이 크게 저하되어 잠재 공간 구조에 대한 REPA 의 필요성이 확인됩니다.
  • 재구성과 생성: 모든 VAE 구성이 거의 완벽한 토큰 재구성 (>97.5%) 을 달성했음에도 불구하고, 생성 품질은 REPA 와 잠재 차원에 따라 크게 달라져 재구성이 아닌 잠재 기하학이 성능을 주도함을 증명합니다.
  • 하이퍼파라미터: logit-normal 타임스텝 스케줄 (std=1.5) 과 CFG 스케일 7 이 최적의 결과를 제공했습니다.

중요성과 주장

이 논문은 TextLDM 이 "시각적 DiT 레시피"가 언어에 효과적으로 이전될 수 있음을 입증했다고 주장합니다. 통합 아키텍처가 시각 합성과 텍스트 생성 모두를 처리할 수 있음을 증명함으로써, 이 연구는 향후 다중 모달 생성 및 이해를 위한 통합 확산 아키텍처에 대한 구체적인 기반을 제공합니다. 저자들은 성공의 열쇠가 확산 백본 자체뿐만 아니라, 이산 텍스트를 연속 확산 프레임워크에서 효과적으로 모델링할 수 있게 해주는 표현 정렬을 통한 잠재 공간의 특정 적응에 있었다고 강조합니다.

저자들은 두 단계 훈련 파이프라인이 복잡성을 추가하고, 도메인 외 데이터 (예: Wikipedia) 에 대한 재구성 정확도가 도메인 내 데이터보다 약간 낮아 오류가 전파될 수 있음을 지적하며 한계에 대해 겸손한 입장을 유지합니다. 그들은 더 다양한 데이터로 훈련 코퍼스를 확장하면 이러한 문제를 완화할 수 있다고 제안합니다.