LLM 수학 03 — 미분과 최적화

02편에서 cross-entropy loss를 만들었다. 이제 질문은 하나다. scalar loss 하나가 수십억 개 parameter 각각을 어느 방향으로 움직이게 하는가?

미분은 local sensitivity를 계산하고, backpropagation은 계산 graph의 chain rule을 재사용하며, optimizer는 그 gradient로 parameter를 갱신한다.

함수와 변화율

함수 $f$ 는 입력을 출력에 대응시킨다.

y=f(x)

$x$ 를 작은 양 $h$ 만큼 바꿨을 때 출력 변화율은

\frac{f(x+h)-f(x)}{h}

이다. $h\to0$ 극한이 존재하면 derivative다.

f'(x) = \lim_{h\to0}\frac{f(x+h)-f(x)}{h}

$f(x)=x^2$ 라면

\frac{(x+h)^2-x^2}{h} =2x+h

이므로 $h\to0$ 에서

f'(x)=2x

이다. $x=3$ 부근에서 $x$ 가 0.01 증가하면 $f$ 는 대략 $6\times0.01=0.06$ 증가한다.

Local linear approximation

미분의 핵심 직관은 복잡한 함수를 현재 점 근처에서 직선으로 근사하는 것이다.

f(x+\Delta x) \approx f(x)+f'(x)\Delta x

여러 parameter에서는

L(\theta+\Delta\theta) \approx L(\theta)+\nabla_\theta L^\top\Delta\theta

다. gradient 반대 방향 $\Delta\theta=-\eta\nabla L$ 을 넣으면

L(\theta+\Delta\theta) \approx L(\theta)-\eta\lVert\nabla L\rVert_2^2

가 되어 충분히 작은 step에서는 loss가 내려갈 것으로 기대한다. 큰 learning rate에서는 local approximation이 깨질 수 있다.

자주 쓰는 derivative

함수	derivative
$c$	$0$
$x^n$	$nx^{n-1}$
$e^x$	$e^x$
$\log x$	$1/x$
$\sigma(x)$	$\sigma(x)(1-\sigma(x))$
$\tanh x$	$1-\tanh^2x$

합과 scalar 배의 derivative는 항별로 계산한다.

\frac{d}{dx}[af(x)+bg(x)] =af'(x)+bg'(x)

곱과 나눗셈은

(fg)'=f'g+fg'

\left(\frac{f}{g}\right)' =\frac{f'g-fg'}{g^2}

이다.

Partial derivative

여러 입력을 가진 함수에서 한 variable만 바꾸고 나머지는 고정한다.

f(x,y)=x^2+3xy

\frac{\partial f}{\partial x}=2x+3y

\frac{\partial f}{\partial y}=3x

parameter가 여러 개이므로 neural network 미분은 대부분 partial derivative다.

Gradient

scalar-valued function $L\colon\mathbb{R}^n\to\mathbb{R}$ 의 partial derivative를 vector로 모은다.

\nabla_\theta L = \begin{bmatrix} \frac{\partial L}{\partial\theta_1}\\ \vdots\\ \frac{\partial L}{\partial\theta_n} \end{bmatrix}

gradient는 가장 빠르게 증가하는 방향이다. gradient descent는 반대 방향으로 움직인다.

각 component의 단위와 scale이 다를 수 있어 raw gradient의 크기만으로 parameter 중요도를 비교하면 안 된다.

Jacobian

vector-valued 함수 $f\colon\mathbb{R}^n\to\mathbb{R}^m$ 의 모든 output-input partial derivative를 matrix로 모은다.

J_{ij}=\frac{\partial f_i}{\partial x_j}

J_f(x) \in\mathbb{R}^{m\times n}

softmax는 vector를 vector로 보내므로 Jacobian이 있다.

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

$i=j$ — $p_i(1-p_i)$
$i\ne j$ — $-p_ip_j$

한 logit을 올리면 normalization 때문에 다른 token probability가 내려간다.

Hessian

scalar loss의 second partial derivative matrix다.

H_{ij} = \frac{\partial^2L}{\partial\theta_i\partial\theta_j}

curvature를 나타낸다. eigenvalue가 큰 방향은 loss가 급하게 휘고 작은 방향은 평평하다. full Hessian은 LLM 규모에서 너무 커 직접 만들기 어렵다. second-order optimizer와 sharpness 분석은 Hessian 또는 그 근사를 사용한다.

