HF Transformers -2-

토큰 번호는 모델 안에서 어떤 길을 지나갈까?

이전 포스트 HF Transformers -1- 에서 모델은 토큰 번호를 받아 라벨 점수 2개를 내놓았다. 그 사이에서 토크나이저가 문장을 바꿔 놓은 토큰 번호(input_ids) 는 세 구간을 지난다.

아래 그림처럼 임베딩이 토큰 번호(input_ids) 를 768개 숫자로 된 벡터로 바꾸고, Transformer 층 6개가 그 벡터를 문맥에 맞게 다듬고, 헤드가 다듬어진 벡터로 점수를 낸다.

이번에서는 구간마다 중간 결과를 꺼내 보며, 토큰 번호가 벡터를 거쳐 점수가 되기까지 모양과 값이 어떻게 바뀌는지 따라가 본다.

1. 모델 구조 다시 보기

먼저 모델이 어떤 부품으로 이루어져 있는지 열어 본다.

import torch
from transformers import AutoTokenizer, AutoModelForSequenceClassification

name = "distilbert-base-uncased-finetuned-sst-2-english"
tokenizer = AutoTokenizer.from_pretrained(name)
model = AutoModelForSequenceClassification.from_pretrained(name)

print(model.distilbert.embeddings)

print(len(model.distilbert.transformer.layer))   # 6

Embeddings(
  (word_embeddings): Embedding(30522, 768, padding_idx=0)
  (position_embeddings): Embedding(512, 768)
  (LayerNorm): LayerNorm((768,), eps=1e-12, elementwise_affine=True)
  (dropout): Dropout(p=0.1, inplace=False)
)

6

본체는 임베딩 하나와 똑같이 생긴 층 6개로 되어 있다. 원래 BERT는 층이 12개인데, DistilBERT는 이걸 절반으로 줄인 모델이다. 이름의 "Distil"이 바로 큰 모델의 지식을 작은 모델로 증류했다는 뜻이다.

1) 임베딩 - 번호를 벡터로

토큰 번호 1045는 그냥 사전 속 위치일 뿐, 그 자체로는 아무 의미가 없다. 그래서 모델은 번호마다 768개 숫자로 된 벡터를 하나씩 배정해 둔다.

위에서 본 Embedding(30522, 768)이 바로 그 표다. 행이 30,522개 (사전 크기), 열이 768개인 표에서 번호에 해당하는 행을 꺼내 오는 것이 임베딩이다.

그런데 단어 벡터만으로는 부족하다. 단어 표는 번호만 보고 벡터를 꺼내므로, 같은 단어는 문장 어디에 있든 똑같은 벡터가 된다. 예를 들어 "dog bites man"과 "man bites dog"을 넣으면 두 문장 모두 dog, bites, man 세 벡터를 받는다. 놓인 순서만 다를 뿐 재료는 같다.

문제는 뒤에 오는 어텐션도 순서를 모른다는 것이다. 어텐션은 토큰끼리 Q와 K를 비교해 섞을 비율을 정하는데, 이 계산에는 "몇 번째 자리인가"가 들어가지 않는다. 입력 벡터의 순서를 뒤섞으면 출력도 똑같이 자리만 뒤섞여 나온다. 이대로라면 모델은 누가 누구를 물었는지 가릴 단서가 없다.

그래서 순서 정보를 벡터 안에 직접 넣어 준다. 자리마다 벡터를 하나씩 정해 두고 단어 벡터에 더하는 것이다. 이게 위치 임베딩(position_embeddings)이다. 단어 표가 "번호 > 벡터"라면 위치 표는 "자리 > 벡터"다. 0번 자리, 1번 자리, 2번 자리마다 768개 숫자로 된 벡터가 한 행씩 들어 있고, 이 값도 학습으로 정해졌다.

이제 dog는 0번 자리에 있을 때와 2번 자리에 있을 때 서로 다른 벡터가 된다. 각각 "dog 벡터 + 0번 자리 벡터", "dog 벡터 + 2번 자리 벡터"이기 때문이다. 단어의 뜻과 자리 정보가 한 벡터 안에 함께 실려 층으로 들어간다. 옆에 이어 붙이지 않고 더하는 이유는 모양을 768로 유지하기 위해서다. 그래야 뒤의 층이 받는 입력 크기가 바뀌지 않는다.

