LLM: KV Cache 이해하기 (feat. Self Attention)

LLM에 관련된 글이나 이야기들을 보다보면 빠지지 않는 기술 중 하나가 KV Cache 인데요. 그만큼 널리 쓰이고 또 중요하다는 뜻일 겁니다. 오늘은 KV Cache가 무엇이고 왜 중요한지 이야기해보려고 합니다.

LLM은 느리다

네, LLM은 느립니다. 덩치가 크니까 느린게 당연하지 않냐고 생각할 수도 있습니다. 물론 LLM이 느린 첫번째 이유는 너무 크기 때문입니다. 하지만 또다른 이유들도 있습니다.

첫번째는 Self Attention 구조 때문이고, 두번째는 Auto-regressive라고 불리는 생성 방식 때문입니다. KV Cache는 이 두 가지 조합에서 생기는 문제를 해결해보려는 시도입니다. 그렇다면 이 두 가지 문제에 대해서 먼저 이해해봐야겠죠?

문맥 파악하기: Self Attention

Transformer가 기가 막히게 잘 작동하는 이유의 중심에는 Self Attention이 있습니다. 이름 그대로 스스로에게 집중한다는 뜻인데요. 여기에서 Self는 입력을 의미합니다.

예를 들어보겠습니다. "I went to the bank to deposit money"라는 문장이 있다고 해봅시다. "돈을 저축하기 위해서 은행에 갔다."라는 뜻인데요. 문제는 bank가 "은행"이라는 뜻이 될 수도 있고 "강둑"이라는 뜻이 될 수도 있다는 점입니다. 하지만 우리는 아무 문제없이 이 문장에서 bank가 은행이라는 것을 이해할 수 있습니다. 어떻게 이런 것이 가능할까요?

아무리 사람이라고 하더라도 "bank" 라는 단어만 주고 무슨 뜻인지 물어보면 은행인지 강둑인지 알 수 없습니다. 일반적으로 bank를 은행이라는 의미로 사용하니 강둑보다는 은행일 가능성이 높다고 찍을 수는 있겠지만요.

하지만 "I went to the bank to deposit money" 라고 문장을 주면 이야기가 달라집니다. 여기에서 bank는 은행이라는 것을 알 수 있습니다. bank 외의 단어들까지 고려해서 문맥을 이용할 수 있기 때문입니다.

Self Attention의 핵심 아이디어가 바로 이것입니다. Transformer가 bank라는 단어를 처리하고 있다고 가정해보죠. Transformer도 사람과 비슷하게 bank만 보고는 bank의 뜻을 정확히 이해할 수 없습니다. 그럼 어떻게 할까요? 문장 내의 다른 단어들을 함께 살펴봅니다. 그리고 아래와 같은 것들을 파악합니다.

• I는 bank와 관련성이 크지 않다.
• went는 bank와 관련성이 조금 있다.
• money는 bank와 관련성이 크다.

이렇게 모든 단어와 bank(자신을 포함해서) 사이의 관련성을 파악하고 나면 Transformer는 입력 문장에 대해 훨씬 풍부한 정보를 가지게 됩니다. bank와 관련성이 높은 단어 중 money가 있으니 이 문장에서 bank는 강둑보다는 은행이 될 가능성이 높아지는 것이죠.

실제로 Self Attention 이 어떻게 계산되나

앞에서는 관련성이 크다, 작다고 추상적으로 표현했는데요. 실제로 Self Attention이 어떻게 계산될까요? 이 부분을 이해하기 위해서는 Self Attention을 계산하는 수식을 봐야하지만 오늘도 수식은 생략하고 개념적으로 꼭 필요한 부분은 살펴보겠습니다.

먼저 입력의 각 단어(엄밀히는 token이지만 이 글에서는 논의의 편의를 위해 단어라고 하겠습니다)에 대해 Query (Q), Key (K), Value (V) 라는 세 가지 값을 계산합니다. 예를 들어 "I went to the bank to deposit money"라는 문장에서는 Q_I, K_I, V_I, Q_went, K_went, V_went, ..., Q_money, K_money, V_money 를 계산합니다. 8단어니까 총 8 * 3 = 24 개의 값을 계산합니다.

