기술 요약: SparseSAM

문제 정의

Segment Anything Model(SAM) 은 최첨단 오픈-보카불러리 분할을 달성하지만, 주로 비전 트랜스포머 (ViT) 이미지 인코더에 기인한 상당한 추론 지연 시간 및 메모리 병목 현상에 직면해 있습니다. 인코더는 모델의 파라미터 및 연산의 99% 이상을 차지합니다. 기존 최적화 전략은 다음과 같은 치명적인 한계를 겪고 있습니다:

• 토큰 병합 (Token Merging): Token Merging(ToMe) 과 같은 방법은 시퀀스 길이를 줄이지만 분할 작업에는 적합하지 않습니다. 디코더가 전체 공간 해상도를 요구하기 때문에, 모든 병합 작업은 이에 상응하는 언병합 (unmerge) 단계를 필요로 합니다. 이러한 레이어별 관리 오버헤드 (산포 - 수집 연산) 는 종종 연산 절감분을 상회하여 지연 시간을 증가시키고, 높은 압축률에서 치명적인 품질 저하를 초래합니다.
• 희소 어텐션 (Sparse Attention): 어텐션에만 초점을 맞춘 접근법 (예: SpargeAttention, PISA) 은 SAM 의 로컬 블록에서 특히 두드러지는 MLP 레이어가 기여하는 상당한 지연 시간을 해결하지 못합니다. 또한, 동적 희소성 메커니즘은 마스크 구성 및 온라인 중요도 추정을 위한 상당한 런타임 오버헤드를 수반하는데, 이는 토큰 당 상대적 개수가 적은 (200 미만) 많은 수의 로컬 윈도우를 가진 SAM 의 아키텍처를 고려할 때 금지적입니다.
• 양자화 (Quantization): 훈련 후 양자화 (예: W4A4) 는 SAM 의 작은 행렬 차원으로 인해 텐서 코어 처리량을 포화시키지 못해 GPU 속도 향상을 거의 제공하지 않습니다.

고품질 분할에 필요한 토큰 수준의 충실도를 유지하면서 어텐션과 MLP 레이어를 모두 가속화할 수 있는 훈련이 불필요한 방법의 공백이 존재합니다.

방법론

SparseSAM 은 재학습이나 동적 런타임 검색 없이 계산을 줄이기 위해 두 가지 상호 보완적인 기법을 도입하는 훈련 불필요 구조화 희소화 프레임워크입니다.

1. 스트라이프-소트 어텐션 (Stripe-Sort Attention)

이 메커니즘은 결정론적 Z-순서 (Morton) 치환을 사용하여 밀집 어텐션을 정적이며 하드웨어 친화적인 희소 패턴으로 변환합니다.

• 기반-기반 순서 (Gradient-Based Ordering): 이 방법은 경량 파라미터 없는 Sobel 연산자를 사용하여 국소 기울기 크기를 계산하며, 이는 정보 밀도 (객체 경계 대 매끄러운 영역) 의 대리 지표로 작용합니다.
• 스트라이프-소트 치환 (Stripe-Sort Permutation): 단순히 토큰을 중요도대로 정렬하는 것 (전경 토큰을 시작 부분에 모으는 결과) 대신, 이 방법은 "스트라이프-소트" 전략을 적용합니다. 정렬된 토큰 리스트를 행렬로 간주하고 평탄화하기 전에 전치 (transpose) 합니다. 이는 고중요도 토큰을 시퀀스 전체에 교차 배치하여, 잘리거나 희소화된 부분집합이 이미지의 광범위한 공간적 커버리지를 유지하도록 보장합니다.
• 정적 희소 패턴: 이 치환은 $N \times N$ 어텐션 행렬을 $4 \times 4$ 서브맵 모자이크로 분해합니다. 결과 패턴은 (국소 상호작용을 포착하는) 대역 대각 구조와 소규모 글로벌 유지 집합이 결합된 형태를 보입니다. 이 정적 패턴은 동적 마스크 구성의 필요성을 제거하여 런타임 오버헤드 없이 커스텀 CUDA 커널 (예: FlashAttention-2 스타일) 을 통한 효율적인 실행을 가능하게 합니다.

2. 잔류-일관성 MLP (Residual-Consistency MLP)

이 기법은 SAM 의 디코더가 정밀한 토큰별 표현보다 영역 간 대비에 더 의존한다는 관찰을 활용하여 MLP 병목 현상을 해결합니다.

