Model Tells You What to Discard Adaptive KV Cache Compression for LLMs

논문 : https://arxiv.org/abs/2310.01801

아래글은 Scholar GPT와 대화를 나눈 내용입니다.

Q : 논문의 핵심 내용과 강점, 알고리즘 설명, 그리고 한계점

A :

논문은 FastGen이라는 적응형 KV 캐시 압축 기법을 제안합니다. 이 방법은 대규모 언어 모델(LLM)의 생성 추론에서 메모리 사용량을 줄이고, 추가적인 모델 재훈련 없이도 고성능을 유지할 수 있도록 설계되었습니다.

논문의 강점 및 독창성

  1. 적응형 KV 캐시 압축: FastGen은 LLM의 어텐션 헤드별로 서로 다른 구조를 활용하여, 메모리를 최적화하는 다양한 압축 정책을 적용합니다. 이는 메모리 사용량을 줄이면서도 중요한 정보를 유지하는 방식으로, 각 어텐션 헤드의 특성을 반영해 최적의 압축 방식을 선택하는 독창적인 접근법입니다.

  2. 재훈련 불필요: FastGen은 가벼운 프로파일링을 통해 모델을 다시 훈련하지 않고도 사용할 수 있습니다. 이는 모델을 손쉽게 적용할 수 있도록 하여, LLM을 사용하는 연구실이나 개발자들에게 큰 장점이 됩니다.

  3. 메모리 절약과 성능 유지: FastGen은 실험에서 GPU 메모리 사용량을 크게 줄이면서도 생성 품질을 거의 유지함을 보여주었습니다. 다양한 벤치마크에서 최대 40% 이상의 메모리 감소와 미미한 성능 저하로 LLM을 효율적으로 사용할 수 있음을 입증했습니다.

핵심 알고리즘 및 예시

FastGen의 핵심은 Prompt Encoding 단계와 Token Generation 단계로 이루어집니다.

  1. Prompt Encoding 단계: 모델의 프로파일링을 수행하여 각 어텐션 헤드의 구조를 분석합니다. 그런 다음, 각 헤드에 맞는 압축 정책을 선택하여 KV 캐시를 구성합니다. 이를 통해 메모리 사용량을 최소화하고, 필요한 정보만 캐시에 유지합니다.

  2. Token Generation 단계: 선택된 압축 정책에 따라 새로 생성된 토큰의 KV 벡터를 캐시에 추가합니다. 캐시 용량이 가득 차면, 기존 정책에 따라 오래된 정보를 삭제하고, 중요한 정보를 유지합니다. 예를 들어, 특수 토큰, 문장부호, 그리고 높은 빈도수의 토큰은 유지하며, 거리 기준을 초과하는 오래된 정보는 삭제하는 식입니다.

예시

  • 주어진 텍스트가 “고양이가 집 안을 걸어다니고 있다.”고 할 때, 첫 번째 몇 토큰만 초기 프롬프트로 유지하고, 이후 “걸어다니고 있다”와 같은 최신 정보만 캐시에 유지합니다. 각 어텐션 헤드의 구조에 따라 ‘고양이’와 같은 특수 토큰을 유지하는 등의 정책을 선택할 수 있습니다.

논문의 한계점

  1. 정밀한 문맥 처리의 한계: FastGen은 메모리 효율성을 위해 일부 문맥 정보를 삭제하기 때문에, 장기 문맥 유지가 필요한 경우 한계가 있을 수 있습니다. 따라서 중요한 정보를 전부 보관해야 하는 작업에는 적합하지 않을 수 있습니다.

  2. 정책 적용의 복잡성: 각 어텐션 헤드에 대해 별도의 압축 정책을 적용하기 때문에, 비정형적인 어텐션 구조를 갖는 LLM에 적용하기가 까다로울 수 있습니다.

이 논문은 FastGen을 통해 LLM의 메모리 사용 문제를 효과적으로 해결하며, 추가적인 학습 없이도 간단히 적용할 수 있는 방안을 제안한 점에서 중요한 기여를 하고 있습니다.

Q : adaptive KV Cache Compression에 대해서 자세하게 설명해 주겠니?

A :

Adaptive KV Cache Compression에 대해 자세히 설명드리겠습니다. 이 기법은 대규모 언어 모델(LLM)의 메모리 사용량을 줄이기 위해 KV 캐시를 효율적으로 관리하는 방법으로, 각 어텐션 헤드의 특성을 반영하여 다양한 압축 정책을 적용합니다.

KV 캐시의 역할

Transformer 모델은 각 토큰의 키(Key)와 값(Value) 벡터를 어텐션 계산에 사용합니다. 일반적으로, 생성 모델이 새로운 토큰을 생성할 때마다 모든 이전 토큰의 KV 벡터를 캐시에 저장해 두고 재사용하는데, 이를 통해 이전 토큰의 계산을 반복하지 않고도 최신 토큰을 생성할 수 있습니다. 그러나 시퀀스 길이가 길어질수록 이 캐시의 메모리 사용량이 급격히 증가하여, 큰 모델에서는 메모리 사용 효율이 중요한 문제가 됩니다.

Adaptive KV Cache Compression의 개념

