"LLM 핵심 학습 (5/6) — 수학 직관: Softmax·CE·KL·Gradient·LayerNorm"

신경망의 학습은 다섯 식이 차례차례 깨지지 않고 흐르는 과정이다. softmax가 포화되지 않고, cross-entropy 그래디언트가 살아 있으며, KL이 의미 있는 거리 신호를 주고, 그래디언트가 소실/폭주하지 않고, LayerNorm이 분산을 잡아 주어야 한다. 이번 편은 각 식의 직관·유도·실패 양상을 끝까지 풀어 본다.


0. 학습 목표

  • Softmax의 형태와 온도가 분포에 미치는 영향을 도식으로 그린다.
  • Cross-Entropy의 그래디언트가 \(p - y\)로 깔끔하게 떨어지는 이유를 유도한다.
  • KL divergence가 비대칭인 이유와 그것이 학습에 미치는 영향을 안다.
  • Perplexity가 cross-entropy의 지수라는 점을 이해한다.
  • Vanishing/Exploding Gradient의 원인과 잔차·정규화·초기화 세 가지 처방을 설명한다.
  • LayerNorm·RMSNorm의 식을 직접 적는다.

1. 핵심 요약

  • Softmax: \(\mathrm{softmax}(z)_i = e^{z_i}/\sum_j e^{z_j}\). 온도 \(T\)는 \(z/T\)로 들어간다.
  • Cross-Entropy + softmax의 그래디언트 = \(p - y\) (정답과 예측 분포의 차).
  • KL\((p \| q)\) = \(\sum p \log(p/q)\). \(p \to q\)의 질문 분포 정렬.
  • Vanishing/Exploding: 곱이 길게 누적될 때 분산이 \(0\) 또는 \(\infty\)로 쏠리는 문제. 잔차 + LayerNorm + He 초기화가 함께 해결.
  • LayerNorm: 토큰 별로 평균·분산 정규화. RMSNorm: 평균을 빼지 않는 더 단순한 변형.

2. Softmax — 분포를 만드는 한 줄 식

2.1 정의

벡터 \(z \in \mathbb{R}^{K}\)에 대해

$$ \mathrm{softmax}(z)_i = \frac{e^{z_i}}{\sum_{j=1}^{K} e^{z_j}}. $$

  • 출력은 항상 \([0, 1]\) 사이, 합 1 → 확률 분포.
  • 단조 변환 불변: \(z_i \to z_i + c\)에 대해 결과 불변. 수치 안정 트릭으로 \(z\)에서 \(\max\)를 빼는 것이 표준.

2.2 온도와 포화

\(z\)를 \(z/T\)로 바꾸면

  • \(T \to 0^+\): 분포가 \(\arg\max\)에 집중 → 결정론(원-핫).
  • \(T \to \infty\): 모든 클래스가 \(1/K\)로 균등 → 정보 없음.
  • 그래디언트 관점: 출력이 너무 1 또는 0에 가까우면 포화되어 그래디언트가 0에 수렴. 이것이 1편에서 \(\sqrt{d_k}\) 스케일링이 필요했던 이유.

2.3 그래디언트

표준 결과: \(p = \mathrm{softmax}(z)\)일 때

$$ \frac{\partial p_i}{\partial z_j} = p_i (\delta_{ij} - p_j). $$

곱셈 안정성: 입력 분산이 너무 크면(\(z_i\)가 매우 크면) 한 \(p_i\)가 1에 가까워져 \(p_i (1 - p_i) \to 0\). 포화. softmax 직전에 분산을 잡는 정규화/초기화가 중요한 이유.

2.4 코드

import torch

def softmax_stable(z: torch.Tensor, dim: int = -1) -> torch.Tensor:
    z = z - z.max(dim=dim, keepdim=True).values
    e = z.exp()
    return e / e.sum(dim=dim, keepdim=True)

3. Cross-Entropy — 분류·언어 모델링의 표준 손실

3.1 정의

정답이 원-핫 \(y\), 예측이 분포 \(p\)일 때

$$ \mathrm{CE}(y, p) = -\sum_i y_i \log p_i = -\log p_{i^*}. $$

(여기서 \(i^*\)는 정답 인덱스.)

3.2 Softmax + CE의 그래디언트

\(z\)에 대한 손실 \(\mathcal{L} = -\log p_{i^*}\)의 그래디언트는

$$ \frac{\partial \mathcal{L}}{\partial z_j} = p_j - y_j. $$

깨끗한 결과 덕분에 신경망은 softmax-CE 조합을 표준으로 채택한다. 유도 과정:

$$ \mathcal{L} = -\log \frac{e^{z_{i^*}}}{\sum_j e^{z_j}} = -z_{i^*} + \log \sum_j e^{z_j}. $$

