초기화 분산 전파: Xavier·He 유도와 폭 스케일링

50층짜리 MLP를 만들고 가중치를 표준편차 0.01짜리 정규분포로 채운 뒤 순전파를 한 번 돌리면 출력이 전부 0입니다. 표준편차를 1.0으로 올리면 이번에는 nan이 나옵니다. 코드는 한 글자도 안 바꿨고 바뀐 것은 초기화 숫자 하나뿐입니다.

nn.init.kaiming_normal_(layer.weight, mode="fan_in", nonlinearity="relu")

라이브러리는 이 한 줄로 답을 알려 주지만, fan_in이 무엇이고 왜 relu를 따로 적어야 하는지는 알려 주지 않습니다. 이 글은 층 하나를 지날 때 값의 크기가 어떻게 변하는지를 계산하고 그 계산에서 저 한 줄을 끌어냅니다.

지난 글에서 워밍업이 필요한 이유 중 하나로 "초기화 근처는 곡률이 크다"를 그냥 두고 왔습니다. 애초에 그 출발점을 어디에 놓느냐가 여기서 계산할 문제입니다.

한 층의 이득

팬인과 분산의 합

완전연결층 하나를 봅니다. 입력이 $n_{\text{in}}$ 개, 출력이 $n_{\text{out}}$ 개일 때

$y_i = \sum_{j=1}^{n_{\text{in}}} W_{ij}\,x_j$

입니다. 가중치 $W_{ij}$ 를 평균 0, 분산 $\sigma_w^2$ 인 분포에서 서로 독립으로 뽑았다 하고, 입력 $x_j$ 도 서로 독립이며 가중치와 독립이라고 두겠습니다. 그러면 각 항의 평균이 0이므로 $E[y_i] = 0$ 이고, 독립인 것들의 합이므로 분산이 더해집니다.

$\operatorname{Var}(y_i) = \sum_{j=1}^{n_{\text{in}}} \operatorname{Var}(W_{ij}x_j) = \sum_{j=1}^{n_{\text{in}}} \sigma_w^2\,E[x_j^2] = n_{\text{in}}\,\sigma_w^2\,E[x^2]$

한 층을 지날 때 값의 제곱평균이 $n_{\text{in}}\sigma_w^2$ 배가 된다는 뜻입니다. 이 배수를 그 층의 이득이라고 부르겠습니다.

$g = n_{\text{in}}\,\sigma_w^2$

들어오는 쪽 연결 수 $n_{\text{in}}$ 을 팬인, 나가는 쪽 연결 수 $n_{\text{out}}$ 을 팬아웃이라고 합니다. PyTorch의 mode="fan_in"이 가리키는 것이 이 $n_{\text{in}}$ 입니다. 이득이 팬인에 비례하는 까닭은 단순합니다. 출력 하나는 입력 $n_{\text{in}}$ 개를 저마다 무작위 가중치로 곱해 더한 것이고, 서로 무관한 무작위 값을 더하면 분산이 항 수만큼 쌓입니다. 폭이 512인 층에 표준편차 0.05를 쓰면 $g = 512 \times 0.0025 = 1.28$ 이고, 표준편차를 절반인 0.025로 줄이면 이득은 넷으로 나뉘어 0.32가 됩니다. 표준편차가 아니라 분산이 이득에 들어가므로 손잡이를 조금만 돌려도 제곱으로 움직입니다.

편향과 세 가정

실제 층에는 편향, 곧 가중치 합에 따로 더하는 상수 $b_i$ 가 붙어 $y_i = \sum_j W_{ij}x_j + b_i$ 입니다. 초기화 함수를 직접 부르는 코드는 대개 편향을 0으로 함께 채우므로 위의 계산은 한 글자도 안 바뀝니다. 편향을 분산 $\sigma_b^2$ 인 무작위 값으로 채우면 층마다 $\sigma_b^2$ 가 더해지는데, 이 항은 입력과 무관해서 서로 다른 입력을 똑같은 쪽으로 미는 성분이 됩니다. 층을 쌓을수록 입력끼리 구별되던 부분이 이 공통 성분에 묻히므로, 0에서 시작하는 편이 안전합니다.

유도가 기댄 가정은 셋입니다.

• 가중치는 평균 0이고 서로 독립이다
• 입력 성분끼리는 서로 상관이 없다
• 가중치와 입력이 서로 독립이다

첫째와 셋째가 교차항 $E[W_{ij}x_jW_{ik}x_k]$ 를 0으로 만들어 분산을 단순한 합으로 바꿔 줍니다. 둘째는 ReLU를 한 번 지나면 곧바로 흔들립니다 — ReLU의 출력은 전부 0 이상이라 평균이 양수이기 때문입니다. 그래도 계산은 무너지지 않는데, 식에 들어간 것이 분산이 아니라 제곱평균 $E[x^2]$ 이어서 입력의 평균이 0이 아니어도 그대로 성립하기 때문입니다. 가장 먼저 진짜로 깨지는 것은 셋째입니다. 초기화 순간에는 가중치를 입력과 무관하게 새로 뽑았으니 정확히 참이지만, 첫 갱신부터 가중치가 입력을 보고 움직입니다. 이 이야기는 마지막 절에서 다시 꺼냅니다.

