멀티헤드 잠재 어텐션(MLA)이란? KV 캐시 절감의 원리 정리

핵심 메시지

• 언어모델이 긴 문맥을 처리할 때 메모리를 채우는 주범은 계산 자체가 아니라 토큰마다 쌓이는 KV(키·값) 캐시다.
• MQA·GQA는 헤드가 키·값을 공유해 캐시를 줄이지만, 그 대가로 헤드들이 토큰을 똑같은 방식으로만 보게 되어 다양성을 잃는다.
• MLA는 헤드를 공유하는 대신 은닉 상태를 작은 잠재 벡터 하나로 압축해 저장하고, 필요할 때 헤드마다 다른 업프로젝션으로 되살린다.
• 행렬 곱을 재배치하면 복원 과정이 기존 가중치 속으로 흡수되어, 캐시는 MQA만큼 작으면서 헤드 다양성은 풀 어텐션 수준으로 유지된다.
• 위치 정보를 넣는 회전 위치 인코딩(RoPE)은 흡수를 깨므로, 키를 내용 레인과 위치 레인으로 나눠 이 문제를 실용적으로 해결한다.

쉽게 이해하기

언어모델은 새 단어를 만들 때마다 앞서 나온 모든 단어를 다시 참조해야 한다. 같은 계산을 반복하지 않으려고 각 토큰마다 키와 값을 저장해 두는데, 이 저장소가 바로 KV 캐시다. 문제는 캐시가 토큰당 숫자 하나가 아니라 헤드 수와 각 헤드 차원의 두 배만큼을 담는다는 점이다. 예컨대 차원 64짜리 헤드가 8개면 토큰 하나에 1,000개가 넘는 숫자가 쌓이고, 긴 문맥에서는 이 캐시가 메모리를 가득 채운다.

그래서 "키·값 헤드를 정말 몇 개나 유지해야 하는가"라는 질문이 어텐션 변형들을 가른다. 풀 멀티헤드 어텐션은 8개 헤드 각각의 키·값을 따로 두어 다양성을 최대로 확보하지만 캐시 비용도 가장 크다. 그룹 쿼리 어텐션(GQA)은 헤드들이 몇 개 그룹으로 키·값을 공유하고, 멀티쿼리 어텐션(MQA)은 극단적으로 키 하나·값 하나를 전부가 공유한다. 캐시는 줄지만 모든 헤드가 같은 키를 읽으면 각 토큰을 똑같이 보게 되어 헤드를 다르게 만든 이유가 사라진다.

MLA(멀티헤드 잠재 어텐션)는 이 맞바꿈을 거부한다. 헤드를 공유하는 대신 저장할 내용을 압축한다. 은닉 상태를 다운프로젝션으로 눌러 작은 잠재 벡터 하나로 만들고, 캐시에는 그 벡터만 담는다. 필요할 때 두 개의 업프로젝션이 이를 다시 펼치는데, 각 헤드가 업프로젝션의 자기 몫을 소유하므로 공유된 잠재 하나에서 서로 다른 8개의 키와 8개의 값이 나온다. 캐시 자체는 헤드를 전혀 모르고, 헤드별 정체성은 그때그때 다시 만들어진다.

매 단계마다 키·값을 복원하면 메모리 문제를 계산 문제로 바꾼 것 아니냐고 물을 수 있다. 여기서 핵심 트릭이 등장한다. 행렬 곱은 묶는 방식을 바꿀 수 있어서, 쿼리 가중치와 업프로젝션을 가운데의 고정된 한 행렬로 합칠 수 있다. 이 행렬은 변하지 않으므로 로드 시점에 한 번만 계산해 두면 된다. 그러면 쿼리는 곧바로 잠재 공간에 도착해 캐시된 벡터와 내적하고, 값 쪽 업프로젝션도 모든 토큰에 동일하므로 출력 행렬로 흡수된다. 복원 과정이 사실상 가중치 속으로 사라지는 것이다.

그 결과 MLA는 MQA의 캐시 크기로 풀 어텐션의 헤드 다양성을 얻는다. MQA는 문자 그대로 키가 하나뿐이라 모든 헤드가 같은 대상에 내적하지만, MLA는 공유된 잠재 벡터에 헤드마다 자기 쿼리로 접근하므로 3번 헤드와 5번 헤드는 여전히 각자의 키로 재는 셈이다. 다만 흡수는 가운데 행렬이 고정이어야 성립하는데, RoPE는 위치에 따라 달라지는 회전을 넣어 이를 깨뜨린다. 해법은 키를 잠재에서 복원해 흡수하는 내용 레인과, 회전을 담아 헤드 간 공유하며 따로 캐시하는 작은 위치 레인으로 나누는 것이다. 최종 점수는 내용 내적과 위치 내적의 합이 된다.

주요 인사이트

• 메모리 병목의 진짜 원인은 어텐션 연산이 아니라 토큰마다 누적되는 KV 캐시라는 점을 이해하면, 왜 어텐션 변형들이 "헤드를 몇 개나 저장하느냐"로 갈리는지 보인다.
• MQA/GQA의 캐시 절감은 헤드 다양성 상실이라는 숨은 비용을 치른다. 캐시를 줄이는 것과 표현력을 지키는 것은 서로 당기는 축이다.
• MLA의 통찰은 "다양성을 캐시에서 지우지 말고 쿼리·출력 가중치로 옮기라"는 것이다. 저장은 압축하되 정체성은 복원 시점에 재구성한다.
• 행렬 곱의 결합 법칙 덕분에 복원 비용이 사전 계산 가능한 고정 행렬로 흡수된다 — 메모리 절약이 곧바로 계산 폭증으로 이어지지 않는 이유다.
• RoPE처럼 위치 의존적인 연산은 사전 계산을 막으므로, 내용과 위치를 분리하는 설계가 실용적 타협점이 된다.

자주 묻는 질문

관련 AI 소식