멀리 보기: 200줄이 안 되는 GPT

이 장의 물음

라운지 창가 자리에서 민준은 「Once upon a time, there was a girl named Luna.」로 시작한 샘플을 다시 띄워 놓고 있었다. 첫 문장에서 이름을 붙인 아이는 분명 Luna 였는데, 두 문장 뒤에는 「Jimmy smiled and nodded.」가 나왔다.

루나라고 해 놓고 두 문장 만에 Jimmy 야. 창이 여섯 칸이니까.

창을 넓히면 되잖아. 여섯 칸 말고 이백 칸.

넓히면 자리마다 가중치가 따로 늘어나. 같은 「 Luna」도 여섯 칸 앞이랑 백 칸 앞을 따로 배워야 하고.

서연은 공책의 설정서 둘째 줄 「앞에서 한 말을 기억한다」에 펜 끝을 대고 있다가, 지난 학기 디퓨전 강의 노트를 꺼냈다. 이 장이 끝나면 루나는 앞 256 토큰을 모두 보고 다음 토큰을 맞히고, 같은 학습 조건에서 바이트당 비트를 MLP 의 0.9326 에서 0.6129 로 낮춘다.

앞 토큰 여섯 개를 이어 붙이는 MLP 언어 모델은 창을 넓힐수록 은닉층의 입력 가중치가 창 길이에 비례해 늘었고, 같은 토큰이라도 몇 칸 앞에 있느냐에 따라 따로 배워야 했다. 이야기 첫 문장의 이름을 수백 토큰 뒤에서 찾아 쓰려면 자리에 묶이지 않고 내용으로 골라 읽는 장치가 필요한데, 그 장치는 이미 다른 강의에서 배웠다. 다만 그때는 그림을 다듬는 모델이었고, 지금은 다음 토큰을 맞히는 모델이다.

이 장은 다음 물음에 차례로 답한다.

내용으로 골라 읽는 어텐션을 언어 모델에 쓰려면 무엇을 더해야 할까? 학습 중에 다음 토큰을 미리 훔쳐보지 않게 하려면?
어텐션, MLP, 정규화, 잔차 연결을 엮은 GPT 한 벌은 파라미터가 몇 개이고, 파라미터 수는 무엇이 정할까?
같은 데이터, 같은 스텝으로 학습하면 GPT 는 MLP 보다 얼마나, 그리고 이야기의 어느 위치에서 더 잘 맞힐까?

마스크를 씌운 셀프 어텐션: 앞 토큰 전부를 골라 읽되 뒤는 보지 않기

앞 토큰 여섯 개를 이어 붙이는 방식으로 창을 256 토큰까지 넓히면, 은닉층의 입력 가중치만 해도 256 × 64 × 512 에 바이어스 512 를 더한 8,389,120 개가 된다(sources/gpt/position_and_mask.out 3절). 여섯 칸일 때의 197,120 개보다 40배 정도 많은데, 문제는 수가 아니라 그 가중치가 「세 칸 앞의 Lily」와 「백 칸 앞의 Lily」를 서로 다른 입력으로 다룬다는 데 있다. 이야기 속 이름은 몇 칸 앞에 있든 같은 이름이므로, 자리마다 따로 배우는 대신 토큰의 내용을 보고 지금 필요한 토큰을 골라 읽는 계산이 있으면 좋겠다.

디퓨전 강의의 「질의·열쇠·값」을 이 책의 코드에서는 쿼리(q)·키(k)·값(v)이라 부른다. 셀프 어텐션의 가중치는 쿼리·키·값을 만드는 행렬 몇 개뿐이고, 그 행렬은 모든 위치가 함께 쓴다. 창이 256 토큰이든 1,024 토큰이든 가중치 수는 그대로이고, 세 칸 앞의 Lily 와 백 칸 앞의 Lily 는 같은 키를 만든다.

