Key and Value Weights Are Probably All You Need: On the Necessity of the Query, Key, Value weight Triplet in Self-Attention Transformers

이 논문은 자기주의 (Self-Attention) 메커니즘에서 쿼리 (Query) 가중치가 불필요할 수 있음을 이론적으로 증명하고, 이를 제거하여 매개변수를 25% 줄이면서도 GPT 스타일 모델의 성능을 유지하고 학습 안정성을 높일 수 있음을 실험적으로 입증했습니다.

핵심

🧠 "쿼리 (Query) 가 없어도 될까요?" - 트랜스포머의 숨겨진 비밀

이 논문은 최근 인공지능 (LLM) 을 만드는 데 쓰이는 **'트랜스포머 (Transformer)'**라는 구조에 대해 아주 흥미로운 질문을 던집니다.

"정말 우리가 3 개의 키 (Key), 쿼리 (Query), 밸류 (Value) 가 모두 필요할까? 아니면 그중 하나는 버려도 되지 않을까?"

저자들은 이 중 **쿼리 (Query) 가 사실은 '불필요한 장난감'**일 수 있음을 수학적으로 증명하고, 실제로 실험해 보았습니다. 결과는 놀랍습니다. 쿼리 부분을 없애고 '1' (정체성 행렬) 로만 대체해도 성능은 그대로 유지되면서, 모델 크기는 25% 줄어듭니다.

이 복잡한 내용을 일상적인 비유로 쉽게 설명해 드릴게요.

---

1. 비유: 도서관 사서와 책 찾기 📚

트랜스포머 모델이 문장을 이해하는 과정을 거대한 도서관에 비유해 봅시다.

• 입력 (Input): 도서관에 들어온 손님들이 찾는 주제입니다.
• 키 (Key): 책들의 제목이나 목차입니다. (무엇이 있는지 알려줌)
• 밸류 (Value): 책의 실제 내용입니다. (정보의 원천)
• 쿼리 (Query): 손님이 사서에게 "저기, 이 주제에 맞는 책이 어디 있나요?"라고 묻는 질문입니다.

기존의 방식은 사서가 손님의 질문 (쿼리) 을 받아서, 키 (제목) 와 비교한 뒤, 가장 적합한 책 (밸류) 을 찾아주는 복잡한 과정을 거칩니다.

이 논문의 핵심 발견은 다음과 같습니다:

"사실 사서가 질문 (쿼리) 을 다시 적어줄 필요는 없습니다. 손님이 원래 가진 질문을 그대로 가져가서 키 (제목) 와 비교해도, 결국 같은 책을 찾아낼 수 있어요!"

즉, 쿼리라는 '변환 과정'을 생략하고, 손님의 질문을 그대로 (Identity, 즉 1) 사용해도 도서관의 업무 흐름이 전혀 깨지지 않는다는 것입니다.

2. 왜 이것이 중요할까요? (효율성의 마법) ✨

지금까지의 AI 모델은 쿼리, 키, 밸류 세 가지 '가중치 (Weight)'를 모두 학습시켰습니다. 마치 사서가 질문을 변형하는 데도 별도의 훈련을 시킨 셈이죠.

하지만 이 논문에 따르면:

1. 쿼리 가중치를 없애고 '1'로 고정하면?
   • 모델이 배워야 할 파라미터 (학습 데이터) 가 25% 줄어듭니다.
   • 이는 마치 도서관 사서 한 명을 줄여도 업무가 똑같이 돌아가는 것과 같습니다.
2. 성능은 어떨까요?
   • 실험 결과, 쿼리를 없앤 모델이 기존 모델과 똑같은 성능을 냈습니다.
   • 오히려 줄어든 파라미터를 다른 부분 (MLP, 즉 문장 생성을 담당하는 뇌) 에 더 투자하면, 기존 모델보다 더 좋은 성능을 내기도 했습니다.

3. 왜 우리가 몰랐을까요? (과도한 설계) 🏗️

