기술 요약: InstructMoLE

문제 제기

다중 조건부 이미지 생성 (예: 다중 주제 구성, 이미지 편집) 을 위한 확산 트랜스포머 (DiT) 의 파라미터 효율적 미세 조정 (PEFT) 은 단일 모놀리식 어댑터 (LoRA 등) 를 사용할 때 작업 간 간섭과 파국적 망각으로 인해 상당한 어려움에 직면합니다. 전문가 혼합 (MoE) 아키텍처는 다양한 작업을 위해 전문가를 다양화함으로써 모듈식 솔루션을 제공하지만, 언어 모델에서 유래된 라우팅 정책으로 인해 확산 모델에 적용되는 데 한계가 있었습니다.

현재 접근 방식은 일반적으로 각 이미지 패치가 로컬 은닉 상태를 기반으로 전문가를 독립적으로 선택하는 토큰 수준 라우팅을 사용합니다. 본 논문은 이러한 로컬 의사 결정 과정이 사용자 명령의 전역적 성격과 근본적으로 불일치한다고 주장합니다. 복잡한 생성 작업 (예: "잔디 위를 기어가는 아기, 근처에 있는 흰 말, 그리고 풋볼 헬멧") 에서 명령은 전역적 의미 관계를 설정합니다. 토큰 수준 라우팅은 이러한 관계를 유지하지 못해 다음과 같은 문제를 초래합니다:

1. 공간적 분열: 스타일과 색조에서 영역별 불연속성 발생.
2. 라우팅 지터: 고주파 노이즈의 증폭.
3. 의미적 편향: 명령의 전체적 의도를 따르지 못해 비일관된 주제나 무시된 속성 발생.

방법론

저자들은 토큰 수준 라우팅을 전역적으로 일관되고 명령에 기반한 정책으로 대체하는 InstructMoLE 프레임워크를 제안합니다. 이 아키텍처는 Flux.1 Kontext 백본을 기반으로 구축되었으며, 저랭크 어댑테이션 (LoRA) 전문가를 활용합니다.

1. 명령 기반 라우팅 (IGR)

토큰별로 전문가를 선택하는 대신, InstructMoLE 는 인스턴스 수준 라우팅 정책을 사용합니다.

• 전역 신호 증류 ($Z_{global}$): 사용자 텍스트 명령 ($I_c$) 에서 전역 라우팅 신호를 증류합니다. 이 신호는 두 가지 표현을 융합합니다:
  • CLIP 임베딩: 견고하고 전체적인 의미적 앵커를 제공합니다.
  • Perceiver 스타일 증류: T5 인코더의 토큰 수준 특징에 대해 학습 가능한 잠재 쿼리가 반복적으로 주의를 기울입니다. 이는 단순한 풀링된 CLIP 임베딩이 평균화할 수 있는 정교한 구성 세부 사항 (예: 특정 색상 - 주제 연관성) 을 포착합니다.
• 통합 전문가 위원회: 모델의 각 레이어에서 라우터는 $Z_{global}$을 기반으로 단일 게이트 가중치 세트를 계산합니다. 이 "전문가 위원회" (Top-$k$ 전문가 집합) 는 해당 레이어 내의 모든 공간 토큰에 동일하게 브로드캐스트됩니다. 이를 통해 특정 추상화 수준에서의 전체 이미지 생성 과정은 사용자의 의도와 정렬된 일관된 전문가 집합에 의해 지배받게 됩니다.

2. 출력 공간 직교성 손실

전문가 붕괴(서로 다른 전문가가 중복된 기능을 학습하는 현상) 를 방지하기 위해 저자들은 새로운 정규화 항을 도입했습니다.

• 표준 부하 균형 손실은 균일한 활용 빈도만 보장할 뿐 기능적 다양성을 보장하지는 않습니다.
• **직교성 손실 ($L_{ortho}$)**은 모든 전문가의 게이트 전 출력에 작용합니다. 이는 모든 전문가 쌍의 평탄화된 출력 벡터 간의 제곱 코사인 유사도를 계산합니다.
• 이 손실을 최소화함으로써 모델은 각 전문가에 대해 상호 직교하는 기능적 역할을 학습하도록 강제되며, 선택된 위원회가 중복된 솔루션이 아닌 다양한 기능 세트를 제공하도록 보장합니다.

3. 훈련 목적 함수

총 손실 함수는 세 가지 구성 요소를 결합합니다:
$$L_{total} = L_{flow} + \lambda_{aux}L_{aux} + \lambda_{ortho}L_{ortho}$$

