목차

Overview

이번 과제에서는 Batch Normalization 의 개념과 BatchNorm 레이어의 forward, backward 를 구현한다.

또한 Batch Norm과 Weight Initialization 의 관계, Batch Norm과 Batch Size와의 관계를 Baseline 기반으로 비교한다.

이후 Layer Normalization 의 개념과 LayerNorm 레이어의 forward, backward 구현 후

Layer Norm 과 Batch Size 의 관계를 baseline 기반으로 비교한다.

핵심 아이디어

정규화 계층은 ‘선택한 축’ 을 기준으로 평균(mean)와 분산(variance) 을 계산하고, 입력을 표준화 한다.
매 layer로 전달되는 activation 의 분포를 일정하게 유지해 학습을 더 안정적이고 빠르게 만들기 위해서 수행한다.

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

입력 데이터를 정규화 한다.

$$\hat{x}_i = \frac{x_i-\mu} {\sqrt{\sigma^2+\epsilon}}$$

scale param과 shift를 덧붙인다.
정규화만 수행할 경우 출력이 항상 평균 0, 분산 1 근처로 강제되기 때문에 표현력이 강제된다.
그래서 학습 가능한 파라미터인 scale $\gamma$와 shift $\beta$ 를 추가한다.

$$y_i = \gamma \hat{x}_i + \beta$$

Batch Normalization

Batch Normalization은 pre-activation 의 채널 C를 기준으로 분리하여 통계를 도출한다.
직관적으로, 채널 C를 고정 후 n,h,w 축을 합친다. (N, C, H, W) to (C, )

평균

$$\mu_c = \frac{1}{NHW} \sum_{n,h,w} X_{n,c,h,w}$$

분산

$$\sigma_c^2 = \frac{1}{NHW} \sum_{n,h,w} (X_{n,c,h,w}-\mu_c)^2$$

Layer Normalization

Layer Normalization은 pre-activation 의 샘플 N를 기준으로 분리하여 통계를 도출한다.
직관적으로, 샘플 N을 고정 후 c,h,w 축을 합친다. (N, C, H, W) to (N, )

평균

$$\mu_n = \frac{1}{CHW} \sum_{c,h,w} X_{n,c,h,w}$$

분산

$$\sigma_n^2 = \frac{1}{CHW} \sum_{c,h,w} \left( X_{n,c,h,w}-\mu_n \right)^2$$

구현

Batch Norm

forward

def batchnorm_forward(x, gamma, beta, bn_param):
    N, D = x.shape
    running_mean = bn_param.get("running_mean", np.zeros(D, dtype=x.dtype))
    running_var = bn_param.get("running_var", np.zeros(D, dtype=x.dtype))

    out, cache = None, None

    if mode == "train":
        mean = x.mean(axis=0)
        var = x.var(axis=0)
        
        x_hat = (x - mean) / np.sqrt(var + eps)

        out = gamma * x_hat + beta
        cache = (x, x_hat, mean, var, gamma, eps)

        running_mean = momentum * running_mean + (1 - momentum) * mean
        running_var = momentum * running_var + (1 - momentum) * var

    elif mode == "test":
        x_hat = (x - running_mean) / np.sqrt(running_var + eps)
        out = gamma * x_hat + beta
        cache = None

    bn_param["running_mean"] = running_mean
    bn_param["running_var"] = running_var

    return out, cache

running_mean, running_var는 momentum optim처럼 학습 중 여러 미니배치의 평균과 분산을 누적해서 전체 데이터 분포의 평균/분산을 근사한 값이다.

Batch Norm forward 에서 train 과 test 분기가 발생하는 이유는, train 때는 미니배치 통계로 정규화가 가능하지만, test 때는 현재 배치 통계를 쓰면 출력이 배치 구성에 따라 달라지기 때문이다.

