트랜스포머 QKV 투영 생략 연구
트랜스포머 모델의 핵심인 어텐션(Attention)에서 반드시 필요한 Query, Key, Value(QKV) 세 개의 투영(Projection)을 공유하거나 생략해도 모델 성능을 유지할 수 있다는 연구 결과입니다. 특히 Key와 Value를 하나로 합치는 방식은 추론 시 메모리 캐시를 최대 96.9%까지 획기적으로 줄이면서도 언어 모델의 성능 저하를 최소화하여, 스마트폰 등 엣지(Edge) 기기에서의 AI 추론 활용을 크게 앞당길 수 있는 실용적인 기여를 합니다.
제목: 트랜스포머는 세 개의 투영(Projections)이 필요한가? QKV 변형에 대한 체계적인 연구
초록: 트랜스포머는 다양한 AI 작업을 위한 표준 솔루션이 되었으며, 쿼리(Query), 키(Key), 밸류(Value)(QKV) 어텐션 구조가 그 중심 역할을 하고 있습니다. 하지만 이 세 가지 투영이 각각 개별적으로 기여하는 바와 이 중 일부를 생략했을 때 미치는 영향에 대해서는 아직 제대로 이해되지 않았습니다. 본 연구에서는 세 가지 투영 공유 제약 조건을 체계적으로 평가했습니다: a) Q-K=V (키-밸류 공유), b) Q=K-V (쿼리-키 공유), c) Q=K=V (단일 투영). 마지막 두 변형은 대칭적인 어텐션 맵(Attention map)을 생성하며, 이를 해결하기 위해 2D 위치 인코딩(Positional encoding)을 통한 비대칭 어텐션도 함께 탐구했습니다.
합성 작업, 비전(MNIST, CIFAR, TinyImageNet, 이상 탐지) 및 언어 모델링(100억 토큰에 대한 3억 및 12억 파라미터 모델)에 걸친 실험을 통해, 이러한 변형 트랜스포머들이 기존 QKV 트랜스포머와 동등하거나 때로는 더 나은 성능을 보여준다는 것을 발견했습니다. 언어 모델링의 경우, Q-K=V 투영 공유 방식은 KV 캐시를 50% 줄이면서도 퍼플렉시티(Perplexity) 저하를 단 3.1%로 억제했습니다. 특히 주목할 점은, 투영 공유가 헤드 공유(GQA/MQA)와 상호 보완적이라는 것입니다. Q-K=V를 GQA-4와 결합하면 캐시를 87.5% 줄일 수 있고, Q-K=V + MQA를 적용하면 96.9%를 감소시켜 실제 온디바이스(On-device) 추론을 가능하게 합니다.
Q-K=V 방식이 품질을 유지하는 이유는 키와 밸류가 유사한 표현 공간을 차지할 수 있고 어텐션이 저계수(Low-rank) 체계에서 작동하기 때문인 반면, Q=K-V는 어텐션의 방향성을 훼손한다는 것을 보여주었습니다. 본 연구 결과는 투영 공유를 어텐션 내 가중치 공유(Weight tying)의 탐구 부족한 사례로 체계적으로 특성화하며, 특히 엣지 배포(Edge deployment)에 매우 가치 있는 직접적이고 정량화 가능한 추론 메모리 이점을 제공합니다.
코드는 공개적으로 사용 가능합니다.