기술 요약: 트랜스포머는 세 개의 프로젝션이 필요한가?

문제 제기

트랜스포머 아키텍처가 도입된 이래, Query, Key, Value (QKV) 어텐션 메커니즘은 시퀀스 전반의 정보를 집계하는 표준으로 자리 잡았습니다. 그러나 이 세 가지 구성 요소를 위해 세 개의 서로 독립적인 선형 프로젝션을 유지해야 하는 필요성은 여전히 미해결 과제로 남아 있습니다. CNN이나 상태 공간 모델(SSM)과 같은 다른 아키텍처들이 종종 더 통합된 내부 표현을 사용하는 것과 달리, 트랜스포머는 삼분법적 구조를 고수하고 있습니다. 컨텍스트 윈도우가 확장되고 실시간 추론에 대한 수요가 증가함에 따라, 연구 커뮤니티는 셀프 어텐션의 이차적 병목 현상을 완화하기 위한 아키텍처 효율성에 집중해 왔습니다. 본 논문은 세 개의 프로젝션이 진정으로 필요한 것인지, 아니면 성능 저하 없이 파라미터 수를 줄이고 특히 더 중요한 추론 메모리 오버헤드를 줄일 수 있는 중복성인지를 조사합니다.

방법론

저자들은 학습된 변환의 수를 점진적으로 줄이는 세 가지 "프로젝션 공유(Projection-Sharing)" 제약 조건을 제안하고 체계적으로 평가합니다:

1. Q=K-V (통합 Query-Key): 쿼리와 키 프로젝션을 통합($Q=K$)하되, 밸류 프로젝션은 별도로 유지합니다. 이는 대칭적인 어텐션 행렬($KK^T$)을 생성합니다. 순차적 작업에서의 대칭성 한계를 해결하기 위해, 저자들은 2D 위치 인코딩을 통해 비대칭성을 주입하는 변형 모델인 $(Q=K-V)^+$를 도입합니다.
2. Q-K=V (통합 Key-Value): 키와 밸류 프로젝션을 통합($K=V$)하고, 쿼리는 별도로 유지합니다. 이는 비대칭적 어텐션 맵($QK^T$)을 보존하며, 키와 밸류 사이에 일종의 가중치 타이잉(weight tying)을 부과합니다.
3. Q=K=V (단일 프로젝션): 세 가지 역할 모두가 단일 프로젝션 행렬을 공유합니다. 이는 첫 번째 변형의 대칭성과 두 번째 변형의 표현력 병목 현상을 결가합한 형태입니다.

본 연구는 또한 이러한 프로젝션 공유 전략이 그룹 쿼리 어텐션(GQA) 및 멀티 쿼리 어텐션(MQA)과 같은 기존의 "헤드 공유(Head-Sharing)" 기술과 어떻게 직교하는지 탐구합니다. 저자들은 프로젝션 공유(K=V 제약)를 헤드 공유(KV 헤드 수 감소)와 결합하면 메모리 효율성 측면에서 복합적인 이득을 얻을 수 있을 것이라고 가설을 세웠습니다.

실험은 12가지 다양한 작업에 대해 수행되었습니다:

• 합성 추론(Synthetic Reasoning): 리스트 조작(Reverse, Sort, Sub, Swap, Copy)과 관련된 5가지 작업.
• 컴퓨터 비전: MNIST, FashionMNIST, CIFAR-10, CIFAR-100, TinyImageNet에 대한 이미지 분류, 그리고 세트 이상 탐지(set anomaly detection) 및 시맨틱 세그멘테이션.
• 언어 모델링: SlimPajama 데이터셋의 100억 개 토큰을 사용한 300M 및 1.2B 파라미터 GPT 스타일 모델의 사전 학습.

주요 결과

1. 성능 대 효율성 트레이드오프

