ViT-AdaLA: 선형 어텐션을 활용한 비전 트랜스포머 적응 프레임워크

고해상도 이미지나 비디오 처리를 위해 비전 트랜스포머의 연산 효율성을 높이는 것은 필수적이다. 이 논문은 기존의 무거운 소프트맥스 기반 모델을 성능 저하 없이 가벼운 선형 어텐션 모델로 변환하는 3단계 적응 전략을 제시하여, 대규모 재학습 없이도 효율적인 비전 모델 구축을 가능하게 한다.

Transformer의 Self-Attention은 모든 토큰 간의 관계를 계산하기 위해 시퀀스 길이 N의 제곱(N²)에 비례하는 연산량을 요구한다. 이는 고해상도 이미지처럼 토큰 수가 많은 경우 메모리 부족과 속도 저하의 주범이 된다. 기존에는 이를 해결하기 위해 연산량을 N에 비례하도록 줄이는 '선형 어텐션'을 처음부터 다시 학습했으나, 이는 막대한 비용이 들고 기존의 강력한 사전 학습 모델의 지식을 활용하기 어렵다는 한계가 있었다.

ViT-AdaLA는 기존 소프트맥스 기반 모델의 지식을 '복제'하는 대신, 점진적으로 '닮아가게' 만드는 방식을 취한다. 먼저 각 블록 내부의 어텐션 결과물을 맞추고, 이후 모델 전체의 최종 출력값이 원본과 일치하도록 미세 조정하여 소프트맥스의 정교함과 선형 어텐션의 효율성을 동시에 확보한다.

이 접근법은 특히 공간적 정보가 중요한 비전 작업에서 효과적이다. 단순히 어텐션 메커니즘만 바꾸는 것이 아니라, 전체적인 특징 맵의 일관성을 유지하도록 유도함으로써 고해상도 이미지에서도 원본 모델의 뛰어난 인식 능력을 그대로 유지하면서 연산 속도만 2배 이상 끌어올리는 결과를 낳는다.

Stage 1 (Attention Alignment)에서는 각 트랜스포머 블록 내의 소프트맥스 어텐션을 선형 어텐션 모듈로 근사한다. 원본 모델의 가중치는 고정하고, 추가된 선형 어텐션의 가중치 행렬만 업데이트하여 소프트맥스 출력값과의 MSE(Mean Squared Error)를 최소화한다. 이때 활성화 함수로 ELU(x)+1을 사용하여 비음수성을 확보한다.

Stage 2 (Feature Alignment)에서는 블록별 근사 오차가 쌓여 발생하는 성능 저하를 방지하기 위해 모델 전체를 미세 조정한다. 원본 소프트맥스 모델(Teacher)과 선형화된 모델(Student)에 동일한 이미지를 입력하고, 최종 레이어의 피처 맵 간의 MSE 손실을 계산하여 Student 모델의 전체 가중치를 갱신한다. 이 과정에서 [두 모델의 출력 피처 차이 계산 → 오차 역전파 → Student 가중치 업데이트] 순으로 연산이 수행되어 전역적인 시맨틱 정보를 보존한다.

Stage 3 (Supervised Fine-tuning)에서는 정렬된 모델에 특정 작업용 헤드(예: 분류용 Linear 레이어 또는 세그멘테이션용 Mask2Former)를 부착하고 하위 데이터셋으로 최종 학습한다. 이 단계에서 [이미지 입력 → 선형 어텐션 기반 특징 추출 → 작업별 헤드 출력 → 정답 레이블과의 손실 계산] 과정을 거쳐 모델을 최종 최적화한다.

ImageNet-1K 분류 실험에서 DINOv2-L 기반 ViT-AdaLA는 원본 소프트맥스 모델(86.8%) 대비 단 0.8%p 차이인 86.0%의 정확도를 기록했다. 이는 기존 선형화 기법인 Hedgehog(58.8%)나 LoLCATS(61.6%)보다 월등히 높은 수치이며, 추론 속도는 36.52 imgs/s에서 41.56 imgs/s로 향상되었다.

ADE20K 세그멘테이션 작업에서는 mIoU 55.55%를 달성하여 원본(56.73%)에 근접한 성능을 보였다. 특히 고해상도 확장성 분석에서 1024x1024 해상도 기준 메모리 사용량을 3.28GB에서 1.37GB로 50% 이상 절감했으며, 연산량(GFLOPS) 또한 1241.0에서 429.35로 대폭 감소시켜 고해상도 작업에서의 실용성을 입증했다.

Ablation Study 결과, Stage 1(어텐션 정렬) 없이 Stage 2(피처 정렬)만 수행했을 때보다 두 단계를 모두 거쳤을 때 수렴 속도가 더 빠르고 최종 성능이 높게 나타났다. 이는 점진적인 정렬 전략이 비전 모델의 복잡한 특징 공간을 보존하는 데 효과적임을 시사한다.

본 연구는 표준 ViT의 어텐션 메커니즘을 Softmax(QK^T)V에서 φ(Q)(φ(K)^T V) 형태로 변환하는 선형 근사 방식을 채택한다. 여기서 φ(x) = ELU(x) + 1 커널 함수를 사용하여 연산 순서를 변경함으로써 복잡도를 O(N²D)에서 O(ND²)로 낮춘다. [입력 토큰 X → Q, K, V 투영 → φ 함수 적용 → (K^T V) 선행 계산 → Q와 곱셈] 순으로 연산하여 시퀀스 길이 N에 대한 의존성을 제거한다.

LLM 선형화 기법들이 비전 모델에서 성능이 급락하는 원인을 '오차 누적 패턴의 차이'로 규명한다. LLM은 시간축(Temporal)으로 오차가 쌓이는 반면, ViT는 공간적(Spatial) 및 계층적(Hierarchical)으로 오차가 누적되어 전역적인 시맨틱 매니폴드가 왜곡된다. 이를 해결하기 위해 각 블록의 국소적 정렬(Stage 1)과 최종 출력의 전역적 정렬(Stage 2)을 결합한 이중 정렬 구조를 설계했다.

학습 전략 측면에서 Stage 1은 COCO 데이터셋을 사용하여 4에폭의 짧은 학습으로도 충분한 근사가 가능함을 보여주었으며, Stage 2는 ImageNet-22K를 활용해 10~30에폭 동안 전체 네트워크를 미세 조정함으로써 원본 모델의 표현력을 복원한다. 이러한 단계적 접근은 학습 효율성을 극대화하면서도 소프트맥스 어텐션의 정교한 특징 추출 능력을 선형 모델로 성공적으로 전이시킨다.

객체 탐지(Object Detection)나 이미지 생성(Image Generation) 작업에 대한 일반화 성능은 아직 실험적으로 검증되지 않았으며, 저해상도 이미지 처리 시에는 소프트맥스 어텐션 대비 연산 오버헤드가 소폭 발생할 수 있다.