Deep Networks with Stochastic Depth (ECCV2016)

안녕하세요. 지난 포스팅의 [IC2D] Identity Mappings in Deep Residual Networks (ECCV2016)에서는 ResNet에서 Identity Mapping의 중요성과 activation function의 위치에 따른 성능 변화를 분석하였습니다. 결과적으로 Batch Normalization과 ReLU를 Skip connection에서 제일 먼저 적용하는 것이 가장 높은 성능을 얻었음을 확인하였죠. 오늘은 ResNet의 변형 구조인 Stochastic ResNet에 대해서 소개해드리도록 하겠습니다. 

 

Deep Networks with Stochastic Depth

 

Background

ResNet과 PreAct ResNet의 실험결과로 네트워크의 깊이는 성능을 향상한다는 것을 알게 되었습니다. 하지만, 문제는 네트워크의 깊이는 커질수록 그만큼 메모리를 많이 차지하고 학습 속도가 굉장히 늘어나게 됩니다. 본 논문에서는 이를 해결하기 위해서 학습 동안에는 ResNet의 특정 블록을 확률적으로 제거하고 Identity Mapping으로 치환하여 사용하는 것을 제안합니다. 물론, 평가 단계에서는 전체 네트워크를 사용하게 되죠. 이론상으로는 학습할 때는 기존의 네트워크보다 더 적은 개수의 블록이 포함되기 때문에 학습속도를 감소하게 되지만 평가할 때는 전체 네트워크를 사용하기 때문에 성능은 보존됩니다. 

 

뿐만 아니라, 깊은 네트워크를 학습할 때 Gradient vanishing problem 말고도 feature diminishing problem도 존재합니다. gradient vanishing은 다들 아시다싶이 backward 과정에서 손실함수에 대한 각 레이어에 대한 기울기가 얕아질수록 소실되는 문제입니다. 이를 해결하기 위해 Highway Network 및 ResNet에서는 skip connection 구조를 차용하였습니다. feature diminishing은 이와는 반대로 forward 과정에서 발생하는 문제로 feature가 여러 계층의 합성곱 및 풀링이 적용되면서 feature가 소실되는 현상입니다. 물론 skip connection을 통해 어느 정도 해결은 할 수 있지만, 두 문제 모두 핵심은 네트워크를 얕게 만들어야 해결할 수 있죠. 그래서 본 논문에서는 학습에서는 얕게, 평가에서는 깊게 네트워크를 구성하게 됩니다. 

 

Residual Network

본 논문의 방법은 ResNet을 중심으로 적용하기 때문에 먼저 ResNet에 대해서 간략하게 설명하도록 하겠습니다. 기존의 VGGNet과 GoogLeNet에서는 네트워크가 깊어질 수록 오히려 성능이 낮아지는 문제가 있었기 때문에 ResNet에서는 이를 타파하고자 Skip connection을 추가하였습니다. 그림으로 그리면 다음과 같습니다. 

 

 

이를 통해 다음과 같은 수식으로 나타낼 수 있죠. 

 

$$H_{l} = \text{ReLU}(f_{l}(H_{l - 1}) + \text {id}(H_{l - 1}))$$

 

여기서 $H_{l}$은 $l$번째 계층의 출력, $f_{l}$은 $l$번째 계층의 합성곱 계층,  $\text {id}$는 identity function, $\text {ReLU}$는 $ReLU$ 합성곱을 의미합니다. 이때, $\text {id}$는 skip connection의 output과 Residual Block 사이의 shape이 맞지 않으면 단순한 $1 \times 1$ 합성곱 계층을 이용해서 resolution과 channel의 개수를 맞추게 됩니다. 

 

Stochastic Depth에서는 해당 Residual Block을 통째로 날려서 사용할지, 아니면 그대로 사용할지에 대한 확률을 적용하여 학습을 진행하기 때문에 보다 빠르게 학습할 수 있습니다. 

 

DropOut

