nn.Embedding에 이해하기

토큰 번호로 바뀐 문장이 어떻게 벡터가 되는지, nn.Embedding의 동작을 간단하게 정리해봤다.

(vocab_size, model_dim) 가중치 행렬#

nn.Embedding은 결국 (vocab_size, model_dimension) 크기의 가중치 행렬 하나다. 이 행렬의 행벡터 하나가 사전의 단어 하나를 model_dim 차원의 벡터로 나타낸 것이다. 즉 단어 하나에서 model_dim 벡터로 가는 매핑이 행렬 형태로 들어 있는 셈이다.

token index로 행을 뽑아온다#

토큰 번호(index)를 지닌 source 문장이 들어오면, nn.Embedding은 weight[token index]에 해당하는 행을 그대로 가져온다. 문장의 모든 단어를 각자의 index로 조회해서 model_dim 벡터로 바꾸는 것이다. 곱셈이라기보다 조회하는 느낌에 가깝다.

차원이 어떻게 바뀌는가#

예를 들어 길이 10짜리 문장 하나가 토큰 번호로 src = (1, 10) 형태로 들어온다고 하자. 이걸 (vocab_size, model_d) 크기의 embedding에 통과시키면, 각 토큰 번호가 model_d 벡터로 바뀌면서

$$(1,10)\stackrel{\text{nn.Embedding}}{\to }(1,10,model\_d)$$

가 된다. 문장의 토큰 개수는 그대로 두고, 각 토큰 자리에 model_d 차원 벡터가 채워지는 것이다.

정리#

• nn.Embedding은 (vocab_size, model_dim) 가중치 행렬이고, 행벡터 하나가 단어 하나의 model_dim 표현이다.
• token index로 해당 행을 lookup해서 가져온다.
• (batch, length) 토큰이 (batch, length, model_dim)으로 바뀐다.