Beyond Self-Attention: Sub-Quadratic Vision Transformers for Fast Image Captioning

기술 요약: 셀프 어텐션을 넘어: 빠른 이미지 캡셔닝을 위한 아임계(Sub-Quadratic) 비전 트랜스포머

문제 정의
이미지 캡셔닝은 시각적 콘텐츠로부터 일관되고 의미 있는 텍란적 설명을 생성하는 과정으로, 컴퓨터 비전과 자연어 처리 사이의 가교 역할을 수행한다. 트랜스포머 기반 아키텍처는 셀프 어텐션(self-attention)을 통해 장거리 문맥 의존성을 포착함으로써 이 분야를 크게 발전시켰으나, 두 가지 주요 한계에 직면해 있다. 첫째, 표준 셀프 어텐션 메커니즘은 모든 이미지 패치 간의 쌍별 상호작용을 계산하므로 이차 복잡도($O(n^2)$)를 갖는다. 이는 고해상도 이미지 및 대규모 애플리케이션 적용 시 상당한 메모리와 계산 오버헤드를 발생시킨다. 둘째, 기존 방식들은 풍부한 국소 특징 표현이 부족하거나 계산적 중복이 발생하는 경우가 많은데, 이는 많은 패치가 의미론적으로 중요한 정보를 포함하지 않음에도 불구하고 모든 패치를 동일하게 처리하기 때문이다.

방법론
제안된 아키텍처는 기존의 셀프 어텐션 메커니즘을 가우시안 혼합 모델(Gaussian Mixture Model, GMM)에 기반한 확률적 클러스터링으로 대체하여 비전 트랜스포머(ViT) 프레임워크를 재구조화함으로써 이러한 한계를 해결한다. 시스템은 두 가지 주요 단계로 작동한다:

1. 이미지 인코딩 (GMM-트랜스포머):

   • 패치 임베딩 (Patch Embedding): 입력 이미지는 겹치지 않는 패치로 나뉘고, $D$-차원 임베딩으로 투영된 후 위치 인코딩(positional encodings)이 추가된다.
   • 문맥 가중치 적용 GMM 클러스터링 (Context-Weighted GMM Clustering): 패치 임베딩 간의 쌍별 어텐션 점수를 계산하는 대신, 모델은 패치 임베딩을 $K$개의 다변량 가우시안 분포($K \ll N$)의 혼합에 의해 생성되는 연속 확률 변수로 취급한다.
   • 학습 가능한 문맥 가중치 (Learnable Context Weights): 하위 레이어는 통계적 특성(평균 및 분산)에 의해 조절되는 학습 가능한 공간 가중치를 패치에 할당하여, 모델이 의미론적으로 중요한 영역을 우선시할 수 있도록 한다.
   • EM 알고리즘: 소프트 클러스터링을 수행하기 위해 기대값 최대화(Expectation-Maximization, EM) 알고리즘이 사용된다. E-단계에서 모델은 각 패치가 특정 클러스터에 속하는 정도를 나타내는 '책임도(responsibilities, 소프트 할당 확률)'를 계산한다. M-단계에서는 클러스터 파라미터(평균, 공분산, 혼합 계수)가 업데이트된다.
   • 복잡도 감소: 이 과정은 $N \times N$ 어텐션 행렬을 $N \times K$ 책임도 행렬로 대체하여, 계산 복잡도를 이차 복잡도 $O(n^2)$에서 근사 선형 복잡도 $O(nK)$로 줄인다.
   • 집계 (Aggregation): 정제된 패치 특징은 잔차 연결(residual connections) 및 레이어 정규화(layer normalization)를 포함하는 피드포워드 네트워크를 통해 처리된다. 마지막으로, 전역 평균 풀링(Global Average Pooling) 연산을 통해 클러스터 중심점들을 압축된 전역 특징 벡터로 집계한다.

