XQuant란 무엇인가?

XQuant는 일반적인 KV 캐시를 저장하는 대신, 'X'라고 부르는 ‘레이어 입력 활성화(Layer Input Activation)만 양자화해서 저장합니다. 그리고 필요할 때 X를 가지고 Keys와 Values를 추론 중에 즉시 재구성(Recalculate)한다는 겁니다. X라고 부르는 레이어 입력 활성화가 바로 원래 K와 V를 가중치와 함께 계산하게 해 주는 요소거든요.

요리사가 요리하는 과정으로 비유를 해 볼께요:

트랜스포머 모델을 요리사가 요리를 만드는 주방이라고 생각해 봅시다. 각 레이어(층)는 요리사가 특정 요리 단계(예: 재료 손질, 양념 만들기, 굽기)를 담당하는 작업대라고 할게요.

• 레이어 입력 활성화(X): 요리사가 작업대에서 요리를 시작하기 전에 준비된 재료 세트(예: 썰린 채소, 고기, 양념 재료)입니다. 이 재료 세트는 이전 작업대(이전 레이어)에서 넘겨준 결과물이거나, 처음 입력(예: 사용자가 입력한 문장)에서 변환된 것이겠죠.

• Key(K)와 Value(V): 이 재료 세트(X)를 가지고 요리사가 특정 작업(예: 양념 버무리기)을 해서 만든 중간 요리(K와 V)입니다. 이 중간 요리는 나중에 최종 요리(모델 출력)를 만들 때 필요한 정보로 사용됩니다.

• 일반 KV 캐시 방식: 요리사가 중간 요리(K, V)를 전부 완성해서 냉장고에 저장해둡니다. 하지만 냉장고 공간이 작아서 금방 꽉 차버리죠(메모리 소모가 큼).

• XQuant 방식: XQuant는 냉장고에 중간 요리(K, V)를 저장하는 대신, 재료 세트(X)만 간단히 메모장에 적어둡니다(양자화로 압축). 나중에 필요할 때 이 메모를 보고 K와 V를 다시 만듭니다. 메모장은 냉장고보다 공간을 훨씬 덜 차지하죠. 대신 요리를 다시 만드는 데 약간의 시간(계산 비용)이 들지만, 공간을 엄청나게 절약할 수 있습니다.

이 과정이 실제로 어떻게 동작하는지 한 번 살펴보겠습니다:

XQuant는 어떻게 작동하나?

XQuant는 결국 '저장은 적게 하고, 계산은 더 많이 다시 하는' 방식으로 작동하는 셈이라고 말씀드렸죠. 그 과정은 다음과 같습니다:

1. 레이어 정규화(Layer Normalization)를 거친 표현인 입력 활성화 X를 저장합니다.

2. X를 양자화해서 더 적은 비트로 압축합니다.

3. Keys와 Values가 필요할 때, Key와 Value 투영 행렬(Wk, Wv)을 X와 곱해서 즉석에서 다시 계산합니다.

생각해 보면, 꽤 괜찮은 절충안이죠. X를 저장하는 데 표준 KV 캐싱에 비해서는 절반 정도의 메모리만 필요하다고 해요. 하지만 연구자들은 이 버전에서 멈추지 않고 더 흥미로운 것을 개발해 냈습니다.

XQuant-CL: 더 나은 버전

XQuant 개발자들은 여러 레이어에 걸쳐 관찰을 해 보면 ‘X가 서로 비슷하게 보인다’는 사실을 발견했습니다 - 각각의 레이어가 입력을 다음으로 전달하기 전에 수정을 아주 조금만 하기 때문에 발생하는 현상인데, '잔여 스트림(Residual Stream)'이라고 불립니다.

이런 단서를 가지고 연구자들은 이미 만들어 놓은 접근 방식을 좀 더 개선하기 위해 노력했고, 결과적으로 ‘각 레이어의 전체 X를 저장하는 대신, X들 간의 '차이(델타)'만을 저장’하는 XQuant-CL을 개발했습니다. 이 기법을 사용하면, 아래와 같이 아주 극단적인 압축을 할 수가 있습니다:

