가중치 초기화와 학습 안정성

잘못 초기화하면 학습(training)은 시작도 못 합니다. 제대로 초기화하면 50개 층(layer)도 3개 층처럼 부드럽게 학습됩니다.

유형: Build

언어: Python

선수 지식: Lesson 03.04 활성화 함수(Activation Functions), Lesson 03.07 정규화(Regularization)

소요 시간: 약 90분

학습 목표

• 0 초기화(zero initialization), 무작위 초기화(random initialization), Xavier/Glorot 초기화, Kaiming/He 초기화 전략(initialization strategy)을 구현하고 50개 층을 지날 때 활성화 크기(activation magnitude)가 어떻게 달라지는지 측정합니다.
• Xavier 초기화가 Var(w) = 2/(fan_in + fan_out)을 쓰고 Kaiming 초기화가 Var(w) = 2/fan_in을 쓰는 이유를 유도합니다.
• 0 초기화의 대칭 문제(symmetry problem)를 보여 주고 무작위 스케일(random scale)만으로는 충분하지 않은 이유를 설명합니다.
• 활성화 함수(activation function)에 맞는 초기화 전략을 고릅니다. sigmoid/tanh에는 Xavier, ReLU/GELU에는 Kaiming을 사용합니다.

문제

모든 가중치(weight)를 0으로 초기화해 보면 아무것도 학습되지 않습니다. 모든 뉴런(neuron)이 같은 함수(function)를 계산하고, 같은 기울기(gradient)를 받으며, 동일하게 갱신(update)됩니다. 10,000 에폭(epoch)을 돌려도 512개 뉴런으로 구성된 은닉층(hidden layer)은 결국 같은 뉴런 512개의 복사본일 뿐입니다. 파라미터(parameter) 512개에 해당하는 비용을 지불했지만 실제로 얻은 것은 하나뿐인 셈입니다.

가중치를 지나치게 크게 초기화하면 활성화(activation) 값이 네트워크(network)를 지나며 폭발합니다. 10번째 층에서 1e15에 도달하고, 20번째 층에서는 무한대(infinity)로 넘쳐(overflow) 버릴 수 있습니다. 기울기 역시 반대 방향으로 같은 문제를 겪습니다.

표준 정규 분포(standard normal distribution)에서 무작위 초기화(random initialization)를 하면 3개 층에서는 작동할 수 있습니다. 하지만 50개 층에서는 무작위 스케일(random scale)이 조금만 작아도 신호(signal)가 0으로 사라지고(collapse), 조금만 커도 무한대로 폭발합니다. "작동"과 "고장"의 경계가 면도날처럼 얇습니다.

가중치 초기화는 딥러닝(deep learning)에서 과소평가되는 결정입니다. 아키텍처(architecture)와 옵티마이저(optimizer)는 더 많이 이야기되지만, 초기화를 잘못하면 네트워크는 학습이 시작되기 전부터 죽어 있습니다.

개념

대칭 문제(The Symmetry Problem)

한 층(layer) 안의 모든 뉴런은 같은 구조를 가집니다. 입력(input)에 가중치를 곱하고 편향(bias)을 더한 뒤 활성화 함수를 적용합니다. 모든 가중치가 같은 값으로 시작한다면, 극단적으로 0으로 시작한다면 모든 뉴런이 같은 출력(output)을 계산합니다. 역전파(backpropagation) 중에도 같은 기울기를 받고, 갱신도 같은 양만큼 일어납니다.

네트워크에 수백 개의 파라미터가 있어도 모두가 보조를 맞춰(lockstep) 움직입니다. 이것이 바로 대칭(symmetry)입니다. 무작위 초기화는 각 뉴런을 가중치 공간(weight space)의 서로 다른 지점에서 시작하게 만들어 대칭을 깨뜨리고, 뉴런마다 서로 다른 특징(feature)을 학습하게 합니다.

하지만 단순히 "무작위(random)"인 것만으로는 부족합니다. 무작위성의 스케일(scale)이 네트워크가 학습되는지 여부를 결정합니다.