ReLU의 절반

선형층 뒤에는 활성함수가 붙습니다. 먼저 ReLU를 봅니다.

들어오는 $z$ 가 평균 0인 대칭 분포이면 절반이 음수이고 그것은 전부 0이 됩니다. 남는 절반만 제곱해서 더하므로

$E[\text{relu}(z)^2] = \int_0^\infty z^2 p(z)\,dz = \frac{1}{2}\int_{-\infty}^{\infty} z^2 p(z)\,dz = \frac{\sigma^2}{2}$

입니다. 정확히 절반입니다. 대칭이라는 조건 하나만 썼으므로 정규분포든 균등분포든 같습니다. 그래서 ReLU가 붙은 층의 이득은

$g = \tfrac{1}{2}\,n_{\text{in}}\,\sigma_w^2$

가 됩니다. nonlinearity="relu"를 적어야 하는 이유가 이 $\tfrac12$ 입니다.

tanh의 이득 계수

tanh는 원점 근처에서 기울기가 1이라 값이 작을 때는 거의 그대로 통과시킵니다. 하지만 입력이 커질수록 기울기가 1보다 작아지고 출력이 ±1 안에 갇히므로, 실제로는 제곱평균을 줄입니다. 표준정규분포를 넣으면 $E[\tanh(z)^2] \approx 0.39$ 로, 이득이 1이어야 할 자리에서 값의 60%를 잃습니다.

이 손실이 층마다 쌓이는 모양은 ReLU와 다릅니다. 제곱평균을 $q$ 라 하면 $q$ 가 작을 때 $\tanh(x)^2 \approx x^2 - \tfrac23 x^4$ 이므로

$q_{l+1} \approx q_l - 2q_l^2$

이고, 양변의 역수를 취하면 $1/q_{l+1} \approx 1/q_l + 2$ 입니다. 곧 $q_L \approx 1/(2L)$ 이라 50층 뒤에 0.01 근처로, 지수가 아니라 층 수에 반비례해 천천히 줄어듭니다. 수치로 따라가도 0.0101이 나옵니다.

라이브러리가 tanh에 붙이는 이득 계수 5/3은 이 손실을 되돌리려고 표준편차에 곱하는 값입니다. 분산으로는 $25/9 \approx 2.78$ 배이고, 표준정규 입력에서 잃는 몫의 역수 $1/0.39 \approx 2.5$ 와 비슷합니다. 유도로 딱 떨어지는 값이 아니라 경험값이지만, 이 계수를 쓰면 층을 지나도 제곱평균이 1.18 근처의 고정점에 머뭅니다. PyTorch의 calculate_gain("tanh")가 돌려주는 것이 이 5/3이고, ReLU에 대해서는 $\sqrt2$ 를 돌려줍니다 — 분산으로 2배, 앞의 $\tfrac12$ 을 정확히 되돌리는 값입니다.

이득의 거듭제곱

50층의 계산

층을 $L$ 개 쌓으면 이득이 $L$ 번 곱해집니다.

$E[x_L^2] = g^L\,E[x_0^2]$

$g$ 가 1에서 조금만 벗어나도 $L$ 이 커지면 답이 없습니다. $g = 0.9$ 면 50층 뒤에 $0.9^{50} = 0.005$ 배, $g = 1.1$ 이면 $1.1^{50} = 117$ 배입니다. 처음의 표준편차 0.01은 얼마나 어긋난 걸까요. 폭이 512이므로

$g = \tfrac12 \times 512 \times 0.01^2 = 0.0256$

이고, 50층이면 $0.0256^{50} \approx 10^{-80}$ 입니다. 제곱근을 취한 크기로도 $10^{-40}$ 이라, float32가 담을 수 있는 가장 작은 양수(약 $10^{-45}$)의 문턱에 닿아 있습니다. 출력이 0인 것이 당연합니다.

import numpy as np
rms = lambda a: float(np.sqrt((a ** 2).mean()))

def forward(sig_w, depth=50, width=512, seed=0):
    r = np.random.default_rng(seed)
    x = r.normal(0, 1, (256, width))
    out = []
    for _ in range(depth):
        W = r.normal(0, sig_w, (width, width))
        x = np.maximum(x @ W, 0.0)          # 선형층 + ReLU
        out.append(rms(x))
    return out

w = 512
for name, sw in (("std 0.01     ", 0.01),
                 ("Xavier √(1/n)", np.sqrt(1.0 / w)),
                 ("He √(2/n)    ", np.sqrt(2.0 / w))):
    s = forward(sw)
    print(f"{name} 1층 {s[0]:.3e}  10층 {s[9]:.3e}  25층 {s[24]:.3e}  50층 {s[49]:.3e}")

# std 0.01      1층 1.604e-01  10층 1.078e-08  25층 1.131e-20  50층 1.184e-40
# Xavier √(1/n) 1층 7.091e-01  10층 3.064e-02  25층 1.541e-04  50층 2.195e-08
# He √(2/n)     1층 1.003e+00  10층 9.806e-01  25층 8.925e-01  50층 7.365e-01