• 3비트 양자화 시 → 10배 메모리 절약, 정확도 손실은 거의 없음

• 2비트 양자화 시 → 12.5배 메모리 절약, 원래 모델의 정확도에 매우 근접함

좀 더 자세히 XQuant-CL의 단계별 작업 흐름을 살펴보면 이렇습니다.

• 첫 번째 레이어(X0​)는 베이스라인으로 삼기 때문에 더 높은 정밀도로 저장하구요.

• 다음 레이어들(X1​, X2​, ...)은 아래와 같이 처리됩니다.

  • 이전 레이어와의 작은 변화량(ΔX=Xi​−Xi−1​)을 계산합니다.

  • 이 변화량들은 값의 범위가 작기 때문에, 훨씬 더 공격적으로 양자화(압축)할 수 있습니다. 따라서 각 변화량(Δ)이 실제로 양자화되어 캐시됩니다.

• 재구성(Reconstruction): 특정 레이어의 X를 복구하려면, XQuant-CL은 단순히 변화량들을 모두 더하기만 하면 됩니다, 이렇게요: X^i​=X0​+ΔX^1​+ΔX^2​+⋯+ΔX^i​

이렇게 단계적으로 정제된 임베딩을 다시 만들어냅니다.

지속적인 재로딩을 피하고 재구성 비용을 낮추기 위해서, XQuant-CL은 과거의 모든 변화량과 X0​ 레이어의 누적 합을 항상 담고 있는 '누산기(Accumulator)'를 사용합니다. 그래서, 새로운 레이어를 위해서는 누산기 값에 최신의 변화량만 더하면 되죠.

XQuant-CL은 간단한 균일 양자화 기법을 사용함에도 불구하고, 더 복잡한 KV 압축 방법들보다도 더 나은 성능을 보여줍니다.

연구자들은 다양한 모델 유형에서도 최고로 높은 효율성을 달성하기 위해서 XQuant와 XQuant-CL의 능력을 더욱 확장했구요. 그 다음으로 중요한 목표로 삼은 건, GQA(그룹화된 쿼리 어텐션)을 지원하는 것이었다고 해요 - 이게 왜 그렇게 중요했을까요?

GQA(그룹화된 쿼리 어텐션)와 함께 작동하는 XQuant

XQuant와 GQA 기반 모델

아마 많이 들어보셨겠지만, LLM의 표준 어텐션 방식은 멀티 헤드 어텐션(MHA)입니다. 모든 쿼리 헤드가 자신만의 Keys와 Values를 가지죠. 반면에, Llama, Mistral, Gemma, Qwen 등 다양한, 효율적인 모델들은 여러 헤드가 동일한 Keys와 Values를 공유하는 GQA(그룹화된 쿼리 어텐션)을 사용하죠 - 이걸로만 해도 벌써 KV 캐시 메모리 크기가 꽤 줄어듭니다.

GQA는 AI 모델에 큰 이점을 제공하는데, 일반적인 XQuant 기법이 GQA에서 몇 가지 이슈가 있습니다. 입력 활성화인 'X'는 크지만, GQA의 Keys와 Values는 훨씬 작습니다. KV 대신 단순히 X나 XQuant-CL의 변화량(델타)을 캐시하는 게 오히려 일반적인 KV 캐싱보다 더 많은 메모리를 사용할 수가 있게 되는 겁니다. 이렇게 되면, 당연히 GQA 기반 모델에서 XQuant가 가지는 장점을 상쇄시키게 될 테고, 전혀 연구자들이 원하는 바가 아니었던 거죠.

이 문제를 해결하기 위해서, XQuant와 XQuant-CL은 Keys와 Values 투영 행렬(Wk​,Wv​)에 SVD(Singular Value Decomposition)를 오프라인으로 적용합니다. 그 원리는 이렇습니다.

• 첫 번째로, 추론 전(오프라인)에 Key와 Value의 가중치 행렬인 Wk​와 Wv​를 SVD로 분해해서 저차원의 부분 공간인 Uk​와 Uv​를 만듭니다.

• 추론 중에는 입력값 X를 이 부분 공간에 투영해서 차원을 낮추고, 양자화해서 캐시합니다.

