기술 요약: FullFlow

문제 정의

현대적인 텍스트-이미지 (T2I) 확산 및 정류 흐름 (rectified-flow) 모델은 풍부한 시각적 사전 지식과 강력한 생성 능력을 보유하고 있지만, 엄격히 단방향 방식으로 작동합니다: 텍스트 입력, 이미지 출력. 통합된 비전 - 언어 모델 (VLM) 이 존재하지만, 이러한 모델들은 일반적으로 사전 훈련된 언어 모델을 시각 인코더로 조건부 설정하거나 처음부터 고비용의 결합 사전 훈련을 수행하여 구축됩니다. 이러한 접근 방식들은 종종 T2I 백본에 이미 인코딩된 강력한 이미지 사전 지식을 폐기하거나, 계산적으로 금지적인 상당한 재훈련을 요구합니다. 본 연구가 다루는 핵심 질문은 사전 훈련된 T2I 흐름 모델을 완전한 멀티모달 사전 훈련 없이도 원래 이미지 사전 지식을 희생하지 않고 양방향 비전 - 언어 생성기로 업그레이드할 수 있는지 여부입니다.

방법론

저자들은 사전 훈련된 정류 흐름 T2I 모델 (구체적으로 Stable Diffusion 3 및 FLUX.1-dev) 을 양방향 생성기로 업그레이드하는 파라미터 효율적인 레시피인 FullFlow를 제안합니다. 핵심 혁신은 단일 공유 백본 내에서 두 모달리티를 비대칭적이지만 호환되게 처리하는 데 있습니다.

1. 네이티브 공간 멀티모달 확장

• 이미지 모달리티: 이미지는 네이티브 연속 정류 흐름 잠재 공간에 유지됩니다. 모델은 사전 훈련된 연속 이미지 과정을 보존합니다.
• 텍스트 모달리티: 텍스트는 Edit Flows에 기반한 이산 삽입 과정을 통해 처리됩니다. frozen 인코더와 호환되지 않는 어휘 외 토큰을 도입하는 마스크 기반 이산 확산과 달리, FullFlow 는 삭제 기반의 전방 과정을 사용합니다. 이는 임의의 노이즈 레벨에서 손상된 텍스트 상태가 보조 인코더 (예: CLIP) 를 위해 디코딩 및 재 토큰화될 수 있는 유효한 토큰 부분 시퀀스가 되도록 보장합니다.
• 아키텍처: 단일 공유 백본 (MM-DiT) 이 두 모달리티를 모두 처리합니다. 모델은 이미지 시간을 위한 $t$와 텍스트 시간을 위한 $\tau$라는 두 개의 분리된 시간 단계에 조건부 설정됩니다.
  • 고정 구성 요소: 백본 가중치와 모든 보조 텍스트 인코더 (T5, CLIP) 는 고정된 상태로 유지됩니다.
  • 학습 가능한 구성 요소: 트랜스포머 선형 레이어에 적용된 LoRA 어댑터, 삽입 수 및 토큰 식별자를 위한 경량 텍스트 헤드, 새로운 시간 단계 조건부 경로만 학습됩니다. 이는 SD3 의 경우 파라미터의 약 5%, FLUX.1 의 경우 약 2.5% 만 학습하는 것에 해당합니다.

2. 이중 시간 단계 학습 목적 함수

학습 목적 함수는 $(t, \tau)$ 쌍의 시간 단계에 기반하여 이미지 속도와 누락된 텍스트 삽입을 공동으로 예측하는 것을 포함합니다.

• 손상 일정: 저자들은 엔드포인트 조건부 작업에서만 학습하는 "교대 - 클린 (alternating-clean)" 일정과 $(t, \tau) \sim \text{Unif}([0, 1]^2)$인 "혼합 - 손상 (mixed-corruption)" 일정을 비교합니다. 그들은 교차 모달 대응을 학습하기 위해 혼합 손상으로 사전 훈련한 후, 교대 - 클린으로 정제하는 하이브리드 접근 방식이 최상의 결과를 산출한다는 것을 발견합니다.
• 안정화:
  • 적응형 그래디언트 균형: MSE 기반 이미지 손실과 교차 엔트로피/수치 기반 텍스트 손실 간의 규모 불일치를 해결하기 위해, 저자들은 그래디언트 비율의 지수 이동 평균 (EMA) 을 사용하여 텍스트 손실 가중치 $\lambda_{txt}$를 동적으로 조정합니다.
  • 교사 매칭: 새로운 데이터 분포에 대한 미세 조정 중 사전 훈련된 이미지 사전 지식의 저하를 방지하기 위해, 이미지 속도 회귀는 "동일 노이즈 교사 매칭" 목적 함수로 대체됩니다. 교사는 LoRA 레이어가 비활성화된 동일한 백본이며 추가 메모리가 필요하지 않습니다.

