트랜스포머 — 어텐션을 손으로 계산한다 · 교차 엔트로피와 퍼플렉시티 · 실습

손실 하나를 손으로 만든다

목표

로짓에서 출발해 "이 모델이 얼마나 잘한다" 를 숫자 하나로 만드는 과정을 표준 라이브러리만으로 직접 짠다. 안정적인 log-softmax, 한 자리의 손실, 여러 자리의 평균인 교차 엔트로피, 지수로 되돌린 퍼플렉시티까지 올라간 뒤, 라벨을 한 칸 미는 일과 패딩 자리를 빼는 일이 그 숫자를 얼마나 바꾸는지 나란히 놓고 잰다. 마지막으로 밑이 2 인 로그로 옮겨 비트/토큰으로도 읽는다.

왜 중요한가

학습도 평가도 이 숫자 하나를 보고 움직인다. 그런데 이 숫자를 만드는 과정에는 오류를 내지 않으면서 조용히 틀리는 자리가 네 곳 있다. 로그를 언제 취했는가, 라벨을 한 칸 밀었는가, 패딩을 뺐는가, 로그의 밑이 무엇인가. 넷 다 코드로는 한 줄이고, 틀려도 예외가 나지 않으며, 대개 숫자가 좋아지는 쪽으로 틀린다. 그래서 의심할 계기가 없다.
이 실습은 실제 모델을 부르지 않는다. 이 파드의 시스템 파이썬에는 numpy·torch·transformers 가 없다. 대신 자리마다 어휘 전체에 대한 점수 한 줄을 결정적으로 만들어 두고 그 위에서 같은 계산을 손으로 짠다. 그래서 "어느 모델의 퍼플렉시티는 얼마" 같은 말은 여기서 하지 않는다. 나오는 숫자는 전부 여러분이 만든 자료로 잰 것이다.
이웃 모듈이 분포에서 하나를 고르는 법(온도·top-k·top-p)이라면 여기는 그 분포가 얼마나 틀렸는지 재는 법이다. 고르기 전에 재는 일이 먼저다.
채점기는 여러분이 적어 둔 설명을 믿지 않는다. 여러분의 모듈을 실제로 불러 매번 다른 로짓으로 함수를 직접 두드려 보고, 채점기가 따로 계산한 값과 대조한다. 입력은 실행마다 바뀌므로 값을 외워 넣을 수 없다.

단계

1. /root/work/tf-loss/loss.py 에 VOCAB·PAD_ID·SEQ·dataset() 과 log_softmax(xs) 를 만드세요. 확률을 거치지 않고 로짓에서 바로 로그 확률로 갑니다.
2. NEG_INF 와 naive_log_softmax(xs) 를 더해 일부러 틀린 순서의 판을 만드세요. 확률을 먼저 구한 뒤 로그를 취해 바닥에서 -inf 가 나는 것을 재현합니다.
3. token_loss(logits, target) 을 더해 한 자리의 손실을 재게 하세요. 정답 토큰의 로그 확률을 뒤집은 값입니다.
4. cross_entropy(rows, targets) 를 더해 여러 자리의 손실을 평균 내게 하세요.
5. perplexity(rows, targets) 와 uniform_perplexity(vocab_size, length) 를 만드세요. 균등 분포에서 퍼플렉시티가 어휘 크기와 같아지는 것을 확인합니다.
6. shift_pairs(rows, ids)·shifted_loss(rows, ids)·unshifted_loss(rows, ids) 를 만들어 라벨을 한 칸 밀었을 때와 안 밀었을 때를 나란히 재세요.
7. kept_positions(targets, pad_id) 와 masked_cross_entropy(rows, targets, pad_id) 를 만들어 패딩 자리를 빼고 재게 하세요. 분자와 분모 둘 다에서 뺍니다.
8. bits_per_token(loss)·nats_per_token(bits) 를 더하고, /root/work/tf-loss/loss_report.json 과 /root/work/tf-loss/loss_report.md 에 잰 값을 기록하세요.

참고