Query, Key, Value가 Self Attention에서 무슨 역할을 하는지 알아보겠습니다.

• Query (Q): 현재 단어의 '질문' 또는 '궁금증'입니다. "얘들아, 나랑 관련된 정보 가진 사람 누구니?" 하고 외치는 역할이죠.
• Key (K): 다른 단어들이 들고 있는 '이름표' 또는 '키워드'입니다. Query의 질문에 "어? 나 너랑 관련 있을지도? 내 이름표 한번 봐봐!" 하고 대답하는 역할입니다.
• Value (V): 그 단어가 가진 진짜 '의미' 또는 '내용물'입니다. Key가 Query와 잘 맞는다고 판단되면, "자, 여기 내 진짜 정보를 줄게!" 하며 건네주는 값이죠.

도서관에서 책 찾기

아직 모호합니다. 도서관에서 책을 찾는 것에 비유해보겠습니다. 책 제목은 정확히 모르지만 "AI의 역사"에 대한 책을 찾고 싶다고 해보겠습니다. 컴퓨터 공학 관련 책들이 모여있는 곳으로 가서 꽂혀있는 책들을 살펴봅니다. 책을 하나하나 펼쳐서 찾아보는 것은 너무 시간이 걸리고 어려운 일이니 먼저 책 제목을 살표봅니다. "운영체제 이론"이나 "컴파일러 개론"같은 제목은 관련성이 낮을 거고요. "패턴인식"은 좀 더 관련성이 높아 보이고, "AI 개론" 은 매우 관련성이 높아 보입니다.

하지만 책을 표지만으로 판단하지 말라고 했죠. 이제 마음에 드는 책들을 꺼내서 실제로 책을 펼치고 내용을 살펴봅니다.

• "운영체제 이론": 운영체제를 어떻게 만드는지에 대한 내용입니다
• "컴파일러 개론": 컴파일러에 대한 이론을 다루지만 일부 Rule-based 나 Logic-based AI에 적용할만한 내용도 있습니다
• "패턴인식": AI의 한 분야인 패턴인식에 대한 내용을 담고있습니다
• "AI 개론": 초기 AI 기술부터 최신 AI 기술까지 설명이 적혀있습니다

"운영체제 이론"은 제목으로 유추했을 때와 비슷하게 실제 내용도 큰 관련성은 없고, 반대로 "AI 개론"은 제목으로 유추했을 때와 비슷하게 실제 내용도 관련성이 있습니다. "컴파일러 개론"은 관련성이 낮아보였는데 실제로 내용을 보니 약간은 관련이 있어보이고요.

지금까지 이야기한 도서관에서 책을 찾는 비유를 Query, Key, Value와 연결해보겠습니다.

• "AI의 역사"가 Query에 해당합니다. "AI의 역사"와 관련도가 높은 책들을 찾는 것이 우리의 목적이고요. 다르게 말하면 우리는 도서관에 있는 각 책들이 "AI의 역사"라는 Query와 얼마나 관련이 있는지 찾고자 합니다.
• "운영체제 이론", "컴파일러 개론", "패턴인식", "AI 개론" 같은 제목이 Key에 해당합니다. 각 책의 내용을 함축적으로 표현하는 키워드인 셈이죠. 우리가 관심있는 내용(Q, "AI의 역사")을 각 책의 제목(K)와 비교해서 얼마나 관련성이 있는지를 1차로 판단했습니다.
• 책의 실제 내용이 Value에 해당합니다.

그럼 각 책이 얼마나 우리의 관심사인 Q("AI의 역사")와 관련있는지 계산할 수 있을까요? "AI의 역사"(Q)와 책의 제목(K)만을 고려하기는 뭔가 부족해보입니다. 아무리 제목을 잘 지었더라고 하더라도 책의 모든 내용을 담기는 어려우니까요. 좀 더 정확히 관련성을 표현하기 위해서는 책의 실제 내용(V)도 함께 고려해야합니다.

다시 말하면,

