효율적인 마스크드 어텐션 트랜스포머를 통한 소수샷 분류·세그멘테이션

초록

EMAT은 메모리 효율적인 마스크드 어텐션, 학습 가능한 다운스케일링, 파라미터 절감 기법을 도입해 기존 최첨단 모델보다 적은 파라미터로 작은 객체에 대한 분류·세그멘테이션 정확도를 크게 향상시킨다. 또한 기존 평가에서 무시되던 지원 이미지의 모든 라벨을 활용하는 두 가지 새로운 평가 설정을 제안한다.

상세 분석

본 논문은 소수샷 분류와 세그멘테이션(FS‑CS)이라는 복합 과제를 다루면서, 특히 이미지 내 작은 객체에 대한 성능 저하 문제를 집중적으로 해결한다. 기존 최첨단 모델인 CST는 고해상도 상관 토큰을 처리하기 위해 지원 이미지 토큰을 과도하게 다운샘플링하고, 그 결과 공간적 디테일이 손실돼 작은 객체의 마스크가 흐릿해지는 한계를 보였다. EMAT은 이 문제를 세 가지 핵심 개선으로 극복한다. 첫째, 마스크드 어텐션 연산을 ‘마스크된 엔트리만 남기고 연산’하도록 재정의함으로써 메모리 사용량을 크게 줄인다. 구체적으로, 마스크 M_f에 의해 0인 위치의 키·밸류를 완전히 제외하고, 남은 토큰에만 소프트맥스와 어텐션 가중치를 적용한다. 이 설계는 동일한 하드웨어에서 t_s(지원 토큰 수)를 2.7배~11배 확대해도 메모리 초과 없이 고해상도 상관 맵을 유지하게 만든다. 둘째, 다운스케일링을 고정 평균 풀링이 아닌 가벼운 2D·3D 컨볼루션과 풀링을 결합한 학습 가능한 모듈로 대체한다. 이는 큰 풀링 커널 없이도 지원 토큰의 공간 해상도를 효과적으로 축소해 연산량을 조절하면서도 중요한 특징을 보존한다. 셋째, 전체 트랜스포머와 헤드의 채널 수를 대폭 감소시켜 파라미터 수를 기존 대비 최소 ¼ 수준으로 줄인다. 파라미터가 적을수록 소수샷 상황에서 과적합 위험이 낮아지고, 학습이 더 안정된다. 이러한 설계는 모두 사전 학습된 ViT‑S(DINOv2) 백본을 고정하고, 1‑way‑1‑shot 설정에서 학습한 뒤 N‑way‑K‑shot 테스트 시 각 클래스를 독립적인 1‑way‑K‑shot 과제로 변환해 추론한다. 또한, 마스크가 전혀 존재하지 않을 경우(클래스가 쿼리 이미지에 없을 때) 빈 마스크를 출력하도록 설계돼 실제 현장 적용성을 높였다. 평가 측면에서 저자는 기존 FS‑CS 벤치마크가 지원 이미지의 다중 라벨을 무시하는 문제를 지적하고, ‘부분 증강’과 ‘전체 증강’ 두 가지 새로운 설정을 제안한다. 부분 증강은 지원 이미지에 포함된 모든 라벨을 유지하고, 전체 증강은 학습/테스트 전체 클래스 라벨을 모두 포함해 모델이 보다 풍부한 라벨 정보를 활용하도록 만든다. 실험 결과, EMAT은 PASCAL‑5i와 COCO‑20i 모두에서 mIoU와 mAP 모두 기존 최고 성능을 앞서며, 특히 작은 객체에 대한 IoU가 현저히 개선되었다. 파라미터 수가 4배 적음에도 불구하고 연산 시간은 기존과 동등하거나 약간 빠른 수준을 유지한다. 전체적으로 EMAT은 메모리·파라미터 효율성을 크게 높이면서도 작은 객체에 대한 정밀한 세그멘테이션을 가능하게 하는 실용적인 솔루션이라 할 수 있다.