층을 지나는 분산 전파(Variance Propagation Through Layers)

fan_in개의 입력을 받는 단일 층(single layer)을 생각해 봅시다.

z = w1*x1 + w2*x2 + ... + w_n*x_n

가중치 w_i의 분산(variance)이 Var(w)이고 입력 x_i의 분산이 Var(x)라면 출력 분산은 다음과 같습니다.

Var(z) = fan_in * Var(w) * Var(x)

Var(w) = 1이고 fan_in = 512라면 출력 분산은 입력 분산의 512배입니다. 10개 층을 지나면 512^10 = 1.2e27이 되어 신호가 폭발합니다.

반대로 Var(w) = 0.001이면 층마다 0.001 * 512 = 0.512배가 되고, 10개 층 뒤에는 0.512^10 = 0.00013이 됩니다. 신호가 사라지는 것입니다.

목표는 Var(z) = Var(x)가 되도록 Var(w)를 고르는 것입니다. 신호 크기(signal magnitude)가 층을 지나도 일정하게 유지되어야 합니다.

Xavier/Glorot 초기화(Xavier/Glorot Initialization)

Glorot와 Bengio(2010)는 sigmoid와 tanh 활성화 함수를 위한 해를 유도했습니다. 순전파(forward pass)와 역전파(backward pass) 모두에서 분산을 일정하게 유지하기 위해 다음 공식을 사용합니다.

Var(w) = 2 / (fan_in + fan_out)

실제로는 아래처럼 표본을 뽑아(sampling) 사용합니다.

w ~ Uniform(-limit, limit)  where limit = sqrt(6 / (fan_in + fan_out))
w ~ Normal(0, sqrt(2 / (fan_in + fan_out)))

sigmoid와 tanh는 0 근처에서 대략 선형(roughly linear)이고, 제대로 초기화하면 활성화 값이 그 근처에 머뭅니다. 그래서 분산이 여러 층에 걸쳐 안정적으로 유지됩니다.

Kaiming/He 초기화(Kaiming/He Initialization)

ReLU는 출력의 절반을 죽입니다. 음수(negative value)는 0이 됩니다. 평균적으로 활성화의 절반이 0으로 처리되므로 유효 fan_in이 반으로 줄어든 것과 같습니다. Xavier 초기화는 이 점을 반영하지 못해서 ReLU에서 필요한 분산을 과소평가합니다.

He 등(He et al., 2015)은 다음 공식을 제안했습니다.

Var(w) = 2 / fan_in

가중치는 보통 아래와 같이 표본을 뽑습니다.

w ~ Normal(0, sqrt(2 / fan_in))

곱해진 2는 ReLU가 활성화의 절반을 0으로 만드는 효과를 보정하는 계수(factor)입니다. 이 보정이 없으면 신호는 층마다 약 0.5배씩 줄어들고, 50개 층 뒤에는 0.5^50 = 8.8e-16이 됩니다.

트랜스포머 초기화(Transformer Initialization)

GPT-2는 잔차 연결(residual connection) 때문에 다른 패턴을 사용했습니다.

x = x + sublayer(x)

덧셈은 분산을 증가시킵니다. N개의 잔차 층(residual layer)이 있으면 분산은 N에 비례해 커집니다. GPT-2는 잔차 층의 가중치를 1/sqrt(2N)로 스케일링해 누적된 신호 크기를 안정화했습니다. Llama 3(파라미터 405B, 126개 층) 같은 대규모 모델(model)도 비슷한 방식(scheme)을 사용합니다. 이 스케일링이 없으면 잔차 스트림(residual stream)은 126개의 어텐션(attention)과 피드포워드(feedforward) 블록을 지나며 한없이(unbounded) 커질 수 있습니다.

