StructSAM: Structure- and Spectrum-Preserving Token Merging for Segment Anything Models

이 논문은 SAM 모델의 구조와 스펙트럼을 보존하면서 토큰 병합 시 발생하는 경계 손실과 프롬프트 정보 누출 문제를 해결하기 위해, 그라디언트 기반 토큰 에너지 점수와 그리드 기반 평탄도 선별을 활용한 'StructSAM'을 제안하여 계산 비용을 대폭 줄이면서도 정밀한 분할 성능을 유지하는 방법을 제시합니다.

핵심

이 논문은 **'Segment Anything Model(SAM)'**이라는 인공지능의 속도를 획기적으로 높이면서도, 그 성능을 떨어뜨리지 않는 새로운 방법을 제안합니다.

이 복잡한 내용을 일상적인 비유로 쉽게 설명해 드릴게요.

🎨 비유: 거대한 모자이크 그림과 바쁜 화가

1. 문제 상황: 너무 많은 타일 (Token)
SAM 이라는 AI 는 이미지를 이해할 때, 이미지를 작은 정사각형 타일 (Token) 수백만 개로 나누어 봅니다. 마치 거대한 모자이크 그림을 그리는 화가처럼요.

• 기존 방식: 화가는 그림의 모든 타일 (배경의 하늘, 나무, 사람 얼굴 등) 을 똑같이 자세히 보고 분석합니다.
• 문제점: 이 방식은 너무 느립니다. 특히 '사람'이나 '물체'의 경계 (가장자리) 를 정확히 그리는 게 중요한데, 배경 같은 평범한 부분까지 다 똑같이 세세하게 보면 시간이 너무 오래 걸립니다.

2. 기존 해결책의 실패: 무작위 줄이기
기존에는 "타일 수를 줄여서 속도를 내자"는 방법들이 있었습니다. 하지만 이는 마치 무작위로 타일을 떼어내는 것과 비슷했습니다.

• 결과: 배경 타일은 줄여도 좋지만, **물체의 경계 (예: 강아지의 귀 끝, 자동차 바퀴)**에 있는 중요한 타일까지 실수로 지워버리거나 뭉개버리는 경우가 많았습니다. 그래서 AI 가 그리는 그림의 테두리가 뭉개지거나, "어? 이게 뭐지?" 하는 실수가 생겼습니다.

3. 이 논문의 해결책: StructSAM (똑똑한 정리꾼)
이 논문에서 제안한 StructSAM은 "무작정 줄이는 게 아니라, 어디를 줄이고 어디는 아껴야 할지를 아는 똑똑한 정리꾼"입니다.

• 비유 1: 에너지 지도 (Energy Score)
  StructSAM 은 그림을 볼 때 "어디가 가장 활발하게 움직이는가?"를 감지합니다.

  • 경계선 (물체의 테두리): 색이 급격히 변하고 선이 뚜렷한 곳 (예: 강아지의 귀, 건물의 모서리). 여기는 에너지가 높은 곳입니다. 이 타일들은 절대 건드리지 않고 보호 구역으로 지정합니다.
  • 평평한 곳 (배경): 하늘, 벽, 풀밭처럼 색이 일정하고 변화가 없는 곳. 여기는 에너지가 낮은 곳입니다. 이 타일들은 서로 비슷하니까 하나로 합쳐도 괜찮습니다.

• 비유 2: 그리드 (Grid) 와 평탄도 검사
  그림을 작은 사각형 (방) 으로 나눕니다.

  • "이 방은 벽이 너무 평평하네? (Flatness)" → 합치기 (Merging): 이 방에 있는 타일 10 개를 대표하는 타일 1 개만 남기고 나머지는 없앱니다.
  • "이 방은 문과 창문이 복잡하네?" → 보호 (Protection): 이 방의 타일들은 모두 그대로 둡니다.

• 비유 3: 다시 원래대로 복구 (Unmerging)
  AI 가 계산을 할 때는 타일 수를 줄여서 빠르게 계산하지만, 결론을 내릴 때는 다시 원래의 타일 개수로 되돌립니다.

  • 마치 "빠르게 요약해서 생각했다가, 발표할 때는 다시 상세한 내용을 다 말해주는" 것과 같습니다. 덕분에 AI 는 빠르면서도 세부적인 경계선을 잃지 않습니다.

