🤖 machine learning

QK-Normed MLA: QK normalization without full key caching

본 논문은 QK 정규화가 전체 키 캐싱을 피하도록 연산을 수학적으로 재구성함으로써 Multi-head Latent Attention(MLA)과 원활하게 통합될 수 있음을 입증하며, 이를 통해 무시할 만한 수준의 지연 시간 오버헤드로 훈련 안정성과 다운스트림 정확도를 향상시킨다.

원저자: Yizhou Han, Yao Zhao, Jun Zhou, Longfei Li, Ruoyu Sun

게시일 2026-06-16

📖 3 분 읽기☕ 가벼운 읽기

CC BY 4.0

원저자: Yizhou Han, Yao Zhao, Jun Zhou, Longfei Li, Ruoyu Sun

원본 논문은 CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/) 라이선스로 제공됩니다. ✨ 이것은 아래 논문에 대한 AI 생성 설명입니다. 저자가 작성하거나 승인한 것이 아닙니다. 기술적 정확성을 위해서는 원본 논문을 참조하세요. 전체 면책 조항 읽기

당신이 거대한 도서관을 운영하고 있다고 상상해 보세요. 여기서 사서(AI)는 질문에 답하기 위해 완벽한 책(정보)을 찾아야 합니다. 이를 효율적으로 수행하기 위해, 사서는 두 가지 주요 기술을 사용합니다.

"잠재 요약" 기술 (MLA): 세상의 모든 책을 방 안에 거대하고 무거운 형태로 보관하는 대신(공간을 너무 많이 차지하니까요), 사서는 각 책에 대한 아주 작고 압축된 "요약 카드"를 보관합니다. 질문이 들어오면, 사서는 답을 찾기 위해 그 요약 카드를 즉석에서 확장합니다. 이는 엄청난 양의 공간을 절약해 줍니다.
"볼륨 조절" 기술 (QK Normalization): 때때로 사서가 너무 흥분할 때가 있습니다. 질문과 책 사이의 연결 "볼륨"이 너무 커져서 다른 모든 것을 집어삼키고, 이로 인해 사서가 당황하거나 실수를 하게 됩니다. 이를 해결하기 위해, 우리는 보통 소리 크기를 일정하게 유지하도록 책에 "볼륨 조절 노브"를 달아둡니다.

문제점:
오랫동안 이 두 가지 기술은 서로 잘 맞지 않았습니다. "볼륨 조절" 기술은 볼륨을 조절하기 위해 사서가 반드시 책의 전체적이고 무거운 복사본을 손에 쥐고 있어야 하는 것처럼 보였습니다. 하지만 "잠재 요약" 기술은 바로 그 무거운 복사본들을 피하기 위해 설계된 것이었습니다. 즉, 공간을 아낄 것인지(잠재 요약), 아니면 볼륨을 일정하게 유지할 것인지(볼륨 조절) 둘 중 하나를 선택해야만 하는 것처럼 보였습니다.

해결책 (QK-Normed MLA):
저자들은 두 마리 토끼를 모두 잡을 수 있는 영리한 수학적 해킹 방법을 발견했습니다. 그들은 "볼륨 조절" 노브가 사실 두 부분으로 구성되어 있다는 것을 깨달았습니다:

정적 설정 (변하지 않는 고정된 다이얼).
동적 읽기 (현재 책의 소리가 얼마나 큰지 빠르게 훑어보는 것).

그들은 다음을 입증했습니다:

정적 설정은 사서가 요약 카드를 확인하기도 전에 미리 옮겨질 수 있습니다. 이는 질문 자체에 미리 녹여낼 수 있습니다.
동적 읽기는 매우 단순해서, 전체 책을 필요로 하는 대신 사서는 각 요약 카드에 대해 단 하나의 작은 숫자(스칼라 값 하나)만을 기록하면 됩니다.

결과:
이제 사서는 (엄청난 공간을 절약하는) 작은 "요약 카드"를 계속 사용하면서도, 동시에 완벽한 볼륨 조절을 할 수 있습니다. 더 이상 무거운 책들을 들고 다닐 필요가 없습니다.

실험이 보여준 것:
연구팀은 400만 개의 파라미터를 가진 모델(중형 규모의 AI)을 사용하여 방대한 양의 텍스트(1,000억 단어)로 학습시킨 후 이 새로운 방법을 테스트했습니다.