직관적으로, 채널 C를 고정하기 때문에, 출력 자체가 현재 배치의 평균과 분산에 직접 의존한다. (샘플들이 달라지면 평균, 분산이 달라진다. batch-dependent 하다.)

반면 아래에서 나올 Layer Norm의 경우 샘플 N을 고정하고 그 샘플 내부의 특징들로 통계를 계산하므로, 다른 샘플이나 배치 구성에 영향을 받지 않는다. (batch-independent 하다.)

---

backward 1. 그래프 따라가며 계산

참고자료: bprop_batch_norm

%%{init: {"flowchart": {"useMaxWidth": false, "nodeSpacing": 55, "rankSpacing": 22, "padding": 7, "diagramPadding": 3, "curve": "linear"}, "themeVariables": {"fontSize": "11px"}}}%%

flowchart TB
    x1(("x₁"))
    x2(("x₂"))
    mu(("μ"))
    variance(("σ²"))
    xhat1(("x̂₁"))
    xhat2(("x̂₂"))
    y1(("y₁"))
    y2(("y₂"))

    x1 --> mu
    x2 --> mu

    x1 --> variance
    x2 --> variance
    mu --> variance

    x1 --> xhat1
    x2 --> xhat2

    mu --> xhat1
    mu --> xhat2

    variance --> xhat1
    variance --> xhat2

    xhat1 --> y1
    xhat2 --> y2

    classDef nodeStyle fill:#ffffff,stroke:#333333,color:#111111;
    class x1,x2,mu,variance,xhat1,xhat2,y1,y2 nodeStyle;
    
    classDef nodeStyle fill:#ffffff,stroke:#333333,color:#111111,font-size:11px;
    class x1,x2,mu,variance,xhat1,xhat2,y1,y2 nodeStyle;

    linkStyle default interpolate basis;

위 참고자료는 두 개의 입력이 $x_1, x_2$만 있다고 가정한다.

1. $\beta$ graident

$$\frac{\partial L}{\partial \beta} = \sum_{i=1}^{2} \frac{\partial L}{\partial y_i}$$

2. $\gamma$ gradient

$$\frac{\partial L}{\partial \gamma} = \sum_{i=1}^{2} \frac{\partial L}{\partial y_i} \hat{x}_i$$

3. normalized input gradient

$$\frac{\partial L}{\partial \hat{x}_i} = \frac{\partial L}{\partial y_i} \gamma$$

4. variance gradient

$$\frac{\partial L}{\partial \sigma^2} = \sum_{i=1}^{2} \frac{\partial L}{\partial \hat{x}_i} (x_i-\mu) \left(-\frac{1}{2}\right) (\sigma^2+\epsilon)^{-3/2}$$

5. mean gradient

$$\frac{\partial L}{\partial \mu} = \sum_{i=1}^{2} \frac{\partial L}{\partial \hat{x}_i} \frac{-1}{\sqrt{\sigma^2+\epsilon}} + \frac{\partial L}{\partial \sigma^2} \frac{\sum_{i=1}^{2}-2(x_i-\mu)}{2}$$

• 사실상 두번째 항은 0이 되지만 전체 수식 자체를 구현했다. 0이 되는 이유는 $\mu$가 배치들의 평균이기 때문에, 각 값에서 평균을 뺀 편차들의 합은 0이 되기 때문이다.

def batchnorm_backward(dout, cache):

    dx, dgamma, dbeta = None, None, None

    x, x_hat, mean, var, gamma, eps = cache
    N, D = dout.shape

    # beta graident
    dbeta = dout.sum(axis=0) #(D,)

    # gamma gradient
    dgamma = np.sum(dout * x_hat, axis=0)

    # normalized input gradient(dx_hat)
    dx_hat = dout * gamma

    # variance gradient
    dvar = np.sum(
      dx_hat 
      * (x - mean) 
      * (-0.5) 
      * (var + eps)
      **(-1.5), axis=0)
    
    # mean gradient
    dmean = (np.sum(
      dx_hat 
      * (-1 / (np.sqrt(var + eps))), axis=0) 
      + dvar * np.sum(-2 * (x - mean), axis=0) / N
    )

    # input gradient
    dx = (
      dx_hat 
      * (1 / np.sqrt(var + eps)) 
      + dvar 
      * (2 * (x - mean)) / N 
      + dmean * 1 / N
    )

    return dx, dgamma, dbeta

