GenMask: 직접적인 마스크 생성을 통한 세그멘테이션용 DiT 적응 기술

생성 모델을 단순히 특징 추출기로 쓰는 대신, 모델이 직접 분할 마스크를 생성하도록 학습시켜 성능을 높였습니다. 특히 이진 마스크가 노이즈에 강하다는 물리적 특성을 이용해 학습 효율을 극대화한 것이 핵심입니다.

기존의 Diffusion 모델은 이미지의 미세한 질감을 복원하기 위해 모든 노이즈 단계에서 정교한 학습이 필요하다. 하지만 세그멘테이션 마스크는 흑백의 단순한 구조를 가지며, VAE 잠재 공간에서 노이즈에 매우 강한 특성을 보인다. 기존 방식은 이러한 차이를 무시하고 생성 모델을 단순히 특징 추출기로만 활용하여 성능 한계가 존재했다.

이 논문은 마스크가 노이즈가 가득 찬 상태에서도 형태 정보를 유지한다는 점에 주목했다. 따라서 학습 시 노이즈가 아주 많은 단계에 집중하는 새로운 샘플링 전략을 도입했다. 이는 모델이 이미지의 세부 사항보다는 전체적인 윤곽과 텍스트 지시문 사이의 관계를 더 명확히 학습하도록 유도한다.

결과적으로 DiT 아키텍처를 전혀 수정하지 않고도 텍스트에 맞는 마스크를 직접 생성할 수 있게 되었다. 추론 시에도 여러 단계를 거치지 않고 단 한 번의 계산만으로 정확한 마스크를 얻을 수 있어 실용성이 크게 향상되었다.

전체 시스템은 사전 학습된 WAN-2.1 DiT와 Qwen2.5-VL-7B VLM을 기반으로 구축된다. VLM은 텍스트 지시문과 입력 이미지를 인코딩하여 DiT의 조건부 입력으로 전달하며, DiT는 이를 바탕으로 마스크를 생성한다.

핵심 메커니즘은 세그멘테이션 전용 타임스텝 샘플링 함수이다. 타임스텝 t와 하이퍼파라미터 a를 입력으로 p(t) = 2a^2t / (t^2 + a^2)^2 연산을 수행해 특정 시점이 선택될 확률을 계산한다. a가 0.05와 같이 작을수록 t=1 근처의 높은 노이즈 영역이 선택될 확률이 높아지며, 이는 마스크의 견고한 특성을 학습에 반영하는 지표가 된다.

학습 시에는 입력 이미지의 VAE 잠재 표현을 DiT 입력에 직접 주입하는 Shortcut 구조를 사용한다. 또한 추론 시에는 노이즈 ε과 모델의 예측 속도 v를 입력으로 x_mask = ε + v(ε, 1) 연산을 수행해 최종 마스크의 잠재 표현을 얻는다. 이는 여러 번의 반복 없이 단 한 번의 연산으로 분할 결과를 도출함을 의미한다.

RefCOCO 벤치마크에서 oIoU 83.3%, mIoU 83.7%를 기록하며 기존 SOTA 모델인 LISA(79.1%)를 능가했다. RefCOCO+와 RefCOCO-g에서도 각각 78.7%, 75.6%의 oIoU를 달성하며 일관된 성능 우위를 입증했다.

Ablation Study를 통해 샘플링 하이퍼파라미터 a=0.05일 때 가장 높은 성능이 나타남을 확인했다. a=0.5로 설정하여 균등하게 샘플링할 경우 mIoU가 66.0%까지 하락하여, 고농도 노이즈 영역 학습의 중요성이 증명됐다.

입력 이미지의 VAE 잠재 표현을 DiT에 직접 주입하는 Shortcut 구조를 제거할 경우 mIoU가 82.2%에서 74.1%로 급감했다. 이는 정확한 픽셀 단위 예측을 위해 저수준의 시각 정보가 필수적임을 나타낸다.

전체 시스템은 1.3B 파라미터의 WAN-2.1 DiT와 Qwen2.5-VL-7B VLM으로 구성된다. VLM은 텍스트와 이미지를 동시에 처리하여 DiT의 Cross-Attention 레이어에 조건부 정보를 제공한다. VAE 인코더와 디코더는 고정된 상태로 유지되며 DiT 모델만 엔드투엔드로 미세 조정된다.

이진 마스크의 VAE 잠재 표현이 선형 분리 가능하다는 사실을 PCA 분석으로 확인했다. 단 하나의 주성분(Principal Component)만으로도 원래 마스크의 형태를 거의 완벽하게 복원할 수 있었으며, 이는 마스크 생성이 일반 이미지 생성보다 훨씬 단순한 매핑 문제임을 수학적으로 뒷받침한다.

학습 과정에서 이미지 생성 데이터와 세그멘테이션 데이터를 1:1 비율로 혼합하여 학습하는 Mix Training 기법을 사용했다. 이를 통해 생성 모델의 일반화 능력을 유지하면서도 세그멘테이션 작업에 특화된 성능 향상을 이끌어냈다.

더 큰 규모의 DiT 백본에 대한 확장성 검증이 필요하며, 의료 영상이나 오디오-비주얼 분할과 같은 특수 도메인으로의 적용 연구가 추가로 요구됨.