더 나은 학습: 이 새로운 "볼륨 조절"을 사용한 모델은 "클리핑(clipping)"이라 불리는 다른 투박한 방식(소리가 너무 커지면 그냥 "멈춰!"라고 외치는 것과 같은 방식)을 사용한 모델보다 더 빠르게 학습하고 실수를 덜 했습니다.
더 나은 답변: 다양한 과업을 테스트했을 때, 이 새로운 방식은 더 정확한 답변을 제공했습니다.
속도: 유일한 단점은 그 하나의 작은 숫자를 확인하는 과정이었는데, 이는 긴 이야기를 읽는 데 걸리는 시간에 2% 미만의 시간만을 추가했습니다. 이는 너무 작아서 거의 느껴지지 않는 수준입니다.

요약하자면: 이 논문은 메모리 효율성을 위해 AI의 안정성을 희생할 필요가 없다는 것을 증명합니다. 당신은 원래 피하려고 했던 무거운 전체 데이터를 저장하지 않고도, "볼륨"을 제어할 수 있는 작고 영리한 수학적 기술을 사용할 수 있습니다.

기술 요약: QK-Normed MLA

문제 정의
현대적인 대규모 언어 모델은 두 가지 상충하는 압력에 직면해 있습니다: 긴 문맥 시나리오에서의 메모리 효율적인 디코딩 필요성과 학습 안정성의 필요성입니다.

메모리 효율성: 표준 멀티 헤드 어텐션(MHA)은 모든 토큰에 대해 전체 키(key) 및 값(value) 벡터를 캐싱하므로 높은 메모리 및 대역폭 비용을 발생시킵니다. 멀티 헤드 잠재 어텐션(MLA)은 전체 키 대신 저차원 잠재 상태( $c_t$ )를 캐싱하여 이를 해결하며, 디코딩 중에 업-프로젝션( $W^K_g$ )을 통해 전체 키 효과를 복구합니다.
학습 안정성: 쿼리(query)와 키(key)의 통제되지 않은 스케일은 어텐션 포화 및 학습 불안정성을 초축할 수 있습니다. 쿼리-키(QK) 정규화(특히 RMSNorm)는 도트 곱(dot product) 전 벡터를 정규화함으로써 학습을 안정화합니다. 그러나 표준적인 포스트-프로젝션(post-projection) QK RMSNorm은 MLA와 호환되지 않는 것처럼 보이는데, 이는 모든 캐싱된 토큰에 대해 정규화 통계량을 계산하기 위해 완전히 투영된 키( $k_{t,g} = c_t W^K_g$ )가 필요해 보이기 때문입니다. 단순한 구현 방식은 전체 키를 실체화(materialize)해야 하므로 MLA의 메모리 이점을 상쇄하게 됩니다.

방법론
본 논문은 QK-Normed MLA를 제안하며, 이는 포스트-업-프로젝션(post-up-projection) QK RMSNorm을 전체 키를 실체화하지 않고도 MLA의 잠재 디코드 경로에 통합하는 방식입니다. 핵심 통찰은 RMSNorm의 대수적 분해에 있습니다:
$\text{RMSNorm}(z) = z \odot \frac{\gamma}{S(z)}$
여기서 $\gamma$ 는 정적 아핀 가중치(static affine weight)이고 $S(z)$ 는 동적 스칼라 RMS 통계량입니다.

저자들은 이러한 구성 요소들을 분리하는 "흡수된(absorbed)" 공식을 유도합니다:

정적 가중치 흡수(Static Weight Absorption): 정적 키 측 가중치 $\gamma_k$ 는 쿼리 측 프로젝션 행렬에 흡수됩니다. 쿼리 $q$ 를 투영한 후 키를 정규화하는 대신, 쿼리를 잠재 도트 곱 전에 $\gamma_k$ (및 정적 가중치 $\gamma_q$ )로 사전 곱합니다.
동적 스칼라 캐싱(Dynamic Scalar Caching): 동적 통계량 $S(z)$ 는 토큰당 및 KV 그룹당 하나의 역-RMS 스칼라( $\alpha_{t,g} = 1/S_{t,g}$ )로 축소됩니다.
캐시 레이아웃(Cache Layout): 디코드 캐시는 원래의 저차원 잠재 상태 $c_t$ 를 유지하면서, 역-RMS 값을 포함하는 작은 스칼라 캐시 $A^k_{cache}$ 를 추가합니다. 이는 정규화된 전체 키를 저장하지 않습니다.
블록 단위 정규화(Blockwise Normalization): 효율성을 유지하기 위해, "콘텐츠(content)"와 "RoPE" 구성 요소를 별도로 정규화합니다. 콘텐츠 경로는 흡수된 잠재 공식을 사용하고, RoPE 경로는 표준 실체화된 벡터를 사용합니다. 이는 융합된 커널(fused kernel) 구현을 복잡하게 만들 수 있는 동기화 지점을 방지합니다.

