ML

[Normalization Techniques] 2탄 Decomposing Normalization

ML_MJSHIN 2022. 5. 30. 18:30

   앞선 포스팅에서 다루었던 NORMALIZING ACTIVATIONS AS FUNCTIONS 의 Batch Normalization 은 굉장히 많이 사용되고 있는 Normalization techniques 중 하나입니다. 

 

  하지만 아래의 2가지 문제점이 크게 있다고 말씀드렸었는데, 이를 해결하기 위한 새로운 normalization techniques 가 등장하였고 좋은 성능을 보여주고 있습니다. 

  1. inference시에 operation이 다르게 작동하여 performance에 불안정성이 존재
  2. small batch size를 사용하는 경우 학습과정 및 inference시에 error가 큼 (input data에 대한 mean & std 가 mini-batch 단위로 수행되어 학습에 사용됨)

  실제로 최근에 주목받고 있는 Denoising Diffusion Probabilistic Model (DDPM)같은 경우 github 정식 코드에서는 GroupNormalization을 사용하는 Unet 구조 기반의 모델을 사용하는 모습을 볼 수 있습니다. 

 

  이번 포스팅에서는 다양한 BN을 대체하기 위한 방법들에 대해서 살펴보도록 하겠습니다.

 

   NORMALIZING ACTIVATIONS AS FUNCTIONS 방법을 크게 3 단계로 나눌 수 있습니다.

  1. Normalization Area Partitioning (NAP)
  2. Normalization Operation (NOP)
  3. Normalization Representation Recovery (NRR) 

   

    BN을 예로 위의 단계들에 대해서 설명을 하면 다음과 같습니다. BN에서 1단계의 NAP은 BN이 batch 와 데이터의 heigth와 width 에 대해서 Area로 묶게 됩니다. 즉 NAP은 아래의 그림의 파란색 영역과 같이 Channel 별로는 다르게 수행되며 NHW 에 대해서는 하나의 Area가 되어 operation이 이루어 집니다. 

    BN의 NOP은 파란색 Area에서 정규화를 수행하는 것을 의미합니다. 

    마지막으로 NRR은 γ 와 β 라는 learnable parameter를 통해 정규화 된 activation 값을 다시 affine transform을 수행하여 분포를 변형시키는 과정을 의미합니다. 정규화는 activation의 분포를 제한해서 optimization에 도움이 될 수 있지만 이러한 제약은 ML 모델이 가지는 feature에 대한 activation의 representation ability를 떨어뜨릴 수 있습니다. 그래서 NRR과 같은 affine transform은 이런 representation ability를 어느 정도 복구시켜주기 위해서 존재한다고 생각할 수 있습니다.

  

    아래의 코드는 pytorch에서 작동할 수 있도록 구현한 BatchNorm 2d 의 코드입니다. 각 코드를 보면서 위의 3단계가 어떤 부분에 속하는지에 대해서 주석을 달아두었으니 살펴보시면 좋을 것 같습니다.

class BatchNorm2d(nn.BatchNorm2d):
    def __init__(self, num_features, eps=1e-5, momentum=0.1,
                 affine=True, track_running_stats=True):
        super().__init__(num_features, eps, momentum, affine, track_running_stats)

    def forward(self, input):
        exponential_average_factor = 0.0

        if self.training and self.track_running_stats:
        	# cumulative mean과 std를 계산할 때 moving average의 weighted factor를 계산
            # 균일하게 과거와 현재를 더해 나아갈 지 or 과거나 현재에 가중치를 줄 것인지 ...
            if self.num_batches_tracked is not None:
                self.num_batches_tracked += 1
                if self.momentum is None:  # use cumulative moving average
                    exponential_average_factor = 1.0 / float(self.num_batches_tracked)
                else:  # use exponential moving average
                    exponential_average_factor = self.momentum

        if self.training:
        	# Normalization Area Partitioning (NAP) 수행 
       		# (B, C, H, W) 가 pytorch image input
            # (B, H, W)에 대해서 mean 과 variance를 계산하여 standardization 수행이 필요 
            mean = input.mean([0, 2, 3]) 
            var = input.var([0, 2, 3], unbiased=False)
            
            n = input.numel() / input.size(1)
            with torch.no_grad():
                self.running_mean = exponential_average_factor * mean\
                    + (1 - exponential_average_factor) * self.running_mean
                # update running_var with unbiased var
                self.running_var = exponential_average_factor * var * n / (n - 1)\
                    + (1 - exponential_average_factor) * self.running_var
        else:
        	# Inference 시에는 학습 중 누적 된 mean 과 var을 사용하여 값을 추정
            mean = self.running_mean
            var = self.running_var

		# Normalization Operation (NOP) 수행 
        input = (input - mean[None, :, None, None]) / (torch.sqrt(var[None, :, None, None] + self.eps))
        
        if self.affine:
        	# gamma와 beta를 통해서 NRR (Normalization Representation Recovery) 수행
            input = input * self.weight[None, :, None, None] + self.bias[None, :, None, None]

        return input

   그러면, 본격적으로 다른 Normalization에 대해서 알아보도록 하겠습니다.

 

1. Layer Normalization (LN) 

    LN은 layer 의 입력에 대해서 standardization을 수행합니다. 이때, batch 데이터를 묶어서 정규화를 수행했던 BN과 다르게 각각의 데이터 별로 정규화를 따로 수행합니다. 이를 통해서 mini-batch 단위로 묶어서 정규화를 수행하면서 발생했던 문제들을 해결할 수 있습니다.  ((1) small batch에 따른 추정 문제 (2) inference 시에 batch 단위로 training시에 만났던 mean과 std 를 누적해서 사용)

   즉, LN은 training과 inference 시에 동일한 operation이 작동할 수 있습니다. LN은 현재 NLP task에서 상대적으로 좋은 성능을 보여주며, transformer 의 attention layer에서 채택되어 중요하게 사용되고 있습니다. 

2. Group Normalization (GN)

    GN은 LN을 더 일반화 한 방법이라고 생각할 수 있습니다. LN에서 묶여있던 전체 Channel 을 몇 개의 group 으로 나누어서 해당 group 별로 정규화를 수행합니다. 만약 1개의 group으로 나눈다고 한다면 LN과 GN은 같은 NOP를 가지게 됩니다. 

    GN은 주로 small-batch size를 가지고 학습하는 image segmentation이나 object detection 분야에서 좋은 성능을 보여주고 있습니다. 

3. Instance Normalization (IN)

    IN은 LN에서 모든 채널을 하나의 정규화 과정에서 포함시켰다면, 각각의 채널이 독립적으로 정규화에 참여하게 됩니다. 주로 StyleTransfer 같은 분야에서 주로 사용되며 StyleGAN과 StyleGAN2에서도 주로 사용됩니다. 

 

 더 많은 normalization 방법들이 소개되고 있지만 실제로 자주 사용되는 3개의 normalization 방법에 대해서 간단하게 알아보았습니다. 다음 포스팅에서는 NRR에 대해서 자세히 다루고 weight normalziation에 대해서 살펴보도록 하겠습니다.