위에서 출력한 구조를 다시 보면 위치 표(position_embeddings)는 Embedding(512, 768)이다. 행이 512개, 즉 0번부터 511번 자리까지만 벡터가 준비되어 있다. 513번째 토큰에게 줄 자리 벡터는 표에 없다. 이 모델이 최대 512 토큰까지만 읽을 수 있는 이유이고, truncation 한계도 여기서 나온다. [CLS]와 [SEP]도 자리를 하나씩 차지하므로, 실제 문장에 쓸 수 있는 자리는 510개다.

"movie"(3185번)의 단어 벡터와 다섯 번째 자리(4번)의 자리 벡터를 더한다. 단어 표는 문장과 상관없이 늘 같은 줄을 내주므로, 순서 정보는 자리 표가 맡는다.

inputs = tokenizer("I loved this movie!", return_tensors="pt")
ids = inputs["input_ids"]                       # torch.Size([1, 7])
emb = model.distilbert.embeddings

with torch.no_grad():
    word = emb.word_embeddings(ids)                            # 단어 벡터
    pos = emb.position_embeddings(torch.arange(ids.shape[1]))  # 자리 벡터
    mine = emb.LayerNorm(word + pos)

    outputs = model(**inputs, output_hidden_states=True)

print(mine.shape)       # torch.Size([1, 7, 768])
print(torch.allclose(mine, outputs.hidden_states[0], atol=1e-6))   # True

word + pos: 단어 벡터와 자리 벡터를 그냥 더한다. 이어 붙이는 게 아니라 더하기라서 크기는 768 그대로다.
LayerNorm: 벡터 하나하나의 숫자 분포를 고르게 맞춰 주는 정리 단계다. 값이 너무 크거나 작게 튀지 않게 해서 다음 계산을 안정적으로 만든다.
output_hidden_states=True: 모델이 중간 결과를 모두 돌려주게 하는 옵션이다. hidden_states[0]이 임베딩을 지난 직후의 값이다.

결국 코드를 보면 모델의 임베딩 단계가 하는 일은 "LayerNorm(단어벡터 + 자리벡터)" 이 전부다 라는 것이다.

간단히 다시 설명하면, 순차적으로 토크나이저로 문장을 임베딩 한 것이 아무런 은닉층을 거치지 않은 것이 hidden_state[0] 이 되고 각 층을 거치면서 해당 모델의 작업에 맞게 다듬어져 나가는 과정이라고 보면 된다.

2) 층 6개를 지나며 다듬기

hidden_states에는 몇 개가 들어 있을까.

print(len(outputs.hidden_states))   # 7
for i, h in enumerate(outputs.hidden_states):
    print(i, h.shape)

0 torch.Size([1, 7, 768])
1 torch.Size([1, 7, 768])
...
6 torch.Size([1, 7, 768])

임베딩 직후 1개와 층마다 1개씩 6개, 합쳐서 7개다.

눈여겨볼 점은 모양이 처음부터 끝까지 [1, 7, 768]로 (3차원 텐서에 768차원) 같다는 것이다. 층을 지나도 토큰 수도, 벡터 크기도 바뀌지 않는다. 바뀌는 건 벡터 안의 숫자뿐이다.

여기서 위 텐서의 숫자는 각각 다른 것을 의미한다.

1: 문장 개수(배치 크기)
지금은 “I loved this movie!” 한 문장만 넣었으니 1이고 문장 3개를 한 번에 넣으면 3이 됩니다.
7: 토큰 개수
[CLS], i, loved, this, movie, !, [SEP]로 7개. 문장이 길어지면 늘어남 (최대 512).
768: 토큰 하나를 나타내는 벡터의 크기
DistilBERT가 정한 값(config.dim)이라 어떤 문장을 넣어도 바뀌지 않음

다시 돌아가서, 그러면 층마다 무엇이 바뀔까. 임베딩 직후의 "movie" 벡터는 문장과 상관없이 늘 같다. "I loved this movie"에서도, "The movie was awful"에서도 똑같은 행을 꺼내 오기 때문이다. 층을 지나면서 각 토큰은 주변 토큰들의 정보를 조금씩 받아 와서, 점점 "이 문장 속의 movie"가 된다. 층마다 벡터가 얼마나 바뀌는지 숫자로 확인해 보자.