1. "AI의 역사"(Q)와 각 책의 제목(K) 사이의 관련성을 먼저 파악하고,
2. 이를 기반으로 각 책의 내용(V)을 얼마나 고려할지 결정합니다.

1단계에서 결정한 관련성이 높을수록 책의 내용이 최종 관련성에 더 많이 반영되는 셈입니다.

다시 Transformer로

원래 우리 예제로 돌아가보겠습니다. "I went to the bank to deposit money"를 입력으로 받았습니다.

첫번째 단어인 "I"와 다른 단어들 사이의 관련성을 확인하려고 합니다. "I"가 Q가 됩니다. 정확히는 "I"를 그대로 쓰는 것이 아니라 계산을 통해 I에 해당하는 Q 값을 계산합니다. "AI의 역사"에 관련된 책을 찾고 싶을 때 "AI의 역사"를 글자 그대로 해석하는게 아니라 우리 뇌에서 나름의 방식으로 의미를 해석하는 것과 비슷합니다.

그리고 모든 단어들의 K값과 V값을 계산합니다. 도서관에서 책찾기 예제에서는 모든 책의 K(제목)과 V(내용)이 정해져있었지만, 우리는 어떤 입력이 들어올지 모르기 때문에 K와 V를 동적으로 계산합니다. 도서관 예제에서는 책의 작가가 이미 K와 V를 미리 계산해두었다고 볼 수 있습니다.

이제 우리는 "I"에 해당하는 Q 값과 모든 단어에 해당하는 K와 V을 가지고 있습니다. 다음으로 아래 과정을 진행합니다.

1. Q_I 와 K_I가 얼마나 관련성이 높은지 봅니다. 같은 단어이니 아무래도 관련성이 꽤 높을 것 같습니다.
2. Q_I와 K_went가 얼마나 관련성이 높은지 봅니다. 주어, 동사 관계이니 어느 정도 관련성이 있을 것 같습니다.
3. Q_I와 K_to가 얼마나 관련성이 높은지 봅니다. 의미적으로나 문법적으로나 관련성이 크게 높아 보이지 않습니다.
4. 이 과정을 K_money까지 수행합니다.

앞에서 말한 Q와 각 K 사이의 관련성을 파악하는 단계가 끝났습니다. 이렇게 결정된 각 Q와 K 사이의 관련성을 바탕으로 각 단어의 V의 값을 조합해 최종적인 단어간 관련성을 결정합니다.

• (Q_I, K_I, V_I): "I"와 "I" 사이의 관련성
• (Q_I, K_went, V_went): "I"와 "went"사이의 관련성
• (Q_I, K_to, V_to): "I"와 "to"사이의 관련성
• ...

여기까지 하면 첫번째 단어인 "I"의 Self Attention 계산이 끝났습니다. 남은 일은 같은 과정을 "went", "to", "the", "bank", "to", "deposit", "money"에 대해서 수행하면 됩니다. "bank"라면 아래와 같은 계산을 하겠죠.

• (Q_bank, K_I, V_I): "bank"와 "I"사이의 관련성
• (Q_bank, K_went, V_went): "bank"와 "went"사이의 관련성
• ....
• (Q_bank, K_money, V_money): ""bank"와 "money"사이의 관련성

여기에서 각 단어 사이의 관련성이라고 표현한 것이 바로 Self Attention 값입니다. Self Attention을 계산한다고 하는 것은 입력 문장에 있은 각 단어 간의 관련성을 계산한다는 의미입니다.

Self Attention 값은 각 단어 간의 관련성을 표현하고 있기 때문에 문맥을 표현한다고도 볼 수 있습니다. Self Attention을 사용하면 "I went to the bank to deposit money" 에서 bank가 강둑이 아닌 은행이라는 것을 알 수 있게되는 것입니다.

Q, K, V를 계산하다는 것이 실제로 무슨 의미인가

앞에서 간 단어에 해당하는 Q, K, V 값을 계산하다는 표현을 했는데요. 실제로 이 말이 무슨 의미일까요?

이 부분은 약간의 수학 이야기를 담고 있습니다. 전체 논의에 관계없으니 다음으로 넘어가도 됩니다.

