VHA가 KV 캐시와 학습 연산량을 줄이면서 모델 품질을 높인다

무엇인가

긴 문맥을 다루는 대규모 언어 모델에서 어텐션 메커니즘은 KV 캐시 크기, 학습 및 프리필 FLOPs, 모델 품질이라는 세 가지 목표가 서로 충돌하는 '불가능한 삼각형'에 직면해 있다. GQA-8은 품질이 좋지만 BF16 기준 레이어당 토큰당 4KB의 KV 캐시를 사용한다. KV 헤드를 2개로 줄인 GQA-2는 캐시를 1KB로 압축하지만, 1.7B 규모에서 품질이 1.00% 하락한다. MLA는 1.12KB로 압축하지만 여러 규모에서 가장 낮은 품질을 보이고, MFA는 1KB로 줄이면서 GQA-8에 가까운 품질을 내지만 1.7B, S=4096에서 학습 및 프리필 FLOPs가 49.8% 증가한다. 기존 방법 중 세 가지를 동시에 만족하는 것은 없다.

어떻게 동작하나

이 논문은 VHA(Virtual Head Attention)를 제안한다. VHA는 물리적 쿼리 헤드 수를 절반으로 줄이고 KV 헤드를 2개로 설정한 뒤, 두 가지 경량 선형 연산으로 표현력을 복원한다. Q Premix는 각 KV 그룹 내 쿼리에 근사 항등 선형 변환을 적용해 H_q'개의 물리적 쿼리 헤드를 H_v = H_kv × H_q'개의 가상 쿼리 헤드로 확장한다. 변환 행렬 W_g^premix는 그룹별로 존재하며 I + α·ε (ε ~ N(0, σ^2))로 초기화되어 초기에는 항등에 가깝지만 학습이 진행되면서 그룹별로 다른 어텐션 패턴을 만들어낸다. Linear Postmix는 어텐션 출력에 저랭크 항등-잔차 변환 I + BA^T를 적용해 그룹 간 헤드 특징을 융합한다. A, B는 H_v × r 크기이고 r = H_q'/2 (1.7B에서 r=4)이다. A는 작은 랜덤, B는 0으로 초기화되어 학습 초기에는 항등으로 작동한다. 두 연산 모두 비선형성이 없어 추론 시 Premix는 W_Q에, Postmix는 W_O에 흡수되어 모델이 표준 GQA-2로 축소된다. 따라서 vLLM, SGLang 같은 기존 GQA 추론 최적화를 그대로 재사용할 수 있다. KV 헤드가 2개뿐이라 대형 모델에서 그룹 크기가 64에 달해 디코딩 시 Tensor Core 병렬성을 최대한 활용한다.

무엇과 다른가

실험은 0.6B, 1.7B, 4B dense, 8B-A1B MoE 네 규모에서 4K 문맥, 50B 토큰으로 FineWeb-Edu를 사전학습해 진행했다. VHA는 모든 규모에서 GQA 베이스라인을 +0.05%에서 +0.81%까지 앞서면서 KV 캐시를 dense 규모에서 GQA-8 대비 4배, MoE 규모에서 GQA-4 대비 2배 줄였다. 1.7B, S=4096에서 학습 및 프리필 FLOPs를 7.9% 낮췄다. 동일한 KV 캐시(1.00KB)에서 VHA는 MFA를 모든 규모에서 +0.10%에서 +0.77%까지 앞서면서 FLOPs를 7.9% 줄인 반면, MFA는 49.8% 증가했다. MHA는 더 큰 캐시를 쓰고도 GQA-8보다 낮은 품질을 보였고, MLA는 모든 규모에서 뒤처졌다. 이중 압축 형태에서의 회복도 주목할 만하다. GQA의 KV 헤드를 8에서 2로 줄인 GQA 16-2-2는 1.7B에서 -1.00%를 기록했고, 물리적 쿼리 헤드까지 절반으로 줄인 GQA 8-2-2는 추가로 -0.71%를 잃어 총 -1.71%가 됐다. 같은 (8,2,2) 형태에서 VHA는 이 손실을 모두 회복하고 GQA-8보다 +0.41% 높은 성능을 냈다. GQA-8의 KV 캐시 25%, 물리적 쿼리 헤드 50%만 사용하고도 그렇다.