코드의 Xavier와 He는 다음 절에서 유도할 두 초기화의 이름이고, 지금은 분산이 각각 $1/n$ 과 $2/n$ 이라는 것만 보면 됩니다. 세 줄이 50층에서 40자릿수 넘게 벌어집니다. 첫 줄의 측정값 $1.2\times10^{-40}$ 은 위의 $10^{-40}$ 과 맞습니다. 가운데 줄이 재미있는데, Xavier는 $n\sigma_w^2 = 1$ 이라 이득이 딱 $\tfrac12$ 입니다 — ReLU의 절반이 그대로 남으므로 층마다 $\sqrt{2}$ 씩 줄어들고 50층이면 $2^{-25} \approx 3\times 10^{-8}$, 측정값 $2.2\times10^{-8}$ 과 맞습니다. 틀린 초기화가 아니라 다른 활성함수를 위한 초기화입니다.

유한 폭의 흔들림

셋째 줄 He는 이득이 정확히 1인데도 50층에서 0.74로 내려앉았습니다. 이득 1은 평균에 대한 약속일 뿐이고, 한 번 뽑은 가중치 행렬이 주는 배수는 그 평균 둘레에서 흔들리기 때문입니다.

활성함수가 없는 층부터 보면, 출력 성분 $y_i$ 는 정규분포이고 $y_i^2$ 의 분산은 평균의 제곱의 두 배입니다. 그런 성분 $n$ 개를 평균 내면 상대 분산이 $2/n$ 이 되어, 한 층이 주는 제곱평균의 배수가 1 둘레에서 상대오차 $\sqrt{2/n}$ 만큼 흔들립니다. ReLU가 붙으면 $E[\text{relu}(z)^4] = \tfrac32\sigma^4$ 이라 상대 분산이 $5/n$ 으로 커집니다.

층마다의 오차는 서로 독립이고 배수로 곱해지므로 로그를 취하면 더하기가 됩니다. 매 걸음 무작위 방향으로 한 칸씩 움직이는 랜덤워크와 같은 모양이라, $L$ 걸음 뒤의 흔들림은 걸음 수가 아니라 그 제곱근에 비례합니다. 크기(제곱근)의 로그로 적으면 표준편차가

$\tfrac12\sqrt{5L/n}$

입니다. 폭 512, 50층이면 $\tfrac12\sqrt{250/512} \approx 0.35$ 이고, 크기로 바꾸면 $e^{\pm0.35}$, 곧 0.70배에서 1.42배 사이가 흔한 범위입니다. 입력 하나로 씨앗을 200번 바꿔 재면 0.33이 나옵니다. 로그의 평균도 층마다 조금씩 아래로 끌려가서, 50층에 $e^{-5L/(4n)} \approx 0.89$ 배쯤 됩니다.

그림은 위 코드의 He 설정 그대로 씨앗만 스무 번 바꾼 것입니다. 폭 512의 띠는 50층에서 0.60~1.51이고, 앞에서 본 0.74는 그 띠 안의 한 줄일 뿐입니다. 폭을 64로 줄이면 식의 $n$ 이 8분의 1이 되어 띠가 $\sqrt8$ 배쯤 넓어지고 0.12까지 내려갑니다. 이득을 정확히 맞춰도 좁은 망을 깊게 쌓으면 운에 기대게 된다는 뜻이고, 다음 조건 하나를 더 맞춰도 이 흔들림은 남습니다.

순전파와 역전파의 두 조건

팬아웃과 역전파

$g = 1$ 로 두면 $\sigma_w^2 = 1/n_{\text{in}}$ 이 나옵니다. 그런데 이것으로 끝이 아닙니다. 순전파에서 활성값이 지나가는 것과 똑같이 역전파에서는 그래디언트가 지나가고, 그쪽 방향은 전치행렬이 곱해집니다.

$\frac{\partial L}{\partial x_j} = \sum_{i=1}^{n_{\text{out}}} W_{ij}\,\frac{\partial L}{\partial y_i}$

합의 개수가 $n_{\text{in}}$ 이 아니라 $n_{\text{out}}$ 입니다. 같은 계산을 다시 하면

$\operatorname{Var}\!\left(\frac{\partial L}{\partial x}\right) = n_{\text{out}}\,\sigma_w^2\,\operatorname{Var}\!\left(\frac{\partial L}{\partial y}\right)$

이므로 그래디언트가 유지되려면 $\sigma_w^2 = 1/n_{\text{out}}$ 이어야 합니다. ReLU가 붙으면 역전파에서는 그 미분이 곱해지는데, 미분은 절반의 자리에서 1이고 나머지에서 0이라 여기에도 똑같이 $\tfrac12$ 이 붙습니다.

조화평균 절충

$n_{\text{in}} \ne n_{\text{out}}$ 인 층에서는 두 조건이 서로 다른 $\sigma_w$ 를 요구합니다. 하나를 고르면 다른 쪽이 어긋나므로 절충합니다.

$\sigma_w^2 = \frac{2}{n_{\text{in}} + n_{\text{out}}}$