3. 추론 모드

$t$와 $\tau$를 분리함으로써 추론은 2 차원 생성 공간에서의 궤적 선택으로 변환됩니다:

• 텍스트 $\to$ 이미지: $\tau=1$(클린 텍스트) 을 고정하고 $t$를 0 에서 1 로 진화시킵니다.
• 이미지 $\to$ 텍스트: $t=1$(클린 이미지) 을 고정하고 $\tau$를 0 에서 1 로 진화시킵니다.
• 공동 샘플링: $t$와 $\tau$를 모두 동시에 진화시킵니다.
• 부분 텍스트 생성 (VQA): 이미지와 질문 토큰을 고정하고 삽입 과정을 답변 영역으로 제한합니다. 이는 별도의 디코더 없이 시각적 질문 응답 (VQA) 을 가능하게 합니다.

주요 기여

1. C1: LoRA 어댑터와 경량 텍스트 헤드만 학습하여 사전 훈련된 T2I 정류 흐름 모델을 양방향 생성기로 변환하고, 사전 훈련된 이미지 사전 지식을 성공적으로 보존합니다.
2. C2: 이미지는 연속 흐름에 머무르고 텍스트는 이산 삽입 과정을 사용하는 이중 시간 단계 프레임워크를 도입하여, 조건부, 공동, 그리고 부분 텍스트 생성을 통합된 $(t, \tau)$ 공간에서의 궤적으로 가능하게 합니다.
3. C3: SD3 및 FLUX.1-dev 에서 텍스트 - 이미지 보존과 이미지 - 텍스트 품질 모두에서 매칭된 Dual Diffusion-LoRA 베이스라인보다 우수한 성능을 입증하며, 최소한의 계산 자원으로 하류 VQA 를 지원합니다.

실험 결과

저자들은 매칭된 파라미터 수와 LoRA 랭크를 사용하여 재구현된 Dual Diffusion 베이스라인에 대해 SD3 및 FLUX.1-dev 에서 FullFlow 를 평가합니다.

• 효율성: FullFlow 는 두 대의 RTX A5000 GPU 에서 24 시간 미만으로 학습하며, 피크 VRAM 을 Dual Diffusion 의 약 84GB 에서 약 38GB 로 줄이고 처리량을 약 8 배 증가시킵니다.
• 텍스트 - 이미지 (SD3): 매칭된 벽시계 시간에서 FullFlow 는 FID 를 62.7 에서 31.6 으로 (매칭된 스텝에서는 29.3) 개선하고 CMMD 를 0.870 에서 0.121 로 개선하여 강력한 이미지 사전 지식 보존을 입증합니다.
• 이미지 - 텍스트 (SD3): FullFlow 는 매칭된 벽시계 시간에서 CIDEr 점수 99.4(베이스라인의 2.0 대비) 를 달성하며, 매칭된 스텝에서는 13.4 를 달성합니다.
• 하류 작업: VQA 벤치마크 (VQAv2, VizWiz, OKVQA) 에서 FullFlow(학습 가능한 파라미터 130M 만 사용) 는 경쟁력 있는 정확도를 달성하며, 특히 VizWiz 에서 Dual Diffusion 베이스라인을 크게 능가합니다 (50.7 대 29.9).
• 이식성: 이 레시피는 FLUX.1-dev 로 성공적으로 이전되어 36 시간의 학습 후 CIDEr 73.6 을 달성합니다.

중요성 및 주장

본 논문은 FullFlow가 완전한 멀티모달 사전 훈련 없이도 사전 훈련된 T2I 흐름 모델에서 강력한 양방향 비전 - 언어 능력을 해제할 수 있음을 보여준다고 주장합니다. 결과는 사전 훈련된 T2I 사전 지식이 일반적으로 가정된 것보다 더 풍부한 시각 - 의미 구조를 인코딩하여 경량 적응을 통해 역작업 (이미지 - 텍스트) 과 공동 생성을 지원할 수 있음을 시사합니다.

저자들은 이 작업을 피크 엔드 - 투 - 엔드 성능에 대한 주장보다는 "계산 가능한 개념 증명"으로 위치시킵니다. 그들은 이 방법이 통합 모델에 일반적으로 필요한 파라미터의 작은 부분과 상용 GPU 만 요구함으로써 양방향 정류 흐름 모델을 더 넓은 연구 커뮤니티에 접근 가능하게 만든다고 강조합니다. 그들은 현재 256$\times$256 해상도에 초점을 맞추고 있으며 대규모 데이터셋으로의 확장이나 복잡한 대화를 위한 지시 튜닝이 부재하다는 한계를 인정하며, 이를 향후 과제로 지적합니다.