왜 그동안没人이 이걸 발견하지 못했을까요?

• 과도한 설계 (Over-parameterization): 우리는 "무조건 3 개가 필요할 거야"라고 생각하며 모델을 설계해 왔습니다. 하지만 실제로는 그중 하나가 redundant(중복된) 역할을 하고 있었을 뿐입니다.
• 수학적 장벽: 쿼리를 없애면 수학적으로 '기저 변환 (Basis Transformation)'이라는 복잡한 문제가 생깁니다. 마치 도서관의 책 정렬 방식을 바꿀 때, 모든 책을 다시 정리해야 하는 것처럼요. 하지만 저자들은 **"MLP(신경망) 가 이 정렬 작업을 자동으로 해낼 수 있다"**는 것을 증명했습니다.

4. 실험 결과: 실제로 작동합니다! 🚀

저자들은 1 억 1,700 만 개의 파라미터를 가진 작은 GPT 스타일 모델을 처음부터 다시 훈련시켰습니다.

• 조건: 쿼리 가중치를 없애고, 학습 방식만 약간 조정했습니다.
• 결과:
  • 성능: 기존 모델과 똑같은 점수 (Validation Loss) 를 기록했습니다.
  • 안정성: 학습이 훨씬 더 안정적으로 이루어졌습니다. (가중치 감소라는 '규제'를 3 배나 줄여도 모델이 망가지지 않음)
  • 효율: 파라미터를 줄인 덕분에 계산 비용과 메모리 사용량이 크게 감소했습니다.

5. 요약: 무엇을 얻었나요? 🎁

이 논문은 AI 의 미래를 위한 두 가지 큰 선물을 줍니다.

1. 비용 절감: 더 적은 전력과 메모리로 똑똑한 AI 를 만들 수 있습니다. (쿼리 가중치 제거 = 25% 효율 향상)
2. 새로운 설계 철학: "무조건 복잡한 게 좋은 게 아니다." 우리가 당연시 여겼던 구조 (쿼리, 키, 밸류의 3 중주) 가 사실은 과도한 설계일 수 있음을 깨닫게 해줍니다.

한 줄 요약:

"AI 가 문장을 이해할 때, '질문 (쿼리) 을 변형하는 과정'은 사실 불필요한 수작업이었습니다. 이걸 없애고 그대로 두면, 모델은 더 가볍고 빠르면서도 똑똑해집니다!"

이 발견은 앞으로 더 크고 복잡한 AI 모델을 만들 때, 불필요한 장비를 덜어내고 핵심 기능에 집중할 수 있는 길을 열어줄 것입니다.

기술 요약

이 논문은 트랜스포머 (Transformer) 아키텍처의 자기 주의 (Self-Attention) 메커니즘에서 Query(Q), Key(K), Value(V) 가중치 triplet 의 필요성을 이론적으로 분석하고 실험적으로 검증한 연구입니다. 저자들은 이 중 Query 가중치 ($W_Q$) 가 불필요하며, 이를 단위 행렬 (Identity Matrix) 로 대체할 수 있음을 증명하여 모델의 파라미터 수를 25% 줄이면서도 성능을 유지하거나 향상시킬 수 있음을 보여줍니다.

다음은 논문의 주요 내용을 기술적으로 요약한 것입니다.

---

1. 연구 배경 및 문제 제기 (Problem)

• 배경: 트랜스포머 기반 언어 모델의 학습 및 배포는 계산 비용이 매우 높습니다. 이를 해결하기 위해 양자화, 효율적 어텐션, 가중치 공유 등 다양한 최적화 기법이 제안되었습니다.
• 문제: 자기 주의 메커니즘 내부의 Query-Key-Value 가중치 triplet 은 실제로 모두 필요한 것일까요? 특히, 입력 $X$가 어텐션에 영향을 미치는 방식이 $XW_Q$, $XW_K$, $XW_V$의 곱을 통해서만 이루어진다는 점에 착안하여, 이 중 하나를 제거하거나 단순화할 수 있는지 의문을 제기합니다.
• 목표: Query 가중치 ($W_Q$) 를 단위 행렬 ($I$) 로 대체하여 파라미터 수를 줄이고, 모델의 구조적 중복성을 규명하며, 이를 통해 효율성을 높이는 것입니다.