Chain rule

합성함수 $y=f(g(x))$ 의 derivative는

\frac{dy}{dx} = \frac{dy}{dg}\frac{dg}{dx}

다.

예를 들어

g(x)=2x+1,qquad y=g^2

이면

\frac{dy}{dg}=2g,qquad \frac{dg}{dx}=2

따라서

\frac{dy}{dx}=4g=4(2x+1)

이다.

neural network는 함수의 합성이다.

L=f_n(f_{n-1}(\cdots f_1(x)))

backpropagation은 output에서 input 방향으로 local derivative를 곱한다.

Computation graph 손계산

간단한 model을 둔다.

\hat{y}=wx+b

L=(\hat{y}-y)^2

$x=2$ , $w=3$ , $b=1$ , $y=5$ 라면 forward는

\hat{y}=3\cdot2+1=7

e=\hat{y}-y=2

L=e^2=4

이다.

backward는 역순이다.

\frac{\partial L}{\partial e}=2e=4

\frac{\partial e}{\partial\hat{y}}=1

\frac{\partial\hat{y}}{\partial w}=x=2

\frac{\partial\hat{y}}{\partial b}=1

따라서

\frac{\partial L}{\partial w}=4\cdot1\cdot2=8

\frac{\partial L}{\partial b}=4\cdot1\cdot1=4

learning rate $\eta=0.1$ 인 gradient descent 한 step은

w\leftarrow3-0.1\cdot8=2.2

b\leftarrow1-0.1\cdot4=0.6

새 prediction은 $2.2\cdot2+0.6=5$ 로 이 sample loss는 0이 된다. 다른 sample까지 좋아진다는 보장은 없다.

Backpropagation이 효율적인 이유

naive하게 parameter마다 forward를 다시 계산하면 공유되는 중간 계산을 반복한다. reverse-mode automatic differentiation은 각 node에 도착한 upstream gradient를 한 번 모으고 local derivative를 곱해 부모에게 보낸다.

node $u$ 가 여러 경로로 loss에 영향을 주면 gradient를 더한다.

\frac{\partial L}{\partial u} = \sum_k \frac{\partial L}{\partial v_k} \frac{\partial v_k}{\partial u}

residual connection $y=f(x)+x$ 에서는

\frac{\partial L}{\partial x} = \frac{\partial L}{\partial y} \left(\frac{\partial f}{\partial x}+I\right)

가 된다. identity 경로가 gradient 전달을 돕는다.

Linear layer gradient와 shape

row-vector convention으로

Y=XW

X\in\mathbb{R}^{B\times d_{in}},quad W\in\mathbb{R}^{d_{in}\times d_{out}},quad Y\in\mathbb{R}^{B\times d_{out}}

라고 하자. upstream gradient $G=\partial L/\partial Y$ 의 shape는 $Y$ 와 같다.

\frac{\partial L}{\partial W}=X^\top G \in\mathbb{R}^{d_{in}\times d_{out}}

\frac{\partial L}{\partial X}=GW^\top \in\mathbb{R}^{B\times d_{in}}

gradient shape는 미분 대상 shape와 같아야 한다. autodiff 결과가 왜 transpose를 포함하는지 matrix multiplication shape로 확인한다.

Cross-entropy backward

softmax와 one-hot cross-entropy를 합치면 logit gradient는

\frac{\partial L}{\partial z}=p-y

다. 마지막 hidden $h$ 와 vocabulary weight $W$ 에 대해 $z=hW$ 라면

\frac{\partial L}{\partial W}=h^\top(p-y)

\frac{\partial L}{\partial h}=(p-y)W^\top

오차가 vocabulary projection을 거쳐 hidden representation으로 돌아간다.

Gradient descent

\theta_{t+1} = \theta_t-\eta\nabla_\theta L(\theta_t)

$\eta$ 가 너무 작으면 느리다.
너무 크면 minimum을 넘거나 발산한다.
non-convex neural loss에서는 global minimum 보장이 없다.
training loss 감소와 validation 성능 개선은 다른 문제다.

Full-batch, stochastic, mini-batch

전체 dataset gradient는

g=\frac{1}{N}\sum_{i=1}^{N}\nabla_\theta L_i

다. LLM dataset 전체를 매 step 계산할 수 없으므로 mini-batch $B_t$ 를 사용한다.