Transformer 구조는 다른 Machine Learning 알고리즘처럼 숫자만을 이해할 수 있습니다. "went", "bank" 같은 표현을 이해하지 못하죠. 그렇기 때문에 각 단어를 Transformer가 이해할 수 있는 숫자, 정확히는 숫자의 나열인 vector 형태로 먼저 변환합니다. 예를 들어 아래과 같은 변환됩니다.

• I: [213, 92, 42, 89, 0, 21]
• went: [5, 2, 43, 99, 92, 111]

이렇게 변환된 vector를 embedding이라고 부릅니다.

Q, K, V 계산은 각 단어의 embedding을 다른 vector로 변환하는 과정입니다. 예를 들어보겠습니다. I의 embedding을 E_I라고 해보겠습니다.

• E_I: [213, 92, 42, 89, 0, 21]
• Q_I: [234, 1, 53, 111, 5, 4]
• K_I: [4, 23, 62, 13, 34, 93]
• V_I: [24, 12, 32, 51, 112, 34]

위와 같이 변환될 수 있습니다.

이 예시에서실제 숫자는 제가 임의로 적었는데요 실제로 임베딩을 Q, K, V로 변환하기 위해 입력 embedding에 W_Q, W_K, W_V라는 행렬을 각각 곱합니다.

• E_I * W_Q = Q_I
• E_I * W_K = K_I
• E_I * W_V = V_I

그럼 W_Q, W_K, W_V 는 어떻게 만들까요? Transformer의 다른 모든 파라미터들처럼 W_Q, W_K, W_V도 학습 대상입니다. 처음에는 의미없는 값을 가지고 시작하지만, 학습 과정을 통해 의미있는 값을 가진 W_Q, W_K, W_V로 변합니다.

계산량: N^2

Self Attention은 각 단어 사이의 관련성입니다. "I"는 "I", "went", "to", ..., "money" 각각과 관련성을 계산해야하고, "money"도 "I", "went", "to", ..., "money" 각각관 관련성을 계산해야합니다.

입력 문장이 N개의 단어로 이루어져있다고하면, "I"는 N 단어와 관련성을 계산해야합니다. 이 과정을 모든 단어 (N개)에 대해 수행해야하기 때문에 입력 문장이 N개의 단어로 이루어져있다면 Self Attention 계산은 N*N 즉, N^2 의 계산 복잡도를 가집니다.

N^2 계산 복잡도라는 것의 의미가 무엇일까요? 입력이 10단어라면 Self Attention 계산은 10 * 10 = 100 번을 하면됩니다. 입력이 100단어라면 100 * 100 = 10000 번을 하면 되고요. 문제가 보이시나요? 입력은 10->100으로 10배가 커졌는데 필요한 계산은 100->10000으로 100배가 커졌습니다. 즉, 입력이 커질수록 필요한 계산은 제곱으로 커집니다.

Transformer에서 Self Attention을 입력의 문맥을 파악하게 해주는 매우 중요한 기술이지만 한편으로는 매우 비싸고 시간이 오래 걸리는 작업입니다.

생성 모델의 눈물겨운 반복 작업: Auto-Regressive

KV Cache를 이야기하기 전에 또다른 중요한 LLM의 특성인 Auto Regressive에 대해서 이야기해보겠습니다. Auto Regressive 라는 표현이 익숙치 않을 수 있는데요. LLM이 어떻게 문장을 생성하는지를 먼저 이야기해보겠습니다.

LLM에게 Prompt로 "What is the tallest mountain?"을 주었다고 해보겠습니다. LLM은 아마도 "It is Mountain Everest"라고 답할 겁니다.

한단계 더 들어가보겠습니다.

1. LLM은 "What is the tallest mountain?"에 이어질 가장 가능성이 높은 다음 단어를 고릅니다. -> "It"
2. LLM은 기존 입력에 새로 만든 "It"을 더합니다. LLM은 "What is the tallest mountain? It"에 이어질 가장 가능성 높은 다음 단어를 고릅니다. -> "is"
3. LLM은 기존 입력에 새로 만든 "is"를 더합니다. LLM은 "What is the tallest mountain? It is"에 이어질 가장 가능성 높은 다음 단어를 고릅니다. -> "Mountain"
4. LLM은 기존 입력에 새로 만든 "Mountain"를 더합니다. LLM은 "What is the tallest mountain? It is Mountain"에 이어질 가장 가능성 높은 다음 단어를 고릅니다. -> "Everest"
5. LLM은 기존 입력에 새로 만든 "Everest"를 더합니다. LLM은 "What is the tallest mountain? It is Mountain Everest"에 이어질 가장 가능성 높은 다음 단어를 고릅니다. -> "<END>"