flowchart TD
    subgraph "Zero Init"
        Z1["Layer 1<br/>All weights = 0"] --> Z2["Layer 2<br/>All neurons identical"]
        Z2 --> Z3["Layer 3<br/>Still identical"]
        Z3 --> ZR["Result: 1 effective neuron<br/>regardless of width"]
    end

    subgraph "Xavier Init"
        X1["Layer 1<br/>Var = 2/(fan_in+fan_out)"] --> X2["Layer 2<br/>Signal stable"]
        X2 --> X3["Layer 50<br/>Signal stable"]
        X3 --> XR["Result: Trains with<br/>sigmoid/tanh"]
    end

    subgraph "Kaiming Init"
        K1["Layer 1<br/>Var = 2/fan_in"] --> K2["Layer 2<br/>Signal stable"]
        K2 --> K3["Layer 50<br/>Signal stable"]
        K3 --> KR["Result: Trains with<br/>ReLU/GELU"]
    end

50층을 지나는 활성화 크기(Activation Magnitude Through 50 Layers)

graph LR
    subgraph "Mean Activation Magnitude"
        direction LR
        L1["Layer 1"] --> L10["Layer 10"] --> L25["Layer 25"] --> L50["Layer 50"]
    end

    subgraph "Results"
        R1["Random N(0,1): EXPLODES by layer 5"]
        R2["Random N(0,0.01): Vanishes by layer 10"]
        R3["Xavier + Sigmoid: ~1.0 at layer 50"]
        R4["Kaiming + ReLU: ~1.0 at layer 50"]
    end

올바른 초기화(init) 고르기

flowchart TD
    Start["Activation은 무엇인가요?"] --> Act{"Activation type?"}

    Act -->|"Sigmoid / Tanh"| Xavier["Xavier/Glorot<br/>Var = 2/(fan_in + fan_out)"]
    Act -->|"ReLU / Leaky ReLU"| Kaiming["Kaiming/He<br/>Var = 2/fan_in"]
    Act -->|"GELU / Swish"| Kaiming2["Kaiming/He<br/>(ReLU와 동일하게 시작)"]
    Act -->|"Transformer residual"| GPT["1/sqrt(2N)로 scale<br/>N = number of layers"]

    Xavier --> Check["activation magnitude가<br/>모든 layer에서 0.5~2.0 근처인지 확인"]
    Kaiming --> Check
    Kaiming2 --> Check
    GPT --> Check

만들어 보기

Step 1: 초기화 전략(initialization strategy)

가중치 행렬(weight matrix)을 초기화하는 네 가지 방법입니다. 각 함수는 fan_in개의 열(column)과 fan_out개의 행(row)을 가진 2차원 행렬(2D matrix), 즉 리스트의 리스트(list of lists)를 반환합니다.

import math
import random


def zero_init(fan_in, fan_out):
    return [[0.0 for _ in range(fan_in)] for _ in range(fan_out)]

def random_init(fan_in, fan_out, scale=1.0):
    return [[random.gauss(0, scale) for _ in range(fan_in)] for _ in range(fan_out)]

def xavier_init(fan_in, fan_out):
    std = math.sqrt(2.0 / (fan_in + fan_out))
    return [[random.gauss(0, std) for _ in range(fan_in)] for _ in range(fan_out)]

def kaiming_init(fan_in, fan_out):
    std = math.sqrt(2.0 / fan_in)
    return [[random.gauss(0, std) for _ in range(fan_in)] for _ in range(fan_out)]

Step 2: 활성화 함수(activation function)

sigmoid, tanh, ReLU를 사용해 각 초기화 전략이 의도한 활성화 함수에서 작동하는지 비교합니다.

def sigmoid(x):
    x = max(-500, min(500, x))
    return 1.0 / (1.0 + math.exp(-x))


def tanh_act(x):
    return math.tanh(x)


def relu(x):
    return max(0.0, x)

Step 3: 50층 순전파(forward pass)

무작위 데이터를 깊은 네트워크(deep network)에 통과시키고 각 층의 평균 활성화 크기(mean activation magnitude)를 측정합니다. 0 초기화, 무작위 N(0,1), 무작위 N(0,0.01), Xavier, Kaiming을 비교하면 신호가 사라지거나 폭발하는 경우를 볼 수 있습니다.

