어텐션

왜 어텐션인가: 뜻은 문맥이 정한다

4장에서 각 토큰을 하나의 벡터(임베딩)로 바꿨다. 그런데 임베딩 테이블은 토큰마다 벡터를 딱 하나만 준다. 다음 세 문장의 '배'는 모두 같은 토큰이니 처음에는 같은 벡터에서 출발한다.

사람은 '시장·달다', '항구·들어왔다', '먹었더니·아프다'를 보고 즉시 뜻을 고른다. 모델도 그렇게 하려면 각 토큰의 벡터를 주변 토큰의 정보로 고쳐 쓰는 장치가 필요하다. 영어의 고전적인 예도 있다.

두 문장은 마지막 단어 하나만 다르다. 그런데 'it'이 가리키는 대상이 'animal'에서 'street'로 바뀐다. 'it'의 의미를 정하려면 문장 끝의 'tired/wide'와 문장 앞의 'animal/street'를 동시에 참고해야 한다. 거리가 먼 두 토큰 사이에 정보가 직접 오가야 한다는 뜻이다.

RNN의 병목: 모든 것을 벡터 하나에

Transformer 이전의 대표 구조인 RNN(Recurrent Neural Network)은 토큰을 왼쪽부터 하나씩 읽으며 은닉 상태 \(h_t = f(h_{t-1}, x_t)\)를 갱신했다. 번역용 인코더-디코더 RNN(2014)은 원문 전체를 읽은 마지막 은닉 상태 하나, 즉 고정 길이 벡터 c에 문장 전체를 압축한 뒤 디코더가 그 벡터만 보고 번역문을 만들었다. 문장이 길어질수록 앞부분 정보가 희미해지고, 번역 품질이 문장 길이에 따라 뚝 떨어졌다.

어텐션의 장점은 세 가지로 정리된다. 첫째, 거리와 무관한 직접 연결: 1번 토큰과 1000번 토큰 사이도 한 번의 연산으로 정보가 오간다(RNN이라면 1000단계를 거쳐야 한다). 둘째, 병렬성: 모든 토큰 쌍의 유사도는 행렬곱 한 번으로 동시에 계산되므로 GPU에 꼭 맞는다. 셋째, 내용 기반 라우팅: 어디를 볼지가 고정된 규칙이 아니라 토큰의 내용에 따라 매번 달라진다. 대가도 있다. 모든 쌍을 보므로 비용이 길이의 제곱에 비례한다(9절).

직관: 부드러운 사전 검색

파이썬 딕셔너리 d["animal"]는 키가 정확히 일치하는 항목 하나의 값을 돌려준다. 어텐션은 이것을 미분 가능하게 '부드럽게' 만든 것이다. 세 가지 역할이 있다.

• 쿼리(query, q) — "나는 지금 이런 정보를 찾는다." 질문하는 쪽 토큰이 낸다.
• 키(key, k) — "나는 이런 정보를 갖고 있다." 각 토큰이 내건 색인표.
• 값(value, v) — "선택되면 실제로 건네줄 내용." 키와 짝을 이루는 내용물.

쿼리를 모든 키와 비교해 유사도 점수 \(s_j = q\cdot k_j\)를 얻고, 소프트맥스로 합이 1인 가중치 \(a_j\)로 바꾼 뒤, 값을 그 가중치로 평균한다. 정확히 하나만 고르는 대신 '조금씩 다' 고르는 셈이다.

이 시뮬레이터의 값 \(v_j\)는 색(RGB 3차원 벡터)이다. 실제 어텐션에서 값은 수십~수백 차원 벡터이고, 출력은 그 벡터들의 가중 평균이다. 원리는 똑같다. 한 가지 짚어 둘 것은 내적은 방향뿐 아니라 길이에도 반응한다는 점이다. 그래서 학습된 모델은 키·쿼리의 길이(노름)를 이용해 "이 토큰은 웬만하면 주목받게" 같은 효과도 낸다.

Scaled dot-product attention

문장에 토큰이 n개 있고 각각 d차원 벡터라고 하자. 이를 행으로 쌓은 것이 입력 행렬 \(X\in\mathbb{R}^{n\times d}\)다. 모든 토큰이 동시에 쿼리이자 키이자 값이므로(셀프 어텐션), 행렬 세 개 \(W_Q, W_K\in\mathbb{R}^{d\times d_k}\), \(W_V\in\mathbb{R}^{d\times d_v}\)로 각 역할의 벡터를 만든다.

왜 \(\sqrt{d_k}\)로 나누는가