<END>라는 특별한 단어가 나오고 LLM은 생성을 멈춥니다. 결과적으로 "It is Mountain Everest"를 출력합니다.

결국 LLM이 하는 것은 입력 Prompt에 이어질 가장 그럴듯한 단어를 고르고(생성하고), 이 단어를 원래 입력에 더한 다음, 다시 이 새로운 입력에 이어질 가장 그럴듯한 단어를 고르는 것입니다. 이 과정을 생성 끝을 나타내는 특별한 단어가 나올 때까지 반복하고요.

regressive는 이전 값에 의존해서 새로운 값을 예측한다는 것을 의미합니다. 어제 날씨를 기반해서 오늘 날씨를 예측한다면 이 일기 예보 시스템은 regressive입니다.

auto는 자기 자신이 예측한 값이 다시 입력으로 쓰인다는 의미입니다. 위 예에서 보면 처음 생성된 It이 다음 차례에서는 입력이 된 것을 볼 수 있습니다.

이 둘을 합쳐 LLM이 문장을 생성하는 방식을 Auto Regressive라고 부릅니다.

다음 단어 예측하기

위 단계 중에서 1단계만 먼저 조금 더 자세히 살펴보겠습니다.

1. LLM이 "What is the tallest mountain?"을 입력으로 받습니다.
2. LLM은 이 입력에 대해 Self Attention을 계산합니다.
3. 마지막 단어인 mountain? 의 self attention 값을 기반으로 다음 단어를 예측합니다. (좀 더 정확히는 이 값을 기반으로 fully connected layer 거쳐 hidden state가 만들어지고 이를 기반으로 다음 단어를 예측합니다)
4. 결과적으로 "It" 을 생성합니다.

3단계에서 mountain? 의 값만 봐도 충분하냐고 생각할 수 있지만, 이 값에는 Self Attention 계산 방식에 따라 What, is, the, tallest 에 대한 정보와 관계가 모두 담겨있습니다. Self Attention 계산 과정을 되새겨보면 Query인 mountains?과 모든 단어의 Key, Value 값을 종합적으로 고려한다는 것을 알 수 있습니다.

그리고 각 단어의 Q, K, V는 어떻게 계산되나요? 앞에 언급한 것처럼 미리 학습한 W_Q, W_K, W_V를 embedding에 곱해서 만듭니다.
다시 말하면 1단계에서 What, is, the, tallest, mountain?에 대한 Q, K, V를 한번 계산한 것입니다.

이제 오늘의 주인공 KV Cach가 등장할 모든 배경이 갖추어졌습니다.

KV Cache 등장

2단계로 가보겠습니다. 이제 입력은 "What is the tallest mountains? It"이 됐습니다. LLM은 It 다음 단어를 예측해야하기 때문에 It을 Query로하고 Self Attention을 계산합니다. It의 Self Attention을 계산하기 위해서는 It을 Query로 변환하고 다른 모든 단어 (What, is, the, tallest, mountain?)들의 K, V와 조합합니다.

그런데! What의 K와 V를 계산하려고 보니, 이미 It을 생성했던 이전 단계에서 K, V를 계산했습니다. is, the, tallest, mountain?도 마찬가지입니다. 굳이 같은 값을 다시 계산할 필요가 없습니다. 그래서 우리는 what, is, the, tallest, mountain?의 K, V 값을 계산하지 않고, 앞에서 이미 계산한 K, V 값을 재사용하기로 합니다.

이전 단계에서 각 단어의 K, V를 계산한 후에 버리지 않고 저장(Cache)해두고 사용하면 됩니다. 복잡한 행렬곱(W_K, W_V를 embedding에 곱하기)을 할 필요가 없기 때문에 시간을 엄청나게 절약할 수 있습니다.

