attention의 직관: 임베딩만으로는 왜 부족한가

오늘의 질문

지난 글에서 단어를 벡터로 바꾸면 “가까움”을 내적으로 계산할 수 있다는 것을 봤다. 개와 고양이는 0.98, 개와 자동차는 0.46이었다. 곱셈과 덧셈만으로 의미의 유사도가 나온다는 게 핵심이었다.

그런데 여기까지만으로는 풀리지 않는 문제가 있다.

“고양이가 자동차를 보았다. 그것은 배가 고팠다.” 모델은 “그것”이 고양이라는 걸 어떻게 알아내는가?

임베딩 표를 아무리 뒤져도 답이 안 나온다. “그것”은 어느 문장에 있든 표에서 같은 행을 가져오기 때문이다. 이 글은 그 간극을 메우는 계산을 다룬다. 미리 말해두면 새로운 수학은 나오지 않는다. 지난 글의 내적을 다른 벡터에 적용하는 것이 전부다.

먼저 답하기

두 문장으로 줄이면 이렇다.

임베딩 벡터를 그대로 쓰지 않고, 학습된 행렬 세 개를 통과시켜 하나의 임베딩에서 세 종류의 벡터를 만든다. 그런 다음 지난 글에서 배운 내적으로 관련도를 재고, 그 관련도만큼 섞는다.

세 종류의 벡터에 붙은 이름이 query, key, value다. 이름이 셋이라 복잡해 보이지만, 정체는 같은 임베딩이 갈아입은 세 벌의 옷이다. 하나는 “내가 찾는 것”, 하나는 “내가 어떤 존재인지 내거는 이름표”, 하나는 “찾아졌을 때 건네줄 내용”이다.

주제에 맞는 핵심 설명

▎지난 글에서 그대로 가져오는 것

먼저 이미 아는 것을 확인하고 가자. 지난 글의 세 단어를 그대로 쓴다.

“그것”은 새로 추가한 것인데, 아직 무엇도 가리키지 않는 중립적인 위치에 놓았다. 이 세 벡터로 지난 글의 코사인 유사도를 계산하면 이렇다.

여기서 문제가 드러난다. 0.949와 0.838은 별로 차이가 안 난다. 임베딩만으로는 “그것”이 둘 중 무엇을 가리키는지 판단할 근거가 약하다.

그리고 더 근본적인 문제가 둘 있다.

첫째, 값이 고정되어 있다. 두 번째 문장을 “그것은 빨간색이었다”로 바꾸면 답은 자동차가 되어야 한다. 그런데 임베딩은 문장이 바뀌어도 그대로다. 계산 결과도 그대로 0.949와 0.838이다. 문맥을 반영할 통로가 없다.

둘째, 대칭이다. 코사인 유사도는 cos(a, b) = cos(b, a)다. 순서를 바꿔도 값이 같다. 그런데 우리가 원하는 관계는 방향이 있다. “그것”이 고양이를 찾는 것과 고양이가 “그것”을 찾는 것은 다른 일이다.

▎해법: 같은 임베딩을 세 벌로 갈아입힌다

위 두 문제를 한 번에 푸는 방법이 임베딩을 그대로 쓰지 않는 것이다.

임베딩 벡터 하나를 서로 다른 행렬 세 개에 각각 통과시킨다. 행렬을 곱한다는 건 벡터를 다른 방향으로 늘이고 돌린다는 뜻이다. 같은 임베딩에서 출발했지만 통과한 행렬이 다르므로 결과 벡터도 서로 다르다.

• 임베딩 × W_Q → query: 이 token이 무엇을 찾고 있는가
• 임베딩 × W_K → key: 이 token이 어떤 존재로 검색되길 원하는가
• 임베딩 × W_V → value: 검색되었을 때 무엇을 건넬 것인가

이 세 행렬은 학습으로 정해진다. 사람이 값을 정하는 게 아니라, 모델이 데이터를 보면서 조정한다.

이렇게 하면 앞의 두 문제가 어떻게 풀리는지 보자. query와 key가 서로 다른 행렬에서 나오므로 q(A)·k(B)와 q(B)·k(A)가 다른 값이 된다. 대칭성이 깨진다. 그리고 이 벡터들은 실제로는 앞선 층들을 거치면서 문맥 정보가 섞인 표현에서 만들어지므로, 같은 token이라도 문장에 따라 다른 query와 key를 갖게 된다.

이 구조는 트랜스포머를 제시한 Vaswani 외의 논문 “Attention Is All You Need”(2017)에 그대로 기술되어 있다. 논문은 attention을 query와 key-value 쌍의 집합을 출력으로 대응시키는 함수로 정의하고, 출력은 value들의 가중합이며 각 가중치는 query와 그에 대응하는 key의 호환성 함수로 계산된다고 설명한다.

