[논문 리뷰] Rethinking Query, Key, and Value Embedding in Vision Transformer under Tiny Model Constraints

Jaesin Ahn, Jiuk Hong외 2명|arXiv (Cornell University)|2021. 11. 18.

Machine Learning and ELM인용 수 2

한 줄 요약

이 논문은 초소형 모델 제약 조건 하에서 비선형, 공유, 코드 기반의 쿼리(Q), 키(K), 값(V) 임베딩 구조를 제안한다. 표준 선형 투영 대신 학습 가능한 비선형 매핑(특히 공유 레이어와 학습 가능한 코드 파라미터를 통해)을 도입함으로써, 파rameter 수를 크게 줄이며 이미지 분류 정확도를 향상시킨다. 이는 오직 310만 개의 파라미터로 ImageNet-1k에서 71.4%의 top-1 정확도를 달성하며, 원본 XCiT-N12 모델을 능가한다.

ABSTRACT

A vision transformer (ViT) is the dominant model in the computer vision field. Despite numerous studies that mainly focus on dealing with inductive bias and complexity, there remains the problem of finding better transformer networks. For example, conventional transformer-based models usually use a projection layer for each query (Q), key (K), and value (V) embedding before multi-head self-attention. Insufficient consideration of semantic $Q, K$, and $V$ embedding may lead to a performance drop. In this paper, we propose three types of structures for $Q$, $K$, and $V$ embedding. The first structure utilizes two layers with ReLU, which is a non-linear embedding for $Q, K$, and $V$. The second involves sharing one of the non-linear layers to share knowledge among $Q, K$, and $V$. The third proposed structure shares all non-linear layers with code parameters. The codes are trainable, and the values determine the embedding process to be performed among $Q$, $K$, and $V$. Hence, we demonstrate the superior image classification performance of the proposed approaches in experiments compared to several state-of-the-art approaches. The proposed method achieved $71.4\%$ with a few parameters (of $3.1M$) on the ImageNet-1k dataset compared to that required by the original transformer model of XCiT-N12 ($69.9\%$). Additionally, the method achieved $93.3\%$ with only $2.9M$ parameters in transfer learning on average for the CIFAR-10, CIFAR-100, Stanford Cars datasets, and STL-10 datasets, which is better than the accuracy of $92.2\%$ obtained via the original XCiT-N12 model.

연구 동기 및 목표

• 비선형성과 의미적 Q, K, V 임베딩 고려 부족으로 인한 초소형 비전 트랜스포머의 성능 저하 문제를 해결하기 위해.
• Q, K, V 임베딩 메커니즘의 재고를 통해 엄격한 파라미터 제약 조건 하에서 ViT 성능을 향상시키기 위해.
• 공유 비선형 레이어와 학습 가능한 코드 파라미터를 통해 Q, K, V 간 지식 공유를 탐색하기 위해.
• 비선형 및 공유 QKV 임베딩 구조가 ImageNet 분류 및 전이 학습 성능을 향상시킬 수 있음을 입증하기 위해.

제안 방법

• 입력 토큰을 별개의 비선형 공간으로 변환하기 위해, Q, K, V에 대해 두 계층의 ReLU 기반 비선형 임베딩을 도입한다.
• Q, K, V 전반에 걸쳐 단일 비선형 레이어를 공유하는 한 계층의 공유 구조를 제안하여 지식 전이를 장려한다.
• 학습 가능한 코드 파라미터(Cq, Ck, Cv)를 갖춘 두 계층의 공유 구조를 개발하여 Q, K, V의 임베딩 변환을 정의한다.
• 백프로파게이션을 사용하여 분류 손실을 최소화할 수 있도록 코드 파라미터를 공동으로 학습시킨다.
• 코드 유사도와 직교성을 분석하기 위해 F-SNE 시각화를 활용하여, 코드가 구분되며 임의의 작업에 관계없이 특징을 학습하고 있음을 확인한다.
• 파라미터 제약 조건 하에서 ImageNet-1k 및 전이 학습 벤치마크(CIFAR-10, CIFAR-100, Stanford Cars, STL-10)에서 성능을 평가한다.

실험 결과

연구 질문

• RQ1 초소형 모델 제약 조건 하에서 선형 QKV 투영을 비선형 매핑으로 대체하면 ViT 성능이 향상되는가?
• RQ2 Q, K, V 간 비선형 레이어를 공유하면 특징 학습과 분류 정확도가 향상되는가?
• RQ3 Q, K, V 임베딩을 동시에 정의하는 학습 가능한 코드 파라미터는 독립적 투영보다 더 나은 성능을 낼 수 있는가?
• RQ4 제안된 QKV 임베딩 구조는 ImageNet 및 전이 학습 과제에서 XCiT-N12와 같은 최신 기술 모델과 비교해 어떻게 성능을 내는가?

주요 결과

• 제안된 방법은 오직 310만 개의 파라미터로 ImageNet-1k에서 71.4%의 top-1 정확도를 달성하였으며, 원본 XCiT-N12 모델(69.9%)을 능가하였다.
• 전이 학습에서, CIFAR-10, CIFAR-100, Stanford Cars, STL-10 평균 정확도가 93.3%로, 원본 XCiT-N12 모델(92.2%)을 초월하였다.
• 학습 가능한 파라미터를 갖춘 코드 기반 공유 구조는 특히 CIFAR-100과 STL-10에서 뛰어난 성능을 보이며, 다양한 작업 간 효과적인 특징 학습이 이루어지고 있음을 시사한다.
• F-SNE 시각화 결과, 학습된 코드(Cq, Ck, Cv)가 거의 직교성을 보이며, 서로 다른 임의의 작업에 관계없는 표현을 학습하고 있음을 확인하였다.
• 코드의 l2-노름은 ImageNet, Cars, STL-10에서 일관되게 유지되었지만, CIFAR-10과 CIFAR-100에서는 다름을 보이며, 이는 데이터셋에 특화된 코드 적응이 이루어지고 있음을 나타낸다.
• 나노 모델의 경우 코드 크기 8이 최적의 성능을 보였고, 티니 모델의 경우 16이 최적의 성능을 보였다. 이는 코드 크기가 임베딩 차원에 따라 조정되어야 함을 시사한다.