4. 추가 기능: "이 부분을 봐!" (Prompt-Aware)
사용자가 "이 박스 안의 강아지를 찾아줘"라고 지시하면 (Prompt), StructSAM 은 강아지가 있는 박스 안은 절대 줄이지 않고, 박스 밖의 배경만 과감하게 줄입니다. 이렇게 하면 속도는 더 빨라지고 정확도는 유지됩니다.

🚀 이 방법의 성과

• 속도: AI 가 계산해야 하는 양 (FLOPs) 을 25~40% 이상 줄였습니다. (기존보다 훨씬 빠름)
• 정확도: 물체의 경계가 뭉개지지 않아서, 의료 영상 (유방암 종양 찾기) 이나 정밀한 객체 인식에서도 기존 방법들보다 훨씬 좋은 결과를 냈습니다.
• 장점: AI 모델을 처음부터 다시 훈련시킬 필요 없이, 기존에 만들어진 모델을 그대로 쓰면서 속도를 높일 수 있습니다.

💡 한 줄 요약

"StructSAM 은 AI 가 그림을 볼 때, '중요한 경계선'은 지켜주면서 '평범한 배경'만 지혜롭게 줄여서, 똑똑하면서도 엄청나게 빠른 AI 를 만든 방법입니다."

기술 요약

1. 연구 배경 및 문제 정의 (Problem)

• 배경: 'Segment Anything Model'(SAM) 은 다양한 시각 도메인에서 뛰어난 분할 성능을 보이는 기반 모델 (Foundation Model) 이지만, 계산 비용이 매우 높습니다. 특히 ViT-L, ViT-H 와 같은 대형 변형의 경우 이미지 인코더가 전체 모델 파라미터와 FLOPs 의 98% 이상을 차지하여 추론 속도가 느립니다.
• 기존 방법의 한계: 최근 Vision Transformer(ViT) 에서 토큰 병합 (Token Merging) 기술이 추론 속도를 높이기 위해 도입되었으나, 이를 SAM 에 직접 적용하는 것은 비자명합니다.
  • 구조적 차이: SAM 은 윈도우 기반 어텐션과 전역 어텐션을 혼합하며, 마스크 디코더가 정밀한 경계 예측을 위해 밀집된 (dense) 특징 맵을 필요로 합니다.
  • 성능 저하: 기존 토큰 병합 방법 (ToMe, PiToMe 등) 은 무작위 또는 단순한 유사도 기반의 병합을 수행하여, 객체 경계 (boundary) 나 프롬프트 (프롬프트) 관련 영역의 정보가 손실됩니다. 이는 특히 의료 영상이나 얇은 구조물을 분할해야 하는 작업에서 성능이 급격히 떨어지는 원인이 됩니다.
• 핵심 문제: 재학습 (Retraining) 없이 SAM 을 그대로 사용하면서 (Off-the-shelf), 계산 비용을 줄이되 객체 경계와 프롬프트 정보를 보존할 수 있는 효율적인 토큰 병합 전략이 필요합니다.

2. 제안 방법: StructSAM (Methodology)

저자들은 SAM 의 구조와 스펙트럼 특성을 보존하는 StructSAM을 제안합니다. 이는 '병합 - 계산 - 복원 (Merge-Compute-Unmerge)' 프레임워크를 기반으로 합니다.

2.1. 핵심 아이디어

1. 경계 인식 에너지 추정 (Gradient-based Energy Estimation):
   • 인코더의 특징 맵 (Feature Map) 에서 1 차 미분 (Sobel 연산자 또는 유한 차분) 을 사용하여 국소적인 특징 변화 (기울기) 를 계산합니다.
   • 기울기 크기가 큰 토큰은 객체 경계에 해당하므로 **보호 (Protected)**하고, 기울기가 작아 시각적으로 평탄한 (Flat) 영역의 토큰만 병합 대상으로 선정합니다.
2. 셀 기반 평탄도 스크리닝 (Grid-based Flatness Screening):
   • 토큰 그리드를 $s \times s$ 크기의 셀로 분할합니다.
   • 각 셀의 최대 기울기 값을 '평탄도 점수'로 사용하여, 경계가 없는 평탄한 셀을 병합 가능 영역으로 선택합니다.
   • 이를 통해 공간적 일관성을 유지하며 경계 영역을 보호합니다.
3. 지향성 병합 및 복원 (Destination Selection & Recovery):
   • 병합 가능한 셀 내에서 기울기가 가장 낮은 (가장 안정적인) 토큰을 '목적지 (Destination)'로 선택하고, 나머지 토큰을 평균화하여 병합합니다.
   • 토큰 복원 (Unmerging): 어텐션 연산 후, 병합된 특징을 다시 원래의 토큰 수와 공간 해상도로 복원합니다. 이는 SAM 의 마스크 디코더가 밀집된 특징 그리드를 필요로 하기 때문에 필수적입니다.