def forward_deep(init_fn, activation_fn, n_layers=50, width=64, n_samples=100):
    random.seed(42)
    layer_magnitudes = []

    inputs = [[random.gauss(0, 1) for _ in range(width)] for _ in range(n_samples)]

    for layer_idx in range(n_layers):
        weights = init_fn(width, width)
        biases = [0.0] * width

        new_inputs = []
        for sample in inputs:
            output = []
            for neuron_idx in range(width):
                z = sum(weights[neuron_idx][j] * sample[j] for j in range(width)) + biases[neuron_idx]
                output.append(activation_fn(z))
            new_inputs.append(output)
        inputs = new_inputs

        magnitudes = []
        for sample in inputs:
            magnitudes.append(sum(abs(v) for v in sample) / width)
        mean_mag = sum(magnitudes) / len(magnitudes)
        layer_magnitudes.append(mean_mag)

    return layer_magnitudes

Step 4: 실험 실행(The Experiment)

0 초기화, 무작위 N(0,1), 무작위 N(0,0.01), Xavier + sigmoid, Xavier + tanh, Kaiming + ReLU 조합을 모두 실행합니다. 주요 층에서 크기를 출력합니다.

def run_experiment():
    configs = [
        ("Zero init + Sigmoid", lambda fi, fo: zero_init(fi, fo), sigmoid),
        ("Random N(0,1) + ReLU", lambda fi, fo: random_init(fi, fo, 1.0), relu),
        ("Random N(0,0.01) + ReLU", lambda fi, fo: random_init(fi, fo, 0.01), relu),
        ("Xavier + Sigmoid", xavier_init, sigmoid),
        ("Xavier + Tanh", xavier_init, tanh_act),
        ("Kaiming + ReLU", kaiming_init, relu),
    ]

    print(f"{'Strategy':<30} {'L1':>10} {'L5':>10} {'L10':>10} {'L25':>10} {'L50':>10}")
    print("-" * 80)

    for name, init_fn, act_fn in configs:
        mags = forward_deep(init_fn, act_fn)
        row = f"{name:<30}"
        for idx in [0, 4, 9, 24, 49]:
            val = mags[idx]
            if val > 1e6:
                row += f" {'EXPLODED':>10}"
            elif val < 1e-6:
                row += f" {'VANISHED':>10}"
            else:
                row += f" {val:>10.4f}"
        print(row)

Step 5: 대칭 시연(Symmetry Demonstration)

0 초기화에서는 뉴런의 출력이 모두 동일합니다. 너비(width)가 4여도 유효 파라미터는 1개처럼 움직입니다. 이 실험은 무작위 초기화가 왜 대칭 깨기(symmetry breaking)에 필요한지 보여 줍니다.

def symmetry_demo():
    random.seed(42)
    weights = zero_init(2, 4)
    biases = [0.0] * 4

    inputs = [0.5, -0.3]
    outputs = []
    for neuron_idx in range(4):
        z = sum(weights[neuron_idx][j] * inputs[j] for j in range(2)) + biases[neuron_idx]
        outputs.append(sigmoid(z))

    print("\n대칭 시연(4 neurons, zero init):")
    for i, out in enumerate(outputs):
        print(f"  Neuron {i}: output = {out:.6f}")
    all_same = all(abs(outputs[i] - outputs[0]) < 1e-10 for i in range(len(outputs)))
    print(f"  모두 동일한가: {all_same}")
    print(f"  유효 파라미터: 1개({len(weights) * len(weights[0])}개가 아님)")

Step 6: 층별 크기 리포트(Layer-by-Layer Magnitude Report)

50층을 지나는 활성화 크기를 시각적인 막대 차트로 출력합니다.