2. 방법론 (Methodology)

가. 이론적 분석 (Theoretical Analysis)

저자는 재매개변수화 (Reparametrization) 레마를 기반으로 한 이론적 증명을 통해 $W_Q$ 제거의 가능성을 규명했습니다.

1. 재매개변수화 레마 (Reparametrization Lemma):

   • 어텐션 메커니즘은 입력 $X$를 $XW_Q, XW_K, XW_V$의 곱을 통해서만 관측합니다.
   • 따라서, 가역 행렬 $\Theta$에 대해 $(X, W_Q, W_K, W_V) \to (X\Theta, \Theta^{-1}W_Q, \Theta^{-1}W_K, \Theta^{-1}W_V)$로 변환해도 출력은 동일하게 유지됩니다.
   • 이를 통해 $W_Q$를 $\Theta = W_Q$로 선택하여 단위 행렬 ($I$) 로 만들 수 있으며, 그 반대의 효과는 $W_K$와 $W_V$를 적절히 조정하여 흡수할 수 있음을 보였습니다.

2. 단일 레이어 제거 (Single-Layer Elimination):

   • 정규화 (Normalization) 레이어가 없는 단일 레이어 트랜스포머에서는 $W_Q$를 $I$로 변경하는 것이 항상 가능합니다 (Theorem 4.1). 이는 기존 학습된 모델에서 레이어 노름을 제거한 후 적용 가능한 "무료 점심 (Free Lunch)"입니다.

3. 다중 레이어 제거 조건 (Multi-Layer Elimination):

   • 모든 레이어에서 동시에 $W_Q$를 제거하려면 다음과 같은 조건 중 하나가 필요합니다 (Theorem 4.2, 4.3):
     • Attention Skip Connection Only: MLP 주변이 아닌 어텐션 주변에만 스킵 연결이 있는 경우. 이 경우 각 블록이 서로 다른 기저 (basis) 를 변환하며 연결될 수 있습니다.
     • Weight Sharing: 모든 레이어가 동일한 가중치를 공유하는 경우 (예: ALBERT, TRM).
   • LayerNorm 의 역할: LayerNorm 이 존재할 경우 정확한 기저 변환이 불가능할 수 있으나, 근사적으로 해결 가능하거나 (Appendix 8.3), 실험적으로 조정된 하이퍼파라미터로 극복 가능함을 보였습니다.

4. 스킵 연결과 MLP 의 기하학적 분석:

   • ReLU 기반 MLP 에서 스킵 연결이 포함된 함수 ($W_2 \text{ReLU}(W_1 x) + x$) 가 동일한 폭의 일반 ReLU MLP 로 정확히 표현될 수 있는 조건을 수학적으로 규명했습니다 (Theorem 8.4). 이는 스킵 연결이 단순히 표현력을 높이는 것이 아니라 함수 공간의 다른 영역에 접근하게 함을 시사합니다.

나. 실험적 검증 (Empirical Validation)

• 모델 설정: OpenWebText 데이터셋에서 0 부터 학습 (Pretraining) 하는 GPT 스타일 모델 (117M ~ 124M 파라미터) 을 사용했습니다.
• 비교 대상:
  • Baseline: 표준 GPT-2 아키텍처 (모든 $W_Q, W_K, W_V, W_O$ 포함).
  • Reduced ($W_Q=I$): $W_Q$를 단위 행렬로 대체한 모델.