▎숫자로 끝까지 따라가기

이제 실제로 계산해본다. 위 세 단어를 그대로 쓴다.

가정 — 아래 행렬 값은 설명을 위해 내가 임의로 정한 것이며 학습된 값이 아니다. 실제 차원은 수백 이상이지만 여기서는 2차원으로 줄였다.

W_Q = [[2,  0],        W_K = [[2, 0.3],
       [0, -2]]                [0,  1  ]]

0단계 — 변환. 각 임베딩에 행렬을 곱한다.

q("그것")   = (2×0.60 + 0×0.60,  0×0.60 + (−2)×0.60) = ( 1.20, −1.20)
k("고양이") = (2×0.80 + 0.3×0.40, 0×0.80 + 1×0.40)   = ( 1.72,  0.40)
k("자동차") = (2×0.20 + 0.3×0.95, 0×0.20 + 1×0.95)   = ( 0.69,  0.95)

여기서 W_Q가 두 번째 성분에 −2를 곱한다는 점을 보자. “그것”의 query가 두 번째 축에 대해 음의 방향을 향하게 됐다. 두 번째 축 값이 큰 단어(자동차는 0.95)를 밀어내는 질문이 된 셈이다.

1단계 — 내적. 지난 글에서 한 그 계산이다.

그것 → 고양이 : 1.20 × 1.72 + (−1.20) × 0.40 = 2.064 − 0.480 =  1.584
그것 → 자동차 : 1.20 × 0.69 + (−1.20) × 0.95 = 0.828 − 1.140 = −0.318

임베딩끼리 쟀을 때는 0.949 대 0.838로 거의 붙어 있었는데, 변환 후에는 1.584 대 −0.318로 뚜렷하게 갈렸다. 자동차 쪽은 음수까지 내려갔다. 이게 학습된 행렬이 하는 일이다.

2단계 — √d로 나눈다. d는 key 벡터의 차원이고 여기서는 2다.

고양이 :  1.584 ÷ √2 =  1.120
자동차 : −0.318 ÷ √2 = −0.225

3단계 — softmax. 지수를 취하고 합으로 나눈다. 음수여도 지수를 취하면 양수가 되므로 가중치는 항상 0보다 크다.

exp( 1.120) = 3.065
exp(−0.225) = 0.799
              합 = 3.864

고양이 : 3.065 ÷ 3.864 = 0.793
자동차 : 0.799 ÷ 3.864 = 0.207

4단계 — 섞는다. 여기서는 설명을 단순하게 하려고 W_V를 항등행렬로 두어 value가 임베딩과 같다고 하자.

0.793 × (0.80, 0.40) = (0.635, 0.317)
0.207 × (0.20, 0.95) = (0.041, 0.196)
                합계 = (0.676, 0.514)

결과 (0.676, 0.514)가 무엇을 뜻하는지는 지난 글의 도구로 확인할 수 있다. 이 벡터와 고양이 임베딩의 코사인 유사도를 재면 0.983이다. 입력이었던 “그것”의 임베딩은 고양이와 0.949였으니, 고양이 쪽으로 더 끌려갔다.

정리하면 “그것” 자리의 표현이 문맥에 맞게 갱신된 것이다. 그리고 이 갱신은 규칙을 심어서가 아니라 곱셈과 덧셈으로 이루어졌다.

적용 한계: 위 행렬은 결과가 직관과 맞도록 내가 고른 값이다. 실제 학습된 행렬의 각 성분이 무엇을 뜻하는지는 사람이 지정할 수 없다. 또 실제로는 이 계산이 모든 token에 대해 동시에, 여러 층에 걸쳐 반복된다.

▎√d로 나누는 이유

2단계의 나눗셈만 따로 보자. 논문은 이 방식에 Scaled Dot-Product Attention이라는 이름을 붙였고, 나누는 이유도 밝혀두었다. d가 커지면 내적 값의 크기도 함께 커지는데, 그러면 softmax가 기울기가 극도로 작은 영역으로 밀려난다는 것이다.

무슨 뜻인지 위 숫자로 확인할 수 있다. 나눗셈을 생략하면 이렇게 된다.

나누지 않으면 분포가 더 뾰족해진다. 차원이 2라서 차이가 작아 보이지만, 실제처럼 d가 64나 512라면 √d가 8이나 22.6이 되므로 효과가 훨씬 크다. 가중치가 0과 1에 완전히 붙어버리면 학습 중에 조금씩 조정할 여지가 사라진다.