$$ \frac{\partial \mathcal{L}}{\partial z_j} = -\delta_{ji^*} + \frac{e^{z_j}}{\sum_k e^{z_k}} = p_j - y_j. \ \square $$

3.3 LLM에서

언어 모델은 위치마다 분류기다. 각 위치 \(t\)에서 다음 토큰 분포 \(p_t\)와 정답 \(y_t\) (다음 토큰 원-핫)의 CE를 합산.

$$ \mathcal{L}_{\text{LM}} = -\frac{1}{T} \sum_{t=1}^{T} \log p_t(y_t). $$

평균이 토큰당 평균 음의 로그 우도이고, 이것이 perplexity의 입력이다.

3.4 Label Smoothing

원-핫 정답이 너무 날카로워 과적합을 부를 수 있다. \(\epsilon = 0.1\) 정도로 부드럽게:

$$ \tilde{y}_j = (1-\epsilon) y_j + \epsilon/K. $$

CE는 \(\mathrm{CE}(\tilde{y}, p)\)로 변경. 출력 분포의 과신을 억제. 단점은 분포 보정(calibration)이 약간 어긋날 수 있음.


4. Perplexity

4.1 정의

$$ \mathrm{PPL} = \exp\!\left( -\frac{1}{T} \sum_{t=1}^{T} \log p_t(y_t) \right) = \exp(\mathcal{L}_{\text{LM}}). $$

해석: 모델이 정답 토큰을 고를 때 평균적으로 몇 개의 선택지 사이에서 헤매는가. PPL 10이면 모델이 10개 후보 중에서 정답을 고른 것과 같다.

4.2 한계

  • 토큰화에 종속. 어휘가 다르면 비교 불가.
  • 외부 사실성·논리성을 평가하지 않음. 분포 적합도일 뿐.

5. KL Divergence

5.1 정의

두 분포 \(p, q\)에 대해

$$ \mathrm{KL}(p \,\|\, q) = \sum_i p_i \log \frac{p_i}{q_i} = \mathbb{E}_{p}\!\left[ \log \frac{p}{q} \right]. $$

  • 비음수: \(\ge 0\), 등호는 \(p = q\)에서.
  • 비대칭: \(\mathrm{KL}(p\|q) \neq \mathrm{KL}(q\|p)\).
  • \(p\)가 0이 아닌 곳에서 \(q\)가 0이면 발산.

5.2 CE와의 관계

$$ \mathrm{CE}(p, q) = H(p) + \mathrm{KL}(p \,\|\, q). $$

\(H(p)\)가 정답 분포의 엔트로피로 상수(원-핫이면 0)이므로, CE 최소화는 KL 최소화와 동등.

5.3 LLM 응용

  • Distillation (2편): KL(teacher\(\|\)student)을 손실로 사용.
  • RLHF·DPO: 정책 \(\pi_\theta\)가 기준 정책 \(\pi_{\text{ref}}\)에서 너무 멀어지지 않도록 KL 패널티를 더한다.

$$ \mathcal{L}_{\text{RLHF}} = -\mathbb{E}[r(y)] + \beta \, \mathrm{KL}(\pi_\theta \,\|\, \pi_{\text{ref}}). $$

5.4 Reverse KL vs Forward KL

  • Forward \(\mathrm{KL}(p \| q)\): \(q\)가 \(p\)의 모든 모드를 덮어야 패널티 작음. mode-covering.
  • Reverse \(\mathrm{KL}(q \| p)\): \(q\)가 한 모드에만 집중해도 됨. mode-seeking.
  • 분포가 멀티모달인 경우 둘이 매우 다른 최적해를 준다.

6. Vanishing & Exploding Gradients

6.1 정의

깊은 네트워크에서 손실의 그래디언트가 층을 거슬러 갈수록 0에 가까워지거나(소실), 폭주(exploding)하는 현상.

6.2 원인

\(L\)층 네트워크의 그래디언트는 야코비 행렬의 곱으로 표현된다.

$$ \frac{\partial \mathcal{L}}{\partial h^{(0)}} = \prod_{\ell=1}^{L} \frac{\partial h^{(\ell)}}{\partial h^{(\ell-1)}} \cdot \frac{\partial \mathcal{L}}{\partial h^{(L)}}. $$

각 야코비 행렬의 특이값이 1 미만이면 곱은 \(0\)으로, 1 초과면 \(\infty\)로.

6.3 처방 1 — 잔차 연결 (He 2016)

