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

왜 중요한가

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

핵심 기여

직접적인 생성 방식의 세그멘테이션 도입

DiT 아키텍처를 수정하지 않고 생성 목적 함수를 그대로 사용하여 흑백 세그멘테이션 마스크를 RGB 공간에서 직접 생성하는 방식을 확립함.

이진 마스크 특화 타임스텝 샘플링 전략

이진 마스크의 VAE 잠재 표현이 노이즈에 매우 견고하다는 점을 발견하고, 고농도 노이즈 영역에 집중하는 롱테일 샘플링 기법을 통해 학습 효율을 개선함.

통합 학습 및 단일 단계 추론 구현

이미지 생성과 세그멘테이션을 동일 구조 내에서 동시에 학습하며, 추론 시에는 여러 단계의 반복 없이 단 한 번의 연산으로 마스크를 생성하는 효율성을 확보함.

핵심 아이디어 이해하기

기존의 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 백본에 대한 확장성 검증이 필요하며, 의료 영상이나 오디오-비주얼 분할과 같은 특수 도메인으로의 적용 연구가 추가로 요구됨.

실무 활용

텍스트 지시문에 따라 이미지 내 특정 객체를 정확히 분할하는 도구로 활용 가능함. 별도의 복잡한 디코더 없이 기존 생성 모델 구조를 그대로 사용할 수 있어 구현이 용이함.

자연어 기반 이미지 편집 도구의 객체 선택 기능
자율주행 시스템의 실시간 장애물 및 도로 분할
의료 영상 내 특정 병변 부위 자동 추출 및 분석

코드 공개 여부: 미확인

키워드

DiT(확산 트랜스포머)Segmentation(세그멘테이션)VAE Latents(VAE 잠재 표현)Flow Matching(플로우 매칭)Referring Segmentation(참조 세그멘테이션)

용어 해설

확산 트랜스포머(Diffusion Transformer): — U-Net 대신 트랜스포머 아키텍처를 노이즈 제거 백본으로 사용하는 확산 모델이다. 데이터의 장거리 의존성을 잘 포착하며 모델 크기 확장에 유리하여 고품질 이미지 및 영상 생성에 널리 쓰인다.
VAE 잠재 표현(VAE Latent): — 변이형 오토인코더(VAE)를 통해 고해상도 이미지를 압축한 저차원 벡터 공간이다. 확산 모델은 이 잠재 공간에서 노이즈를 제거하며 학습하므로 연산 효율성이 높고 데이터의 핵심 특징을 효과적으로 다룬다.
플로우 매칭(Flow Matching): — 데이터 분포와 노이즈 분포 사이의 연속적인 변환 경로를 학습하는 생성 모델 프레임워크이다. 기존 확산 모델보다 학습이 안정적이고 추론 속도가 빠르며, 직선적인 경로를 통해 데이터를 생성하는 것이 특징이다.
평균 교집합 대비 합집합 비율(mIoU): — 세그멘테이션 모델의 정확도를 측정하는 대표적인 지표이다. 모델이 예측한 영역과 실제 정답 영역의 교집합을 합집합으로 나눈 값의 평균이며, 1에 가까울수록 분할 성능이 완벽함을 의미한다.
분류기 없는 가이드(Classifier-Free Guidance): — 별도의 분류기 없이 조건부 모델과 무조건부 모델의 점수를 조합하여 생성 결과가 입력 조건(텍스트 등)에 더 잘 부합하도록 만드는 기법이다. 생성의 다양성과 정확도 사이의 균형을 조절하는 데 필수적이다.

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

왜 중요한가

핵심 기여

핵심 아이디어 이해하기

방법론

주요 결과

기술 상세

한계점

실무 활용

키워드

용어 해설

관련 토론

댓글

관련 기사

이미지 세그멘테이션의 역사와 Segment Anything Model의 등장

왜 중요한가

핵심 기여

핵심 아이디어 이해하기

방법론

주요 결과

기술 상세

한계점

실무 활용

키워드

관련 토론

댓글

관련 기사

이미지 세그멘테이션의 역사와 Segment Anything Model의 등장