hs = outputs.hidden_states
for i in range(1, len(hs)):
    sim = torch.cosine_similarity(hs[i][0, 4], hs[i-1][0, 4], dim=0)
    print(f"{i-1}층 → {i}층  movie 벡터 유사도: {sim:.3f}")

0층 → 1층  movie 벡터 유사도: 0.782
1층 → 2층  movie 벡터 유사도: 0.829
2층 → 3층  movie 벡터 유사도: 0.873
3층 → 4층  movie 벡터 유사도: 0.893
4층 → 5층  movie 벡터 유사도: 0.531
5층 → 6층  movie 벡터 유사도: 0.555

여섯 값 모두 1보다 작다. 층을 지날 때마다 movie 벡터가 실제로 바뀐다는 뜻이다. 흐름은 두 갈래로 나뉜다. 층 1~4에서는 유사도가 0.782에서 0.893으로 조금씩 올라간다. 뒤로 갈수록 한 층이 고치는 양이 줄어든다는 뜻이다. 그런데 층 5와 층 6에서는 0.531, 0.555로 크게 떨어진다. 마지막 두 층이 movie 벡터를 가장 크게 바꾸는 셈이다. 이 모델은 감정 분류로 미세조정한 모델이라, 그 작업에 맞춰 벡터를 다듬는 일이 뒤쪽 층에 몰려 있는 것으로 보인다.

층 하나가 안에서 정확히 무슨 계산을 하는지는 "주변 토큰을 참고해 내 벡터를 고친다"는 정도로만 기억해 두기로 한다.

3) 헤드 - 마지막 벡터에서 점수로

마지막 층을 지난 결과 hidden_states[-1]이 last_hidden_state다. 분류 헤드는 이 중 첫 토큰 [CLS]의 벡터 하나만 받아 점수 2개로 줄인다. 이 과정도 직접 해서 logits와 맞춰 보자.

with torch.no_grad():
    cls = outputs.hidden_states[-1][:, 0]   # [CLS] 벡터
                                            # torch.Size([1, 768])

    x = model.pre_classifier(cls)           # 768 → 768
    x = torch.relu(x)

    logits = model.classifier(x)            # 768 → 2

print(logits)
print(outputs.logits)                       # 같은 값

[:, 0]
모든 문장에서 0번 자리, 즉 [CLS]만 꺼낸다. 나머지 6개 토큰의 벡터는 헤드에서 쓰이지 않는다.
pre_classifier → relu → classifier
헤드를 이루는 두 층이다. 중간의 relu는 음수를 0으로 바꾸는 함수다.

다시 위 코드를 시각화 해보면 위와 같다.

[CLS] 벡터(768차원)는 먼저 pre_classifier를 지난다. 여기서 768개 숫자를 서로 다른 비율로 섞어, 감정을 가르기 좋은 새 숫자 768개를 만든다. 새 숫자 하나를 만들 때 입력 768개에 비율을 하나씩 곱하므로, 비율은 모두 768 × 768개다. 이렇게 만든 768차원 벡터는 relu를 지나며 음수가 0으로 바뀐다.

그런데 [CLS]는 문장 맨 앞에 붙인 특수 토큰일 뿐인데, 어떻게 이 하나로 문장 전체의 감정을 판단할 수 있을까.

층 6개를 지나는 동안 [CLS]도 다른 토큰들의 정보를 계속 받아 왔기 때문에, 마지막에는 문장 전체를 요약한 벡터가 되어 있다. [CLS]가 문장의 "대표"가 되는 이유다.

결국 전체 모델 구조를 보면 세 구간이다. 임베딩이 번호를 768개짜리 벡터로 바꾸고, 층 6개가 모양은 그대로 둔 채 벡터 안의 숫자를 문맥에 맞게 고치고, 헤드가 [CLS] 벡터 하나로 점수를 낸다.

다음 포스트에는 층안에서 주변 토큰을 참고하는 어텐션을 알아보기로 한다.