g_t = \frac{1}{|B_t|}\sum_{i\in B_t}\nabla_\theta L_i

sample이 적절히 뽑혔다면 full gradient의 noisy estimate다. batch size는 다음을 동시에 바꾼다.

gradient variance
activation memory
accelerator utilization
optimizer step 수
learning-rate tuning 조건

gradient accumulation은 여러 micro-batch gradient를 더해 큰 effective batch를 흉내 낸다. BatchNorm 같은 batch-dependent 연산, dropout random state, loss reduction 방식 때문에 항상 물리적으로 큰 batch와 완전히 같지는 않다. Transformer의 LayerNorm은 sample 내부 feature를 normalize해 이 차이를 줄이는 편이다.

Momentum

현재 gradient만 쓰지 않고 과거 방향의 moving average를 사용한다.

v_t=\beta v_{t-1}+(1-\beta)g_t

\theta_{t+1}=\theta_t-\eta v_t

oscillation을 줄이고 일관된 방향의 이동을 누적한다. library마다 $(1-\beta)$ factor와 update 표기가 다를 수 있다.

Adam

Adam은 gradient의 first moment와 squared gradient의 second raw moment를 지수이동평균으로 추정한다.

m_t=\beta_1m_{t-1}+(1-\beta_1)g_t

v_t=\beta_2v_{t-1}+(1-\beta_2)g_t^2

초기 $m_0=v_0=0$ 의 bias를 보정한다.

\hat{m}_t=\frac{m_t}{1-\beta_1^t}

\hat{v}_t=\frac{v_t}{1-\beta_2^t}

update는 element-wise다.

\theta_{t+1} = \theta_t-eta\frac{\hat{m}_t}{\sqrt{\hat{v}_t}+\epsilon}

$m_t$ — gradient 방향을 부드럽게 한다.
$v_t$ — 최근 squared gradient가 큰 coordinate의 step을 줄인다.
$\epsilon$ — 0으로 나누는 것을 막고 numerical behavior에도 영향을 준다.

“second moment”는 variance와 같지 않다. $v_t$ 는 centered variance가 아니라 $E[g^2]$ 의 moving estimate다.

AdamW와 weight decay

L2 regularization을 loss에 넣으면

L'(\theta)=L(\theta)+\frac{\lambda}{2}\lVert\theta\rVert_2^2

\nabla L'=\nabla L+\lambda\theta

가 된다. plain SGD에서는 이 항이 multiplicative weight decay와 대응한다.

adaptive optimizer에서는 gradient에 L2 항을 넣으면 coordinate별 normalization의 영향을 받아 단순 weight decay와 같지 않다. AdamW는 decay를 gradient update와 분리한다.

\theta_{t+1} =(1-\eta\lambda)\theta_t -eta\frac{\hat{m}_t}{\sqrt{\hat{v}_t}+\epsilon}

실제 optimizer 구현의 update 순서와 parameter group 제외 항목을 확인해야 한다. bias와 normalization parameter에 weight decay를 적용하지 않는 설정이 흔하지만 보편 법칙은 아니다.

Learning-rate schedule와 warmup

constant learning rate만 쓰지 않는다.

Warmup

초기 몇 step 동안 learning rate를 0 근처에서 target까지 올린다. optimizer moment와 activation scale이 불안정한 초기에 큰 update를 피한다.

linear warmup은

\eta_t=\eta_{max}\frac{t}{T_{warmup}}

처럼 쓸 수 있다.

Cosine decay

\eta_t = \eta_{min} +\frac{1}{2}(\eta_{max}-\eta_{min}) \left(1+\cos\frac{\pi t}{T}\right)

schedule 이름만 같아도 warmup 포함 여부, final learning rate, restart 여부가 다를 수 있다.

Gradient clipping

global gradient norm이 threshold $c$ 보다 크면 scale을 줄인다.

g' = g\cdot\min\left(1,\frac{c}{\lVert g\rVert_2}\right)

방향은 유지하고 norm만 제한한다. element-wise value clipping과 다르다. exploding gradient로 한 step이 training을 망치는 것을 막지만 근본 원인을 숨길 수도 있다.

Vanishing과 exploding gradient

깊은 합성함수의 gradient는 Jacobian product다.

\frac{\partial h_L}{\partial h_0} = \prod_{l=1}^{L}\frac{\partial h_l}{\partial h_{l-1}}