이것이 Xavier 초기화입니다(Glorot 초기화라고도 합니다). 두 조건 $1/n_{\text{in}}$ 과 $1/n_{\text{out}}$ 의 조화평균입니다. 산술평균이 아니라 조화평균이 나오는 까닭은 식을 $n_{\text{in}}\sigma_w^2 + n_{\text{out}}\sigma_w^2 = 2$ 로 고쳐 쓰면 보입니다 — 순전파 이득과 역전파 이득의 평균을 1로 맞춘 것입니다.

폭이 크게 다른 층에서 이 절충이 얼마나 어긋나는지 수로 보겠습니다. 트랜스포머의 MLP처럼 512에서 2048로 넓혔다가 다시 512로 좁히는 두 층을 활성함수 없이 생각합니다. 첫 층의 Xavier 분산은 $2/2560 = 1/1280$ 이라 순전파 이득이 $512/1280 = 0.4$, 역전파 이득이 $2048/1280 = 1.6$ 입니다. 둘의 평균은 정확히 1이지만 어느 쪽도 1이 아닙니다. 두 조건을 산술평균했다면 분산이 $5/4096$ 이 되어 이득이 0.625와 2.5로, 평균까지 1.56으로 어긋났을 것입니다. 다행히 둘째 층이 거꾸로 1.6과 0.4를 주므로 한 쌍을 지나면 양쪽 모두 0.64로 돌아옵니다.

He 초기화와 fan_out

ReLU면 순전파 조건에 $\tfrac12$ 이 붙어 $\tfrac12 n_{\text{in}}\sigma_w^2 = 1$ 이 되고

$\sigma_w^2 = \frac{2}{n_{\text{in}}}$

이 나옵니다. 이것이 He 초기화(Kaiming 초기화)입니다. 폭 512면 분산 $1/256$, 표준편차 $1/16 = 0.0625$ 입니다. 같은 폭의 Xavier는 표준편차 $\sqrt{1/512} \approx 0.044$ 이니, 차이가 $\sqrt2$ 배뿐인데 50층 뒤에는 앞의 코드처럼 일곱 자릿수가 벌어집니다.

He는 절충하지 않고 한쪽만 고릅니다. 그래도 되는 까닭은 역전파 쪽 이득을 층마다 곱해 보면 드러납니다. 팬인으로 맞추면 한 층의 역전파 이득이 $\tfrac12 n_{\text{out}} \cdot 2/n_{\text{in}} = n_{\text{out}}/n_{\text{in}}$ 이고, 앞 층의 팬아웃이 다음 층의 팬인이므로 전체를 곱하면 가운데가 다 지워지고 맨 끝 폭과 맨 앞 폭의 비 하나만 남습니다. 512→2048→512 한 쌍이면 $4 \times 0.25 = 1$ 입니다. 한쪽을 정확히 맞추면 다른 쪽의 어긋남은 층 수에 따라 거듭제곱되지 않는다는 뜻입니다.

그래서 mode="fan_in"과 mode="fan_out"의 차이는 이 양 끝 폭의 비 한 번뿐이고, 어느 쪽이든 지수로 벌어지지 않습니다. 어느 쪽을 고를지는 어느 신호를 정확히 지키고 싶은가의 문제입니다. torchvision의 ResNet 구현은 합성곱에 fan_out을 씁니다 — 역전파 쪽 크기를 정확히 지키는 쪽을 택한 셈입니다.

활성함수	이득에 붙는 계수	권장 $\sigma_w^2$	이름
없음(항등)	1	$2/(n_{\text{in}}+n_{\text{out}})$	Xavier
tanh	1보다 작음	$(5/3)^2 \cdot 2/(n_{\text{in}}+n_{\text{out}})$	Xavier + 계수 5/3
ReLU	$1/2$	$2/n_{\text{in}}$	He
LeakyReLU($a$)	$(1+a^2)/2$	$2/((1+a^2)n_{\text{in}})$	He 변형

분포와 층별 예외

균등분포 환산

지금까지의 계산은 가중치 분포의 분산만 썼으므로 정규분포 대신 균등분포로 뽑아도 됩니다. 구간 $[-a, a]$ 의 균등분포는 분산이 $a^2/3$ 이므로, 분산을 $\sigma_w^2$ 에 맞추려면 $a = \sqrt3\,\sigma_w$ 입니다. Xavier는 $\pm\sqrt{6/(n_{\text{in}}+n_{\text{out}})}$, He는 $\pm\sqrt{6/n_{\text{in}}}$ 이 됩니다. 폭 512끼리의 층이면 Xavier 경계가 $\sqrt{6/1024} \approx 0.077$, He 경계가 $\sqrt{6/512} \approx 0.108$ 입니다.

두 분포의 실제 차이는 꼬리입니다. 균등분포에는 경계 밖의 값이 아예 없고, 정규분포는 드물게 표준편차의 서너 배짜리 가중치를 만듭니다. 하지만 출력 하나는 수백 개 항의 합이라 중심극한정리에 따라 어느 쪽으로 뽑든 거의 정규분포가 되고, 이득 계산도 두 분포에서 같습니다. 극단값이 걱정되면 표준편차의 두 배 밖을 잘라 다시 뽑는 잘린 정규분포를 쓰기도 하는데, 이때는 잘린 만큼 분산이 줄어드는 것을 보정해야 합니다.

직교 초기화