\(h^{(\ell)} = h^{(\ell-1)} + f^{(\ell)}(h^{(\ell-1)})\)로 두면 야코비가 \(I + f'\)이 되어 기본값이 항등. 깊이가 늘어도 그래디언트가 통과한다. ResNet의 핵심 아이디어이자 트랜스포머의 표준 구성.

6.4 처방 2 — 정규화 (Layer/RMSNorm)

활성화의 분산을 \(\mathbb{E}[h^2] = 1\) 부근으로 강제. 다음 절.

6.5 처방 3 — 초기화

He 초기화(He 2015): \(W \sim \mathcal{N}(0, 2/n_{\text{in}})\) (ReLU 가정). Xavier 초기화: \(\mathcal{N}(0, 1/n_{\text{in}})\). 트랜스포머는 보통 \(\sigma = 0.02\)의 작은 분산을 쓰고, 추가로 깊이에 따라 출력 가중치를 \(1/\sqrt{2L}\)로 다운스케일하는 LLM 트릭이 흔하다.


7. LayerNorm & RMSNorm

7.1 LayerNorm (Ba 2016)

토큰 단위로 평균·분산을 정규화한 뒤 학습된 스케일·시프트 \(\gamma, \beta\)를 적용.

$$ \mu = \frac{1}{d}\sum_i h_i,\ \ \sigma^2 = \frac{1}{d}\sum_i (h_i - \mu)^2, $$ $$ \hat{h}_i = \frac{h_i - \mu}{\sqrt{\sigma^2 + \epsilon}},\ \ \mathrm{LN}(h)_i = \gamma_i \hat{h}_i + \beta_i. $$

배치 통계가 아니라 각 샘플의 마지막 축에서 통계를 잡으므로 배치 크기와 무관. 토큰 단위 모델인 트랜스포머에 자연스러움.

7.2 Pre-Norm vs Post-Norm

  • Post-Norm (원논문): \(\mathrm{LN}(h + f(h))\). 큰 모델에서 학습 불안정.
  • Pre-Norm: \(h + f(\mathrm{LN}(h))\). 현재 표준. 그래디언트가 잔차 경로로 안정적으로 통과.

7.3 RMSNorm (Zhang & Sennrich 2019)

평균을 빼지 않고 RMS로만 정규화.

$$ \mathrm{RMS}(h) = \sqrt{\frac{1}{d}\sum_i h_i^2 + \epsilon},\ \ \mathrm{RMSNorm}(h)_i = \gamma_i \frac{h_i}{\mathrm{RMS}(h)}. $$

연산이 약 30% 적고, LayerNorm 대비 성능 손실이 거의 없어 LLaMA·Mistral·Qwen 등 대다수 최신 LLM의 표준.

7.4 다이어그램 — 트랜스포머 한 블록 (Pre-Norm)

diagram-1

8. AdamW 옵티마이저 — 짧은 직관

표준 식(2편 3장 참조):

$$ m_t = \beta_1 m_{t-1} + (1-\beta_1) g_t,\ \ v_t = \beta_2 v_{t-1} + (1-\beta_2) g_t^2. $$ $$ \hat{m}_t = m_t/(1-\beta_1^t),\ \ \hat{v}_t = v_t/(1-\beta_2^t). $$ $$ \theta_t = \theta_{t-1} - \eta \frac{\hat{m}_t}{\sqrt{\hat{v}_t} + \epsilon} - \eta \lambda \theta_{t-1}. $$

  • 모멘텀 \(m\): 그래디언트의 지수이동평균. 노이즈 감쇠.
  • 분산 \(v\): 그래디언트 제곱의 이동평균. 큰 그래디언트는 작게, 작은 그래디언트는 크게 → 축별 스케일 보정.
  • W: weight decay를 그래디언트에 합치지 않고 분리 적용. L2 정규화와 미세하게 다르고, 트랜스포머 학습에서 더 안정.

LLM 학습에서 흔한 설정: \(\beta_1 = 0.9, \beta_2 = 0.95, \epsilon = 10^{-8}, \lambda = 0.1, \eta = 3 \times 10^{-4}\) (cosine warmup-decay).


9. 정리된 결론 — 보지 않고 답해 보라

이 글이 답한 핵심 질문 5개. 먼저 답을 가린 채로 한 줄로 답해 본 뒤 정답과 비교.

Q1. softmax + cross-entropy의 그래디언트가 왜 \(\hat{p} - y\) (one-hot)인지?

정답 softmax의 도함수와 cross-entropy의 \(\log\)가 합성될 때 깔끔하게 cancel out — 체인 룰 결과가 정확히 \(\hat{p} - y\). 미분 시 분모의 \(\sum e^{z_j}\) 가 사라지는 마법. 이 깔끔함이 softmax + CE 페어 가 분류 표준이 된 이유. PyTorch의 cross_entropy()log + softmax + NLL 을 한 번에 하는 것도 수치 안정성 + 같은 단순화 활용. (본문 §2·§3)

Q2. KL Divergence가 비대칭 이라는 사실의 직관은?

정답 \(\mathrm{KL}(P \| Q) \neq \mathrm{KL}(Q \| P)\). \(\mathrm{KL}(P \| Q)\)는 P 기준 으로 측정 — P가 큰 곳에서 Q가 작으면 비용 큼. 반대는 다름. Forward KL (\(P\) data, \(Q\) model): mass-covering — Q가 P의 모든 모드를 덮으려 함. Reverse KL (\(Q\) data, \(P\) model, VI에서 사용): mode-seeking — Q가 P의 한 모드 에 집중. RLHF의 KL penalty는 reverse 방향 사용. (본문 §5)

Q3. Perplexity와 cross-entropy의 관계를 한 줄로?

정답 Perplexity = \(e^{\mathrm{CE}}\) 또는 \(2^{\mathrm{CE}}\) (밑 log 따라). cross-entropy가 낮을수록 perplexity 낮을수록 모델 우수. Perplexity의 직관: "모델이 다음 토큰을 몇 개 중에서 고민하는가". PPL=10이면 평균 10개 후보 중에서 고민. LLM 평가의 기본 메트릭. (본문 §4)

Q4. 잔차 연결 (residual connection) 이 그래디언트 흐름을 어떻게 살리나? LayerNorm·RMSNorm 식 차이는?

정답 Residual: \(y = F(x) + x\). 역전파 시 \(\frac{\partial y}{\partial x} = \frac{\partial F}{\partial x} + I\) → 항등 경로가 항상 살아있어 vanishing 차단. LayerNorm: \(\frac{x - \mu}{\sqrt{\sigma^2 + \epsilon}} \cdot \gamma + \beta\). RMSNorm: \(\frac{x}{\sqrt{\text{RMS}(x)^2 + \epsilon}} \cdot \gamma\) — 평균 제거 단계 생략. Residual: ResNet (He 2016)의 핵심 발견 + 모든 트랜스포머 블록의 표준. RMSNorm은 LayerNorm에서 평균 제거를 빼서 계산 절감 — 분포 가정에 따른 단순화 (LLaMA 채택). (본문 §6·§7)

Q5. AdamW의 weight decay와 L2 정규화의 차이는?

정답 L2: 손실에 \(\lambda \|w\|^2\) 항을 더해 그래디언트와 함께 적응 됨. Weight decay: \(w \leftarrow w - \eta \lambda w\) 를 옵티마이저 단계에서 직접 적용 — 적응 모멘트와 분리. Adam의 적응 학습률이 L2 정규화를 희석 하는 문제 발견 (Loshchilov & Hutter 2017). AdamW는 weight decay를 분리하여 진짜 정규화 효과를 보장. 트랜스포머 학습의 표준. (본문 §8)


5개를 한 줄로 답할 수 있었다면 LLM 수학의 softmax-CE · KL · Perplexity · 잔차 · 정규화 핵심이 자리잡은 것이다.


10. 실습 과제

  1. NumPy로 \(K=5\) softmax-CE 분류기를 만들고, 정답이 인덱스 2일 때 그래디언트가 \(p - y\)임을 확인한다.
  2. 같은 \(p, q\)에 대해 \(\mathrm{KL}(p\|q)\)와 \(\mathrm{KL}(q\|p)\)를 계산해 비대칭을 수치로 본다.
  3. 단순 30층 MLP를 Pre-Norm·Post-Norm으로 학습하고 그래디언트 노름의 층별 분포를 비교한다.

11. 추가 학습

  • Goodfellow, Bengio, Courville, 2016, Deep Learning. 6장(Optimization), 8장(Regularization).
  • He et al., 2015, Delving Deep into Rectifiers. arXiv:1502.01852. He 초기화.
  • He et al., 2016, Deep Residual Learning. arXiv:1512.03385. 잔차 연결.
  • Ba et al., 2016, Layer Normalization. arXiv:1607.06450.
  • Zhang & Sennrich, 2019, RMSNorm. arXiv:1910.07467.
  • Loshchilov & Hutter, 2019, Decoupled Weight Decay. arXiv:1711.05101. AdamW.
  • Kullback & Leibler, 1951, On Information and Sufficiency. KL divergence 원전.
  • 3Blue1Brown, Neural networks 시리즈: 시각적 직관에 가장 좋은 자료.

이 글은 LLM 핵심 학습 시리즈의 5/6 편이다. 다음 편(6/6)은 학습 로드맵 — 단계별 코스·논문 10편·저장소 10선으로 마무리된다.

시리즈 전체 안내: 시리즈 목차

댓글

이 블로그의 인기 게시물

"LLM 핵심 학습 (1/6) — 기본: 토큰화·임베딩·어텐션·위치 인코딩"

"LLM 핵심 학습 (2/6) — 파인튜닝: LoRA·QLoRA·증류·Adapter"

"ML 기초 학습 (1/9) — 머신러닝과 sklearn: 학습의 좌표계"