비주얼 인컨텍스트 학습을 위한 아날로지 기반 디퓨전 트랜스포머

초록

VIRAL은 사전 학습된 Diffusion Transformer(DiT)를 활용해, 예시 이미지 쌍을 시각적 아날로지(xₛ:xₜ :: x_q:y_q) 형태로 입력함으로써 다양한 비전 작업을 하나의 모델에서 인컨텍스트 방식으로 수행하도록 설계된 프레임워크이다. 역할 인식 멀티이미지 토큰 컨디셔닝과 MoE‑LoRA 어댑터를 도입해 작업 간 그래디언트 간섭을 완화하고, 인식·복원·편집을 아우르는 대규모 데이터셋을 구축해 학습하였다. 실험 결과, 기존 V‑ICL 방법들을 크게 능가하며, 오픈‑도메인 이미지 편집까지 포괄하는 통합 비전 인컨텍스트 학습이 가능함을 입증한다.

상세 분석

VIRAL은 “시각적 아날로지”라는 새로운 관점을 도입해 V‑ICL 문제를 조건부 생성으로 재정의한다. 기존 연구가 이미지‑이미지 변환을 그리드 형태로 결합하거나, 특정 태스크에 맞춘 별도 헤드를 설계하는 반면, 본 논문은 모든 변환을 동일한 RGB 공간에서 xₛ→xₜ 형태의 함수 T 로 모델링하고, 이를 xₛ:xₜ :: x_q:y_q 라는 비율 관계로 해석한다. 핵심 기술은 크게 세 가지로 나뉜다.

1. 역할 인식 멀티이미지 토큰 컨디셔닝

   • 사전 학습된 VAE‑Encoder와 패치화 과정을 통해 각 이미지(예시 소스, 예시 타깃, 쿼리 소스)를 동일 차원의 토큰 시퀀스로 변환한다.
   • 3‑D MSRoPE(Position Encoding) 를 확장해 이미지별 역할(소스, 타깃, 쿼리)을 명시적으로 구분한다. 이는 토큰 간 공간적 관계는 유지하면서, 서로 다른 이미지 간의 정체성을 구분하게 해 주어 DiT의 전역 교차‑어텐션이 변환 로직을 정확히 추론하도록 돕는다.

2. Mixture‑of‑Experts LoRA (MoE‑LoRA)

   • DiT는 파라미터를 고정한 채 어댑터만 학습한다는 점에서 파라미터 효율성이 뛰어나지만, 다양한 태스크가 섞이면 그래디언트 간섭이 발생한다. 이를 해결하기 위해 LoRA에 MoE 구조를 도입, 여러 저‑랭크 어댑터 전문가들을 라우팅 네트워크를 통해 동적으로 선택한다.
   • Top‑k 라우팅과 로드‑밸런싱 손실(L_aux)을 적용해 전문가 사용을 균등화하고, 특정 태스크에 과도히 편향되는 현상을 방지한다. 결과적으로 하나의 어댑터 모듈이 다중 태스크를 동시에 학습하면서도 성능 저하를 최소화한다.

3. 대규모 인컨텍스트 데이터셋 구축

   • 기존 V‑ICL 데이터는 이미지 쌍이 제한적이거나, 그리드 형태로 구성돼 해상도·다양성에 한계가 있었다. 저자는 두 가지 스트림으로 데이터셋을 확장한다.
     • 표준 비전 태스크: DiffusionDB에서 텍스트 프롬프트를 추출하고 Qwen‑Image 로 고해상도 소스 이미지를 생성, ControlNet‑Aux 로 깊이·에지·노멀 등 라벨을 자동 생성한다. 객체 검출, 키포인트, 세그멘테이션 등은 기존 COCO·SAM·EntityV2 등 공개 데이터와 결합해 시각적 타깃(마스크·컬러맵)으로 변환한다.
     • 오픈‑도메인 편집: 텍스트‑이미지 쌍을 이용해 “스타일 변환”, “배경 교체” 등 자유로운 편집을 수행하고, 동일 변환을 적용한 두 이미지 쌍을 아날로지 형태로 매칭한다.
   • 이렇게 구축된 1천만 규모의 (xₛ, xₜ, x_q, y_q) 쿼드러플은 다양한 시각 변환을 포괄하며, 모델이 변환 로직을 일반화하도록 충분한 컨텍스트를 제공한다.

학습 및 추론
학습 단계에서는 노이즈가 섞인 y_q(쿼리 타깃)의 라티스 z_{y,t}를 DiT에 입력하고, 고정된 Z_cond(예시·쿼리 토큰)와 함께 노이즈 예측 ϵ̂_θ를 학습한다. 손실은 전통적인 MSE 기반 디퓨전 손실이며, MoE‑LoRA 파라미터만 업데이트한다. 추론 시에는 예시 쌍과 쿼리 소스만 제공하고, DiT가 노이즈를 점진적으로 제거하면서 최종 ŷ_q 를 생성한다.

실험 결과
다양한 베이스라인(VisualPrompt, Painter, SD‑VICL 등)과 비교했을 때, VIRAL은 이미지 복원, 세그멘테이션, 키포인트 검출, 색채 변환, 텍스트‑이미지 편집 등 12개 이상의 태스크에서 평균 4.2 % 이상의 mAP/PSNR 향상을 기록한다. 특히 오픈‑도메인 편집에서는 기존 방법이 실패하던 복합적인 스타일·구도 변환을 성공적으로 수행한다. Ablation 실험은 MoE‑LoRA가 없을 경우 성능이 급격히 저하되고, 3‑D MSRoPE가 없으면 역할 구분이 모호해져 변환 추론 정확도가 감소함을 보여준다.

의의와 한계
VIRAL은 “시각적 아날로지”라는 통합 프레임워크를 통해 V‑ICL의 범용성을 크게 확장했으며, 사전 학습된 대형 디퓨전 모델을 그대로 활용함으로써 별도 대규모 재학습 비용을 절감한다. 다만 현재는 RGB 공간에 국한돼 있어 멀티모달(텍스트·오디오) 컨텍스트와 고해상도(>1024 px) 실시간 응용에는 추가 최적화가 필요하다. 또한 MoE‑LoRA의 라우팅 비용이 모델 규모에 따라 증가할 수 있어 효율적인 전문가 수와 라우팅 전략에 대한 연구가 뒤따라야 한다.