무작위 행렬의 이득은 평균으로만 1이고, 앞 절에서 본 흔들림이 거기서 나왔습니다. 직교 초기화는 가중치를 직교행렬, 곧 열벡터들이 서로 수직이고 길이가 1인 행렬로 채워 이 흔들림을 없앱니다. 직교행렬은 어떤 벡터를 곱해도 길이를 바꾸지 않으므로 $\lVert Wx\rVert = \lVert x\rVert$ 가 평균이 아니라 모든 입력에서 정확히 성립합니다. ReLU가 붙으면 여기에 $\sqrt2$ 를 곱해 절반을 되돌립니다.

행렬이 정사각이 아니면 제약이 생깁니다. $n_{\text{out}} > n_{\text{in}}$ 이면 열들만 서로 수직으로 만들 수 있고, 그러면 순전파의 길이는 지켜지지만 역전파 쪽은 투영이 되어 길이가 줄어듭니다. 반대로 좁히는 층에서는 순전파가 투영이 됩니다. 두 방향을 모두 정확히 지킬 수 있는 것은 정사각행렬뿐이고, 이것도 앞에서 본 두 조건의 충돌과 같은 모양입니다.

합성곱의 팬인

합성곱층에서 출력 한 칸은 커널이 덮은 자리의 모든 채널을 더해서 만듭니다. 커널이 $k_h \times k_w$ 이고 입력 채널이 $C_{\text{in}}$ 이면 더하는 항이 $k_h k_w C_{\text{in}}$ 개이므로 그것이 팬인입니다.

$n_{\text{in}} = k_h\,k_w\,C_{\text{in}}, \qquad n_{\text{out}} = k_h\,k_w\,C_{\text{out}}$

3×3 커널에 입력 채널 64면 팬인이 576이고 He 표준편차는 $\sqrt{2/576} \approx 0.059$ 입니다. 흔한 실수가 채널 수 64만 넣는 것인데, 그러면 분산이 9배 커져 이득이 층마다 9가 됩니다. 라이브러리의 초기화 함수는 가중치 텐서의 모양에서 팬인을 스스로 계산하므로, 직접 계산할 일은 손으로 초기화를 짤 때뿐입니다.

임베딩과 0 초기화

팬인이라는 말이 뜻을 잃는 자리도 있습니다. 임베딩층은 행렬곱처럼 적을 수 있지만 입력이 원-핫 벡터라 실제로는 한 행을 꺼낼 뿐이고, 더해지는 항이 하나입니다. 어휘 크기를 팬인으로 넣으면 수만 분의 1짜리 분산이 되어 의미가 없습니다. 그래서 임베딩은 폭과 무관하게 표준편차를 정해 두는데, PyTorch의 기본값은 표준정규분포이고 GPT-2는 0.02를 썼습니다.

반대로 일부러 0이나 0에 가까운 값에서 시작하는 층도 있습니다. 잔차 블록의 마지막 층과 출력 투영이 그렇습니다. 이 층이 0이면 블록 전체가 처음에는 아무 일도 하지 않고 입력을 그대로 넘기는데, 왜 그것이 이득인지는 다음 절의 잔차 연결에서 계산이 나옵니다.

잔차 연결과 정규화

잔차의 두 배 성장

$g$ 를 정확히 1에 맞춰도 유한한 폭에서는 흔들림이 남고, 층이 수백 개면 그 흔들림도 쌓입니다. 더 깊게 쌓으려면 곱셈 구조 자체를 바꿔야 합니다.

잔차 연결은 층의 출력을 입력에 더합니다.

$x_{l+1} = x_l + F(x_l)$

두 항이 대체로 독립이면 분산이 더해지므로 $\operatorname{Var}(x_{l+1}) = \operatorname{Var}(x_l) + \operatorname{Var}(F(x_l))$ 입니다. 그런데 $F$ 가 이득 1짜리 층이면 $\operatorname{Var}(F(x_l)) = \operatorname{Var}(x_l)$ 이므로 층마다 정확히 두 배가 됩니다 — 곱셈이 사라진 게 아니라 밑이 $\sqrt2$ 로 바뀌었을 뿐입니다. 활성함수 없는 분기로 50층을 돌리면 크기가 $2^{25} \approx 3.4\times10^7$ 배로 식과 그대로 맞습니다.

정규화의 위치

바뀌는 것은 블록에 정규화를 넣을 때입니다. LayerNorm은 들어온 벡터를 성분끼리 평균 0, 분산 1로 맞춘 뒤 학습 가능한 배율 $\gamma$ 와 이동 $\beta$ 를 적용하는데, 초기값이 $\gamma = 1$, $\beta = 0$ 이라 처음에는 표준화 말고 아무것도 하지 않습니다. 정규화를 어디에 두느냐에 따라 분산이 쌓이는 방식이 갈립니다.

블록 뒤에 두어 $x_{l+1} = \operatorname{LN}(x_l + F(x_l))$ 로 쓰는 것이 Post-LN입니다. 순전파 크기는 층마다 1로 되돌려지므로 아예 쌓이지 않습니다. 대신 역전파의 그래디언트가 주 경로에서 층마다 정규화를 지나야 하고, 초기화 직후 출력 쪽 층의 그래디언트가 크게 잡혀 워밍업 없이는 학습이 불안정하다는 분석이 있습니다. 지난 글에서 두고 온 "초기화 근처의 큰 곡률"이 이 구조에서 특히 두드러집니다.