dx error:  1.7029261167605239e-09
dgamma error:  7.420414216247087e-13
dbeta error:  2.8795057655839487e-12

---

backward 2. 수식을 정리하여 계산

그래프를 따라가며 계산 하는것이 아닌, 수식을 정리하여 빠르게 계산한다. 참고자료: stanford.edu/batch-norm.pdf

포인트는 $\frac{\partial y_j}{\partial x_i}$를 중간 목표로 정하는것이다.

$$\frac{\partial y_j}{\partial x_i} = \frac{\partial}{\partial x_i} \left[ \frac{x_j-\mu}{\sigma} \right] \tag{7}$$

$y_i$가 아니라 $y_j$인 이유는 $x_i$ 하나가 $\mu$와 $\sigma$를 바꾸고, $\mu,\sigma$는 batch의 모든 $y_j$가 공유하기 때문이다.

중요한 것은 $x_i$가 모든 $y_j$에 영향을 주므로 경로가 여러 개라는 것이다.

따라서 $\frac{\partial\ell}{\partial x_i}$을 구하기 위한 연쇄법칙은 다음과 같다.

$$\frac{\partial\ell}{\partial x_i} = \sum_{j=1}^N \frac{\partial\ell}{\partial y_j} \frac{\partial y_j}{\partial x_i} \tag{11}$$

최종적으로 $dx$는 다음과 같이 유도된다.

$$\boxed{ \frac{\partial L}{\partial x_i} = \frac{\gamma}{\sigma} \left[ d_i - \frac1N\sum_jd_j - \hat x_i \frac1N\sum_j\hat x_jd_j \right] }$$

또는 아래처럼 적을 수 있다.

$$\boxed{ dx = \frac{\gamma}{\sigma} \left[ dout - \operatorname{mean}(dout) - \hat x\, \operatorname{mean}(\hat x\odot dout) \right] }$$

def batchnorm_backward_alt(dout, cache):
    dx, dgamma, dbeta = None, None, None

    x, x_hat, mean, var, gamma, eps = cache

    dbeta = dout.sum(axis=0) #(D,)
    dgamma = np.sum(dout * x_hat, axis=0)

    dx = (
      (gamma / np.sqrt(var + eps))
      * (
        dout 
        - np.mean(dout, axis=0)
        - x_hat
        * np.mean(x_hat * dout, axis=0)
      )
    )
    return dx, dgamma, dbeta

dx difference:  7.919015490255652e-13
dgamma difference:  0.0
dbeta difference:  0.0
speedup: 3.84x

그래프를 따라가며 역전파를 계산하는 것 대비 유의미한 속도 향상을 보였다.

PDF에서 유도 순서는 아래 순서였다. 중간에 판별식과 치환이 있어 조금 헷갈렸다.

$$\boxed{ \mu \overset{(1)}{\longrightarrow} v \overset{(2)-(5)}{\longrightarrow} \sigma \overset{(6)}{\longrightarrow} \frac{\partial y_j}{\partial x_i} \overset{(7)-(10)}{\longrightarrow} \frac{\partial L}{\partial x_i} \overset{(11)-(13)}{\longrightarrow} \text{vectorized form }(14) }$$

---

Layer Norm

forward

