크로네커 곱 어텐션(KPA): 키와 값을 컴팩트한 인자의 크로네커 곱의 합으로 나타내 KV 캐시를 최대 10배 압축하고 별도 회전 분기도 필요 없다
이 논문은 키와 값을 토큰에 의존하는 컴팩트한 행렬 인자의 크로네커 곱의 합으로 쓰는 KPA를 제안한다. 텐서 곱 어텐션(TPA)을 엄밀히 일반화하며 TPA의 랭크 1 벡터 외적이 퇴화한 경우다. 이 형태에 맞춘 구조화 회전 연산자가 컴팩트한 인자만 돌리므로 MLA처럼 분리된 RoPE 분기가 필요 없다. 약 1억 2,500만~15억 파라미터의 네 규모와 아홉 벤치마크에서 KPA는 TPA, MHA, MQA, GQA, MLA와 같거나 낫고 KV 압축은 6.4~10배다. 모델은 모두 저자가 처음부터 훈련한 소형이다.
추론 단계에서 대형 모델의 병목으로 점점 두드러지는 것이 KV 캐시입니다. 중복 계산을 피하려고 모델은 모든 층에서 모든 토큰의 키와 값을 저장하는데, 이 캐시는 시퀀스 길이, 배치 크기, 깊이에 비례해 늘어납니다. 컨텍스트가 수십만 토큰에 이르고 처리량을 높이려 배치가 커지면 KV 캐시가 GPU 메모리의 대부분을 차지해 가중치와 활성화에 쓸 여유가 적어지고, 감당할 수 있는 컨텍스트 길이, 배치, 배포 규모가 제한됩니다. 기존 대책은 저마다 무언가를 맞바꿉니다. 멀티쿼리 어텐션(MQA)과 그룹 쿼리 어텐션(GQA)은 여러 헤드가 키와 값을 공유하는 대신 표현의 다양성을 잃습니다. 멀티헤드 잠재 어텐션(MLA)은 압축된 잠재 상태를 캐시하지만 분리된 위치 인코딩 장치가 따로 필요합니다. 텐서 곱 어텐션(TPA)은 저랭크 벡터 외적으로 분해하는데 지나치게 엄격한 저랭크 구조에 묶일 수 있습니다. 논문 2610.12657이 제안하는 크로네커 곱 어텐션(KPA)은 이런 배경에서 나왔습니다.
KPA의 핵심 생각은 전체 키와 값 행렬을 캐시하지 않고, 각 토큰의 키와 값을 그 토큰에 의존하는 컴팩트한 행렬 인자들의 크로네커 곱의 합으로 쓰는 것입니다. 크로네커 곱은 작은 행렬 둘로 큰 행렬을 만들므로 헤드 수준과 헤드 안의 특징 수준 모두에 구조를 둘 수 있습니다. TPA의 벡터 외적은 둘 중 하나만 다룹니다. 논문은 경험적 연구가 어텐션 헤드가 구문, 의미, 위치 같은 서로 다른 역할로 분화함을 보였기 때문에 좋은 KV 분해는 두 수준에서 동시에 작동해야 한다고 말합니다. TPA의 랭크 1 벡터 외적 분해가 KPA의 “퇴화한 크로네커 분할”로 나타남을 보여 KPA가 TPA를 엄밀히 일반화하며, 퇴화하지 않는 행렬 인자는 더 풍부한 블록 구조를 표현할 수 있다고 합니다. 디코딩 때는 인자만 캐시하며 압축률은 헤드 차원을 어떻게 나누느냐에 달려 있고 인자가 비대칭일수록 압축이 강해집니다. 이 “크로네커 헤드 분할”은 설계 때 조절하는 하이퍼파라미터로 다른 구조 변경 없이 압축과 정확도의 균형 위에서 움직일 수 있게 해 줍니다.
위치 인코딩이 또 하나의 핵심입니다. RoPE는 전체 키에 작용하지만 KPA는 인자만 캐시하므로 회전을 어디에 둘지 답해야 합니다. MLA의 답은 별도의 분리된 RoPE 분기입니다. 논문은 크로네커 표현에 맞는 구조화 회전 연산자를 제시합니다. RoPE의 상대 위치 성질을 유지하면서 컴팩트한 인자만 돌리므로 분리된 분기가 더는 필요하지 않습니다. 논문은 KPA가 그런 분기를 요구하지 않으면서도 MLA를 앞선다고 말합니다.
실험에서는 저자들이 약 1억 2,500만, 3억 5,300만, 7억 7,300만, 15억 파라미터의 네 규모에서 SwiGLU와 RoPE를 쓰는 nanoGPT 기반 Llama 풍 모델을 FineWeb-Edu 약 500억 토큰으로 훈련했습니다. 키와 값의 랭크는 RK=RV=2, 쿼리의 랭크는 RQ=6입니다. lm-evaluation-harness로 ARC-Easy, ARC-Challenge, BoolQ, HellaSwag, OpenbookQA, PIQA, WinoGrande, MMLU, SciQ의 아홉 제로샷 벤치마크를 평가했습니다. 논문에 따르면 같은 훈련 조건에서 KPA는 TPA, MHA, MQA, GQA, MLA와 같거나 낫고 KV 캐시를 최대 10배 압축합니다. 한 규모에서 KPA는 TPA와 같은 6.4배 압축비로 가장 좋은 평균을 얻어 정확도가 1.49포인트 높았고 가장 큰 상승은 ARC-Easy입니다. 다른 구성은 8배 압축에 이릅니다. 약 6.7배 압축에서는 다른 기준을 앞서며 ARC-Easy, BoolQ, SciQ에서 차이가 뚜렷합니다. 15억에서는 아홉 중 여덟에서 앞서고, 1억 2,500만에서는 모든 KPA 구성이 TPA의 46.21%를 넘습니다. 헤드 수준 인자를 더 크게 한 비대칭 분할은 대칭 분할과 TPA를 모두 앞섭니다.
이 논문을 어떻게 볼까요? 장점은 생각이 분명하고 엄밀한 일반화를 증명하며, 압축과 회전 위치 인코딩이 공존하는 실제 문제에 답하고, 압축비를 조절할 수 있다는 점입니다. 유보할 점도 분명합니다. 첫째, 모델은 모두 저자가 처음부터 훈련한 작은 것으로 최대 15억 파라미터, 약 500억 토큰이며 기존 대형 모델에 적용한 것이 아닙니다. 둘째, 평가는 제로샷 정확도가 중심입니다. 논문 스스로 맺음말에서 효율적인 분해 디코딩 커널을 만드는 것과 더 긴 컨텍스트에서 KPA를 평가하는 것이 자연스러운 다음 단계라고 하므로, 디코딩 속도, 지연, 긴 컨텍스트 동작은 아직 데이터가 없다는 뜻입니다. 셋째, 결과는 자체 보고이며 독립 재현은 보지 못했습니다. KV 캐시 비용이 걱정이라면 눈여겨볼 설계로 보고, 커널과 긴 컨텍스트 결과가 나온 뒤에 GQA나 MLA를 대체할 수 있는지 판단하세요.
Sources
FAQ
KPA는 KV 캐시를 어떻게 압축하나요?
일반 멀티헤드 어텐션은 토큰마다 전체 키와 값을 캐시합니다. KPA는 대신 각 토큰의 키와 값을 그 토큰의 은닉 상태에서 계산한 컴팩트한 행렬 인자들의 크로네커 곱의 합으로 쓰고, 디코딩 때는 그 인자만 캐시합니다. 크로네커 곱은 헤드 수준과 헤드 안의 특징 수준 모두에 구조를 주므로 TPA의 벡터 외적보다 풍부한 블록 구조를 표현할 수 있습니다. 논문은 TPA의 랭크 1 외적이 KPA의 퇴화한 경우임을 보여 KPA가 엄밀히 더 일반적이라고 합니다.
위치 인코딩은 어떻게 다루나요?
압축된 잠재 형태에 회전 위치 임베딩을 쓰려면 MLA는 별도의 분리된 RoPE 분기가 필요합니다. KPA는 크로네커 표현에 맞는 구조화 회전 연산자를 제시합니다. RoPE의 상대 위치 성질을 유지하면서 컴팩트한 인자만 돌리므로 분리된 분기가 필요 없습니다. 논문은 이 점에서 KPA가 그 분기 없이도 MLA를 앞선다고 말합니다.
결과는 어떻고 유보할 점은 무엇인가요?
저자들은 약 1억 2,500만, 3억 5,300만, 7억 7,300만, 15억 파라미터의 Llama 풍 모델을 FineWeb-Edu 약 500억 토큰으로 처음부터 훈련해 ARC, BoolQ, HellaSwag, PIQA, MMLU 등 아홉 제로샷 벤치마크에서 비교했습니다. 같은 훈련 조건에서 KPA가 기준과 같거나 낫고, 6.4배 압축에서 TPA보다 평균 정확도가 1.49포인트 높으며 최대 10배 압축이고, 15억에서는 아홉 중 여덟에서 앞선다고 합니다. 유보점은 모든 모델이 작고 처음부터 훈련한 것이며 논문 스스로 효율적인 분해 디코딩 커널과 더 긴 컨텍스트 평가를 다음 단계로 꼽는다는 점이어서 실제 디코딩 속도와 긴 컨텍스트 동작은 아직 보이지 않았다는 것입니다.