쿼리와 키의 각 성분이 평균 0, 분산 1인 독립 확률변수라고 하자. 내적 \(q\cdot k = \sum_{i=1}^{d_k} q_i k_i\)는 분산 1짜리 항 \(d_k\)개의 합이므로 분산이 \(d_k\), 표준편차가 \(\sqrt{d_k}\)다. \(d_k = 128\)이면 점수가 대략 ±11 범위로 퍼진다. 소프트맥스는 입력 차이에 지수적으로 반응하므로, 이렇게 큰 점수가 들어가면 가장 큰 항 하나에 가중치가 거의 다 몰린다(포화). 포화된 소프트맥스는 기울기 \(\partial a/\partial s = \mathrm{diag}(a) - aa^\top\)가 거의 0이라 학습 신호가 사라진다. \(\sqrt{d_k}\)로 나누면 분산이 다시 1이 되어 차원이 몇이든 소프트맥스가 '적당히 부드러운' 영역에서 시작한다.

어텐션 행렬을 직접 편집하기

이제 이 장의 핵심 실험이다. 실제 모델에서 점수 행렬 S는 \(QK^\top/\sqrt{d_k}\)로 계산되지만, 여기서는 그 앞단을 잠시 잊고 S의 칸을 손으로 직접 정한다. 칸을 위로 끌면 점수가 올라가고, 아래로 끌면 내려간다. 오른쪽(모바일에서는 아래)의 A는 각 행에 소프트맥스를 적용한 실제 결과이고, 맨 아래 아크 그림은 선택한 행(쿼리 토큰)이 어느 토큰에서 정보를 가져오는지를 선 굵기로 보여 준다.

먼저 관찰해 둘 성질 세 가지가 있다. ① 소프트맥스는 행 안에서의 상대 차이에만 반응한다. 한 행 전체에 같은 값을 더해도 A는 그대로다. ② 점수 차이가 1 늘 때마다 가중치 비율은 \(e \approx 2.7\)배가 된다. 차이가 3이면 약 20배, 5면 약 150배다. ③ 다른 행은 서로 독립이다. 행 i를 고쳐도 다른 행의 A는 변하지 않는다. 각 토큰이 '자기만의 질문'을 따로 던지기 때문이다.

편집기의 숫자는 교육용으로 손으로 넣은 값이지만, 실제로 학습된 Transformer를 시각화하면 비슷한 장면이 보인다. BERT의 일부 헤드는 대명사에서 선행 명사로 강한 가중치를 보내고(Clark 등, 2019), GPT 계열의 많은 헤드는 첫 토큰이나 문장부호에 가중치 대부분을 몰아 둔다. 단, 어텐션 가중치가 크다고 해서 그 토큰이 출력에 '원인'이 된다고 단정할 수는 없다는 점도 기억해 두자. 값 벡터의 크기와 그 뒤의 \(W_O\)가 실제 기여를 함께 정한다.

Q, K, V는 어디서 오나: 같은 X, 다른 W

앞 절에서는 S를 손으로 넣었다. 실제 모델에서 S를 정하는 것은 학습되는 행렬 \(W_Q\)와 \(W_K\)다. 핵심은 \(S_{ij} = (x_i W_Q)\cdot(x_j W_K)/\sqrt{d_k} = x_i\,(W_Q W_K^\top)\,x_j^\top/\sqrt{d_k}\)라는 점이다. 즉 어텐션 점수는 임베딩 공간에 정의된 학습 가능한 쌍선형 유사도다. 그냥 \(x_i\cdot x_j\)를 쓰면 "비슷한 단어끼리" 서로 보게 될 뿐이지만, \(W_Q W_K^\top\)를 거치면 "대명사는 명사를 찾는다", "각 토큰은 바로 앞 토큰을 찾는다"처럼 비대칭적이고 목적이 있는 관계를 표현할 수 있다. 같은 문장에 서로 다른 W를 가진 헤드가 여러 개 있는 이유도 여기에 있다.

값 쪽도 같다. \(V = XW_V\)는 "선택되면 무엇을 넘겨줄지"를 정한다. 키는 '광고'이고 값은 '내용물'이다. 둘을 분리해 두었기 때문에, 어떤 토큰이 "나는 명사다"라는 키로 주목을 받으면서 정작 넘겨주는 값은 "나는 살아 있는 존재이고 피곤할 수 있다" 같은 다른 성질일 수 있다.