대표적인 네트워크 정규화 방법으로 DropOut, DropConnect, MaxOut 등의 방법이 있습니다. 특히, DropOut은 네트워크의 특정 뉴런의 개수를 확률적으로 제거하게 되죠. 이러한 맥락에서 확률적으로 Residual Block을 제거하는 Stochastic Depth 역시 정규화 방법으로 확장할 수 있음을 알 수 있습니다. DropOut은 뉴런의 개수를 줄이기 때문에 더 thin 한 네트워크를 구성하는 반면 Stochastic Depth는 깊이를 줄이기 때문에 더 shallow 한 네트워크를 구성할 수 있게 도와주죠. 본 논문에서는 Stochastic Depth가 정규화 방법으로써도 훌륭함을 보이고자 합니다. 특히, DropOut은 Batch Normalization과 함께 사용하게 되면 성능이 떨어지지만 Stochastic Depth는 성능이 보존되는 것을 통해 좀 더 generic 하게 사용할 수 있다는 것이 장점입니다. 

 

Stochastic Depth

 

위 그림은 Stochastic Depth에서 제안하는 그림입니다. 각 계층별로 확률 (1.0 ~ 0.5)이 부여되어 각 확률에 따라 네트워크의 계층이 활성 (active) 또는 비활성 (inactive) 될지 결정하게 됩니다. 이를 $b_{l}$이라는 베르누이 확률변수로 정의하도록 하겠습니다. 즉, $b_{l} = 1$이면 $l$번째 계층이 활성화되어 기존의 Residual Block을 사용하게 됩니다. 만약, $b_{l} = 0$이라면 $l$번째 계층이 비활성화되어 Identity Function으로 치환되죠. 따라서 다음과 같이 정리할 수 있습니다. 

 

$$\begin{cases} &b_{l} = 1 \Rightarrow H_{l} = \text {ReLU}(f_{l}(H_{l - 1}) + \text {id}(H_{l - 1})) \\ &b_{l} = 0 \Rightarrow H_{l} = \text {id}(H_{l - 1})\end {cases}$$

 

그렇다면 각 계층에 대한 생존확률 $p_{l} = Pr(b_{l} = 1)$을 어떻게 정의할 수 있을까요? 가장 단순한 방법은 모든 계층에 대해서 동일한 확률을 적용하는 것입니다. 하지만, 본 논문에서는 얕은 계층은 최대한 살리고 깊은 계층을 줄이기 위해 1.0에서 시작하여 마지막 계층의 확률을 0.5로 지정하여 각 계층의 확률을 다음과 같이 선형적으로 정의하였습니다. 위 그림에서 회색 점선을 의미하죠. 

 

$$p_{l} = 1 - \frac {l}{L} (1 - p_{L})$$

 

여기서 $L$은 마지막 계층으로 $p_{L} = 0.5$로 고정하였습니다. 이때, 한 가지 더 생각해봐야 할 것은 저희가 확률적으로 특정 계층을 제거했기 때문에 계층의 개수 역시 확률적으로 변하게 됩니다. 이는 계층의 개수 $L$이 확률변수가 됨을 의미하고 저희가 계층의 개수에 대한 기댓값을 구할 수 있음을 알 수 있습니다. 확률변수로 표현된 계층의 개수를 $\tilde {L}$이라고 하면 다음과 같이 구할 수 있습니다. 

 

$$\mathbb {E}(\tilde {L}) = \sum_{i = 1}^{L} p_{l} = \frac {3L - 1}{4} \approx \frac {3}{4} L$$

 

예를 들어, 110개의 계층을 가지는 ResNet-110은 총 54개의 Residual Block을 가지고 있기 때문에 $L = 54$가 됩니다. 따라서, 기대 계층 개수는 $\mathbb {E}(\tilde {L}) \approx 40$임을 알 수 있죠. 

 

마지막으로 고려해야 할 것은 평가단계에서 확률적으로 적용된 계층의 출력을 어떻게 조정해 줄 것인가입니다. 이는 아주 간단하게 기존의 출력값에서 해당 계층의 확률값을 곱해주면 됩니다. 

 

$$H^{\text {test}}_{l} = \text {ReLU}(p_{l} f_{l}(H^{\text {test}}_{l - 1}) + \text {id}(H^{\text {test}}_{l - 1}))$$

 