• Keys와 Values가 필요할 때는 결합된 작은 행렬들을 곱해서 재구성합니다.

흥미로운 점은, 차원을 낮춘 Key 공간에서 모든 '이상점(Outliers)'이 첫 번째 채널에 모이게 되는 효과도 거둘 있다는 점이구요. 반면에, Values는 토큰별로 처리하는 것이 가장 효율적입니다.

이 방법을 통하면, 메모리 사용량은 이제 GQA의 KV 캐시 크기와 동일해집니다. 하지만 차원을 낮춘 X는 KV를 직접 양자화하는 것보다 더 쉬워서, 동일한 비트 폭에서 정확도가 더 높게 유지됩니다.

XQuant-CL과 GQA 기반 모델

• 이 변형 모델의 경우에는, Key와 Value 행렬을 $W_{kv} = [W_k | W_v]$라는 하나의 큰 행렬로 결합하고, 이걸 오프라인에서 SVD로 분해해서 공유되는 저차원 부분 공간을 만듭니다.

• 추론 중에는 XQuant 버전의 X와 비슷하게 ΔX를 이 부분 공간에 투영하고 양자화한 다음 캐시합니다.

• 한 레이어의 X를 재구성할 때, 캐시된 변화량을 차원 복원하고 누적 합산기에 더한 다음, 결합된 Key/Value 행렬을 곱해서 Keys와 Values를 복구합니다.

GQA 기반 모델을 위한 XQuant-CL은 XQuant 버전과 동일한 이점을 보여줍니다. 가장 중요한 SVD 단계는 오프라인에서 진행니까, 추가적인 런타임 비용이 들지는 않구요.

자, 지금까지 XQuant와 XQuant-CL의 내부 작동 방식에 대해서 간단히 살펴보았는데, 결국 어떤 이점을 얻을 수 있을까요?

성능의 향상

XQuant를 도입하는 순간부터 즉각적인 효과를 볼 수 있습니다. 단순히 X(레이어당 하나의 텐서)를 저장하는 것으로 전환하는 것만으로도, 메모리 사용량이 일반적인 KV 캐시에 비해 절반으로 줄어드니까요.

여기에 양자화를 더하면 더 큰 이점을 얻을 수 있습니다. X를 더 적은 비트로 압축해서 메모리 절약 효과는 약 7.7배까지 커지고, 반면에 모델의 품질은 거의 떨어지지 않습니다. FP16과 비교했을 때 퍼플렉서티(Perplexity) 변화는 0.1 미만에 불과한데, 즉 모델은 거의 원본만큼의 정확도를 유지한다는 뜻이죠.

여기서 더 나아가서, XQuant-CL 버전은 여러 레이어에서 X의 값이 서로 비슷하다는 점을 활용, 더 공격적인 압축 전략으로 정확도 손실은 거의 없이 최대 12.5배의 메모리 절약 효과를 달성할 수 있습니다.

GQA 모델(Llama-3.1-8B, Mistral-7B)의 경우에, XQuant의 결과는 다음과 같습니다.

• 4비트: 3.7배 메모리 절약과 함께 퍼플렉서티(Perplexity)는 0.01만 하락했습니다.

• 3비트: 5배 메모리 절약과 함께 퍼플렉서티는 0.1 미만으로 하락했습니다.

• 2비트: XQuant는 여전히 사용 가능하고, 7.1배 더 많은 메모리 절약 효과와 함께 KIVI*보다 퍼플렉서티에서 0.57 낮은 우수한 성능을 보였습니다.

XQuant, XQuant-CL의 장점

XQuant 및 XQuant-CL 기법의 장점들을 종합하면 다음과 같습니다.

다양한 데이터셋에서 XQuant와 XQuant-CL은 정말 미미한 정확도의 희생으로 엄청난 메모리 절약(최대 12.5배) 효과를 꾸준히 보여줍니다.

약간의 추가 연산이 필요하지만 GPU가 충분히 감당할 수 있는 수준이고, 전체적으로는 더 빠르고 효율적인 추론 결과를 만들어냅니다. 어찌보면, 아주 실리적인 접근으로 메모리 병목 현상을 극복하게 해 주는 기법입니다:

• 아주 낮은 정밀도에서도 FP16에 준하는 정확도를 유지합니다.

• XQuant-CL은 기본 XQuant보다 더 효과적입니다. 더 많은 메모리를 절약할 뿐만 아니라, 누산기(Accumulator) 트릭을 사용해서 어떤 레이어의 입력이라도 빠르고 저렴하게 재구성할 수 있습니다.

• SVD 기술 덕분에 두 방법 모두 GQA와 함께 사용할 수 있어서, 다양한 모델에 적용할 수 있습니다.

물론, 이 모든 장점에도 불구하고, 완벽한 방법이란 건 없죠.

XQuant의 한계점

이 방법들은, 어텐션을 계산할 때마다 Keys와 Values를 즉석에서 재구성하기 위해서 추가적인 연산을 더해야 한다고 말씀드렸죠. 이로 인해서 일부 하드웨어나 플랫폼에서는 지연 시간(Latency)이 늘어날 수 있기는 한데, 이걸 단점이라고 말씀드려야 할지는 모르겠어요 - 이 기법의 기본적인 전제가 계산을 더 많이 하고, 메모리는 절약한다는 것이기는 하니까요.

어쨌뜬, 이 말은, XQuant의 효율성이 일반적으로 하드웨어에 따라 달라질 수 있다는 의미이기는 합니다.

XQuant-CL 변형 모델의 경우, 누산기(Accumulator)를 관리해야 하니까 약간의 추가 부담이 있습니다. 그리고, 일반적인 XQuant보다는 구현이나 유지보수가 좀 더 복잡합니다.

GQA 모델의 경우, 두 가지 변형 모두 오프라인 SVD와 다운 투영(Down-Projection) 과정을 필요로 합니다. 그렇지 않으면 메모리 절약이 아니라 거꾸로 메모리 사용량이 늘어날 위험이 있습니다. 재구성이 근사치를 기반으로 하기 때문에 양자화 오류가 누적될 수도 있구요.

그럼에도 불구하고, 메모리 제약이 심하고 정확도가 중요한 시나리오에서 XQuant는 아주 훌륭한, 균형잡힌 해결책이 될 수 있습니다 - 원본에 가까운 정확도를 유지하면서도 엄청난 메모리 절약 효과를 가져다주니까요.

맺으며

오늘 살펴본 UC 버클리, FuriosaAI, ICSI, LBNL 연구진의 이 훌륭한 연구. 메모리 한계를 다루는 새로운, 프레쉬한 관점을 제시해 준다고 할 수 있을 것 같습니다. 기존의 직접적인 KV 캐시 압축 방식을 벗어나 새로운 생각을 해 본 거라고 생각합니다.

• 더 작고 다루기 쉬운 X(입력 활성화)를 압축합니다.

• X의 레이어 간 유사성을 이용해서 메모리 사용량을 더욱 공격적으로 줄입니다.

오래되고 익숙한 KV 캐싱 같은 접근법이 구식이라고 말할 수도 있지만, 여전히 현재와 미래의 혁신을 위한 탄탄한 출발점이자 베이스라인이기는 합니다. XQuant와 XQuant-CL은 여기서 한 걸음 더 나아가서, 단점도 더 적고 메모리는 더 많이 절약하게 해 주는 새로운 시각을 열어줍니다.

물론, 어떤 방법을 사용할지는 사용자의 특정한 필요와 상황에 따라 다르겠지만, 높은 정확도를 추구하고 대규모 메모리 절약이 최우선 목표일 때, 그럴 때 XQuant는 아주 좋은 해답이 될 수 있을 것 같습니다.

보너스: 참고자료

• XQuant: Breaking the Memory Wall for LLM Inference with KV Cache Rematerialization

• Mastering LLM Techniques: Inference Optimization (NVIDIA blog)

• Key, Value, Compress: A Systematic Exploration of KV Cache Compression Techniques

• Reducing Transformer Key-Value Cache Size with Cross-Layer Attention

• Layer-Condensed KV Cache for Efficient Inference of Large Language Models