2. 언어 디코딩 (통합 GPT 기반 디코더):

   • 디코더는 자기회귀적(autoregressive) GPT 스타일의 트랜스포머를 활용한다.
   • 통합 어텐션 (Unified Attention): 마스크드 셀프 어텐션(masked self-attention)과 크로스 어텐션(cross-attention)을 분리하는 전통적인 인코더-디코더 구조와 달리, 이 아키텍처는 이를 단일 멀티 헤드 어텐션 블록으로 통합한다.
   • 메커니즘: 텍스트 토큰은 쿼리(query) 벡터 역할을 하며, 결합된 시퀀스(시각적 토큰인 클러스터 중심점과 이전 텍ual 토큰의 결합)는 키(key) 및 값(value) 벡터 역할을 한다.
   • 인과적 마스킹 (Causal Masking): 특수화된 인과적 마스크는 텍스트 토큰이 모든 시각적 토큰에는 주목하되, 이전에 생성된 텍스트 토큰에만 주목하도록 보장하여 자기회귀적 생성 특성을 유지한다.
   • 생성: 출력은 Softmax 레이어를 통해 투영되어 다음 토스 텍 토큰을 반복적으로 예측한다.

주요 기여

• 복잡도 감소: 본 아키텍처는 셀프 어텐션을 GMM 기반의 소프트 클러스터링으로 대체함으로써, 문맥적으로 일관된 캡션을 생성하는 능력을 유지하면서도 트랜스포머 인코더의 시간 복잡도를 이차 복잡도에서 근사 선형 복잡도($O(nK)$)로 줄인다.
• 효율적인 디코더 설계: 마스크드 셀프 어텐션과 크로스 어텐션을 단일 계산 블록으로 통합한 맞춤형 GPT 기반 디코더를 사용하여 멀티모달 상호작용을 강화하고 중복성을 줄였다.
• 실증적 검증: Flickr30k 데이터셋에 대한 종합적인 평가를 통해, 모델은 기존의 최첨단(SOTA) 연구들과 비교하여 계산 효율성과 캡션 품질 사이의 균형을 맞추며 경쟁력 있는 성능을 보여주었다.

실험 결과
모델은 8개의 트랜스포머 레이어, 4개의 어텐션 헤드, 768의 은닉 차원 설정을 사용하여 Flickr30k 데이터셋(Karpathy split 사용)에서 평가되었다.

• 성능 지표: 제안된 모델은 고급 의미론적 지표인 BLEU-4 (0.364), METEOR (0.368), SPICE (0.199)에서 최고 점수를 기록했다. 이 결과는 TAKG, GNN-Visual, Scene Graph와 같은 여러 SOTA 모델에 비해 우수한 구절 구조 및 장면 그래프 매핑 능력을 나타낸다.
• 베이스라인 비교: 모델은 베이스라인 정밀도 지표(BLEU-1: 0.733 vs. 0.794; ROUGE: 0.568 vs. 0.608; CIDEr: 0.723 vs. 0.946)에서 ADA 모델에 뒤처지지만, TopDown 및 GNN-Visual과 같은 다른 베이스라인보다는 뛰어난 성능을 보인다. 저자들은 이 아키텍처가 단순한 복제보다는 깊은 문맥(deep context)을 선호한다고 언급했다.
• 효율성: 모델은 약 21.17 GFLOPs의 계산 비용으로 작동하며 약 3,280만 개의 학습 가능한 파라미터를 포함한다. 혼동 행렬(confusion matrix) 분석 결과, 높은 참 양성률(45,580)과 낮은 거짓 양성/음성 수를 보여 강력한 분류 능력을 입증했다.

의의 및 주장
본 논문은 제안된 프레임워크가 표준 비전 트랜스포머의 고유한 계산 비효 효율성을 성공적으로 해결한다고 주장한다. 가우시안 혼합 모델과 기대값 최대화(EM) 알고리즘을 활용함으로써, 모델은 의미 있는 특징 관계를 포착하는 능력을 희생하지 않으면서도 계산 복잡도를 이차 복잡도에서 근사 선형 복잡도로 크게 줄였다. 저자들은 모델이 강력한 예측 능력과 확장 가능한 특징 인코딩을 보여주며, 캡션 품질, 계산 효율성 및 확장성 사이의 개선된 균형을 제공한다고 단언한다. 향후 연구는 고급 멀티모달 사전 학습을 통합하고, 도메인 특화 데이터셋에 적응하며, 실시간 에지 애플리케이션을 위한 경량화된 모델을 배포하는 방향으로 진행될 예정이다.