def magnitude_report(name, magnitudes):
    print(f"\n{name}:")
    for i, mag in enumerate(magnitudes):
        if i % 5 == 0 or i == len(magnitudes) - 1:
            if mag > 1e6:
                bar = "X" * 50 + " EXPLODED"
            elif mag < 1e-6:
                bar = "." + " VANISHED"
            else:
                bar_len = min(50, max(1, int(mag * 10)))
                bar = "#" * bar_len
            print(f"  Layer {i+1:3d}: {bar} ({mag:.6f})")

사용하기

PyTorch는 다음과 같은 내장 초기화(built-in initialization) 함수를 제공합니다.

import torch.nn as nn

layer = nn.Linear(512, 256)

nn.init.xavier_uniform_(layer.weight)
nn.init.xavier_normal_(layer.weight)

nn.init.kaiming_uniform_(layer.weight, nonlinearity='relu')
nn.init.kaiming_normal_(layer.weight, nonlinearity='relu')

nn.init.zeros_(layer.bias)

nn.Linear(512, 256)을 호출하면 PyTorch는 기본적으로 Kaiming 균등 초기화(uniform initialization)를 사용합니다. 그래서 단순한 네트워크가 "그냥 작동"하는 경우가 많습니다. 하지만 직접 만든 아키텍처(custom architecture)를 사용하거나 20개 층보다 깊어지면, 초기화가 무엇을 하는지 이해하고 필요할 때는 기본값을 직접 덮어써(override) 야 합니다.

트랜스포머(transformer)에서는 HuggingFace 모델이 보통 _init_weights 메서드(method)에서 초기화를 처리합니다. GPT-2는 잔차 투영(residual projection)을 1/sqrt(N) 계열로 스케일링합니다. 트랜스포머를 처음부터 직접 만든다면 이 부분을 본인이 추가해야 합니다.

산출물 만들기

이 lesson에서 만드는 산출물은 다음과 같습니다.

• outputs/prompt-init-strategy.md: 가중치 초기화 문제를 진단하고 적절한 전략을 추천하는 prompt

연습문제

1. (쉬움) SELU 활성화 함수용 LeCun 초기화(Var = 1/fan_in)를 추가합니다. LeCun + tanh를 Xavier + tanh와 비교합니다.
2. (중간) GPT-2의 잔차 스케일링(residual scaling)을 구현합니다. 잔차 스트림에 더하기 전에 각 층의 출력을 1/sqrt(2*N)로 곱하고, 스케일링 유무에 따른 신호 크기의 증가 속도를 비교합니다.
3. (중간) 네트워크의 층 차원(layer dimension)과 활성화 함수 종류를 받아 알맞은 초기화를 추천하는 초기화 건강 점검(init health check) 함수를 만듭니다.
4. (중간) fan_in = 16과 fan_in = 1024에서 실험합니다. Xavier와 Kaiming은 fan_in에 맞게 조정되지만 단순 무작위 초기화는 그렇지 않습니다. 층이 커질수록 "작동"과 "고장" 사이의 간격이 얼마나 벌어지는지 보여 줍니다.
5. (어려움) 직교 초기화(orthogonal initialization)를 구현합니다. 무작위 행렬의 특이값 분해(SVD)를 계산하고 직교 행렬 U를 사용합니다. 50개 층 ReLU 네트워크에서 Kaiming과 비교합니다.

핵심 용어

더 읽을거리

• Glorot & Bengio, "Understanding the difficulty of training deep feedforward neural networks" (2010) — Xavier 초기화를 처음 제안한 원 논문이며 분산 분석을 함께 다룹니다.
• He et al., "Delving Deep into Rectifiers" (2015) — ReLU 네트워크를 위한 Kaiming 초기화를 도입한 논문입니다.
• Radford et al., "Language Models are Unsupervised Multitask Learners" (2019) — GPT-2 논문과 잔차 스케일링 초기화의 배경을 제공합니다.
• Mishkin & Matas, "All You Need is a Good Init" (2016) — 해석적 공식 대신 경험적으로 단위 분산을 맞추는 LSUV(Layer-Sequential Unit-Variance) 초기화를 다룹니다.