ClustViT: Clustering-based Token Merging for Semantic Segmentation

"ClustViT: Clustering-based Token Merging for Semantic Segmentation" 논문에 대한 상세한 기술 요약은 다음과 같습니다.

1. 문제 제기

비전 트랜스포머 (ViT) 는 컴퓨터 비전 분야에서 최첨단 성능을 달성했지만, 입력 크기에 비례하는 2 차 계산 복잡도로 인해 고통받고 있습니다. 토큰 가지치기 및 병합 기술 (token skimming) 은 이미지 분류에 대한 비용을 성공적으로 줄였으나, 시맨틱 세그멘테이션(밀집 예측 작업) 에서는 어려움을 겪습니다.

• 갈등: 분류는 "관련 없는" 토큰을 폐기할 수 있지만, 세그멘테이션은 표준 디코더 헤드와의 호환성을 보장하기 위해 모든 픽셀에 대한 공간적 세부 정보를 보존해야 합니다.
• 격차: 기존 방법들은 종종 저수준 토큰 유사성이나 구조화된 고정 크기 압축 (예: $2\times2$ 패치만 병합) 에 의존하는데, 이는 이미지의 시맨틱 내용에 적응하지 못하거나 밀집 예측 작업에 대한 일반화 성능이 떨어집니다.

2. 방법론: ClustViT

저자들은 표준 ViT 백본에 학습 가능한 클러스터 모듈과 리제네레이터 (Regenerator) 모듈을 통합한 ClustViT 아키텍처를 제안합니다. 핵심 혁신은 단순한 픽셀 근접성이 아닌 그라운드 트루스 세그멘테이션 마스크에서 유도된 시맨틱 정보를 사용하여 토큰 병합을 안내하는 것입니다.

A. 아키텍처 개요

1. 표준 ViT 백본: 프로세스는 표준 ViT 와 마찬가지로 이미지를 패치로 분할하고 이를 토큰으로 투영하는 것으로 시작합니다.
2. 클러스터 모듈 (주입 지점): 인코더 내의 미리 정의된 레이어 ($ip$) 에서 경량 MLP 가 각 토큰에 대한 클러스터 할당을 예측합니다.
   • 의사 클러스터 (Pseudo-Clusters): 학습 중 모델은 그라운드 트루스 마스크를 사용하여 "의사 클러스터"를 생성합니다. 패치 내 모든 픽셀이 동일한 시맨틱 클래스에 속하면 해당 패치는 해당 클래스 ID 를 할당받고, 그렇지 않으면 혼합 (0) 으로 표시됩니다.
   • 병합: 동일한 클러스터에 할당된 토큰들은 평균 풀링을 통해 단일 대표 토큰으로 집계됩니다. 복잡하거나 혼합된 영역의 토큰들은 세부 정보를 보존하기 위해 클러스터링되지 않은 상태로 남습니다.
   • 결과: 시퀀스 길이가 단축되어 후속 트랜스포머 레이어의 계산 비용이 감소합니다.
3. 리제네레이터 모듈: 최종 인코더 레이어 이후, 압축된 시퀀스는 디코더를 위해 원래 해상도로 복원되어야 합니다.
   • 이 모듈은 저장된 클러스터 할당을 사용하여 대표 토큰을 원래 위치로 확장합니다.
   • 정제 MLP 는 미세한 특징을 재구성하여 출력 시퀀스가 입력 길이와 일치하고 오프더셸 세그멘테이션 헤드 (예: Segmenter, UPerNet) 와 호환되도록 보장합니다.

B. 학습 전략

• 이중 작업 손실 (Dual-Task Loss): 모델은 결합된 손실 함수로 엔드 투 엔드 방식으로 학습됩니다:
  $$L = L_{segm} + \lambda L_{clust}$$
  • $L_{segm}$: 시맨틱 세그멘테이션을 위한 표준 교차 엔트로피 손실.
  • $L_{clust}$: 클러스터 모듈이 그라운드 트루스 마스크에서 유도된 의사 클러스터를 예측하도록 강제하는 교차 엔트로피 손실.
  • $\lambda$: 균형을 맞추는 하이퍼파라미터 (실험에서 0.1 로 설정됨).

3. 주요 기여

