Transformer의 출력을 교체: CE를 사용해도 정확도는 같지만 MSE에서는 0.19 하락
요약
본 글은 Transformer 모델의 출력 레이어에서 손실 함수를 Cross Entropy(CE) 대신 Mean Squared Error(MSE)로 변경했을 때 발생하는 성능 변화를 분석합니다. 정확도는 유사했으나, CE 대비 MSE가 0.19 악화되는 현상을 발견했습니다. 이는 정답 문자에 낮은 확률을 할당하는 위치의 특성 차이 때문으로 추정됩니다.
핵심 포인트
- 손실 함수 변경 시 정확도 하락은 미미함.
- CE에서 MSE로 변경 시 val CE가 0.19 악화됨.
- MSE는 자신이 최소화하던 값에서도 CE에 뒤처짐.
- 출력 레이어의 손실 함수 선택이 모델 성능에 영향을 줄 수 있음.
서론
이 글은 Transformer를 0부터 구축하는 8주 학습 계획인 'Transformer0-1'의 4주차 기록입니다. W2에서는 PyTorch로 문자 단위의 작은 GPT를 구현했고, W3에서는 의도적으로 학습을 망가뜨려 증상과 원인의 대응표를 만들었습니다. W4에서는 모델 본체(Block)는 건드리지 않고 출력 부분만 교체했습니다.
결론
- 손실 함수를 CE에서 MSE로 변경하자, val의 정확도는 46.3% → 45.4%로 거의 같았지만, val CE는 1.82 → 2.01로 0.19 악화되었습니다.
- 그 원인은 '자신 있게 제외한 위치'였습니다. 정답 문자에 0.1% 미만의 확률만 할당했던 위치가 MSE에서는 CE보다 2.3배 많았습니다.
- MSE는 자신이 최소화하던 val MSE에서도 CE에 뒤처졌습니다.
- 학습률을 10배로 늘려도 MSE의 차이는 줄어들지 않았습니다.
- 온도 T는 가중치를 변경하지 않고 정확성과 다양성의 배분만 움직였습니다. T=0.5일 때 실존 단어율은 90%, 단어 중복률은 53%였고, T=2.0일 때는 실존 단어율이 20%, 중복률이 6%가 되었습니다.
변경한 부분
모델은 W2의 GPTMini(어휘 65 문자, d_model=64, block_size=32, 4층, 4 head)를 사용했고, 데이터는 tiny_shakespeare입니다. Block을 통과한 후의 출력은 다음과 같은 형태입니다.
Block × 4 의 출력 (B, T, 64)
↓ 최종 LayerNorm
↓ lm_head : Linear(64 → 65)
...
실험1은 학습 측면이므로 가중치가 변경됩니다. 실험2는 생성 측면이므로 가중치는 변경되지 않습니다.
실험1: 손실 함수를 CE에서 MSE로 변경하기
조건
AdamW(lr=1e-3), batch_size=32, 3000 iter, seed=0으로, 손실 함수만 변경했습니다.
| tag | 손실 함수 | 학습률 |
|---|---|---|
| ce | CE | 1e-3 |
| ... | ||
MSE를 F.mse_loss |
의 상태로 사용하면 65 문자 각각에 대한 평균이 되고, 기울기가 추가로 1/65가 됩니다. 기울기의 '형태' 차이와 '스케일' 차이가 섞이기 때문에, 65 문자는 합계로 했습니다.
ce의 최종 val CE는 1.8223으로, W2의 4 head 모델(1.8205)을 재현했습니다.
### 결과
학습 후 가중치를 val 전체를 나누어 모두 사용했고, 6가지 기준으로 채점했습니다. p는 정답 문자에 할당된 확률입니다.
| val CE | val MSE | 정확도 | p 중앙값 | p<0.01 | p<0.001 |
|---|---|---|---|---|---|
| ce | 1.8157 | 0.6687 | 46.29% | 0.2165 | 5.26% | 0.77% |
| ... |
그림의 세로축은 다음과 같습니다.
- **val CE**: val의 각 위치에서 계산된 −log p의 평균. 보지 못한 문장을 한 글자씩 읽게 했을 때, 모델이 다음 문자에 평균적으로 얼마나 놀랐는지를 나타냅니다. 낮을수록 좋습니다 -
- **val top-1 accuracy**: 1위에 예측한 문장이 정답과 일치하는 위치의 비율. 맞았는지 틀렸지만 보는 것만으로 판단합니다 -
- **lm_head grad norm**: 출력층(lm_head) 가중치의 기울기 크기. 손실이 그 시점에서 출력층을 얼마나 강하게 누르고 있는지를 나타냅니다. ce와 mse는 자릿수가 다르므로 로그 축으로 했습니다.
val CE를 e^(val CE)로 변환하면, ce는 평균 약 6.2 문자, mse는 약 7.5 문자의 후보군에서 고민하는 상태가 됩니다. 학습 전(65 문자가 균등할 때)은 65 문자였습니다.
### 정확도는 거의 같은데 val CE만 차이가 나는 경우
정확도는 '1위가 맞았는지'만 봅니다. val CE는 −log p의 평균이므로, 정답에 할당된 확률의 크기를 봅니다.
p<0.001 위치는 mse가 ce보다 2.3배, p<0.01은 1.5배가 되었습니다. 극단적인 이상치일수록 배율이 큽니다. p<0.001 위치는 하나당 −log(0.001) = 6.9 이상을 CE에 더하기 때문에, 이 약 1%의 차이만으로도 val CE의 차이 0.19 중 적어도 약 0.07을 차지합니다.
mse는 1위를 맞힐 능력은 ce와 거의 같지만, 틀렸을 때 정답 문자에 거의 0에 가까운 확률밖에 남기지 않은 위치가 많습니다.
### 왜 MSE는 자신 있게 놓친 위치를 바로잡지 못하는가
logits에 대한 기울기(gradient)는 CE와 MSE에서 다음과 같은 형태가 된다.
| 손실 | 정답 logit에 대한 기울기 |
|---|---|
| CE | p − 1 |
| MSE | 2p × [(p − 1) − Σ p_i(p_i − y_i)] |
MSE의 기울기에는 softmax 미분에서 오는 p가 곱해져 있다. 하나의 배치에 다음 두 곳이 있다고 가정한다.
| 위치 | 상태 | 정답 확률 p | CE 기울기 | MSE 기울기 |
|---|---|---|---|---|
| A | 적당히 벗어남 | 0.3 | −0.70 | −0.59 |
| B | 자신 있게 벗어남 | 0.001 | −1.00 | −0.004 |
CE는 심하게 벗어난 B 쪽을 강하게 바로잡는다. MSE에서는 B의 기울기가 A의 약 1/150이 되어, B는 거의 방치된다.
학습 전에는 65 글자가 거의 균등(p ≈ 1/65)하므로, MSE의 기울기는 CE의 약 2/65 ≈ 1/32가 되어야 한다. 실제 측정한 iter 0의 lm_head 기울기 노름은 ce 0.69, mse 0.024로, 약 1/29가 되었다.
### val CE는 CE에 유리한 자(尺)가 아닌가
val CE는 ce가 학습 중에 최소화했던 식 그 자체이다. 그래서 mse가 최소화했던 val MSE로도 채점해 보았다. 결과는 ce 0.6687, mse 0.6815로, mse는 자신의 자(尺)로도 ce에 패했다. 차이는 자의 유리함/불리함만으로는 설명할 수 없다.
CE와 MSE(확률에 대한 MSE는 Brier 스코어라고도 불린다) 모두 진정한 확률 분포를 출력했을 때 최소가 된다. 시간과 용량이 무한하다면 같은 분포에 도달해야 한다. 차이가 난 것은, 제한된 학습 시간 속에서 CE 쪽이 진정한 분포에 더 빨리 접근했기 때문이라고 생각된다.
### 학습률을 높이면 MSE도 따라잡을까
초기 기울기가 1/29였다면, 학습률을 높여서 보충할 수 있다는 가설을 세웠다. 학습률을 10배로 한 mse_lr*10의 val CE는 2.0120으로, mse와 거의 같아졌다. 차이는 줄지 않았다.
이유는 AdamW에 있다. AdamW는 기울기의 평균을 기울기의 제곱 평균의 제곱근으로 나누어 업데이트량을 결정한다. 기울기가 모두 c배가 되어도 분자와 분모의 c가 상쇄되기 때문에, 전체 스케일의 작음은 처음부터 보충되어 있었다. iter 0에서 기울기가 29배 달랐는데도, iter 300에서의 val CE 차이가 0.11로 끝난 것은 이 때문이다.
반면 AdamW가 나누는 것은 배치 내 모든 위치의 기울기를 합산한 값의 크기뿐이다. 합산 속에서 A의 목소리가 B의 150배 큰 배분은 변하지 않는다. 학습률을 높여도, A를 바로잡는 방향으로 크게 나아갈 뿐, B는 방치된 채로 남게 된다.
이 결과로부터, AdamW 하에서의 mse_lr*10은 '스케일을 보충하는' 실험이 아니라 '보폭을 크게 하는' 실험이었음도 알 수 있었다.
## 실험 2: 온도 T를 변경하기
### 조건
W2의 학습된 가중치를 불러와, 줄 바꿈부터 1000 글자를 생성했다. T 외에는 모두 같았고, 어떤 T에서도 같은 난수열을 사용했다.
### 결과
| T | 평균 엔트로피(bit) | 실질적 후보 수(2^H) | 단어 중복률 | 실재 단어율 |
|---|---|---|---|---|
| 0.5 | 1.313 | 약 2.5 | 0.527 | 0.896 |
| ... |
단어 중복률은 1 − 다른 단어 수 / 총 단어 수, 실재 단어율은 생성된 단어 중 train에 나오는 단어의 비율이다.
T=0.5의 시작 부분
That make the prester our and hange in the strenge
That you than her shall the see thinks of his stants, and away,
That they are that we mother of well.
T=2.0의 시작 부분
thyremhmadk'swes:? fildHad neamany,
Nding
let, grecourewy his tweetre hath Doig Oureliess. I him;sleat'ory? raius,
### T로 확률이 변하는 이유
softmax는 로짓(logits)의 차이만 본다. 두 문자 확률의 비율은 exp((z_i − z_j) / T)가 된다. 로짓 차이가 2인 두 문자의 경우, 확률 비율은 T=0.5에서 약 55배, T=1.0에서 약 7.4배, T=2.0에서 약 2.7배가 된다. T는 로짓의 차이를 늘리거나 줄이는 조절 장치(knob)이며, 순위는 바꾸지 않고 할당만 바꾼다.
T를 높이면 후보가 약 2.5 문자에서 약 16.7 문자로 넓어지고, 확률이 낮은 문자도 선택된다. 그 결과, 같은 단어를 따라 하는 경향이 사라졌고(중복률 0.53 → 0.06), 존재하지 않는 단어가 늘었다(실재단어율 0.90 → 0.20). 두 수치는 별개의 현상이 아니라, 분포가 넓어졌다는 것의 두 가지 측면이다.
### T로 제거할 수 없는 것들
T=0.5에서도 `tescrands`, `soffather`와 같은 존재하지 않는 단어가 섞였다. T는 확률 할당만 바꿀 뿐이므로, 모델이 정답에 확률을 부여하지 못한 부분은 고칠 수 없다. 생성의 품질 상한선은 가중치(weight)로 결정된다.
## 요약
| 변경한 것 | 가중치 | 알게 된 점 |
|---|---|---|
| 실험1 | 손실 함수 | 달라짐 | CE가 더 좋은 확률 분포를 만든다. MSE는 자신 있게 틀린 위치를 고칠 수 없다 |
| 실험2 | 온도 T | 변하지 않음 | 정확성과 다양성의 할당을 움직일 뿐이다. 상한선은 가중치로 결정된다 |
W2~W3에서 당연하게 사용되던 CE가 표준인 이유를, 기울기(gradient)의 형태와 p의 분포 양쪽에서 확인할 수 있었다.
## 참고 문헌
### 손실 함수 (실험1)
- Goodfellow, Bengio, Courville『Deep Learning』 6.2절
- Golik, Doetsch, Ney (2013)
AI 자동 생성 콘텐츠
본 콘텐츠는 Zenn ML의 원문을 AI가 자동으로 요약·번역·분석한 것입니다. 원 저작권은 원저작자에게 있으며, 정확한 내용은 반드시 원문을 확인해 주세요.
원문 바로가기0