4. 프롬프트 인식 변형 (Prompt-aware Variant):
   • 박스 프롬프트가 제공되는 경우, 프롬프트 영역 내에서는 병합 비율을 낮추고 배경 영역에서는 병합을 더 aggressively 수행하여 추가적인 속도 향상을 꾀합니다.

2.2. 이론적 근거 (Spectral Graph Theory)

• 토큰 병합을 그래프 축소 (Graph Coarsening) 관점에서 분석했습니다.
• 제안한 점수 기반 병합 전략은 라플라시안 (Laplacian) 의 고유값 (Spectral properties) 왜곡을 수학적으로 제한할 수 있음을 증명했습니다.
• 반면, 무작위 또는 구조를 고려하지 않은 병합은 비가역적인 스펙트럼 드리프트 (Spectral Drift) 를 유발하여 분할 성능을 저하시킨다는 것을 이론적으로 설명합니다.

3. 주요 기여 (Key Contributions)

1. 체계적인 평가: SAM 및 Medical SAM 을 재학습 없이 (Off-the-shelf) 다양한 자연 및 의료 벤치마크에서 평가한 최초의 체계적인 연구입니다. 기존 방법들이 경계 민감한 작업에서 실패하는 이유를 규명했습니다.
2. StructSAM 제안: 기울기 기반 에너지와 셀 평탄도를 활용하여 정보-rich 영역 (경계, 프롬프트) 을 보호하고 중복된 배경 토큰을 병합하는 구조 보존 전략을 제시했습니다.
3. 이론적 분석: 그래프 스펙트럼 이론을 통해 제안된 방법의 안정성과 기존 방법의 한계를 수학적으로 증명했습니다.

4. 실험 결과 (Results)

• 실험 설정: DIS5K, HRSOD, ThinObject5K(경계 정밀도), Cityscapes, INbreast(의료 영상) 등 8 개 벤치마크에서 ViT-B 및 ViT-L 백본을 사용하여 평가했습니다.
• 성능 향상:
  • 계산 효율성: 인코더 FLOPs 를 25~30% 감소시켰으며, 프롬프트 인식 모드를 사용하면 40% 이상 감소했습니다.
  • 정확도 유지: 평균 IoU(mIoU) 와 Dice 점수에서 기존 방법 (ToMe, PiToMe, ALGM 등) 보다 훨씬 우수한 성능을 유지했습니다.
    • 예: INbreast 데이터셋에서 MedSAM 기준 FLOPs 28.5% 감소 시 Dice 점수 하락은 0.62 포인트에 불과했으나, 경쟁 방법들은 2~5 포인트 이상 하락했습니다.
  • 경계 보존: 얇은 구조물 (전선, 기둥 등) 이나 날카로운 경계에서 기존 방법들이 흐릿해지거나 배경과 합쳐지는 현상을 방지하고 선명한 분할 결과를 보여주었습니다.
• 비교: ToMe, PiToMe, ToMeSD, VidToMe, ALGM 등 기존 최첨단 방법들을 모두 능가하는 효율성 - 정확도 트레이드오프를 달성했습니다.

5. 의의 및 결론 (Significance)

• 실용적 가치: SAM 과 같은 대형 기반 모델을 재학습이나 아키텍처 변경 없이, 추론 단계에서만 토큰 병합을 통해 경량화할 수 있는 경로를 제시했습니다.
• 응용 분야 확대: 의료 영상 (저해상도 장비, 실시간 수술 지원), 로봇 공학, 임베디드 비전 시스템 등 계산 자원이 제한된 환경에서도 고품질 분할 모델을 배포할 수 있게 합니다.
• 방법론적 통찰: 토큰 병합을 단순한 데이터 압축이 아닌, 구조와 스펙트럼을 보존하는 지역적 (Local) 과정으로 재정의함으로써, 향후 밀집 예측 (Dense Prediction) 태스크를 위한 효율적인 Transformer 설계에 중요한 지침을 제공합니다.

요약하자면, StructSAM은 SAM 의 높은 계산 비용을 줄이면서도 객체 경계와 프롬프트 정보를 잃지 않기 위해 기울기 기반의 지능형 토큰 병합 전략을 도입한 혁신적인 연구입니다.