• $L_{flow}$: 생성 정확도를 위한 표준 흐름 매칭/확산 손실.
• $L_{aux}$: 소수 전문가에 대한 과도한 의존을 방지하기 위한 보조 부하 균형 손실.
• $L_{ortho}$: 기능적 다양성을 위한 제안된 직교성 손실.

주요 기여

1. 라우팅 불일치 식별: 다중 조건부 이미지 생성에서 로컬 토큰 수준 라우팅과 전역 명령 의미 간의 치명적인 불일치를 식별했습니다.
2. 명령 기반 라우팅 (IGR): 사용자 명령에서 유도된 전역 신호에 완전히 의존하여 전문가 선택을 조건화하는 새로운 프레임워크로, 이미지 전반에 걸쳐 구조적 및 의미적 일관성을 보장합니다.
3. 직교성 정규화: 부하 균형만으로는 효과적으로 해결되지 않는 표현적 붕괴를 완화하기 위해 전문가 간의 기능적 다양성을 명시적으로 촉진하는 새로운 손실 함수입니다.
4. 실증적 검증: 다양한 다중 조건부 벤치마크에서 기존 LoRA 어댑터 및 MoE 변형 (토큰 수준 라우팅 베이스라인 포함) 보다 InstructMoLE 가 우수함을 입증하는 광범위한 실험.

실험 결과

이 모델은 합성 및 공개 데이터 (얼굴 교체, 스타일 전환, 다중 주제 구성, 공간 제어 포함) 의 다양한 데이터셋을 사용하여 Flux.1 Kontext 백본에 대해 미세 조정되었습니다.

• 다중 주제 생성 (XVerseBench): InstructMoLE 는 평균 점수 71.07 을 기록하여 강력한 Flux.1 Kontext 베이스라인 (68.04) 과 DreamO-v1.1 (69.57) 을 능가했습니다. 특히 신원 보존 (ID-Sim) 과 객체 충실도 (IP-Sim) 에서 뛰어난 성능을 보였습니다.
• 단일 이미지 편집 (GEdit-EN-full): 비교된 모델 중 평균 점수 6.17 을 기록하여 OmniGen2 (5.45) 를 13% 이상 상회하고 Flux.1 Kontext 베이스라인 (6.06) 을 능가했습니다. 전역 및 로컬 조작이 모두 필요한 작업에서 탁월한 성능을 발휘했습니다.
• 인컨텍스트 생성 (OmniContext): 출시된 모델 중 평균 ID-Sim 점수 38.85 를 기록하여 명령 준수와 신원 보존 사이의 균형을 OmniGen2(명령 추종은 높으나 신원 보존이 낮음) 보다 더 잘 유지했습니다.
• 애블레이션 연구:
  • 라우팅 정책: IGR 은 공간 정렬 작업 (Depth RMSE, Canny F1, Pose F1) 에서 토큰 수준 라우팅 (Standard, Expert Choice, Expert Race) 보다 현저히 우수하여 전역 일관성이 구조적 무결성에 중요함을 확인했습니다.
  • 구성 요소: 직교성 손실이나 Perceiver 스타일 신호 증류를 제거하면 측정 가능한 성능 저하가 발생하여 기능적 다양성과 정교한 명령 인코딩의 필요성을 검증했습니다.

의의 및 주장

본 논문은 **명령 기반 라우팅 (IGR)**이 생성 모델의 명령 기반 미세 조정을 위한 기존 토큰 수준 라우팅보다 더 효과적인 패러다임이라고 주장합니다. 사용자 의의 전체적 의도와 전문가 선택을 정렬함으로써 InstructMoLE 는 다음과 같은 가능성을 제공합니다:

• 우수한 구성 제어: 분열 없이 여러 주제와 특정 공간 관계를 가진 복잡한 장면을 렌더링할 수 있는 능력.
• 사용자 의도에 대한 충실도: 복잡한 체인 오브 씽킹 명령에 대한 준수 개선 및 편집 작업에서의 신원 보존 향상.
• 강건성: 복잡한 명령을 처리할 때에도 높은 충실도를 유지하지만, 저자들은 심각한 텍스트 손상이나 극단적인 참조 부하 (3 개 이상의 엔티티 바인딩) 에서는 한계가 있음을 지적합니다.

이 연구는 단일 모놀리식 어댑터와 로컬 라우팅 MoE 모델이 해결하지 못한 다중 조건부 생성의 특정 과제를 해결하는 확산 트랜스포머의 파라미터 효율적 적응을 위한 견고하고 일반화 가능한 프레임워크를 확립합니다.