Do transformers need three projections? Systematic study of QKV variants

개요

트랜스포머의 쿼리, 키, 값(QKV) 어텐션 구조에서 세 가지 프로젝션의 개별적인 기여와 일부를 생략했을 때의 영향을 체계적으로 분석하여, 프로젝션 공유가 추론 메모리 사용량을 크게 줄이면서도 성능 저하를 최소화할 수 있음을 발견했습니다.

주요 내용

• QKV 어텐션은 다양한 AI 작업의 표준으로 자리 잡았지만, 세 가지 프로젝션(Query, Key, Value) 각각의 역할과 프로젝션 공유 제약 조건(Q-K=V, Q=K-V, Q=K=V)이 모델 성능에 미치는 영향은 아직 충분히 이해되지 않았습니다.
• 세 가지 프로젝션 공유 제약 조건(Q-K=V: 쿼리와 키가 공유, 값은 별도; Q=K-V: 쿼리와 키가 별도, 값이 공유; Q=K=V: 모든 프로젝션이 단일 투영으로 공유)을 체계적으로 평가했습니다.
• Q=K-V 및 Q=K=V 제약 조건은 대칭적인 어텐션 맵을 생성하므로, 이를 해결하기 위해 2D 위치 인코딩을 사용한 비대칭 어텐션도 탐색했습니다.
• 합성 데이터셋, 비전 작업(MNIST, CIFAR, TinyImageNet, anomaly), 그리고 언어 모델링(100억 토큰 데이터셋에서 3억 및 12억 파라미터 모델)에 걸친 실험 결과, 제안된 변형 트랜스포머는 기존 QKV 트랜스포머와 동등하거나 때로는 더 나은 성능을 보였습니다.
• 특히 언어 모델링에서 Q-K=V 프로젝션 공유는 50%의 KV 캐시를 절감하면서도 퍼플렉시티(perplexity) 성능 저하를 3.1% 미만으로 유지했습니다.
• 프로젝션 공유는 헤드 공유(GQA/MQA)와 상호 보완적입니다. Q-K=V와 GQA-4를 결합하면 87.5%의 캐시 감소를, Q-K=V와 MQA를 결합하면 96.9%의 캐시 감소를 달성하여 온디바이스 추론을 실질적으로 가능하게 합니다.
• Q-K=V 방식이 성능을 유지하는 이유는 키와 값 프로젝션이 유사한 표현 공간을 점유할 수 있고, 어텐션이 저랭크(low-rank) 영역에서 작동하기 때문이며, 반면 Q=K-V는 어텐션의 방향성을 깨뜨리는 것으로 나타났습니다.
• 본 연구 결과는 프로젝션 공유를 어텐션 메커니즘 내에서 탐구되지 않은 가중치 공유(weight tying)의 한 형태로 체계적으로 특징화하며, 특히 엣지 디바이스 배포에 중요한 정량적인 추론 메모리 이점을 제공합니다.

시사점

트랜스포머 모델의 프로젝션 공유, 특히 Q-K=V 방식은 모델의 메모리 효율성을 크게 향상시키면서도 성능 저하를 최소화하여, 자원 제약이 있는 환경에서의 딥러닝 모델 배포 가능성을 높이는 실용적인 해결책을 제시합니다.