기술 요약: 고해상도 의미 분해를 위한 STEP (SuperToken 및 Early-Pruning)

1. 문제 제기

비전 트랜스포머 (ViT) 는 장거리 의존성을 포착하는 능력 덕분에 의미 분해 분야에서 최첨단 성능을 달성했습니다. 그러나 이미지 해상도에 비례하는 자기 주의 (self-attention) 메커니즘의 2 차 복잡성으로 인한 높은 계산 비용과 메모리 사용량으로 인해 실제 배포는 방해받고 있습니다. 입력 해상도가 증가할 경우 (예: 1024×1024 로), 계산 비용과 메모리 소비가 크게 증가합니다.

기존 토큰 축소 전략은 밀집 예측 (dense prediction) 작업에서 다음과 같은 특정 한계에 직면해 있습니다:

• 패치 단위 방법 (예: CTS) 은 종종 고정된 하이퍼파라미터 (예: 병합된 패치의 고정된 수) 에 의존하여 복잡한 이미지에서는 과잉 병합 (over-merging) 이나 단순한 이미지에서는 병합 부족 (under-merging) 으로 인해 최적이지 않은 성능을 보입니다.
• 토큰 단위 가지치기 (hard pruning) 는 토큰을 영구적으로 제거하기 때문에 픽셀 단위 출력을 위해 공간 정보를 보존해야 하는 밀집 예측에는 종종 적합하지 않습니다.
• 하이브리드 접근법 은 종종 다양한 이미지 콘텐츠에 대한 적응력이 부족하거나 고해상도 입력의 특정 확장성 문제를 해결하지 못합니다.

이 논문은 표준 그리드 기반 처리가 비효율적인 고해상도 시나리오에서 특히 계산 효율성과 정확도 보존 사이의 균형을 맞춰야 한다고 주장합니다.

2. 방법론: STEP 프레임워크

저자들은 ViT 의 효율성을 크게 저하시키지 않으면서 향상시키기 위해 설계된 하이브리드 토큰 축소 프레임워크인 STEP(SuperToken 및 Early-Pruning) 을 제안합니다. STEP 는 두 가지 보완적인 수준에서 작동합니다:

2.1 동적 CTS(dCTS): 적응형 패치 병합

정적 병합 전략과 달리 STEP 는 EfficientNetLite0 기반의 경량 클래스 무관 정책 네트워크인 dCTS를 도입합니다.

• 메커니즘: 네트워크는 인접한 패치 그룹 ($2\times2$, $4\times4$, $8\times8$, $16\times16$ 윈도우) 을 평가하여 유사도 점수 $S \in [0,1]$을 예측합니다.
• 의사 결정 논리: 유사도 점수가 미리 정의된 임계값 $\tau$를 초과하면 패치들이 **슈퍼패치 (superpatch)**로 병합됩니다. 이 과정은 중첩된 그룹 간의 충돌을 방지하기 위해 coarse-to-fine 계층 구조를 따릅니다.
• 출력: 그 결과 균일하지 않은 크기의 슈퍼패치 그리드가 생성됩니다. 이러한 슈퍼패치는 이선형 보간을 통해 표준 해상도 (16×16) 로 리사이즈되고 **슈퍼토큰 (supertokens)**으로 임베딩 공간에 투영됩니다.
• 이점: 이를 통해 토큰 수를 이미지 콘텐츠의 의미론적 동질성에 따라 동적으로 적응시킬 수 있습니다 (예: 객체 경계에서는 해상도를 보존하면서 큰 배경 영역을 병합).

2.2 조기 가지치기 (Early-Pruning) 메커니즘

STEP 은 "쉬운" 토큰의 처리를 중단하기 위해 DToP(Dynamic Token Pruning) 에서 영감을 받은 조기 종료 전략을 통합합니다.

• 아키텍처: ViT 인코더는 $S$개의 단계로 나뉩니다. 각 단계에는 각 토큰에 대한 신뢰도 점수를 계산하는 보조 디코딩 헤드 (Attention-to-Mask 모듈 사용) 가 포함됩니다.
• 가지치기 논리: 신뢰도 점수가 임계값 $\tau$를 초과하는 토큰은 마스킹되어 후속 레이어로 전파되는 것이 중단됩니다. 낮은 신뢰도 (어려운) 토큰만 더 깊은 레이어를 통과합니다.
• 통합: 이 전략은 dCTS 와 보완적입니다. dCTS 가 초기 시퀀스 길이를 줄이는 동안, 조기 가지치기는 네트워크가 깊어질수록 활성 토큰 집합을 점진적으로 단축시킵니다.

3. 주요 기여

논문은 다음과 같은 구체적인 기여를 제시합니다:

1. 새로운 하이브리드 프레임워크: 의미 분해의 효율성 문제를 해결하기 위해 동적 패치 병합 (dCTS) 과 조기 토큰 가지치기를 결합한 STEP 제안.
2. 고해상도 평가: ViT-Base 및 ViT-Large 백본을 사용하여 고해상도 이미지 (최대 1024×1024) 에서 수행된 광범위한 실험. 저자들은 이것이 의미 분해의 맥락에서 고해상도에 특화된 토큰 가지치기 전략을 체계적으로 평가한 첫 번째 시도라고 주장합니다.
3. 종합적 분석: 병합을 위한 임계값 파라미터, 보조 가지치기 헤드의 최적 배치, 그리고 계산 복잡성 (GFLOPs), 추론 속도 (FPS), 정확도 (mIoU) 간의 트레이드오프에 대한 상세한 절멸 연구 (ablation studies).
4. 확장성 증명: 이미지 해상도와 모델 크기가 증가함에 따라 STEP 가 계산 복잡성과 추론 시간을 크게 줄이면서도 경쟁력 있는 정확도를 유지한다는 증거.

4. 실험 결과

실험은 NVIDIA A100 GPU 를 사용하여 COCOStuff10k, ADE20K, Cityscapes 데이터셋에서 수행되었습니다.

4.1 표준 해상도 (512×512) 성능

• dCTS 단독 적용: ViT-Large 에 적용 시, dCTS 는 표준 16×16 패칭 대비 토큰 수를 2.5 배 감소시킵니다. 이는 최소한의 정확도 저하로 계산 비용 2.6 배 감소 및 처리량 3.4 배 증가를 가져옵니다.
• 완전한 STEP: dCTS 와 조기 가지치기를 모두 적용하면 최대 4 배의 계산 복잡성 감소와 1.7 배의 추론 속도 향상을 달성하며, 최대 정확도 저하는 2.0% 이내로 제한됩니다.
• 구성: ViT-Large 의 16 번째 또는 18 번째 레이어에 단일 가지치기 헤드를 배치하는 것이 멀티 헤드 구성에 비해 효율성과 지연 시간 간의 유리한 트레이드오프를 제공했습니다.

4.2 고해상도 (768×768 및 1024×1024) 성능

• 확장성: 해상도가 증가함에 따라 표준 SegViT 는 GFLOPs 가 급격히 증가하고 FPS 가 떨어지는 반면, STEP 구성은 안정적인 트레이드오프를 유지합니다.
• 1024×1024 해상도의 ViT-Large: dCTS 변형만으로도 GFLOPs 를 2086 에서 802 로 줄여 복잡성을 2.6 배 낮추고, FPS 를 12 에서 41 로 3.4 배 증가시켰으며, mIoU 감소는 미미했습니다.
• 토큰 가지치기 동역학: 최대 40% 의 토큰이 최종 인코더 레이어에 도달하기 전에 확신 있게 예측되어 중단될 수 있습니다. ViT-Large 를 사용하는 고해상도 설정에서는 평균 48% 의 토큰이 조기에 중단됩니다.

4.3 처리량 vs FLOPs에 대한 관찰

저자들은 중요한 발견을 지적합니다: 처리량 (FPS) 은 FLOP 감소와 선형적으로 비례하지 않습니다.

• GFLOPs 가 크게 감소함에도 불구하고 추론 속도 향상은 때때로 제한적입니다 (예: 4 배의 GFLOP 감소가 항상 4 배의 FPS 증가를 가져오지는 않음).
• 원인: 조기 가지치기에 필요한 마스킹 연산은 불규칙한 제어 흐름, 동적 메모리 액세스 패턴, 텐서 모양의 가변성으로 인해 병목 현상을 초래합니다. 논문은 불필요한 GPU-CPU 간 메모리 전송과 같은 하드웨어 비효율성이 예상되는 성능 향상을 감소시킬 수 있다고 제안합니다.

5. 중요성 및 주장

이 논문은 STEP 가 ViT 를 고해상도 밀집 예측 작업에 실용적으로 만드는 중요한 단계라고 주장합니다.

• 유일무이한 첫 시도: 저자들은 이것이 의미 분해에서 서로 다른 해상도 수준에 걸쳐 토큰 가지치기를 체계적으로 평가한 첫 번째 작업이라고 주장합니다.
• 효율성 - 정확도 균형: STEP 는 정확도 저하를 허용 가능한 범위 (<2.0%) 내로 유지하면서 상당한 계산 절감 (최대 4 배) 과 속도 향상을 달성할 수 있음을 보여줍니다.
• 적응성: dCTS 모듈은 특히 견고하여 고정된 하이퍼파라미터 없이 이미지 복잡도에 맞춰 병합율을 적응시키므로 다양한 실제 장면에 적합합니다.

저자들은 STEP 가 효율성과 확장성 측면에서 명확한 장점을 제공하지만, FLOP 감소로 예측되는 잠재적 처리량 향상을 완전히 실현하기 위해서는 가지치기 메커니즘 자체의 하드웨어 수준 비효율성을 해결해야 한다고 결론지었습니다.