이 시뮬레이터의 2차원 임베딩은 교육용으로 직접 설계한 것이다. 7개 토큰을 원 둘레에 문장 순서대로 배치해 두었는데, 이는 실제 모델에서 토큰 벡터에 위치 정보가 섞여 들어가는 것(6장의 위치 인코딩·RoPE)을 아주 거칠게 흉내 낸 것이다. 실제 모델은 \(d = 768\)(GPT-2 small)~\(4096\)(Llama 3 8B) 차원이고 헤드당 \(d_k = 64\)~\(128\)이라, 의미 축과 위치 축이 서로 다른 방향을 차지하며 공존한다. 2차원에서는 둘을 동시에 담기 어렵기 때문에 '이전 토큰 헤드'는 원 위의 회전으로, '명사→대명사 헤드'는 특정 방향으로의 사영(rank-1 W)으로 구현했다.

마스크: 인과 마스크와 패딩 마스크

어텐션은 기본적으로 모든 토큰이 모든 토큰을 본다. 하지만 상황에 따라 보면 안 되는 위치가 있다. 마스크 \(M\)은 그런 위치의 점수에 \(-\infty\)를 더해, 소프트맥스 뒤 가중치가 정확히 0이 되게 한다(\(e^{-\infty} = 0\)). 구현에서는 \(-\infty\) 대신 \(-10^4\)이나 dtype의 최솟값을 쓰기도 한다.

• 인과 마스크(causal mask) — GPT·Llama 같은 디코더 언어 모델은 다음 토큰을 예측하도록 학습된다. 위치 i가 i+1 이후를 보면 정답을 훔쳐보는 셈이므로 \(j > i\)인 칸을 모두 막는다. A는 하삼각 행렬이 된다. 덕분에 길이 n 문장 하나로 n개의 '다음 토큰 예측' 문제를 한 번에 병렬로 학습할 수 있고, 추론 때는 지나간 토큰의 K·V가 변하지 않으므로 8장의 KV 캐시가 가능해진다.
• 패딩 마스크(padding mask) — 길이가 다른 문장을 배치로 묶으려면 짧은 쪽 뒤에 [PAD] 토큰을 채운다. 그 위치는 의미가 없으므로 키 쪽 열을 막는다. 패딩 위치의 쿼리 행은 결과를 버린다.
• 양방향(bidirectional) — BERT 같은 인코더는 마스크 없이 앞뒤를 모두 본다. 번역 원문 이해, 문장 분류, 임베딩 생성(10장)처럼 문장 전체가 이미 주어진 작업에 쓴다.

위의 편집기(4절)에서 인과 마스크를 켜고 끄며 차이를 다시 확인해 보자. 마스크를 켜면 행마다 소프트맥스의 분모에 들어가는 항이 줄어들기 때문에, 남은 칸의 가중치가 커진다. 첫 행은 항이 하나뿐이라 무조건 1이다. 이 '첫 토큰은 항상 보인다'는 성질은 8절의 어텐션 싱크와 연결된다.

멀티헤드 어텐션: 여러 관점을 병렬로

한 헤드는 행마다 하나의 확률분포만 만든다. 그런데 한 토큰은 동시에 여러 종류의 정보가 필요하다. 'it'은 지시 대상(animal)도, 바로 앞 단어(because)도, 문장 주어도 알아야 한다. 하나의 소프트맥스로 이 모든 것을 섞으면 서로 경쟁해 흐려진다. 그래서 Transformer는 어텐션을 h개의 헤드로 나눠 병렬로 계산한다.