어떻게 쓰나

산업 규모 검증도 있다. 30B-A3B MoE를 500B 토큰, 8K 문맥으로 학습한 결과, VHA는 12개 베이스 모델 벤치마크 평균에서 MLA와 GQA-4를 각각 +0.70%, +2.60% 앞섰다. 카테고리별로는 General +3.39/+0.68%, Math&STEM +3.31/+2.91%, Coding +1.38/-1.14%, Multilingual +2.30/+0.36% (VHA vs GQA/MLA)로 이득이 넓게 분포했다. 별도의 10B-A1B MoE 장문맥 중간학습 비교에서는 8K-128K 구간에서 VHA의 평균 RULER 점수가 84.4로 GQA의 73.0보다 11.4점 높았다. 128K에서 LongBench-v2는 24.2 대 18.9, BABILong은 38.2 대 31.1로 앞섰고, InfiniteBench는 31.5 대 31.8로 비슷했다. HELMET 카테고리는 3승 1무 1패였다. 이때 VHA는 전체 어텐션 레이어에서 GQA-4보다 적은 2개의 KV 헤드를 사용했다.

전제와 한계

절제 실험은 1.7B 규모에서 Premix만 사용하면 -0.78%, Postmix만 사용하면 -0.14%의 격차가 남지만, 둘을 함께 쓰면 GQA-8 대비 +0.41% 이득을 얻는다. Postmix 랭크를 r=4, 32, 128로 바꾸면 평균이 각각 54.50, 54.24, 54.15로 랭크가 커질수록 학습 손실은 낮아지지만 다운스트림 평가는 나빠져 저랭크 제약이 정규화 효과를 준다는 것을 보여준다. SiLU 같은 비선형성을 추가하면 품질이 떨어지고 추론 흡수도 불가능해진다. 학습된 행렬 분석에서 Premix는 그룹별로 다른 쿼리 변환을 학습했고(||W_g^premix - I||_F 피크 5.85, 그룹 간 거리 피크 6.13), Postmix는 |AB^T|에서 비대각 흐름을 보이며 그룹 내 재분배와 소수의 그룹 간 연결을 만들어냈다.

개발자에게 VHA는 긴 문맥 서빙에서 메모리 병목을 줄이면서 품질을 유지하거나 높이고 싶을 때 유용하다. 특히 GQA-2 수준의 KV 캐시를 쓰면서 GQA-8에 가까운 품질을 원하는 경우, 추론 시 표준 GQA-2로 흡수되어 vLLM이나 SGLang에서 추가 커널 없이 바로 사용할 수 있다는 점이 실무적으로 크다. 도입 시 확인할 점은 학습 시 Premix와 Postmix가 제대로 학습되는지, 그룹 크기가 충분히 커서 Tensor Core 이득을 보는지, 그리고 Postmix 랭크를 모델 규모에 맞게 설정했는지다. 또한 기존 GQA 체크포인트를 VHA로 변환하는 것이 아니라 처음부터 VHA 구조로 학습해야 한다.

논문은 별도의 한계 절을 두지 않았지만, 몇 가지 전제가 있다. VHA의 추론 흡수는 모든 추가 연산이 선형일 때만 성립하며, 비선형성을 넣으면 품질도 떨어지고 흡수도 깨진다. 실험은 최대 30B-A3B MoE와 128K 문맥까지이며, 50B 또는 500B 토큰 학습 예산에 기반한다. MLA가 모든 규모에서 뒤처진 결과도 이 학습 예산에 의존할 수 있다. 또한 학습된 행렬 분석은 1.7B와 4B 모델에 한정되어 있어 다른 규모에서 같은 구조가 나타나는지는 추가 검증이 필요하다. 이론적 보장 없이 경험적 결과에 의존한다는 점도 염두에 둬야 한다.