• Q-K=V (최적의 균형): 이 변형이 명확한 승자로 나타났습니다. 언어 모델링에서 이 방식은 KV 캐시 크기를 50% 감소시키면서도 (300M 모델에서) 퍼플렉시티(perplexity) 저하는 단 3.1%, (1.2B 모델에서) **2.48%**에 그쳤습니다. 또한 비전 및 합성 작업에서도 경쟁력 있는 성능을 유지했습니다.
• Q=K-V (대칭성의 한계): 이 변형은 베이스라인과 유사한 훈련 품질을 달랐으나, 별도의 K와 V 텐서를 여전히 캐싱해야 하므로 추론 메모리 이득은 전혀 없었습니다. 또한 대칭적인 어텐션 맵이 인과적 언어 모델링(causal language modeling)에서의 효과를 제한하여 4.9%의 퍼플렉시티 저하를 초래했습니다.
• Q=K=V (과도한 제약): 이 공격적인 단순화는 치명적인 성능 손실(언어 모델링에서 25.4% 퍼플렉시티 저하)을 일으켰으며, 이는 세 가지 역할 모두에 단일 프로젝션을 강제하는 것이 복잡한 작업에는 너무 제한적임을 나타냅니다.

2. 헤드 공유와의 시너지

연구는 프로젝션 공유와 헤드 공유가 엄격하게 상호 보완적임을 보여줍니다.

• Q-K=V를 GQA-4(4개 그룹)와 결합했을 때, 3.9%의 퍼플렉시티 저하와 함께 87.5%의 캐시 감소를 달야냈습니다.
• Q-K=V를 MQA(1개 그룹)와 결합했을 때, 4.8%의 퍼플렉시티 저하와 함께 96.9%의 캐시 감소를 달성했습니다.
• 1.2B 규모에서 Q-MQA 구성은 (32k 컨텍스트 기준 사용자당 5.9 GB에서 88 MB로) KV 캐시를 67배 감소시키면서도 다운스트림 작업 정확도에서 베이스라인과 거의 대등한 수준을 유지했습니다.

3. 스케일링 동작

변형 모델들의 상대적 순위는 모델 크기가 300M에서 1.2B 파라미터로 증가함에 따라 안정적으로 유지되었습니다. 주목할 점은, 더 큰 모델이 프로젝션 제약에 대해 약간 더 견고해 보였다는 것입니다. Q-K=V의 퍼플렉시티 저하는 300M에서의 3.1%에서 1.2B에서의 2.48%로 감소했습니다.

의의 및 주장

본 논문은 프로젝션 공유가 직접적이고 정량화 가능한 추론 메모리 이득을 제공하는, 아직 충분히 탐구되지 않은 가중치 타이잉(weight tying)의 사례라고 주장하며, 이는 엣지 배포 및 롱 컨텍스트 애플리케이션에 특히 유용합니다.

• Q-K=V가 작동하는 이유: 저자들은 키(Key)와 밸류(Value)가 유사한 표현 공간을 점유할 수 있으며, 어텐션이 Q와 공유된 K-V 표현 사이의 비대칭성보다 K와 V 사이의 구분이 덜 중요한 저차원(low-rank) 영역에서 효과적으로 작동한다고 주장합니다. 경험적 분석 결과, 표준 K와 V 프로젝션 행렬 사이의 코사인 유사도가 0.73으로 높게 나타났으며, 이는 내재된 중복성을 시사합니다.
• Q=K-V가 실패하는 이유: $Q=K$를 강제하면 대칭적인 어텐션 맵이 생성되어, 언어 모델링에 필요한 방향성을 깨뜨립니다.
• 실질적 영향: 주요 기여는 단순히 파라미터 감소(전체 모델 크기의 약 7~14% 감소로 다소 미미함)가 아니라, 50%의 KV 캐시 메모리 감소에 있습니다. 이를 통해 다음이 가능해집니다:
  • 동일한 메모리 예산 내에서 2배 더 긴 컨텍스트 윈도우.
  • 2배 더 높은 처리량(GPU당 동시 사용자 수).
  • 상당한 비용 절감(예: 코드 완성 모델의 서빙 비용 43% 절감 예상).

저자들은 단일 아키텍처가 모든 도메인에서 승리하는 것은 아니지만, Q-K=V 구성이 순차적 작업에 대해 최적의 품질-효율성 트레이드오프를 제공하며, 헤드 공유 기술과의 결합을 통해 수십억 파라미터 모델의 실용적인 온디바이스 추론을 가능하게 한다고 결론짓습니다. 연구를 촉진하기 위해 코드는 공개되어 있습니다.