def layernorm_forward(x, gamma, beta, ln_param):
    # transpose 를 쓰면 batch norm forward 와 유사하다.
    x_t = x.T #(D, N)

    mean = x_t.mean(axis=0) #(N, )
    var = x_t.var(axis=0) #(N, )
    
    x_hat = (x_t - mean) / np.sqrt(var + eps)
    x_hat = x_hat.T #(N, D)
    out = (gamma * x_hat + beta)
    cache = (x, x_hat, mean, var, gamma, eps)

    return out, cache

Batch Norm 에서 추론 시 쓰던 running_mean, running_var가 필요하지 않다. 샘플별 통계를 내기 때문에 현재 배치 구성에 영향을 받지 않기 때문이다.

backward

이전 Batch Norm backward-alt와 비슷하다.
다만 N을 고정하고 평균을 내는 축이 D이다.
Layer Norm은 샘플 축 $N$을 고정하고 나머지 feature 축에 대해 평균을 내는데, d가 변할때 $\gamma_d$ 도 같이 변하게 된다.
따라서 $\gamma_d$를 sum 밖으로 빼낼 수 없으므로, 먼저 $d\gamma$ 를 구하고 그것을 바로 이용해서 $d\hat{x}$ 을 계산하고 최종적으로 $dx$를 계산하게 된다.

$$\frac{\partial L}{\partial \hat{x}_{n,d}} = \frac{\partial L}{\partial y_{n,d}}\gamma_d$$ $$\frac{\partial L}{\partial x_{n,d}} = \frac{1}{\sqrt{\sigma_n^2+\epsilon}} \left( \frac{\partial L}{\partial \hat{x}_{n,d}} - \frac{1}{D} \sum_{j=1}^{D} \frac{\partial L}{\partial \hat{x}_{n,j}} - \hat{x}_{n,d} \frac{1}{D} \sum_{j=1}^{D} \hat{x}_{n,j} \frac{\partial L}{\partial \hat{x}_{n,j}} \right)$$

def layernorm_backward(dout, cache):
    dx, dgamma, dbeta = None, None, None

    x, x_hat, mean, var, gamma, eps = cache

    dbeta = np.sum(dout, axis=0)
    dgamma = np.sum(dout * x_hat, axis=0)

    dx_hat = dout * gamma # (N, D)

    dx = (
        1 / np.sqrt(var[:, None] + eps)
        * (
            dx_hat
            - np.mean(dx_hat, axis=1, keepdims=True)
            - x_hat
            * np.mean(x_hat * dx_hat, axis=1, keepdims=True)
        )
    )

    return dx, dgamma, dbeta

추가 개념들

배치 정규화와 가중치 초기화 스케일의 관계

베이스라인(without batchnorm), batch norm 관찰된 상관관계

최고 검증 정확도 대비 가중치 초기화 스케일 batch norm의 경우 10^-4부터 10^-1까지 꾸준히 높은 val acc를 보였다. 베이스라인의 경우 10^-1 부근에서만 batch norm과 비슷한 val acc를 보였으며, 평균적으로 batch norm 대비 val acc가 낮았다. 다만 두 경우 모두 10^0근처로 갈수록 검증 정확도가 떨어지는 모습을 보였다.

최고 훈련 정확도 대비 가중치 초기화 스케일 batch norm 모델의 경우 0.3 부근의 train acc를 보였다. 평균적으로 baseline 대비 훈련 정확도가 높은 양상을 보였다. 다만 10^-1 부근에서는 베이스라인의 train acc 최고점이 가장 높았다.

최종 훈련 손실 대비 가중치 초기화 스케일 traing loss의 경우 평균적으로 batch norm 모델이 베이스라인 대비 낮은 로스를 보였으며, loss가 weight scale에 따라 일정하게 유지되다가 10^-1부터 급증하는 양상을 보였다. 베이스라인의 경우 평균적으로 batch norm보다는 Loss가 높았으나 10^-1 부근에서만큼은 가장 loss가 낮았다.

이유
전반적으로 BatchNorm은 weight initialization scale에 대한 모델의 민감도를 크게 감소시켜, 훨씬 넓은 초기화 범위에서 안정적인 학습을 가능하게 한다.