어텐션이 내용만 본다는 성질에는 대가가 따른다. 「Tom saw Lily」와 「Lily saw Tom」은 같은 토큰을 다른 순서로 놓은 글인데, 셀프 어텐션만으로는 둘을 구별하지 못한다. 그래서 토큰 임베딩에 자리마다의 벡터를 더해 순서를 알려 준다.

우리 GPT 는 사인 방식도 회전 방식도 쓰지 않고, 2018~2019년의 GPT 들처럼 자리마다 벡터를 하나씩 학습한다. 위치 임베딩 표는 256 자리 × 폭 256 으로 65,536 개의 파라미터다. 회전 위치 부호를 포함해 요즘 모델이 바꾼 선택들은 모델을 크게 키운 뒤의 부에서 하나씩 켜 가며 잰다. 여기서는 가장 오래된 방식을 읽기 쉽게 그대로 따른다.

다음 토큰을 맞히는 모델에 더할 것: 마스크

그림을 다듬는 모델은 모든 칸이 모든 칸을 읽어도 됐지만, 다음 토큰을 맞히는 언어 모델은 학습 때 정답을 볼 수 있느냐가 문제가 된다. 학습 글 「Once upon a time there was a girl named Lily.」를 한 번에 넣으면, 「 named」 자리에서 다음 토큰을 맞혀야 하는데 바로 옆에 정답 「 Lily」가 들어 있다. 어텐션이 「 Lily」 위치를 읽으면 맞히기는 쉽지만, 글을 쓸 때는 다음 토큰이 아직 없으므로 그렇게 배운 것은 쓸모가 없다.

그래서 위치 t 의 쿼리가 키 0 부터 t 까지만 읽도록 점수 표의 뒤쪽 칸을 소프트맥스 전에 −∞ 로 바꾼다. 이렇게 뒤 토큰을 가리는 장치를 마스크 (뒤 토큰을 보지 못하게 하는 마스크 / causal mask)이라 한다. 우리 GPT 의 코드에서는 마스크 표를 직접 만드는 부분이 이렇다.

# sources/common/gpt.py 의 CausalSelfAttention.forward 에서 발췌
q, k, v = self.c_attn(x).split(C, dim=2)                  # 쿼리·키·값을 한 번에
q, k, v = (t.view(B, T, self.n_head, hd).transpose(1, 2) for t in (q, k, v))
att = (q @ k.transpose(-2, -1)) / math.sqrt(hd)          # 쿼리 위치 x 키 위치 점수
mask = torch.ones(T, Tk, dtype=torch.bool, device=x.device).tril(diagonal=Tk - T)
att = att.masked_fill(~mask, float("-inf")).softmax(dim=-1)
y = att @ v                                               # 비중대로 값을 섞는다

tril 은 아래 삼각만 남기는 함수다. 토큰 8 개일 때의 마스크 표와, 학습 전 모델에 예문 11 토큰을 넣어 뽑은 첫 층 헤드 하나의 어텐션은 다음과 같다(sources/gpt/train_gpt_small.out 2절).

## 2. 어텐션 마스크: 토큰 8 개 (행 = 지금 위치가 보는 쪽, 1 = 볼 수 있음)
  1 0 0 0 0 0 0 0
  1 1 0 0 0 0 0 0
  1 1 1 0 0 0 0 0
  1 1 1 1 0 0 0 0
  1 1 1 1 1 0 0 0
  1 1 1 1 1 1 0 0
  1 1 1 1 1 1 1 0
  1 1 1 1 1 1 1 1