• 실용적 조정 (Practical Adjustments):
  • Attention Scaling: $W_Q=I$일 때 쿼리가 학습된 투영이 아닌 입력의 좌표 슬라이스가 되므로, 초기 어텐션 점수의 분산이 커집니다. 이를 보정하기 위해 스케일링 인자를 $\frac{1}{\sqrt{d_k}}$에서 $\frac{1}{2\sqrt{d_k}}$로 변경했습니다.
  • Weight Decay: 이론적으로 남은 파라미터가 원래 기능과 기저 변환을 모두 인코딩할 수 있으므로, 정규화 강도를 낮추어 학습을 안정화했습니다 (0.1 에서 $2^{-5} \approx 0.03$으로 감소).

3. 주요 결과 (Key Results)

1. 파라미터 효율성:

   • $W_Q$를 제거한 Reduced 모델 (117M 파라미터) 은 표준 Baseline 모델 (124M 파라미터) 과 **동등한 검증 손실 (Validation Loss)**을 달성했습니다.
   • 이는 어텐션 파라미터의 25% 를 줄임 (전체 트랜스포머 블록 파라미터의 약 8% 감소) 으로도 성능이 유지됨을 의미합니다.

2. 파라미터 재배치 효과:

   • $W_Q$ 제거로 절약된 파라미터를 MLP 히든 차원에 재배치한 모델 (Reduced 124M, Larger MLP) 은 표준 Baseline (124M) 보다 **더 낮은 검증 손실 (3.004 vs 3.016)**을 기록하여 성능이 향상되었습니다. 이는 MLP 의 표현력이 트랜스포머의 병목 요소일 수 있음을 시사합니다.

3. 학습 안정성:

   • Reduced 모델은 표준 모델보다 3 배 낮은 Weight Decay에서도 학습이 안정적으로 진행되었습니다. 이는 $W_Q$ 제거가 암시적 정규화 (Implicit Regularization) 역할을 하여 모델의 안정성을 높였음을 의미합니다.

4. 비교 우위:

   • 동일한 파라미터 수 (117M) 를 가진 Baseline 모델 (MLP 크기 축소 또는 $d$ 축소) 보다 Reduced 모델이 성능이 더 우수했습니다. 즉, 파라미터를 단순히 줄이는 것보다 $W_Q$를 제거하는 것이 더 효율적인 구조적 단순화입니다.

4. 기여 및 의의 (Contributions & Significance)

• 이론적 통찰: 트랜스포머의 Query-Key-Value triplet 중 Query 가중치가 구조적으로 불필요할 수 있음을 수학적으로 증명했습니다. 이는 트랜스포머 아키텍처의 과잉 매개변수화 (Overparameterization) 에 대한 새로운 관점을 제시합니다.
• 실용적 효율성:
  • KV Cache 및 GQA 호환성: Query 가중치를 제거하더라도 KV Cache 와 Grouped-Query Attention (GQA) 최적화와 충돌하지 않아, 실제 LLM 추론 및 학습에 바로 적용 가능합니다.
  • 학습 비용 절감: 파라미터 수 감소로 인한 메모리 및 계산량 절감 효과가 있습니다.
• 미래 방향성:
  • 이 연구는 트랜스포머의 어떤 구성 요소가 표현력 (Expressivity) 에 필수적인지, 어떤 것이 설계 역사적 산물인지에 대한 질문을 제기합니다.
  • 비선형 Query 변환 ($Q(X) = X + N(X)$) 과 같은 확장 가능성을 제시하며, 다양한 모달리티와 대규모 모델에서의 적용 가능성을 탐구할 수 있는 기반을 마련했습니다.

5. 결론

이 논문은 "Query 가중치는 단위 행렬로 대체 가능하다"는 명제를 이론과 실험을 통해 입증했습니다. 이를 통해 트랜스포머 모델의 파라미터 수를 25% 줄이면서도 성능을 유지하거나 오히려 향상시킬 수 있으며, 특히 MLP 로 절약된 자원을 재배치할 때 더 큰 이점을 얻을 수 있음을 보여주었습니다. 이는 효율적인 트랜스포머 아키텍처 설계에 중요한 이정표가 될 것입니다.