블록 앞, 곧 $F$ 의 입력에만 두어 $x_{l+1} = x_l + F(\operatorname{LN}(x_l))$ 로 쓰는 것이 Pre-LN이고 요즘 트랜스포머는 이쪽이 많습니다. $F$ 에 들어가는 값이 늘 분산 1이므로 $F$ 의 출력 크기가 $x_l$ 의 크기와 무관해지고

$\operatorname{Var}(x_{l+1}) = \operatorname{Var}(x_l) + c$

로 더하기만 남습니다. $L$ 층 뒤의 분산이 $L$ 에 비례하고 표준편차는 $\sqrt{L}$ 에 비례합니다. 지수가 거듭제곱으로 내려온 것입니다. 덤으로 주 경로가 늘 커지므로 뒤쪽 블록 하나가 흐름에 보태는 몫은 $1/\sqrt{l}$ 로 줄어듭니다.

분기 축소

def residual(mode, depth=50, width=512, seed=0):
    r = np.random.default_rng(seed)
    x = r.normal(0, 1, (256, width))
    out = []
    for _ in range(depth):
        W = r.normal(0, np.sqrt(2.0 / width), (width, width))
        h = x
        if mode.startswith("preln"):                       # 블록 앞에서 정규화
            h = (h - h.mean(1, keepdims=True)) / (h.std(1, keepdims=True) + 1e-6)
        f = np.maximum(h @ W, 0.0)
        if mode == "preln_scaled":                         # 분기를 1/√L 로
            f = f / np.sqrt(depth)
        x = x + f
        out.append(rms(x))
    return out

for m in ("plain", "preln", "preln_scaled"):
    s = residual(m)
    print(f"{m:13s}", " ".join(f"{i+1}층 {s[i]:.3g}" for i in (0, 9, 24, 49)))

# plain         1층 1.41 10층 141 25층 3.59e+05 50층 2.25e+11
# preln         1층 1.41 10층 6.29 25층 14.7 50층 28.8
# preln_scaled  1층 1.01 10층 1.33 25층 2.31 50층 4.19

잔차만 쓴 첫 줄은 50층에서 $2\times10^{11}$ 배로, 위에서 계산한 $3.4\times10^7$ 보다도 큽니다. ReLU의 출력은 평균이 양수라 $x_l$ 과 $F(x_l)$ 이 같은 방향의 치우침을 공유하고, 독립이라 두고 버린 교차항이 살아나 한 층에 두 배보다 더 불어나기 때문입니다. 정규화를 넣은 둘째 줄은 28.8배로 내려오는데, 유도가 준 $\sqrt{50} \approx 7$ 보다는 큽니다. 같은 치우침이 층마다 더해지기 때문이고, 중요한 것은 지수가 사라졌다는 점입니다.

셋째 줄처럼 잔차 분기를 $1/\sqrt{L}$ 로 줄이면 $L$ 개의 $c/L$ 을 더하게 되어 합이 깊이와 무관한 상수로 눌리고, 4.2배가 됩니다. GPT-2가 잔차 경로에 들어가는 투영층의 초기 가중치를 $1/\sqrt{N}$ ($N$ 은 잔차 층 수)로 줄인 것이 이 계산입니다. 같은 자리를 노리는 다른 방법은 분기 끝에 학습 가능한 계수 $\alpha$ 를 두고 $x_l + \alpha F(x_l)$ 로 쓰되 $\alpha = 0$ 에서 시작하는 것입니다. 그러면 처음에는 망 전체가 항등함수라 크기 문제가 원천적으로 없고, 학습이 필요한 만큼만 분기를 켭니다. 블록의 마지막 층이나 마지막 BatchNorm의 $\gamma$ 를 0으로 두는 관행도 같은 효과입니다.

정규화가 들어간 뒤로 초기화에 대한 민감도가 크게 줄어든 것이 이 계산의 결론입니다. 그렇다고 초기화가 필요 없어지는 것은 아니고, 정규화가 없는 자리(임베딩, 출력 투영)와 학습 초반의 곡률은 여전히 초기화가 정합니다. 구현과 실제 학습 곡선 비교는 가중치 초기화 실습이 맡고 이 글은 유도만 했습니다.

폭과 갱신 크기

SGD 한 스텝

지금까지의 조건은 전부 $\sigma_w^2 \propto 1/n$ 이었습니다. 폭 $n$ 을 두 배로 하면 가중치를 $1/\sqrt{2}$ 로 줄이라는 것이고, 그러면 $t=0$ 의 순전파는 폭과 무관해집니다. 그런데 학습이 시작되면 사정이 달라집니다.

한 스텝의 SGD 갱신이 이 층의 출력을 얼마나 바꾸는지 계산해 보겠습니다. 출력 쪽에서 내려온 그래디언트를 $\delta_i$ 라 하면 $\partial L/\partial W_{ij} = \delta_i x_j$ 이므로 $\Delta W_{ij} = -\eta\,\delta_i x_j$ 이고,

