기술 요약: 확장 가능한 비전 트랜스포머를 위한 탄성 어텐션 코어

문제 제기

비전 트랜스포머 (ViT) 는 전역 수용 영역과 유연한 표현 학습을 가능하게 하는 모든-대-모든 (all-to-all) 자기 어텐션을 활용하여 강력한 데이터 기반 확장성을 달성했습니다. 그러나 이러한 유연성은 이미지 해상도에 따라 2 차 ($O(N^2)$) 로 증가하는 계산 비용을 수반하며, 여기서 $N$은 이미지 패치 수입니다. 이러한 2 차 확장성은 고해상도 영역에서 금지적 (prohibitive) 이 되어 표준 ViT 의 적용 범위를 제한합니다.

ViT 설계에 있어 지배적인 가정은 풍부한 시각적 - 의미론적 표현을 학습하기 위해 조밀한 쌍별 토큰 상호작용이 필수적이라는 것입니다. 본 논문은 해당 가정에 도전하여 직접적인 패치 간 상호작용 없이도 효과적인 시각적 표현을 학습할 수 있음을 주장합니다.

방법론: VECA(시각 탄성 코어 어텐션)

저자들은 표준 조밀 자기 어텐션을 블록 희소 코어 - 주변부 어텐션 메커니즘으로 대체하는 비전 트랜스포머 아키텍처인 VECA를 제안합니다.

코어 아키텍처

• 코어 - 주변부 구조: 이 아키텍처는 통신 인터페이스 역할을 하는 소수의 학습된 $C$개의 "코어" 토큰 집합을 도입합니다. 이미지 패치는 "주변부"를 형성합니다.
• 상호작용 메커니즘:
  • 코어 토큰: 완전히 연결된 클리크를 형성하여 모든 다른 코어 토큰과 모든 패치 토큰에 어텐션합니다.
  • 패치 토큰: 오직 활성 코어 토큰에만 어텐션합니다. 다른 패치 토큰에는 직접 어텐션하지 않습니다.
  • 정보 흐름: 패치 간의 모든 정보 교환은 코어 토큰을 통해 exclusively(배타적으로) 라우팅됩니다.
• 복잡도: 어텐션 비용은 $O(N^2)$에서 $O(2NC + C^2)$로 감소합니다. 사전 결정된 코어 수 $C$(여기서 $C \ll N$) 에 대해 복잡도는 이미지 해상도에 따라 선형 ($O(N)$) 으로 확장됩니다.
• 공간 역학: 코어 토큰은 학습 가능한 2D 좌표로 초기화되며, 경량 좌표 헤드를 통해 레이어 간 업데이트됩니다. 이들은 2D 축 방향 회전 위치 인코딩 (RoPE) 을 활용하여 의미론적 표현과 진화하는 공간 위치를 모두 유지합니다.

탄성 학습 및 추론

• 중첩 드롭아웃 (Nested Dropout): 학습 중 모델은 정렬된 코어 토큰 뱅크에 중첩 드롭아웃을 적용합니다. 모델은 분포에서 샘플링된 활성 코어의 가변 예산 ($C$) 으로 학습됩니다.
• 암시적 순위 매기기: 이 절차는 코어 토큰을 중요도에 따라 암시적으로 순위 매깁니다. 초기 코어는 모든 예산에서 활성화되어 가장 필수적이고 광범위하게 유용한 시각적 정보를 인코딩하는 법을 학습합니다. 후기 코어는 더 미세한 공간적 세부 사항이나 특정 의미론적 개념에 특화됩니다.
• 추론 유연성: 테스트 시 재학습 없이 활성 코어 수를 조정할 수 있습니다. 이는 계산 비용 (속도) 과 정확도 간의 탄력적인 트레이드오프를 가능하게 합니다. $C$를 줄이면 더 거친 어텐션 맵을 생성하지만 FLOP 를 크게 절감합니다.

학습 목적

모델은 고정된 DINOv3 교사로부터의 특징 증류 (feature distillation) 를 사용하여 학습됩니다. 손실 함수는 레이블이 없는 Objects365 이미지를 사용하여 첫 번째 코어 토큰에서 추출된 전역 이미지 표현과 조밀한 패치 표현을 모두 교사의 목표와 정렬시킵니다.

주요 기여

1. VECA 아키텍처: 모든 레이어에서 2 차 자기 어텐션을 선형 시간 코어 - 주변부 어텐션으로 대체하여 토큰 상호작용을 해상도 불변의 병목 지점을 통해 강제하는 시각 백본입니다.
2. 성능 대비 비용: 분류 및 조밀 공간 작업에 대한 평가는 VECA 가 계산 비용을 크게 줄이면서도 최첨단 기반 모델 (특히 DINOv3) 과 경쟁력 있는 성능을 달성함을 보여줍니다.
   • 분류: ViT-Base 크기에서 VECA 는 어텐션 상호작용의 일부만 사용함에도 불구하고 ImageNet-1k 에서 81.93% 의 Top-1 정확도를 달성하여 DINOv3(83.56%) 와 약 1.6% 이내의 격차를 보입니다.
   • 조밀 예측: PASCAL Context 분할 벤치마크에서 VECA 는 레이어당 87.1% 적은 어텐션 상호작용을 가짐에도 불구하고 DINOv3 의 57.74 mIoU 대비 57.46 mIoU 를 달성합니다.
3. 발현 행동: 저자들은 명시적인 감독 없이도 코어 토큰이 객체 중심 표현을 개발하며, 초기 레이어에서는 등방성 (구형) 분포에서 깊은 레이어로 갈수록 의미론적으로 클러스터링된 구조로 진화함을 확인했습니다.

결과

• 확장성: VECA 는 다양한 입력 해상도 (256px 에서 1024px) 에서 안정적인 성능을 유지하는 반면, 표준 ViT 는 2 차 FLOP 증가로 고통받습니다. VECA 는 해상도가 증가함에 따라 FLOP 와 지연 시간에서 선형 확장을 보입니다.
• 탄력성: 모델은 정확도를 속도와 탄력적으로 교환할 수 있습니다. 예를 들어, 코어 예산을 $C=64$에서 $C=8$로 줄이면 계산량을 극적으로 줄이면서도 상당한 성능 (예: 분류에서 전체 모델 성능의 96% 이상) 을 유지합니다.
• 조밀 특징: 시각화 (UMAP) 는 VECA 가 DINOv3 와 비교할 수 있는 고해상도에서도 견고한 안정적이고 고품질의 조밀 특징을 생성함을 확인시켜 줍니다.

중요성과 주장

본 논문은 고품질 조밀 시각 표현을 학습하기 위해 직접적인 패치 간 자기 어텐션이 반드시 필요한 것은 아님을 주장합니다. 학습된 코어 토큰의 소규모 집합을 통해 정보를 라우팅함으로써 VECA 는 탄성 코어 - 주변부 어텐션을 비전 트랜스포머를 위한 확장 가능한 대체 구성 요소로 확립합니다.

이 연구는 모든-대-모든 어텐션의 유연성이 공간 정렬과 조밀 특징의 품질을 보존하는 구조화된 선형 시간 인터페이스로 대체될 수 있음을 시사합니다. 이는 계산 효율성이 중요한 고해상도 영역에서 비전 트랜스포머를 배포할 수 있는 유망한 경로를 제공하며, 사용 가능한 리소스나 작업 요구 사항에 따라 추론 비용을 동적으로 조정할 수 있는 모델을 가능하게 합니다.