Adaptive KV Cache Compression은 KV 캐시의 압축을 어텐션 모듈의 특성에 맞게 유연하게 적용하는 기법입니다. FastGen은 LLM의 각 어텐션 헤드가 참조하는 정보의 구조가 다르다는 점에 착안하여, 어텐션 헤드별로 다른 압축 방식을 사용해 메모리 효율성을 최적화합니다.

Adaptive KV Cache Compression의 단계

Adaptive KV Cache Compression은 크게 모델 프로파일링 단계와 적응형 캐시 관리 단계로 이루어집니다.

1. 모델 프로파일링 단계

FastGen은 먼저 Prompt Encoding 단계에서 모델을 프로파일링하여 각 어텐션 헤드가 주로 참조하는 정보 구조를 파악합니다. 이를 통해 각 헤드에 맞는 최적의 압축 정책을 선택할 수 있습니다. 예를 들어, 어떤 헤드는 문장의 특수 토큰(예: 시작/종료 토큰)을 자주 참조하고, 또 다른 헤드는 로컬 문맥에만 집중하는 구조를 가질 수 있습니다.

2. 적응형 캐시 관리 단계

이 단계에서는 앞서 선택한 압축 정책을 바탕으로 KV 캐시를 관리합니다. 새로운 토큰이 생성될 때마다 FastGen은 각 어텐션 헤드별로 최적화된 정책에 따라 캐시에 저장하거나 삭제할 KV 벡터를 결정합니다. 이로 인해 불필요한 정보는 제거되고 중요한 정보만 유지되어 메모리 사용량을 절감하게 됩니다.

압축 정책 유형

FastGen은 어텐션 헤드의 특성에 따라 다양한 압축 정책을 제공합니다. 각 정책은 특정한 유형의 정보를 유지하거나 삭제하도록 설계되었습니다:

  1. Special Tokens (특수 토큰): 특수 토큰(예: 문장의 시작/종료, 명령어 토큰)만 캐시에 유지합니다. 이는 주로 특수 토큰에 집중하는 어텐션 헤드에 적용됩니다.

  2. Punctuation (문장부호): 마침표, 쉼표, 물음표 등과 같은 문장부호만 캐시에 유지합니다. 문장의 구조에 중요한 정보를 제공하는 토큰을 유지하는 방식입니다.

  3. Locality (로컬 문맥): 로컬 문맥에 집중하는 어텐션 헤드에는 거리가 먼 토큰의 정보를 제거하고, 일정 거리 내의 토큰만 캐시에 유지합니다. 로컬 컨텍스트를 중요시하는 어텐션에 유리합니다.

  4. Frequency (빈도 기반): 자주 참조되는 높은 빈도의 토큰만 남기고 나머지 토큰은 제거합니다. 이전 어텐션 점수가 높은 토큰만을 선택하는 방식입니다.

이 외에도 하이브리드 정책을 통해 여러 압축 정책을 조합하여, 다양한 특성을 가진 어텐션 헤드를 더 효율적으로 지원할 수 있습니다.

Adaptive KV Cache Compression의 예시

예를 들어, 다음과 같은 텍스트를 처리한다고 가정해 보겠습니다: “고양이는 나무 위에 올라갔다. 고양이는 다시 내려왔다.”

  1. 모델 프로파일링을 통해 어떤 어텐션 헤드는 ‘고양이’와 같은 특수 토큰에 집중하고, 다른 헤드는 문장부호에 집중하며, 또 다른 헤드는 로컬 문맥에 집중한다는 사실을 파악합니다.
  2. 이후 새로운 토큰을 생성할 때, 각 어텐션 헤드에 적합한 압축 정책을 적용합니다. 특수 토큰에 집중하는 헤드에서는 ‘고양이’를 유지하고, 로컬 문맥에 집중하는 헤드는 최근에 생성된 ‘다시 내려왔다’ 부분만 유지하며, 불필요한 초기 문맥은 삭제합니다.

Adaptive KV Cache Compression의 장점

  • 메모리 효율성: 불필요한 KV 벡터를 제거하여 메모리 사용량을 줄이고, 더 큰 모델을 효율적으로 사용할 수 있습니다.
  • 고성능 유지: 각 어텐션 헤드에 맞춤화된 압축 정책을 사용하여 중요한 정보를 유지하면서도 모델의 성능을 안정적으로 유지할 수 있습니다.
  • 추가 학습 불필요: FastGen은 별도의 재훈련 없이 적용할 수 있어, 기존 모델에 간단히 추가할 수 있습니다.

한계점

  • 문맥 유지 한계: 필요에 따라 삭제된 정보가 장기적인 문맥을 필요로 하는 작업에 불리할 수 있습니다. 예를 들어, 문서 요약처럼 전체 문맥을 참조해야 하는 작업에서는 성능이 저하될 수 있습니다.
  • 어텐션 구조에 의존: 모델의 어텐션 헤드가 정형화되지 않았거나 특정 구조에 맞지 않는 경우, 최적의 압축 정책을 적용하기 어려울 수 있습니다.

Adaptive KV Cache Compression은 LLM이 메모리 효율성을 높이면서도 성능을 유지할 수 있도록 돕는 강력한 솔루션입니다. FastGen은 이를 통해 대규모 모델을 더욱 경제적으로 운영할 수 있게 하여, 특히 자원이 제한된 환경에서 중요한 기여를 할 수 있습니다.