$\Delta y_i = \sum_{j=1}^{n} \Delta W_{ij}\,x_j = -\eta\,\delta_i \sum_{j=1}^{n} x_j^2 = -\eta\,\delta_i\,n\,E[x^2]$

팬인 $n$ 이 그대로 곱해져 나옵니다. 초기화는 $n$ 개 항의 무작위 합이라 $\sqrt{n}$ 스케일로 커지는데, 갱신은 모든 항이 같은 방향으로 정렬되어 있어 $n$ 스케일로 커집니다. 폭을 키우면 첫 스텝이 출력을 훨씬 더 크게 흔든다는 뜻입니다.

def step_effect(n, eta=1e-3, seed=0):
    r = np.random.default_rng(seed)
    x = r.normal(0, 1, n)                       # 입력 성분은 폭과 무관하게 Θ(1)
    W = r.normal(0, np.sqrt(2.0 / n), (n, n))   # He 초기화
    d = r.normal(0, 1, n)                       # 출력 쪽에서 내려온 그래디언트
    y0 = W @ x
    W2 = W - eta * np.outer(d, x)               # SGD 한 스텝
    return float(np.abs(W2 @ x - y0).mean()), float(np.abs(y0).mean())

for n in (128, 512, 2048, 8192):
    dy, y = step_effect(n)
    print(f"폭 {n:>5}  |y| {y:.3f}   한 스텝의 |Δy| {dy:.4f}   Δy/폭 {dy/n:.3e}")

# 폭   128  |y| 1.045   한 스텝의 |Δy| 0.1016   Δy/폭 7.934e-04
# 폭   512  |y| 1.165   한 스텝의 |Δy| 0.3745   Δy/폭 7.315e-04
# 폭  2048  |y| 1.135   한 스텝의 |Δy| 1.6014   Δy/폭 8.004e-04
# 폭  8192  |y| 1.122   한 스텝의 |Δy| 6.5568   Δy/폭 8.004e-04

$|y|$ 는 폭이 64배가 되어도 1.1 근처에 머뭅니다 — 초기화가 제 일을 한 것입니다. 그런데 $|\Delta y|$ 는 0.10에서 6.56으로 정확히 64배가 되고, 폭으로 나눈 값이 상수입니다. 같은 학습률을 쓰면 넓은 모델일수록 첫 스텝이 크게 튑니다.

Adam의 갱신

Adam에서는 그림이 조금 다릅니다. Adam은 그래디언트를 그 크기의 이동 평균으로 나누므로, 첫 걸음은 성분마다 $-\eta\,\operatorname{sign}(\delta_i x_j)$ 로 그래디언트의 크기와 무관하게 $\eta$ 짜리입니다. 이것을 같은 식에 넣으면

$\Delta y_i \approx -\eta\,\operatorname{sign}(\delta_i)\sum_{j=1}^{n}|x_j| = -\eta\,\operatorname{sign}(\delta_i)\,n\,E|x|$

로 여기서도 $n$ 이 곱해집니다. 다른 점은 $\delta$ 의 크기가 식에서 사라졌다는 것입니다. SGD에서는 출력 쪽에서 내려오는 그래디언트를 작게 만들어 폭의 효과를 상쇄할 여지가 있지만, Adam은 그 크기를 나눠 없애므로 학습률 자체를 $1/n$ 로 줄이는 것 말고는 길이 없습니다. 폭을 8배로 키운 모델에서 은닉층의 Adam 학습률을 8분의 1로 줄여야 첫 스텝의 흔들림이 같아진다는 계산입니다.

muP

그래서 폭을 바꿀 때 옮겨야 하는 것은 초기화 분산만이 아니라 층마다의 학습률과 출력의 배율까지입니다. 그 규칙을 층의 종류별로 정리한 것이 muP(maximal update parametrization)입니다. 층을 셋으로 나누는 이유는 폭 $n$ 이 팬인과 팬아웃 중 어디에 들어가느냐가 다르기 때문입니다.

• 입력층(임베딩): 팬인이 입력 차원이나 어휘 크기라 폭과 무관하다. 초기화도 학습률도 폭을 따라 바꾸지 않는다
• 은닉층: 팬인과 팬아웃이 모두 $n$ 이다. 초기화는 $1/n$ 그대로 두고, Adam이면 학습률을 $1/n$ 로 줄인다
• 출력층: 팬인만 $n$ 이고 결과가 곧바로 로짓이 된다. 정렬 효과가 가장 먼저 로짓을 흔드는 자리라 출력에 $1/n$ 꼴의 배율을 곱하거나 초기값을 그만큼 더 작게 둔다

정확한 지수는 옵티마이저와 구현마다 적는 방식이 조금씩 다르므로 여기서는 방향만 봅니다. 이 글의 계산은 그 유도의 첫 걸음입니다.

독립 가정의 수명

첫 절에서 가장 먼저 깨진다고 한 셋째 가정이 바로 이 자리에서 깨집니다. 갱신량 $\Delta W = -\eta\,\delta x^\top$ 은 입력 $x$ 를 그대로 품고 있고, 그래서 $\Delta W x$ 의 합에서 항들이 서로 상쇄되지 않고 같은 방향으로 쌓였습니다. $\sqrt{n}$ 이 $n$ 으로 바뀐 것은 가중치와 입력이 더 이상 무관하지 않다는 것의 수치적 모습입니다.