1. 시맨틱 유도 토큰 병합: 이전 방법들이 시각적 유사성이나 고정 그리드에 기반하여 병합하는 것과 달리, ClustViT 는 시맨틱 내용(세그멘테이션 마스크에서 유도됨) 에 기반하여 토큰을 병합하여 비구조적이고 콘텐츠 적응형 압축을 가능하게 합니다.
2. 엔드 투 엔드 학습 가능한 백본: 클러스터링 구성 요소는 ViT 백본에 직접 통합되어 세그멘테이션 작업과 함께 공동 학습되므로, 별도의 정책 네트워크나 사전 학습 단계가 필요하지 않습니다.
3. 밀집 예측을 위한 리제네레이터: 리제네레이터 모듈의 도입은 호환성 문제를 해결하여, 압축된 토큰 표현이 공간 해상도를 잃지 않고 밀집 예측 헤드를 위해 완전히 복원될 수 있게 합니다.
4. 효율성 - 정확성 트레이드오프: 이 방법은 계산 비용 (GFLOPs) 과 추론 시간을 크게 줄이면서도 경쟁력 있는 정확도를 유지하며, 특히 로봇 공학에서 흔히 발생하는 넓은 배경 영역이 있는 시나리오에서 두드러집니다.

4. 실험 결과

저자들은 ClustViT 를 ADE20K(일반 장면), SUIM(수중 로봇 공학), RumexWeeds(농업 로봇 공학) 의 세 가지 데이터셋에서 평가했습니다. 표준 ViT 와 최첨단 콘텐츠 인식 토큰 공유 (CTS) 방법과 비교했습니다.

• 성능 지표:
  • 계산 효율성: 기준 ViT 대비 최대 2.18 배 적은 GFLOPs와 1.64 배 빠른 추론(초당 이미지) 을 달성했습니다.
  • 정확도: 비교 가능한 세그멘테이션 정확도 (mIoU) 를 유지했습니다. ADE20K 에서 최상의 구성 (ClustViT-bk3,ip3) 은 약 1.66 mIoU 포인트 감소만 있으면서 1.47 배의 속도 향상을 달성했습니다.
  • 로봇 공학 시나리오: 이 방법은 넓은 배경 영역이 있는 데이터셋 (SUIM 및 RumexWeeds) 에서 탁월한 성능을 발휘했습니다. 예를 들어, RumexWeeds 에서 기준선 대비 mIoU 를 약간 개선하면서 1.54 배의 FPS와 2.12 배의 GFLOP 향상을 달성했습니다.
• 애블레이션 연구:
  • 클러스터 수 ($k$): $k$(더 많은 클러스터) 를 증가시키면 속도는 향상되지만 더 많은 세부 사항이 병합됨에 따라 정확도가 감소합니다. $k=3$이 가장 좋은 균형을 제공했습니다.
  • **주입 지점 ($ip$):** 네트워크 초기에 압축 ($ip=2$) 하면 속도가 극대화되지만 정확도가 떨어집니다. 중간 주입 지점 ($ip=3$또는$4$) 이 최적의 트레이드오프를 제공했습니다.
  • 모델 규모: 효율성 향상은 더 큰 모델 (Base/Large) 에서 가장 두드러졌으며, 여기서 클러스터링 MLP 의 오버헤드는 어텐션 계산의 절감량에 비해 무시할 수 있었습니다.

5. 중요성 및 영향

• 로봇 공학 적용성: 이 논문은 계산 자원이 제한되고 장면이 종종 크고 균일한 배경을 포함하는 실제 로봇 시스템 (예: 자율 농업, 수중 탐사) 에 ViT 를 배포하는 데 있어 중요한 병목 현상을 해결합니다.
• 패러다임 전환: 토큰 병합이 단순한 시각적 유사성이 아닌 시맨틱 이해에 의해 안내된다면, 밀집 예측에서 토큰 압축이 공간적 충실도를 희생할 필요가 없음을 보여줍니다.
• 일반화 가능성: 표준 디코더 헤드 (Segmenter 및 UPerNet) 를 지원함으로써 ClustViT 는 기존 ViT 백본을 위한 플러그인 대체제를 제공하여 더 넓은 컴퓨터 비전 커뮤니티에 매우 실용적입니다.

결론적으로, ClustViT 는 시맨틱 인식 가능한 학습 가능한 토큰 병합 메커니즘을 도입함으로써 비전 트랜스포머의 높은 정확도와 실시간 로봇 인식의 효율성 요구 사항 사이의 격차를 성공적으로 메웠습니다.