각 레이어에서 생성되는 pre-activation을 정규화하여 activation의 크기와 분포가 가중치 초기화 스케일에 지나치게 의존하지 않도록 한다. 따라서 다음 레이어로 전달되는 activation의 scale이 보다 안정적으로 유지되고, 전체 네트워크가 넓은 범위의 weight initialization scale에서도 학습될 수 있다.

---

배치 정규화와 배치 사이즈

배치 정규화와 배치 사이즈 간의 관찰된 상관관계
훈련 정확도는 Batch Normalization이 적용된 모델에서 배치 사이즈가 커질수록 높아지는 양상을 보였다. 동일한 배치 사이즈 5에서는 Batch Normalization을 적용한 모델이 baseline보다 낮은 훈련 정확도를 보였다. 반면 Batch Normalization을 적용한 모델들끼리 비교했을 때 배치 사이즈 10, 50으로 증가할수록 더 높은 훈련 정확도를 보였으며, 배치 사이즈 50인 모델이 가장 높은 훈련 정확도를 기록했다.

검증 정확도에서는 Batch Normalization을 적용한 모델 중 배치 사이즈가 큰 케이스가 전반적으로 높은 정확도를 보였으며, 가장 높은 검증 정확도 역시 Batch Normalization을 적용하고 배치 사이즈가 가장 큰 50인 모델에서 관찰되었다. 동일한 배치 사이즈 5를 사용한 경우, Batch Normalization 모델과 baseline의 검증 정확도 차이는 크지 않았으며 epoch에 따라 서로 높은 정확도를 보이는 구간이 달랐다.

이유
배치 사이즈가 커질수록 평균과 분산 계산에 사용되는 샘플 수가 증가하여 Batch Normalization에 사용되는 통계값을 보다 안정적으로 추정할 수 있기 때문이다.

레이어 정규화와 배치 사이즈

baseline 대비 모두 뚜렷한 Training Accuracy와 Validation Accuracy를 보였다.

LayerNorm의 정규화 자체는 batch size에 의존하지 않지만, 모델의 학습 과정(gradient)은 batch size에 영향을 받는다.

$$\mu_n = \frac{1}{D} \sum_{d=1}^{D} x_{n,d}$$

즉 parameter update에 쓰는 gradient는 mini batch 평균이다.

$$g_B = \frac{1}{B} \sum_{i=1}^{B} \nabla_{\theta} L_i$$

따라서 배치 사이즈 $B$ 가 바뀌면 gradient 의 성질이 달라진다. 적은 샘플(n=5)을 평균내면 gradient noise가 커지고, 상대적으로 많은 샘플(n=50)에서는 gradient가 안정적인 모습을 보인다.

gradient의 추정 variance는 대략적으로 다음과 같이 감소한다고도 볼 수 있다.

$$\operatorname{Var}(g_B) \propto \frac{1}{B}$$

그리고 validation accuracy를 보면 세 LayerNorm 모델이 거의 비슷하다.
batch size가 training 에는 꽤 영향을 줬지만, LayerNorm 자체는 batch 통계에 의존하지 않기 때문에 BatchNorm처럼 batch size가 작다고 normalization 통계 자체가 망가지는 현상은 없다.

레이어 정규화와 feature

Layer Normalization이 잘 작동하지 않을 가능성이 높은 경우는?

정답: 2. Feature dimension이 매우 작은 경우

LayerNorm은 한 샘플 내부의 feature들을 이용해 평균과 분산을 계산한다. 따라서 feature 개수가 너무 적으면 평균과 분산을 추정할 값이 부족해 통계가 불안정해지고, 정규화의 효과도 떨어질 수 있다.

https://www.adityaagrawal.net/blog/deep_learning/bprop_batch_norm

Stanford Batch Normalization Notes를 보고 유도.