그러니 이 글의 분산 계산은 전부 $t = 0$ 한 점의 이야기입니다. 학습이 진행된 뒤의 가중치는 무작위 행렬이 아니고, 그때 활성값의 크기를 지키는 것은 초기화가 아니라 정규화와 옵티마이저와 학습률입니다. 초기화가 정하는 것은 출발점이고, muP는 그 출발점에서의 첫 몇 걸음이 폭과 무관하게 같은 크기가 되도록 맞추는 규칙입니다.

실용적인 값어치는 여기 있습니다 — 규칙을 지키면 폭이 작은 모델에서 고른 학습률이 폭이 큰 모델에서도 그대로 맞습니다. 학습률 후보를 열 개 놓고 비교하려면 모델을 열 번 학습시켜야 하는데, 그 열 번을 폭이 8분의 1인 모델로 돌리면 은닉층 행렬곱 비용이 대략 64분의 1로 줄어듭니다. 큰 모델은 고른 값으로 한 번만 돌리면 됩니다.

정리

• 한 층의 이득은 $g = n_{\text{in}}\sigma_w^2$ 이고, ReLU가 붙으면 $E[\text{relu}(z)^2] = \sigma^2/2$ 때문에 $\tfrac12$ 이 곱해진다. tanh는 1보다 조금씩 줄이므로 계수 5/3을 곱해 되돌린다.
• 이득은 층마다 곱해지므로 $L$ 층 뒤에 $g^L$ 이다. 폭 512에 표준편차 0.01이면 $g = 0.0256$, 50층이면 $10^{-80}$ — 측정에서 $10^{-40}$(제곱근 기준)로 확인된다.
• 이득을 정확히 1로 맞춰도 유한한 폭에서는 로그 크기가 랜덤워크로 $\tfrac12\sqrt{5L/n}$ 만큼 흔들린다. 폭 512, 50층이면 0.70~1.42배가 흔한 범위다.
• 순전파 조건은 $n_{\text{in}}\sigma_w^2 = 1$, 역전파 조건은 $n_{\text{out}}\sigma_w^2 = 1$ 로 서로 다르다. 두 이득의 평균을 1로 맞춘 조화평균 $\sigma_w^2 = 2/(n_{\text{in}}+n_{\text{out}})$ 이 Xavier, ReLU의 $\tfrac12$ 을 되돌린 $\sigma_w^2 = 2/n_{\text{in}}$ 이 He다. 한쪽만 맞춰도 다른 쪽의 어긋남은 양 끝 폭의 비 하나로 줄어든다.
• 합성곱의 팬인은 $k_h k_w C_{\text{in}}$ 이고, 임베딩처럼 팬인이 뜻을 갖지 않는 자리는 따로 정한다.
• 잔차 연결만으로는 밑이 $\sqrt2$ 인 지수가 남는다. 블록 앞에 정규화를 넣어야 분산이 곱이 아니라 합으로 쌓여 지수가 사라지고, 분기를 $1/\sqrt L$ 로 줄이거나 0에서 시작하면 그 합까지 눌린다.
• 초기화를 $1/n$ 로 맞추면 $t=0$ 의 순전파는 폭과 무관해지지만, 한 스텝의 갱신이 출력에 주는 변화는 SGD든 Adam이든 폭에 비례한다. 폭을 바꿀 때 학습률까지 함께 옮기는 규칙이 muP이고, 그 덕분에 작은 모델에서 고른 값을 큰 모델에 그대로 쓸 수 있다.

처음의 50층 MLP로 돌아가면 두 증상이 모두 이득 하나로 설명됩니다. 표준편차 0.01은 이득 0.0256이라 크기가 층마다 6분의 1로 줄어 0이 되었고, 표준편차 1.0은 이득이 $\tfrac12 \times 512 = 256$ 이라 크기가 층마다 16배씩 불어납니다. $16^{32} = 2^{128}$ 이 float32의 최댓값 $3.4\times10^{38}$ 이므로 32층 언저리에서 무한대가 되고, 무한대끼리의 연산이 nan을 만듭니다. 맨 앞의 kaiming_normal_ 한 줄은 팬인 512와 ReLU의 $\tfrac12$ 을 받아 표준편차 0.0625를 고르는 것이고, 그것이 이득을 1에 놓는 유일한 값입니다.

지금까지 세 편은 전부 1차 정보만 썼습니다 — 그래디언트를 재고, 그 평균과 분산을 다루고, 출발점을 맞췄습니다. 곡률은 $L$ 이라는 상수 하나로만 등장했고 그마저 우리가 알 수 없는 값이었습니다. 다음 글에서 곡률을 직접 쓰는 방법, 즉 헤세 행렬을 쓰는 뉴턴법을 유도하고 그것이 왜 딥러닝에서 그대로 쓰이지 않는지를 비용으로 계산합니다. 그리고 Adam의 대각 근사가 그 자리에서 무엇을 흉내 내고 있었는지 되짚습니다.

---

읽어주셔서 감사합니다. 😊