Output head — hidden state에서 logit으로 · Transformer Quest

마지막 block을 지난 residual stream, 그러니까 최종 hidden state의 shape는 (seq_len, d_model)이야. 예측을 만들려면 그걸 vocab에 대한 logit shape (seq_len, vocab_size)로 바꿔야 해. 이게 output head.

기계적으로는 linear projection: logits = hidden @ W_lm.T, W_lm shape는 (vocab_size, d_model). 마지막 차원 softmax가 확률. 추론 시엔 보통 마지막 위치 logit만 필요(next-token 예측), 학습 시엔 모든 위치 logit을 병렬 계산.

Weight tying

많은 모델이 weight tying을 써. input embedding 행렬과 output head가 같은 파라미터를 공유하는 거야. 수학적으로는 W_lm = E.T. GPT-2가 그렇고, 요즘 모델 계열도 작은 변종 쪽은 대체로 묶어. 이러면 vocab × d_model 파라미터를 아끼고(Llama 3 8B 크기면 524M쯤이야) 보기 좋은 대칭도 생겨. 비슷한 input embedding을 갖는 토큰이 비슷한 output logit 프로필을 갖게 되거든. 반대로 플래그십 open-weight 모델은 대체로 안 묶어. Llama 3 8B랑 Mistral 7B 둘 다 lm_head를 따로 두는데, Llama 3 8B가 7.5B가 아니라 8.03B인 이유가 바로 그거야. 더 큰 모델(GPT-3, GPT-4)도 안 묶어. 규모가 커지면 파라미터 절감이 전체에서 차지하는 비중이 작아지고, 떼어 놓는 쪽이 품질도 살짝 낫거든.

Code

weight tying을 적용한 output head·python

class TiedLM(nn.Module):
    def __init__(self, vocab_size, d_model):
        super().__init__()
        self.embed = nn.Embedding(vocab_size, d_model)
    def forward(self, ids):
        x = self.embed(ids)             # (B, L, d_model)
        # ... transformer body ...
        x = final_norm(x)
        # Reuse embedding weights as the output projection
        logits = x @ self.embed.weight.T     # (B, L, vocab_size)
        return logits

Output head — hidden state에서 logit으로

Weight tying

Code

External links

Exercise

Progress

댓글 0