이렇게 K, V 값을 매번 저장하지 않고 전 단계에서 계산한 K, V 값을 저장해두고 재사용하는 기법을 KV Cache라고 부릅니다.
KV Cache를 이야기 할 때 고려해야하는 두 가지가 있습니다.

직전에 생성한 마지막 단어의 K, V는 새로 계산해야합니다. 전 단계에서 입력이 아니었기 때문에 K, V를 계산한적이 없어서입니다. 이 경우에는 "It" 에 대한 K, V는 새로 계산합니다. 그리고 이 값은 다음 단계에서 사용하기 위해 Cache에 저장합니다.

Q는 왜 Cache하지 않는지 궁금할 수도 있습니다. 다음 단어를 생성할 때 관심사는 입력 중 마지막 단어의 Self Attention(정확히는 앞에 말한대로 hidden state)입니다. 항상 마지막 단어의 Q만 필요하고 이 Q값은 K, V와 마찬가지로 전 단계에서 계산된 적이 없기 때문에 새로 계산해야합니다. 이 Q값을 저장하더라도 의미가 없습니다. 왜냐면 다음 Auto Regressive 구조에서는 다음 단계에서 이미 이 Q값은 과거가 됐기 때문에 사용하지 않으니까요. 정리하자면 마지막 단어의 Self Attention을 계산하기 위해서는 과거의 K, V가 모두 필요하지만 Q는 마지막 단어만 필요하다입니다.

KV Cache의 단점

복잡하고 반복적인 과정을 없앰으로서 KV Cache는 LLM Inference에 엄청난 속도 향상을 가져옵니다. 하지만 세상에 공짜는 없죠. KV Cache에도 단점이 있습니다.

KV Cache는 한번 계산한 값을 버리지 않고 메모리에 저장했다가 나중에 사용하는 방식입니다. 그렇기 때문에 KV Cache를 저장할 메모리 공간을 사용하게 됩니다.

실제 Self Attention은 여러 개의 Head를 가지는 Multi-Head Self Attention(Self Attention이 여러개 있다고 생각하면 됩니다)으로 구현되고, 또 각 Transformer Layer마다 다른 Self Attention 값을 가집니다. 그래서 KV Cache를 저정하기 위한 메모리 크기는 다음과 같이 계산됩니다.

(단어 수) * (Attention head 수) * (Layer 수) * (K, V vector의 차원) * ...

따라서 입력 단어수가 길어질 수록, 모델이 커질 수록 (Attention head 수, Layer 수 증가) KV Cache의 크기가 급격하게 커집니다. KV Cache는 보통 GPU의 VRAM에 저장하기 때문에 매우 고용량의 VRAM을 가진 GPU가 필요하게 됩니다.

특히 최근에 수 십만에서 수 백만 토큰 context window를 생각하는 모델들을 생각하면 KV Cache 크기는 심각한 문제입니다. 이런 문제를 해결하기 위해서 Attention head 당 K, V를 공유하는 Multi-Query Attention이나 Group Query Attention같은 기법을 사용하기도 합니다.

마무리하며

오늘은 Transformer의 추론 성능을 비약적으로 향상시킨 일등공신, KV Cache에 대해 알아봤습니다.

핵심을 요약해보겠습니다.

• Self-Attention은 문맥을 파악하기 위해 Q, K, V라는 세 요소를 사용하며, 계산 비용이 비쌉니다(O(N²)).
• LLM은 이전에 생성한 내용을 포함하여 다음 단어를 예측하는 Auto-Regressive 방식으로 작동합니다.
• 이 과정에서 매번 모든 단어의 K, V를 다시 계산하는 것은 매우 비효율적입니다.
• KV Cache는 한 번 계산한 K, V를 메모리에 저장하고 재활용하여 반복 계산을 없애는 기술입니다.
• 그 결과 추론 속도는 매우 빨라지지만, 문장이 길어질수록 메모리를 많이 차지한다는 단점이 있습니다.

긴 글 읽어주셔서 감사합니다.