디코딩 중에 어텐션 로짓(logits)은 다음과 같이 계산됩니다:
$\ell_{i,t,h} = (\tilde{q}_{i,h} \cdot c_t) \cdot \alpha^q_{i,h} \cdot \alpha^k_{t,g(h)}$
여기서 $\tilde{q}$ 는 흡수된 가중치를 포함하며, 스칼라 인자들은 소프트맥스 전의 행-열 스케일링 단계로 적용됩니다.

주요 기여

정확한 공식화: 본 논문은 MLA의 잠재 디코드 경로와 호환되는 포스트-업-프로젝션 QK RMSNorm의 대수적으로 정확한 공식을 유도합니다.
효율적인 캐시 설계: 잠재 KV 캐시를 지배적인 상태로 유지하면서, 토큰 및 KV 그룹당 하나의 스칼라 키 RMS 통계량만을 추가하는 캐시 레이아웃을 제안합니다(테스트된 구성에서 캐시 크기를 약 1.56% 증가시킴).
경험적 검증: 저자들은 이 접근 방식이 QK 클리핑(QK Clipping)보다 낮은 학습 손실을 달면서도 다운스트림 정확도를 개선하고, 지연 시간 오버헤드는 미미하다는 것을 입증했습니다.

결과
H800 GPU에서 최대 100B 토큰 동안 학습된 400M 파라미터 MLA 모델을 대상으로 실험을 수행했습니다:

학습 안정성: QK-Normed MLA는 100B 토큰 실행 전반에 걸쳐 QK-Clip보다 일관되게 낮은 학습 손실을 달s했습니다. 높은 학습률 스트레스 테스트( $2 \times 10^{-2}$ )에서 QK-Clip은 그래디언트 NaN과 함께 발산한 반면, QK-Normed MLA는 안정성을 유지했습니다.
다운스트림 성능: 다양한 다운스트림 태스크 세트에서 QK-Normed MLA는 8개 중 7개 태스크에서 정확도를 향상시켰으며, 평균 점수를 44.75에서 46.33으로 높였고, LAMBADA 퍼플렉시티(perplexity)를 16.28에서 14.18로 낮추었습니다.
디코드 효율성: 4k에서 256k 토큰 범위의 벤치마크에서 이 방법은 2% 미만의 지연 시간 오버헤드(평균 1.40%)를 추가했습니다. 스칼라 캐시는 레이어당 총 캐시 크기를 약 2.8%만 증가시켰습니다.

의의 및 주장
본 논문은 QK 정규화와 MLA 사이의 불일치처럼 보이는 것이 구조적 제약이 아니라 구현상의 산물이라고 주장합니다. 정적 가중치를 흡수하고 동적 통계량을 스칼라로 축소할 수 있음을 보여줌으로써, 저자들은 잠재 디코딩의 메모리 효율성을 희생하지 않고도 QK 정규화를 MLA 모델을 위한 실용적인 안정화 메커니즘으로 사용할 수 있게 했습니다.

저자들은 다음과 같은 범위 내에서 주장을 신중하게 제시합니다:

결과는 400M 파라미터 모델을 기반으로 하며, 이것이 십억 단위 파라미터 모델에서도 프런티어 규모의 동작을 확립한다고 주장하지는 않지만, 대수적 등가성은 규모와 관계없이 성립합니다.
로짓 스케일 분석은 완전한 최적화 이론이라기보다 메커니즘 수준의 진단 도구로 제시되었습니다.
실제 속도는 불필요한 메모리 트래픽을 피하기 위해 스칼라 인자를 어텐션 경로에 융합(fusion)하는 커널 통합에 달려 있습니다.

유사한 논문