각 factor의 norm이 반복해서 1보다 작으면 vanishing, 크면 exploding이 생길 수 있다. 실제 matrix product는 방향과 singular value가 관여해 scalar 곱보다 복잡하다.

도움을 주는 요소는 다음과 같다.

scale을 고려한 initialization
residual connection
normalization
적절한 activation
gradient clipping
learning-rate tuning

한 기법이 모든 원인을 해결하지는 않는다.

Activation function

linear layer만 여러 개 쌓으면 전체도 하나의 linear transformation이다.

W_2(W_1x)=W'x

nonlinearity가 있어야 복잡한 함수를 표현한다.

ReLU

\operatorname{ReLU}(x)=\max(0,x)

\operatorname{ReLU}'(x)= \begin{cases} 0,&x<0\\ 1,&x>0 \end{cases}

0에서 derivative convention은 implementation이 정한다.

GELU

Transformer에서 자주 사용한다.

\operatorname{GELU}(x)=x\Phi(x)

$\Phi(x)$ 는 standard normal CDF다. ReLU처럼 hard gate를 두기보다 입력 크기에 따라 부드럽게 통과시킨다. 구현은 tanh approximation을 쓸 수 있다.

SiLU 또는 Swish

\operatorname{SiLU}(x)=x\sigma(x)

SwiGLU 같은 gated FFN에서 변형되어 사용된다.

Weight initialization

$y=Wx$ 에서 input component가 독립이고 mean 0이라고 단순 가정하면 output variance는 input dimension에 따라 커질 수 있다.

\operatorname{Var}(y_j) \approx n\operatorname{Var}(W_{ji})\operatorname{Var}(x_i)

$\operatorname{Var}(W)\propto1/n$ 로 잡으면 layer를 지날 때 variance 폭증을 줄일 수 있다. Xavier와 He initialization은 activation 조건에 맞춰 scale을 정한다.

실제 Transformer는 residual, normalization, depth 때문에 더 세밀한 initialization을 쓴다. $1/\sqrt{n}$ 은 출발 직관이지 모든 architecture의 완성된 recipe가 아니다.

Layer normalization

한 token의 feature $x\in\mathbb{R}^D$ 에서 mean과 variance를 계산한다.

\mu=\frac{1}{D}\sum_{i=1}^{D}x_i

\sigma^2=\frac{1}{D}\sum_{i=1}^{D}(x_i-\mu)^2

\operatorname{LN}(x)_i = \gamma_i\frac{x_i-\mu}{\sqrt{\sigma^2+\epsilon}}+\beta_i

$\gamma,\beta$ 는 trainable parameter다. LayerNorm이 모든 token의 mean을 함께 계산하는 것이 아니라 보통 각 token의 feature axis를 normalize한다.

RMSNorm은 mean subtraction 없이 root mean square로 scale을 정규화한다.

\operatorname{RMSNorm}(x) = \gamma\odot \frac{x}{\sqrt{\frac1D\sum_i x_i^2+\epsilon}}

둘은 같은 연산이 아니다.

Pre-norm과 post-norm

Transformer block을 단순화하면

post-norm은

y=\operatorname{LN}(x+F(x))

pre-norm은

y=x+F(\operatorname{LN}(x))

이다. pre-norm에는 residual identity 경로가 normalization을 거치지 않고 이어져 deep network의 gradient flow에 유리한 면이 있다. architecture마다 안정성과 최종 품질 tradeoff가 있어 이름만 보고 우열을 단정하지 않는다.

Dropout

training 때 mask $m_i\sim\operatorname{Bernoulli}(1-p)$ 를 곱한다. inverted dropout은

y_i=\frac{m_i}{1-p}x_i

로 expectation을 유지한다.

\mathbb{E}[y_i]=x_i

inference 때는 mask를 쓰지 않는다. dropout은 parameter를 직접 줄이지 않고 training 경로에 noise를 넣는 regularization이다.

Numerical gradient check

analytic gradient를 finite difference와 비교한다.

중앙차분은

\frac{\partial L}{\partial\theta_i} \approx \frac{L(\theta_i+h)-L(\theta_i-h)}{2h}

다. relative error를

\frac{|g_{analytic}-g_{numeric}|} {\max(1,|g_{analytic}|,|g_{numeric}|)}

처럼 본다.