예문 'Once upon a time there was a girl named Lily.' -> 토큰 11 개: Once | ␣upon | ␣a | ␣time | ␣there | ␣was | ␣a | ␣girl | ␣named | ␣Lily | .
층마다 어텐션 텐서 모양 (1, 4, 11, 11) = (배치, 헤드, 쿼리 위치, 키 위치), 층 4 개
행 합 (모두 1 이어야): [1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0]
대각선 위(미래 칸)의 최댓값: 0.0
학습 전 1층 헤드 1 의 어텐션 (앞 6 위치, 소수 셋째 자리):
  1.000 0.000 0.000 0.000 0.000 0.000
  0.438 0.562 0.000 0.000 0.000 0.000
  0.358 0.347 0.295 0.000 0.000 0.000
  0.245 0.294 0.194 0.267 0.000 0.000
  0.193 0.156 0.204 0.258 0.189 0.000
  0.162 0.176 0.159 0.177 0.161 0.165

행마다 합은 1 이고, 대각선 위 칸은 모두 0 이다. 첫 위치는 볼 수 있는 토큰이 자기 하나뿐이라 1.000 을 다 주고, 여섯째 위치는 여섯 토큰에 0.159~0.177 을 나눠 준다. 1/6 이 0.167 이니 학습 전의 헤드는 아직 아무것도 고르지 않고 앞을 고르게 훑는다. 학습이 끝난 헤드가 이 비중을 어떻게 바꾸는지는 모델 속을 들여다볼 때 다시 본다.

한 번에 256 문제

마스크에는 덤이 하나 따라온다. 위치 t 의 출력은 위치 0~t 의 토큰만 보고 만들어지므로, 256 토큰짜리 조각 하나를 한 번 통과시키면 「앞 1 토큰을 보고 맞히기」부터 「앞 256 토큰을 보고 맞히기」까지 256 문제의 답이 한꺼번에 나온다. 손실은 이 문제들의 크로스 엔트로피를 평균한 것이다.

