CARE-Edit: 컨텍스트 이미지 편집을 위한 조건 인식 전문가 라우팅

기존 이미지 편집 모델은 여러 수정 조건을 한꺼번에 처리하다가 조건끼리 충돌하여 결과물이 어색해지는 문제가 있었다. CARE-Edit은 각 조건에 특화된 전문가를 동적으로 선택해 처리함으로써, 복잡한 편집 상황에서도 물체의 정체성을 유지하고 경계를 깔끔하게 처리한다.

기존의 Diffusion 기반 편집 모델은 ControlNet처럼 외부 조건을 단순히 결합하거나 어댑터를 추가하는 방식을 사용한다. 하지만 Attention 연산 과정에서 텍스트 정보가 마스크 제약을 무시하거나, 참조 이미지의 스타일이 배경으로 번지는 등 조건 간의 간섭이 발생한다. CARE-Edit은 Mixture-of-Experts(MoE) 개념을 도입하여 이 문제를 해결한다. 모든 토큰을 하나의 경로로 통과시키는 대신, Latent-attention Router가 현재 토큰의 특징과 편집 목적을 분석하여 가장 적합한 전문가(Text, Mask, Reference, Base)에게 작업을 배분한다. 특히 Diffusion의 역과정 단계에 따라 필요한 정보가 달라진다는 점에 주목하여, 초기에는 구도를 잡는 전문가가, 후기에는 세부 경계를 다듬는 전문가가 더 많이 활성화되도록 설계했다.

전체 아키텍처는 Frozen 상태의 DiT(Diffusion Transformer) 백본 위에 LoRA 어댑터와 전문가 라우팅 레이어를 삽입한 구조이다. 입력된 텍스트, 베이스 이미지, 참조 이미지, 마스크는 각각 전용 인코더를 통해 잠재 토큰으로 변환된 후 하나의 시퀀스로 결합된다. 핵심인 Token-wise Top-K Routing은 각 토큰 h'_i에 대해 라우터가 4명의 전문가에 대한 확률 분포를 계산한다. [토큰 특징 k_i와 전역 작업 쿼리 q를 입력으로] → [MLP 연산과 Softmax를 수행해] → [각 전문가의 선택 확률을 얻고] → [상위 K개의 전문가만 활성화하여 연산 효율성과 전문성을 확보한다]. Mask Repaint 모듈은 현재 잠재 특징과 참조 인코딩 Z_r 사이의 기하학적 대응 관계를 활용한다. [이전 단계 마스크와 특징 맵을 입력으로] → [Convolution과 Sigmoid 연산을 거쳐] → [잔차 마스크를 생성하고] → [이를 통해 물체 경계에 딱 맞는 정교한 소프트 마스크를 실시간으로 갱신한다]. Latent Mixture 모듈은 전문가들의 출력을 최종 결합한다. [전문가별 가중치 w_e와 출력 h'_e를 입력으로] → [가중합 연산을 수행해] → [융합된 토큰 h'_fuse를 얻으며], 이는 타임스텝에 따른 게이팅 계수를 통해 베이스 이미지의 구조와 최종적으로 조화를 이룬다.

EMU-Edit 및 MagicBrush 벤치마크에서 기존 통합 모델인 OmniGen2, AnyEdit 등을 능가하는 성능을 기록했다. 특히 EMU-Edit 테스트에서 CLIPout 0.313, DINO 0.835를 달성하며 작업 특화 모델 수준의 품질을 보여주었다. DreamBench++의 다중 객체 설정에서도 DINO-I 0.568, CLIP-I 0.720을 기록하며 복잡한 컨텍스트 내에서의 객체 보존 능력을 입증했다. 이는 120K라는 상대적으로 적은 학습 데이터로도 1M 이상의 데이터를 사용한 모델들보다 우수한 효율성을 보인 결과이다. Ablation Study를 통해 전문가 라우팅을 제거할 경우 성능이 크게 하락함을 확인했으며, 활성화되는 전문가 수 K=3일 때 표현의 다양성과 연산 효율 사이의 최적의 균형을 찾음을 확인했다.

CARE-Edit은 FLUX.1-dev를 백본으로 사용하며, 2~3개의 DiT 블록마다 라우팅 레이어를 배치하여 오버헤드를 최소화했다. 모든 전문가는 백본과 동일한 차원 d를 공유하며, 출력단에 LayerNorm과 Linear Projection을 배치해 잔차 연결의 안정성을 확보했다. 라우팅 붕괴를 방지하기 위해 고정된 비율의 토큰을 항상 공유 전문가로 통과시키는 전략을 채택했다. 또한 라우팅 로짓에 지수 이동 평균(EMA)을 적용해 훈련 중의 변동성을 줄이고 동적 선택 과정을 안정화했다. 학습 손실 함수는 확산 재구성 손실에 부하 분산, 마스크 경계 일관성, 잠재 혼합 매끄러움의 세 가지 보조 정규화 항을 결합하여 구성된다. 각 정규화 계수는 실험적으로 (0.01, 0.1, 0.05)로 설정되었다. 데이터 효율성을 높이기 위해 Subjects200K 데이터셋을 기반으로 GPT-4와 VLM을 활용한 자동화 파이프라인을 구축했다. 이를 통해 일관된 배경을 공유하면서 전경 객체만 달라지는 20K 규모의 고품질 멀티 페어 데이터를 생성하여 학습에 활용했다.

추가적인 하이퍼파라미터(예: top-K) 설정이 필요하며, 현재 전문가 세트가 가장 일반적인 작업들로 구성되어 있어 더 넓은 범위의 편집 유형으로 확장할 여지가 남아 있다.