• 토큰 라우팅: 결정론적 치환 순서에 기반하여 토큰을 "유지 집합 (keep-set, $X_k$)"과 "잔류 집합 (residual set, $X_r$)"으로 분할합니다.
• 선택적 계산: MLP 는 유지 집합의 상위-K 토큰 (높은 특징 불일치 또는 업데이트 크기를 가진 토큰) 에 대해서만 적용됩니다.
• 잔류 전파: 잔류 집합의 나머지 토큰은 MLP 를 완전히 우회하여 변경되지 않은 상태로 잔류 연결을 통과합니다 ($y_i = x_i + \epsilon_i$).
• 기하학 보존: 이 접근법은 디코더가 필요로 하는 표현 기하학을 유지하면서 MLP 레이어의 계산 비용을 크게 줄입니다.

두 기법 모두 모든 모델 레이어에서 재사용되는 단일 사전 계산 치환 집합을 활용하여 레이어별 오버헤드를 무시할 수준으로 유지합니다.

주요 기여

1. 공동 가속화: SparseSAM 은 SAM 의 어텐션과 MLP 레이어를 모두 공동으로 가속화하는 최초의 훈련 불필요 프레임워크로, 단일 구성 요소만 최적화하는 방법의 한계를 극복합니다.
2. 스트라이프-소트 어텐션: 결정론적 Z-순서 치환을 사용하여 정적 블록 대역 희소 패턴을 유도하는 새로운 구조화 희소 어텐션 메커니즘입니다. 이는 동적 마스크 오버헤드를 피하고 글로벌 컨텍스트를 보존합니다.
3. 잔류-일관성 MLP: 의미적으로 중요한 토큰에만 MLP 업데이트를 적용하고 나머지를 잔류 경로를 통해 전파하여 밀집 계산 없이 특징 다양성을 유지하는 토큰 라우팅 메커니즘입니다.
4. 훈련 불필요 효율성: 전체 프레임워크는 재학습, 지식 증류, 또는 미세 조정 없이 작동합니다 (다만 선택적 미세 조정이 작은 정확도 손실을 회복하는 것으로 나타났습니다).

실험 결과

네 가지 분할 벤치마크 (MS-COCO, DIS5K-VD, COIFT, ThinObject5K-TE, HRSOD 을 포함한 HQ-44K 하위 집합) 와 세 가지 SAM 변형 (SAM-B, SAM-L, SAM-H) 에서 평가되었습니다:

• 정확도: 0.4 밀도 (60% 희소성) 에서 SparseSAM 은 mIoU 를 0.004 만 손실합니다. 0.3 밀도에서는 손실이 0.021 mIoU 입니다. 이는 토큰 병합 진보 대비 정확도 손실을 2.10 배 감소시킨 것입니다.
• 가속화: 이 방법은 평균 2.0 배의 추론 속도 향상을 달성하며 (공동 압축 시 25% 밀도에서 SAM-B 기준 최대 2.15 배),
• 메모리: 메모리 사용량을 2.8 배 감소시킵니다.
• 견고성: 높은 압축률에서 치명적인 품질 저하를 겪는 토큰 병합과 달리, SparseSAM 은 극단적인 압축률 (밀도 < 0.4) 에서도 기준선과 거의 동일한 충실도를 유지합니다.
• 하드웨어: NVIDIA A100 GPU 에서 공동 압축 구성은 25% 희소성에서 2.05 배의 최대 속도 향상을 달성하여 다양한 배치 크기에 걸쳐 어텐션 전용 기준선 및 토큰 병합 방법을 능가합니다.

중요성 및 주장

이 논문은 SparseSAM 이 추론 속도와 마스크 품질 간의 트레이드오프를 해결함으로써 효율적인 분할의 새로운 표준을 설정한다고 주장합니다. 그 중요성은 다음과 같습니다:

• 엣지 배포 가능화: 재학습 없이 지연 시간과 메모리를 줄여 대규모 분할 기반 모델을 지연 시간에 민감하고 자원이 제한된 환경 (예: 엣지 장치) 에서 실행 가능하게 만듭니다.
• 아키텍처 통찰: SAM 의 로컬 블록이 이중 병목 (어텐션 + MLP) 에 시달리며, 밀집 예측 작업의 경우 동적 또는 병합 기반 접근법보다 정적이며 구조를 보존하는 희소성이 우월함을 보여줍니다.
• 일반화 가능성: 잔류-일관성 MLP 메커니즘이 SAM2(Hiera 인코더) 로 효과적으로 이전되는 것으로 나타났으며, 이는 토큰 업데이트의 내재적 중복성이 분할 백본의 일반적인 속성임을 시사합니다.

저자들은 현재 접근 방식이 콘텐츠 무관 치환을 사용하지만, 향후 연구는 의미적 중요성을 동적으로 우선시하는 하이브리드 전략을 탐구할 수 있음을 언급합니다. 다만, 현재 정적 접근 방식은 이미 상당한 효율성 향상을 제공합니다.