여기서 짚어둘 것은 이 나눗셈이 결과를 더 정확하게 만들려는 게 아니라 학습이 되게 만들려는 장치라는 점이다.

▎질문을 여러 개 던지기

지금까지 query 하나만 다뤘다. 그런데 “그것”이 필요로 하는 정보가 하나뿐일 리 없다. 가리키는 대상도 찾아야 하고, 문법적 역할도 파악해야 한다.

논문의 해법은 Q/K/V 세트를 여러 벌 병렬로 두는 것이고, 이를 multi-head attention이라 부른다. 논문의 기본 모델은 8개를 병렬로 두고 각 세트의 차원을 512 ÷ 8 = 64로 잡았다. 512차원을 64차원짜리 8개로 쪼개서 각각 다른 관점의 검색을 시키는 것이다.

논문은 이 구조가 서로 다른 위치의 서로 다른 표현 부분공간에서 정보를 함께 참조할 수 있게 한다고 설명한다. 차원을 나눠 쓰기 때문에 전체 계산량은 하나만 쓸 때와 비슷한 수준으로 유지된다.

▎이 계산이 메모리에 남기는 것

마지막으로 시스템 쪽 연결을 하나 짚는다. 위 계산에서 key와 value는 앞선 token들이 만든 것이다. 다음 token을 생성할 때도 같은 key와 value가 또 필요하다. 문장 앞부분이 바뀐 게 아니기 때문이다.

그래서 매번 다시 계산하지 않고 저장해둔다. 이것이 앞선 글에서 트래픽을 계산할 때 등장했던 KV 캐시다. 이 글의 관점에서 다시 정의하면, KV 캐시란 지금까지 등장한 모든 token의 이름표(key)와 내용물(value)을 모아둔 것이다.

이렇게 보면 크기가 왜 그런 식으로 자라는지도 자연스럽다. token이 하나 늘 때마다 key 하나와 value 하나가 추가되고, 이게 층마다 존재한다. 그래서 컨텍스트 길이 × 층 수 × key/value 차원에 비례한다. 앞선 글에서 쓴 계산식이 이 구조에서 나온 것이다.

자주 묻는 질문

임베딩 유사도로 직접 비교하면 안 되나요? 위에서 본 두 가지가 걸린다. 첫째, 임베딩은 문맥과 무관하게 고정이라 같은 단어가 문장마다 다른 것을 가리키는 상황을 구분할 수 없다. 둘째, 코사인 유사도는 대칭이라 “A가 B를 찾는다”와 “B가 A를 찾는다”를 구분할 수 없다. 학습된 행렬을 거치면 두 제약이 모두 풀린다. 위 예시에서 임베딩끼리는 0.949 대 0.838로 붙어 있던 것이 변환 후 1.584 대 −0.318로 갈린 것이 그 차이다.

key와 value를 굳이 나눠야 하나요? 합쳐도 계산은 돌아가지만 표현력이 줄어든다. key는 “어떤 질문에 걸릴 것인가”를, value는 “걸렸을 때 무엇을 줄 것인가”를 결정하는데 이 둘이 항상 같아야 할 이유가 없다. 어떤 token이 “생물 주어”라는 조건으로 검색되기를 원하면서, 정작 전달하고 싶은 내용은 그 대상의 상태 정보일 수 있다. 분리해두면 검색 조건과 전달 내용을 각각 학습할 수 있다.

attention 가중치를 보면 모델이 왜 그렇게 답했는지 알 수 있나요? 단서는 되지만 설명으로 받아들이기엔 조심스럽다. 실제 모델에는 head가 여러 개이고 층도 여러 개라 봐야 할 가중치가 매우 많고, 각 head가 무엇을 찾는지는 학습된 결과라 해석이 잘 붙지 않는다. 또 attention은 정보를 섞는 한 단계일 뿐이고 그 뒤로도 여러 변환을 거친다. 가중치가 높은 token이 크게 기여했으리라는 추측은 가능하지만 판단 근거라고 단정하기는 어렵다.

문장이 길어지면 이 계산이 얼마나 무거워지나요? 전체 문장을 한 번에 처리할 때는 token 하나가 앞선 모든 token과 대조되므로 비교 횟수가 길이의 제곱에 비례한다. 답변을 한 token씩 생성하는 단계에서는 새 token 하나가 기존 전부와 대조되므로 비교 횟수 자체는 길이에 선형이지만, 그러기 위해 저장해둔 KV 캐시 전체를 읽어야 한다. 앞선 글에서 봤듯 긴 컨텍스트에서는 이 읽기 비용이 가중치 읽기 비용을 넘어서는 지점이 생긴다. 연산량 문제이면서 동시에 메모리 대역폭 문제다.