\textcolor{#d62728}{\mathcal{L}}(\textcolor{#665522}{\theta}) = -\frac{1}{\textcolor{#3c4bff}{T}} \sum_{\textcolor{#a67a63}{t}=1}^{\textcolor{#3c4bff}{T}} \log \textcolor{#c2408a}{p_{\theta}}\!\left(\textcolor{#1b9e77}{x}_{\textcolor{#a67a63}{t}+1} \,\middle|\, \textcolor{#1b9e77}{x}_{1}, \ldots, \textcolor{#1b9e77}{x}_{\textcolor{#a67a63}{t}}\right)

\begin{array}{ll} \textcolor{#d62728}{\mathcal{L}}(\textcolor{#665522}{\theta}) & \text{조각 하나의 손실 (토큰당 nat)} \\ \textcolor{#665522}{\theta} & \text{모델의 파라미터 전체} \\ \textcolor{#3c4bff}{T} & \text{조각 길이 (문맥 길이, 여기서는 256)} \\ \textcolor{#a67a63}{t} & \text{위치 번호} \\ \textcolor{#1b9e77}{x}_{\textcolor{#a67a63}{t}} & \text{위치 } \textcolor{#a67a63}{t} \text{ 의 토큰} \\ \textcolor{#c2408a}{p_{\theta}}(\cdot \mid \cdot) & \text{앞 토큰들을 보고 모델이 다음 토큰에 준 확률} \end{array}

MLP 언어 모델도 조각 안의 위치마다 손실을 매겼지만, 위치마다 보는 것은 언제나 앞 여섯 토큰이었다. 마스크를 씌운 GPT 는 위치 t 에서 앞 t 토큰을 모두 본다. 그래서 같은 조각에서도 앞쪽 위치는 짧은 문맥으로, 뒤쪽 위치는 긴 문맥으로 맞히는 연습을 함께 한다.

ML에서: 학습과 생성이 같은 계산이 되는 조건

마스크가 있으면 앞 토큰들의 계산은 뒤에 무엇이 붙든 바뀌지 않는다. 학습 때 위치 t 가 본 것과 생성 때 맨 끝 위치가 보는 것이 둘 다 「자기와 앞 토큰들」이 되므로, 학습한 대로 생성할 수 있다. 같은 성질 덕분에 앞 토큰의 키와 값을 적어 두고 다시 쓰는 KV 캐시도 맞는 계산이 된다. GPT 라는 이름의 G 는 생성(generative)이고, 마스크를 씌운 셀프 어텐션은 생성 모델이 되기 위한 최소한의 장치다.

문제 1. 정답지를 펴 놓고 푼 모의고사

모의고사를 풀면서 바로 옆에 정답지를 펴 놓았다. (가) 그 점수는 실력을 알려 주는가? (나) 작은 GPT 를 마스크 없이, 나머지는 같게 500 스텝 학습하고 검증도 마스크 없이 재면 검증 손실은 어떻게 되겠는가? 마스크를 씌운 같은 모델은 500 스텝에 2.4512 였다. (다) 마스크 없이 학습한 모델로 「Once upon a time」 뒤를 이어 쓰게 하면 어떤 글이 나오겠는가?

(가)는 당연히 아니죠. 정답지만 베껴도 만점인데.

(나)는요? 숫자로 짐작해 볼까요?

위치마다 바로 다음 칸에 정답이 들어 있으니까, 어텐션이 다음 칸 하나만 보면 돼요. 손실이 0 근처까지 내려갈 것 같아요.

돌려 봤어. 100 스텝까지는 3.7958 로 마스크를 씌운 모델이랑 비슷하다가, 200 스텝에 0.0763, 500 스텝에 0.0165 야. 마스크를 씌운 모델의 2.4512 보다 150배쯤 낮아.

그럼 이 모델이 가장 잘 배운 모델이네요?

이어 쓰기를 시켰더니 「Once upon a time time time time upon a time a a a time a time time a upon」 이래요. 끝까지 time, a, upon 세 개만 돌려 써요.

글을 쓸 때는 맨 끝 위치가 다음 토큰을 맞혀야 하는데, 맨 끝 위치 오른쪽에는 아무것도 없잖아. 학습 내내 옆 칸을 베끼기만 했으니 베낄 게 없으면 할 줄 아는 게 없는 거네.

그래요. 손실이 낮다는 말은 「학습 때와 같은 조건에서 잘 맞힌다」는 뜻이에요. 생성할 때의 조건이 학습 때와 다르면 손실은 아무것도 보장하지 않아요. 마스크는 학습 조건을 생성 조건과 같게 맞추는 장치예요.

손실 0.0165 를 보고 신나서 자랑할 뻔했네요. 숫자가 너무 좋으면 먼저 의심해야겠어요.

문제 2. 조각 하나에 든 문제 수

학습 조각은 256 토큰이고, 한 스텝에 조각 64 개를 쓴다. (가) 마스크를 씌운 GPT 는 한 스텝에 다음 토큰 맞히기 문제를 몇 개 푸는가? (나) 그 문제들이 보는 앞 토큰 수의 평균은? (다) 같은 조각으로 학습한 MLP 언어 모델(앞 6 토큰)의 문제들은 평균 몇 토큰을 보는가?

(가)는 256 × 64 = 16,384 개요. 학습 기록에 「스텝당 16,384 토큰」이라고 찍히던 그 숫자예요.

(나)는 위치 1 이 1 개, 위치 2 가 2 개, … 위치 256 이 256 개를 보니까 평균은 (1 + 256) / 2 = 128.5 개예요.

(다)는요? MLP 도 위치마다 문제를 풀었죠.

앞 다섯 위치는 1~5 개, 그 뒤 251 위치는 6 개씩이니까 (15 + 1,506) / 256 = 5.94 개요.

문제 수는 둘이 똑같은데, 문제마다 보는 문맥이 20배 넘게 차이 나네요. 그러면 GPT 가 MLP 보다 잘 맞히는 건 주로 조각의 뒤쪽 위치일 것 같은데요?

그건 학습한 두 모델을 위치마다 재 보면 알 수 있어요. 조금 뒤에 확인해 봐요.