STUDYLLM
‹ Lecture 2 16 / 23

Weight Tying & LM Head

Embedding으로 들어올 때 모델은 ID 하나를 vector 하나로 바꾸었다. 나갈 때는 반대로 마지막 hidden state를 Vocabulary의 모든 ID에 대한 점수로 바꾸어 다음 token을 고른다. 같은 VV가 양쪽에 등장하므로, input Embedding과 output LM head를 서로 다른 weights로 둘지 하나로 공유할지를 선택할 수 있다.

  1. Token IDs
  2. EE
  3. Transformer / Residual Stream
  4. Final Norm
  5. LM head
  6. Logits
  7. Sampling / Argmax
  8. next ID append

Shared Weights

저장 weight를 E,Wout∈RV×dE,W_{\mathrm{out}}\in\mathbb{R}^{V\times d}로 쓰면 input은 EE[i], output은 hWout⊤∈R1×VhW_{\mathrm{out}}^\top\in\mathbb{R}^{1\times V}다. Tying은 Wout=EW_{\mathrm{out}}=E로 두 기능에 같은 matrix를 공유한다. 저장 시 transpose matrix를 따로 유지하는 뜻이 아니다.

입력에서는 ID가 고른 row만 읽지만 출력에서는 각 row와 hidden state를 비교해 모든 후보의 score를 만든다. Tying은 이 두 용도에 같은 learned matrix를 쓰는 것이며, transpose한 사본을 따로 저장해 동기화하는 방식이 아니다. nanoGPT의 wte.weight = lm_head.weight라는 코드도 같은 weight를 연결한다는 뜻이다. 강의 예시는 GPT-2/Gemma tied, Llama 3 untied, Qwen은 size/version별 차이다.

From Hidden State to Scores

슬라이드의 worked example은 d=4d=4, V=5V=5다.

h=[0.9,−1.2,0.4,2.0]h=[0.9,-1.2,0.4,2.0]

LM head로 projection한 결과가 다음과 같다고 하자. 각 coordinate는 hidden feature가 아니라 이제 Vocabulary의 특정 token에 대응한다.

hWout⊤=[3.1,0.5,−0.8,1.7,0.2]hW_{\mathrm{out}}^\top=[3.1,0.5,-0.8,1.7,0.2]

순서가 ▁the, ▁cat, ▁dog, ▁a, ▁is라면 greedy argmax는 가장 큰 3.1에 대응하는 ▁the를 고른다. 이 수들은 아직 probabilities가 아니다. Sampling하려면 Softmax로 정규화한 분포를 사용하고, greedy 선택만 필요하다면 Softmax가 순서를 보존하므로 생략할 수 있다.

Autoregressive Generation

Embedding+position 이후 Attention/MLP가 Residual Stream에 updates를 더하고 final norm/LM head로 logits를 얻는다. Sampling은 Softmax probabilities를 사용하고 greedy argmax는 가장 큰 logit을 선택한다. Softmax는 순서를 보존하므로 greedy는 이를 생략할 수 있다.

Logit Lens

고른 ID는 현재 sequence 뒤에 붙고 다시 다음 입력이 된다. Embedding과 position 처리 → Transformer blocks의 Attention/MLP updates → Final Norm → LM head → Logits → 선택이라는 경로를 한 token마다 반복하는 것이 autoregressive generation이다. Position 처리의 구체적 방법은 architecture에 따라 다르며, RoPE는 Embedding에 position vector를 더하는 방식과 구분한다.

Logit lens는 이 경로를 마지막까지 통과하기 전에 읽어 보는 방법이다. Residual Stream의 중간 state도 같은 dmathrmmodeld_{mathrm{model}} 폭이므로 Final Norm과 unembedding을 적용하여 Vocabulary 방향의 점수를 볼 수 있다. 어떤 token의 점수가 특정 layer 뒤 크게 달라졌다면 그 구간에서 representation이 달라졌다는 분석 단서가 된다. 다만 early-layer state가 마지막 output head에 잘 맞도록 calibrated되어 있다는 보장은 없으므로, 그 점수를 완성된 답이나 정확한 확률 해석으로 곧바로 받아들이지 않는다.

h⏟1×dWout⊤⏟d×V=logits⏟1×V\underbrace{h}_{1\times d}\underbrace{W_{\mathrm{out}}^\top}_{d\times V}=\underbrace{\mathrm{logits}}_{1\times V} Tied:Wout=E\text{Tied:}\quad W_{\mathrm{out}}=E

Logits는 확률이 아니다. RoPE는 Embedding에 position vector를 더하는 방식이 아니라 QQ/KK에 적용하는 별도 방식이다.

입력 lookup과 출력 logits에 같은 VV가 두 번 등장한다.

Check

V=5V=5,d=4d=4일 때 h와 Wout⊤W_{out}^\top, logits의 shape를 말하고 tying의 절약량을 설명하라.

Answer

h는 1×41\times 4, Wout⊤W_{out}^\top는 4×54\times 5, logits는 1×51\times 5다. EE와 WoutW_{out}을 공유하면 V×dV\times d matrix 하나의 parameters를 절약한다.

greedy argmax에서 Softmax를 생략할 수 있는 이유와 sampling에서의 역할 차이는?

Answer

Softmax는 logit 대소 순서를 보존하므로 최대 ID는 동일하다. Sampling에는 점수를 확률 분포로 만드는 단계가 필요하다.

Weight Tying을 해도 input lookup과 output projection의 계산이 같은 것은 아닌 이유는?

Answer

같은 E를 공유하지만 input은 ID로 E의 row를 선택하고, output은 hidden state와 모든 token rows의 내적으로 Vocabulary logits를 만든다. Parameter 공유와 연산 역할의 동일함은 다르다.

Contents