Stochastic Depth의 성능 향상에서 가장 큰 기여를 하는 것은 확률적으로 깊이를 제거함으로써 $2^{L}$개의 독립적인 네트워크의 정보가 하나로 응축되어 최종 평가에 사용되었기 때문입니다. 이를 implicit model ensemble이라고 언급하였습니다. 

 

Experiment Results

본 논문에서는 총 4개의 데이터셋 (CIFAR10, CIFAR100, SVHN, ImageNet)에 대해서 실험을 진행하였습니다. 

 

 

Stochastic Depth와 Constant Depth를 비교해 보면 ImageNet을 제외하고는 3가지 데이터셋에 대해서 모두 높은 성능을 달성하였습니다. 

 

Implementation

 

class StochasticBottleNeck(nn.Module):
    """Residual block for resnet over 50 layers
    """
    expansion = 4
    def __init__(self, prob, in_channels, out_channels, stride=1):
        super(StochasticBottleNeck, self).__init__()

        self.conv1 = nn.Conv2d(in_channels, out_channels, kernel_size=1, bias=False)
        self.bn1   = nn.BatchNorm2d(out_channels)
        self.conv2 = nn.Conv2d(out_channels, out_channels, stride=stride, kernel_size=3, padding=1, bias=False)
        self.bn2   = nn.BatchNorm2d(out_channels)
        self.conv3 = nn.Conv2d(out_channels, out_channels * StochasticBottleNeck.expansion, kernel_size=1, bias=False)
        self.bn3   = nn.BatchNorm2d(out_channels * StochasticBottleNeck.expansion)
        self.relu  = nn.ReLU(inplace=True)

        self.shortcut = nn.Sequential()

        if stride != 1 or in_channels != out_channels * StochasticBottleNeck.expansion:
            self.shortcut = nn.Sequential(
                nn.Conv2d(in_channels, out_channels * StochasticBottleNeck.expansion, stride=stride, kernel_size=1, bias=False),
                nn.BatchNorm2d(out_channels * StochasticBottleNeck.expansion)
            )

        self.prob = prob
        self.m = torch.distributions.bernoulli.Bernoulli(torch.Tensor([self.prob]))

    def forward(self, x):
        if self.training:
            if torch.equal(self.m.sample(), torch.ones(1)):
                self.conv1.weight.requires_grad = True
                self.conv2.weight.requires_grad = True
                self.conv3.weight.requires_grad = True

                output = self.conv1(x)
                output = self.bn1(output)
                output = self.relu(output)

                output = self.conv2(output)
                output = self.bn2(output)
                output = self.relu(output)

                output = self.conv3(output)
                output = self.bn3(output)

                output = output + self.shortcut(x)
            else:
                self.conv1.weight.requires_grad = False
                self.conv2.weight.requires_grad = False
                self.conv3.weight.requires_grad = False

                output = self.shortcut(x)
        else:
            output = self.conv1(x)
            output = self.bn1(output)
            output = self.relu(output)

            output = self.conv2(output)
            output = self.bn2(output)
            output = self.relu(output)

            output = self.conv3(output)
            output = self.bn3(output)

            output = self.prob * output + self.shortcut(x)

        return output

 

기본적으로 ResNet을 기반으로 코드를 만들었기 때문에 쉽게 이해할 수 있습니다. Stochastic BottleNet Block을 먼저 분석해보겠습니다. 전체적인 구조는 BottleNeck Block과 동일하지만 해당 block을 무시할 것인지에 대한 확률인 self.m와 그 기준인 self.prob가 추가되었습니다. 만약, self.m이 1이 된다면 해당 Block은 그대로 사용되고 0이라면 무시되어야하기 때문에 각 계층의 기울기 추적을 False로 바꾸어줍니다. 

 

test 역시 동일하지만 마지막에 해당 Block이 얼마만큼의 확률로 삭제되었지를 보정해주기 위해 self.prob를 곱하여 해당 Block의 최종 output을 얻을 수 있습니다.