STUDYLLM
‹ Lecture 2 15 / 23

Embedding Table

Token ID 자체에는 계산할 의미의 좌표가 없다. Embedding layer는 그 ID를 주소로 삼아 학습 가능한 table에서 vector 하나를 가져온다. VV는 서로 다른 주소의 수, dd는 각 주소에 저장한 vector의 폭이다. IDs가 크다고 vector의 값이 크거나 의미가 더 강해지는 것이 아니다.

  1. Token ID i
  2. EE[i] lookup
  3. dmodeld_{model}-width vector
  4. Transformer

Lookup & Gradients

E∈RV×dE\in\mathbb{R}^{V\times d}는 token당 한 row다. one_hot⁡(i)E\operatorname{one\_hot}(i)E와 EE[i]는 수학적으로 같지만 nn.Embedding 구현은 index lookup이다.

Input lookup 경로에서는 batch가 참조한 rows만 직접 gradient를 받는다. Weight tying이면 output-softmax 경로도 같은 matrix를 갱신한다. rare token은 under-trained일 수 있어도 반드시 untouched는 아니다.

Matrix Shape

one_hot⁡(i)E=E[i]\operatorname{one\_hot}(i)E=E[i]

각 row를 ei∈R1×de_i\in\mathbb{R}^{1\times d}라고 쓰면, ID가 하는 일은 다음 matrix에서 특정 row를 선택하는 것이다.

E=[e0e1⋮e15717⋮eV−1]∈RV×d,x=E[15717]=e15717E= \begin{bmatrix} e_0\\ e_1\\ \vdots\\ e_{15717}\\ \vdots\\ e_{V-1} \end{bmatrix} \in\mathbb{R}^{V\times d}, \qquad x=E[15717]=e_{15717}

슬라이드의 ID 15717은 berry에 대응한다. One-hot vector를 곱해도 나머지 rows의 계수는 0이고 선택한 row의 계수만 1이므로 같은 결과가 나온다. 그러나 코드에서 길이 VV의 one-hot을 만들고 dense matrix product를 할 필요는 없다. nn.Embedding은 index lookup으로 필요한 dd개 값만 가져온다. 논문에서의 수학적 표기와 효율적인 구현을 구분해야 하는 이유다.

Modeldmodeld_{\mathrm{model}}Vocabulary VVOne table VdVd
GPT-2 small76850,257약 38.6M
BERT-base76830,522약 23.4M
Llama-3-8B4,096128,256약 525M
GPT-3 175B12,28850,257약 618M

GPT-2 small의 한 table은 124M 전체의 약 31%지만, GPT-3 175B에서는 약 0.35%다. Llama-3-8B의 untied input/output는 각각 약 525M으로 총 1.05B, 약 13%다. Tying하면 한 table을 공유하여 약 525M, 모델의 약 6.5%를 절약한다.

Input + untied output parameters≈2Vd\text{Input + untied output parameters}\approx2Vd

학습에서는 어떤 경로로 table을 사용했는지가 중요하다. Input lookup만 보면 이번 batch가 고른 rows에 직접 gradient가 온다. 그러나 같은 table을 output head로도 쓰면 모든 Vocabulary의 logits 계산에 참여하므로, 입력에서 선택되지 않은 row도 간접 갱신될 수 있다. 이 점은 앞의 Glitch Tokens와 다음 Weight Tying을 연결한다.

Vector가 dd개 수를 가진다고 해서 첫 dimension은 색, 둘째는 장소처럼 명확한 사전이 생기는 것은 아니다. 교수님은 여러 정보가 같은 표현 안에 섞일 수 있다는 superposition을 언급했다. dd는 representation capacity와 관련되지만, 각 dimension을 단일한 의미 feature로 완전히 해석할 수 있다고 보장하지 않는다.

dmodeld_{model}과 Vocabulary Size VV를 혼동하지 않는다. Embedding layer의 lookup은 Attention이 아니다.

EE의 rows=VV, width=dd. Input lookup과 output gradients를 구분한다.

Check

one-hot×EE가 가능한데도 lookup 구현을 쓰는 이유와 결과 shape는?

Answer

동일한 row를 얻는 데 전체 Vocabulary를 곱하는 대신 index로 선택할 수 있다. 단일 token의 결과는 dd-dimensional vector다.

Input lookup에서 안 읽힌 row도 tied model에서 업데이트될 수 있는 경로는?

Answer

같은 EE가 output LM head로 쓰이며 Vocabulary 전체에 대한 Softmax loss의 gradient가 전달될 수 있다.

Token ID를 다른 숫자로 바꾸고 Embedding weights를 그대로 두면 입력 의미는 유지되는가?

Answer

유지되지 않는다. ID는 해당 Embedding row를 선택하는 index이므로, 다른 ID는 다른 learned vector를 읽는다. ID 자체의 숫자 크기가 의미를 나타내는 것은 아니지만 ID↔row 대응은 학습된 모델과 일치해야 한다.

Contents