실행 계약: 채점기는 /root/work/tf-loss/loss.py 를 파이썬 모듈로 불러 VOCAB·PAD_ID·SEQ·dataset·log_softmax·NEG_INF·naive_log_softmax·token_loss·cross_entropy·perplexity·uniform_perplexity·shift_pairs·shifted_loss·unshifted_loss·kept_positions·masked_cross_entropy·bits_per_token·nats_per_token 을 직접 씁니다. 스크립트로 실행하지 않으므로 if __name__ == "__main__" 은 없어도 됩니다.
VOCAB 은 12 이상, SEQ 는 16개 이상의 토큰 번호 목록이고 모든 번호가 0 부터 VOCAB 미만이어야 합니다. 뒤의 3개 이상이 PAD_ID 이고 그 앞에는 PAD_ID 가 없어야 합니다. PAD_ID 도 0 부터 VOCAB 미만입니다.
dataset() 은 (로짓 줄 목록, 토큰 번호 목록) 을 돌려줍니다. 줄 수는 len(SEQ) 와 같고 줄마다 길이는 VOCAB 입니다. 난수를 쓰지 마세요 — 두 번 불러 같은 값이 나와야 합니다.
dataset() 의 자리 t 의 줄은 자리 t+1 에 실제로 오는 토큰이 가장 높은 점수를 받도록 만드세요. 그래야 라벨을 미는 일이 왜 필요한지가 숫자로 드러납니다. 문장이 끝난 뒤의 패딩 자리에는 점수를 더 크게 얹으세요 — 맞히기 쉬운 자리를 평균에 넣으면 값이 좋아 보인다는 것을 보이기 위해서입니다.
log_softmax(xs) 는 x_i - (max + log sum exp(x - max)) 입니다. 확률을 만들어 나눈 뒤 로그를 취하면 안 됩니다. log_softmax([0.0, -800.0]) 은 두 칸 모두 유한해야 하고 둘째 칸은 -800 근처입니다.
naive_log_softmax(xs) 는 반대로 확률을 먼저 만든 뒤 로그를 취합니다. math.log(0.0) 은 예외를 던지므로 확률이 0.0 인 자리는 직접 NEG_INF 로 채우세요. 같은 입력에서 log_softmax 는 유한한데 이쪽만 -inf 가 나는 것이 이 단계의 요점입니다.
token_loss(logits, target) 은 -log_softmax(logits)[target] 입니다. 부호를 뒤집는 것을 잊지 마세요.
cross_entropy(rows, targets) 는 합이 아니라 평균입니다. targets 가 비면 0.0 을 돌려주세요.
perplexity(rows, targets) 는 exp(평균 손실) 입니다. 자리마다 exp 를 취해 평균 내는 것이 아닙니다 — 균등 분포에서는 두 값이 우연히 같아지므로 그것만으로는 구별되지 않습니다.
uniform_perplexity(vocab_size, length) 는 모든 점수가 같은 로짓 줄 length 개를 만들어 퍼플렉시티를 잽니다. 결과는 vocab_size 와 같아야 합니다.
shift_pairs(rows, ids) 는 (rows[:-1], ids[1:]) 입니다. 로짓은 뒤를 하나 버리고 토큰은 앞을 하나 버립니다. shifted_loss 는 그 짝으로 잰 교차 엔트로피이고, unshifted_loss 는 밀지 않고 rows 와 ids 를 그대로 넘겨 잰 값입니다. 여러분의 자료에서는 민 쪽이 더 작아야 합니다.
masked_cross_entropy(rows, targets, pad_id) 는 targets[i] != pad_id 인 자리만 더하고 그 자리 수로 나눕니다. 전체 길이로 나누면 값이 조용히 작아집니다. 남는 자리가 없으면 0.0 입니다.
8단계 보고서는 dataset() 을 한 번 불러 얻은 한 벌로 잽니다. shift_pairs 로 민 짝에서 masked_loss 를 재고, 밀지 않은 값은 unshifted_loss 로 잽니다. bits_per_token 은 masked_loss 를 기준으로 냅니다. probe_gap 은 800 으로 고정이고, naive_is_inf 는 naive_log_softmax([0.0, -800.0])[1] 이 -inf 인지, stable_logprob 은 log_softmax([0.0, -800.0])[1] 입니다.
이 파드에는 인터넷이 없습니다. pip install 은 되지 않고 시스템 파이썬에는 numpy·torch·transformers 가 없습니다. numpy 는 /opt/onnx-lab/bin/python 안에만 있습니다. import math 하나면 충분합니다.
공식 문서: [Attention Is All You Need](https://arxiv.org/abs/1706.03762) · [Python — math](https://docs.python.org/3/library/math.html) · [Python — statistics](https://docs.python.org/3/library/statistics.html) · [Hugging Face — Text generation](https://huggingface.co/docs/transformers/en/llm_tutorial)
흔한 실수: 확률을 먼저 만들고 로그 취하기, 손실의 부호를 안 뒤집기, 평균 대신 합 쓰기, 자리마다 exp 를 취해 평균 내기, 라벨을 반대로 밀기, 마스킹에서 분모를 안 고치기, 밑이 2 인 로그와 자연로그를 섞어 적기.

단계 8개

로짓에서 바로 로그 확률로
일부러 무너뜨려 본다
한 자리의 손실
여러 자리의 평균
퍼플렉시티의 눈금
라벨은 한 칸 밀려 있다
패딩을 빼고 잰다
네 숫자를 나란히 놓는다