$h$ 가 너무 크면 근사 오차, 너무 작으면 floating-point cancellation이 커진다. dropout 같은 randomness를 끄고 작은 double-precision test에서 확인한다. LLM 전체 parameter에 할 검사가 아니라 새 custom operation의 backward를 검증하는 도구다.

Floating-point precision

FP32, FP16, BF16

FP32 — exponent와 mantissa가 넓어 안정적이지만 memory가 크다.
FP16 — mantissa와 exponent가 작고 overflow 범위가 좁다.
BF16 — FP32와 같은 exponent bit 수를 가져 range는 넓고 precision은 낮다.

mixed precision training은 matmul과 activation을 낮은 precision으로 처리하면서 일부 accumulator, master weight, optimizer state를 높은 precision에 둔다. hardware와 kernel에 따라 세부 동작이 다르다.

Loss scaling

FP16에서 작은 gradient가 underflow되는 것을 막기 위해 loss에 scale $S$ 를 곱해 backward한 뒤 gradient를 $S$ 로 나눈다.

\nabla(SL)=S\nabla L

scale이 너무 크면 overflow가 나므로 dynamic loss scaling은 inf와 NaN을 감지해 조절한다.

Quantization의 최소 수학

real value $x$ 를 scale $s$ 와 zero point $z$ 로 integer $q$ 에 근사한다.

q=\operatorname{clip}\left(\operatorname{round}\left(\frac{x}{s}\right)+z,q_{min},q_{max}\right)

복원은

\hat{x}=s(q-z)

다. quantization error는 $x-\hat{x}$ 다.

per-tensor — tensor 전체가 scale 하나를 공유한다.
per-channel — channel별 scale을 둔다.
symmetric — 보통 zero point를 0 근처로 둔다.
asymmetric — range offset을 표현한다.

bit 수만 같아도 calibration data, group size, outlier 처리와 compute kernel이 다르면 품질과 속도가 다르다. storage가 4-bit라고 모든 중간 계산도 4-bit인 것은 아니다.

Optimization에서 자주 하는 실수

Gradient가 0이면 optimum이다

stationary point일 뿐이다. local minimum, maximum, saddle point일 수 있다.

Training loss가 낮으면 generalization이 좋다

training objective와 unseen data 성능은 다르다. validation split, data leakage, distribution shift를 확인한다.

Adam이 learning rate를 없앤다

coordinate별 scale을 조정해도 global learning rate와 schedule은 여전히 중요하다.

Weight decay와 L2 regularization은 항상 같다

plain SGD에서는 밀접하지만 Adam 같은 adaptive optimizer에서는 decoupling 여부가 중요하다.

Gradient clipping이 exploding 원인을 고친다

증상을 제한한다. data outlier, bad initialization, unstable kernel, excessive learning rate를 별도로 찾아야 한다.

NaN을 작은 epsilon 하나로 해결한다

어느 operation에서 처음 inf 또는 NaN이 생겼는지 찾아야 한다. 잘못된 mask, all-masked softmax, log(0), overflow, invalid label도 원인이다.

Debugging 순서

tiny batch 하나에서 forward loss가 유한한지 본다.
random label보다 의미 있는 loss인지 본다.
각 layer activation의 min, max, mean, standard deviation을 기록한다.
gradient norm과 non-finite 위치를 기록한다.
tiny dataset을 overfit할 수 있는지 확인한다.
analytic gradient가 의심되면 finite difference를 쓴다.
mixed precision을 잠시 끄고 FP32와 비교한다.
optimizer, schedule, clipping을 하나씩 바꾼다.

한 번에 여러 설정을 바꾸면 원인을 잃는다.

완료 체크

다음 글은 앞의 선형대수, 확률, 미분을 decoder-only Transformer 한 바퀴에 합친다. LLM 수학 04 — Transformer를 식으로 한 바퀴.

Reference

David E. Rumelhart, Geoffrey E. Hinton, Ronald J. Williams. Learning representations by back-propagating errors, 1986.
Diederik P. Kingma, Jimmy Ba. Adam — A Method for Stochastic Optimization, 2014.
Ilya Loshchilov, Frank Hutter. Decoupled Weight Decay Regularization, 2017.
Jimmy Lei Ba, Jamie Ryan Kiros, Geoffrey E. Hinton. Layer Normalization, 2016.
Dan Hendrycks, Kevin Gimpel. Gaussian Error Linear Units, 2016.