시뮬레이터의 네 헤드는 실제 모델에서 자주 관찰되는 패턴을 규칙으로 손설계한 것이다(교육용). 아래 '혼합'은 모든 헤드가 같은 V를 쓰고 \(W_O\)가 헤드 출력을 켜진 헤드 수로 나눠 더하는 특별한 경우의 유효 어텐션 \(\tfrac{1}{h'}\sum_h A_h\)다. 이 경우 \(\sum_h g_h A_h V = (\sum_h g_h A_h)V\)이므로 정확히 성립한다. 일반적인 \(W_O\)에서는 헤드마다 다른 값 부분공간을 쓰므로, 행렬 하나로 요약되지 않는다. 그 '요약되지 않음'이 바로 멀티헤드의 표현력이다.

어텐션 패턴 동물원

학습된 언어 모델의 어텐션 행렬을 열어 보면, 몇 가지 '종'이 반복해서 나타난다. 해석 가능성(mechanistic interpretability) 연구는 이런 패턴이 무슨 계산을 하는지 역설계한다. 아래 그림은 대표 패턴을 교육용으로 단순화해 그린 것이다(실제 헤드는 훨씬 지저분하고 여러 패턴이 섞여 있다).

• 어텐션 싱크(attention sink) — 많은 헤드가 첫 토큰(보통 BOS)에 가중치 대부분을 보낸다. 소프트맥스는 행 합이 반드시 1이어야 하므로, "지금은 볼 게 없다"는 헤드도 가중치를 어딘가에 버려야 한다. 모든 위치에서 항상 보이는 첫 토큰이 그 쓰레기통이 된다. Xiao 등(2023, StreamingLLM)은 긴 생성에서 오래된 KV를 버릴 때 첫 몇 토큰만은 남겨 두면 성능이 유지된다는 것을 보였다.
• 이전 토큰 헤드(previous-token head) — 각 위치가 바로 앞 토큰을 본다. 단어가 여러 토큰으로 쪼개졌을 때 조각을 이어 붙이거나, 아래 인덕션 헤드의 첫 단계로 쓰인다.
• 인덕션 헤드(induction head) — "… A B … A" 다음에 B를 예측한다. 앞에서 A 다음에 무엇이 왔는지 찾아 그대로 복사하는 2층 회로다. Olsson 등(2022, Anthropic)은 이 헤드가 학습 중 특정 시점에 갑자기 생기며, 그때 7장의 인컨텍스트 학습 능력이 함께 뛰어오른다고 보고했다.
• 지역/윈도우 패턴 — 가까운 몇 토큰만 본다. 일부 모델은 이를 구조로 강제한다(10절의 슬라이딩 윈도우).

인덕션 헤드는 어떻게 작동하나

두 층의 협업이다. 1층의 이전 토큰 헤드가 각 위치 j에 "내 바로 앞 토큰은 \(t_{j-1}\)이다"라는 정보를 써넣는다. 2층의 인덕션 헤드는 현재 위치 i의 토큰 \(t_i\)를 쿼리로, 각 위치에 적힌 '앞 토큰' 정보를 키로 쓴다. 그러면 \(t_{j-1} = t_i\)인 위치 j, 즉 "예전에 지금과 같은 토큰 바로 다음에 왔던 자리"의 점수가 높아진다. 그 위치의 토큰 \(t_j\)를 값으로 복사하면 다음 토큰 예측이 된다.

이 시뮬레이터의 두 층은 학습된 가중치가 아니라 원-핫 벡터와 위치 규칙으로 손설계한 것이다. 그러나 계산 구조(1층 어텐션의 출력이 2층의 키가 되는 'K-합성')는 Elhage 등(2021)과 Olsson 등(2022)이 실제 모델에서 찾아낸 회로와 같다. 해석 가능성 연구는 이처럼 어텐션 헤드, MLP 뉴런, 그리고 최근에는 희소 오토인코더로 찾은 '특징'들을 회로로 엮어 모델 내부를 설명하려 한다. 아직은 작은 모델과 특정 행동에 대한 부분적 설명이 대부분이라는 점을 함께 기억하자.

계산량과 메모리: O(n²)와 FlashAttention

어텐션의 비용은 두 부분으로 나뉜다. Q·K·V·O 투영은 토큰마다 독립이라 길이에 선형이다(\(4\cdot 2nd^2\) FLOPs). 반면 점수 \(QK^\top\)와 \(AV\)는 모든 쌍을 다루므로 길이의 제곱이다(\(2\cdot 2n^2 d\) FLOPs). 둘이 같아지는 지점은 \(n = 2d\)로, Llama 3 8B(d = 4096)라면 약 8K 토큰부터 어텐션 쌍 계산이 투영보다 비싸진다. 메모리는 더 아프다. 점수 행렬 S를 그대로 만들면 헤드마다 \(n^2\)개 원소가 필요하다.

FlashAttention: 큰 행렬을 만들지 않는다

GPU에서 병목은 연산이 아니라 메모리 이동인 경우가 많다. H100의 HBM3는 약 3.35 TB/s지만, 연산은 BF16 dense 약 989 TFLOPS라 바이트당 수백 번 연산할 수 있다. 순진한 구현은 S(n×n)를 HBM에 쓰고, 다시 읽어 소프트맥스하고, A를 쓰고, 다시 읽어 V와 곱한다. Dao 등(2022)의 FlashAttention은 Q, K, V를 작은 타일로 잘라 칩 위의 빠른 SRAM(공유 메모리)에 올리고, 한 타일 안에서 점수·소프트맥스·V 곱을 모두 끝낸 뒤 출력만 HBM에 쓴다. n×n 행렬은 어디에도 통째로 존재하지 않는다. 연산량은 같지만(오히려 역전파 때 재계산으로 약간 늘지만) 메모리 읽기·쓰기가 크게 줄어 수 배 빨라지고, 추가 메모리는 O(n)이다.

문제는 소프트맥스다. \(a_j = e^{s_j}/\sum_l e^{s_l}\)는 분모에 행 전체의 합이 들어가고, 수치 안정성을 위해 보통 행 최댓값 \(m\)을 빼서 \(e^{s_j - m}\)으로 계산한다. 블록을 하나씩 보면 최댓값도 합도 아직 모른다. 온라인 소프트맥스(online softmax, Milakov & Gimelshein, 2018)는 지금까지의 최댓값 \(m\)과 합 \(\ell\)을 들고 다니다가, 새 블록에서 최댓값이 커지면 기존 합과 누적 출력에 \(e^{m_{old}-m_{new}}\)를 곱해 소급 보정한다.

FlashAttention은 2022년 이후 사실상 표준이 되었고, FlashAttention-2(2023)는 병렬화를 개선했으며, FlashAttention-3(2024)는 H100의 비동기 텐서 코어와 FP8을 활용한다. PyTorch의 scaled_dot_product_attention도 조건이 맞으면 이런 커널을 자동으로 쓴다. 중요한 점은 이것이 정확한 어텐션이라는 것이다. 근사로 비용을 줄이는 선형 어텐션·희소 어텐션과 달리, 결과는 순진한 구현과 같고 메모리 접근 순서만 바꿨다.

셀프·크로스 어텐션, MHA·MQA·GQA, 슬라이딩 윈도우

셀프 어텐션 vs 크로스 어텐션

지금까지는 Q, K, V가 모두 같은 X에서 나왔다. 이것이 셀프 어텐션이다. 크로스 어텐션(cross-attention)에서는 쿼리는 한 시퀀스에서, 키와 값은 다른 시퀀스에서 온다. 원래 Transformer의 번역 디코더는 번역문 토큰의 쿼리로 원문 인코더 출력의 키·값을 조회했다(그림 5-1의 b가 바로 그것). 멀티모달 모델에서는 텍스트 쿼리가 이미지 패치 임베딩을 키·값으로 조회하고, Stable Diffusion 같은 이미지 생성 모델은 이미지 특징이 쿼리, 프롬프트 텍스트 임베딩이 키·값이다. 크로스 어텐션의 점수 행렬은 \(n_q\times n_{kv}\)로 정사각형이 아니며, 인과 마스크도 보통 필요 없다.

MHA, MQA, GQA: KV 헤드를 공유하기

생성 단계에서는 이전 토큰들의 K, V를 저장해 두는 KV 캐시(8장)가 메모리의 주인공이 된다. 표준 멀티헤드(MHA)는 쿼리 헤드마다 자기 K, V 헤드가 있다. MQA(Multi-Query Attention, Shazeer 2019)는 모든 쿼리 헤드가 K, V 헤드 하나를 공유해 캐시를 h배 줄이지만 품질이 약간 떨어질 수 있다. GQA(Grouped-Query Attention, Ainslie 등 2023)는 그 중간으로, 쿼리 헤드를 g개 그룹으로 묶어 그룹마다 K, V 헤드 하나를 둔다. Llama 3 8B는 쿼리 헤드 32개, KV 헤드 8개(그룹당 쿼리 4개)를 쓴다.

슬라이딩 윈도우와 그 밖의 변형

슬라이딩 윈도우 어텐션은 각 토큰이 최근 w개 토큰만 보게 해 비용을 \(O(nw)\)로 줄인다(그림 5-6의 넷째 패턴). Mistral 7B(2023)는 w = 4096을 썼고, 층을 쌓으면 정보가 층마다 w씩 더 멀리 전파되므로 L층이면 이론적으로 약 L·w 거리까지 닿는다. Gemma 2·3처럼 지역(윈도우) 층과 전역 층을 번갈아 쓰는 모델도 많다. 이 밖에 키·값을 저차원 잠재 벡터로 압축해 캐시하는 MLA(DeepSeek-V2, 2024), 소프트맥스를 커널 근사로 바꾼 선형 어텐션 계열 등이 있다. 위치 정보를 Q, K의 회전으로 넣는 RoPE는 6장에서, KV 캐시와 디코딩 속도는 8장에서 자세히 다룬다.