ConvNeXt V1

Paper Review 7 : resnet에서 convnext까지

|Intro

이번에 리뷰할 논문은 기존 ConvNet 모델에 현대적 기법들이 적용됨으로써 재탄생한 'A ConvNet for the 2020s'입니다. 기존의 Convolution 기반 모델들은 Transformer 기반의 모델들이 등장한 이후 상대적으로 부진한 모습을 보여왔는데요, 이는 Convolution의 구조적 문제가 아니라 AI가 발전해오며 적용된 현대적 기법들이 적용되지 않았기 때문이라고 본 논문에서는 주장합니다. 때문에 ViT, CvT, Swin에 이르기까지 적용된 개념들을 ConvNet에 적용함으로써 그 이상의 성능을 달성할 수 있다는 것을 보여주는 것이 본 논문의 핵심입니다.

|Paper Summary

아래 그림에서 볼 수 있다시피 본 논문의 시작은 ResNet입니다. ResNet에서부터 다양한 현대적 기법들이 적용되며 ConvNet 구조를 현대화했을 때의 성능은 Transformer 기반의 swin보다 높다라는 것이 요지입니다.

Modernizing Convnet Roadmap

위 그림은 모델링 관점에서만 설명되었지만, 본 논문에서는 학습 방법에 대한 내용도 당연히 존재합니다. 위 그래프에서 구체적으로 언급되지 않은 사항(Train Recipe)들을 간단히 정리하면 다음과 같습니다.

Convnet vs ConvNext

기존 resnet 모델 및 모델 구조를 이용해 그대로 학습할 경우, 제대로 학습되지 않는 문제가 발생하는데요. 이로 인해 성능이 얼마나 오르는지도 명확하게 판단하기 어렵기 때문에, 위 표에서 Optimizer, Scheduler, Augmentation과 Label Smoothing 및 Gradient Clipping이 적용된 결과를 살펴보고 이를 기준으로 ConvNext 상에서 제시한 기법들을 적용함으로써 어느정도 성능이 올라가는지 확인해보도록 하겠습니다. 이는 본 논문에서 제시한 training recipe의 일부를 선적용한 것입니다.

|Modeling & Training 

새로운 개념들을 고안하고, 모델링하고, 적용했던 이전의 논문들과 달리 본 논문은 우리가 여태까지 포스팅해온 기법들을 ResNet에 적용하는 것입니다. 때문에 본 포스팅도 본 논문의 목적에 충실하게 진행할 것인데요, 본 논문에서 제시한 현대적 기법들은 크게 5가지로 나뉘고, 작게는 12가지로 나뉩니다. 본 포스팅에서는 크게 5가지로 구분해서 sports 데이터셋에 학습해보며 그 결과를 살펴보도록 하겠습니다.

|STEP 00 : ResNet-50

본 논문의 시작은 ResNet입니다. 2015년에 발표되었으며, 현재까지도 많이 사용되는 잔차학습 또는 skip connection 등으로 불리는 기법이 제시된 논문입니다. 해당 논문은 모델의 깊이가 깊어짐에 따라 모델의 성능이 오히려 하락하는 문제를 해결하기 위해 모델이 깊어질수록 성능이 향상될 수 있는 근거를 마련하고자 했고, 이에 기반해 잔차학습이라는 기법을 도입해야 한다 주장했고, 그 효율성을 입증했습니다. 

resnet architecture

참고로 resnet-50의 파라미터는 다음과 같습니다.

resnet-50의 파라미터

swin보다는 작지만 cvt보다는 큰 사이즈의 모델입니다. 여기에 위에서 설명한 Training Recipe 일부(데이터 증강, Optimizer, Scheduler, clip grad norm, Label Smoothing)를 적용해 100epoch 학습시킨 결과는 다음과 같습니다. 이미 ViT에서부터 적용했었던 Stochastic Depth와 LayerScale, 그리고 이전에는 적용하지 않았던 EMA는 본 포스팅에서 설명하면서 마지막에 적용해보도록 하겠습니다. (LLRD는 파인튜닝 과정에서 주로 사용되기 때문에 학습만을 다루는 본 포스팅에선 적용하지 않습니다.)

resnet + train recipe의 성능

f1-score가 0.8879 수준으로 상당히 높게 나왔는데요, 본 데이터셋의 특성상 대형 데이터셋에 사전학습시켜서 파인튜닝하지 않는 이상 이 이상의 성능을 얻기란 쉽지 않습니다(적어도 100epoch의 학습만으로는 어렵습니다). 이후의 성능은 100에포크를 기준으로 했을 때 오히려 하락하게 됩니다. 그 이유에는 여러가지가 있지만, 대표적인 한 가지만 제시하자면 현대적 학습 기법 중 대표적인 것 하나가 바로 더욱 긴 epoch를 통해 학습시키는 것이기 때문입니다. 때문에 모델이 resnet에서 점차 업데이트될수록 학습 epoch를 늘려서 학습하도록 하겠습니다.

이에 대한 모델링 코드는 여기를, 학습 코드는 여기를 눌러 확인할 수 있습니다.

|STEP 01 : Macro Design

Resnet + Macro Design

ResNet을 현대화시키는 첫 번째 과정은 'Macro Design'을 추가하는 것입니다. 이는 Swin의 구조의 떠올리면 이해하기 쉽습니다. Macro Design은 크게 Stage Ratio와 Patchify Stem으로 구성돼 있습니다. 간단히 설명하면 Stage Ratio는 Swin과 같이 4개의 계층적 구조로 학습하고, 그 중 3번째 계층에 더욱 집중하는 것입니다. Patchify Stem은 ViT 이후로 진행했던 바와 같이 이미지를 겹치지 않는 패치로 분할해 임베딩하는 방법입니다. 이에 대해 하나씩 살펴보도록 하겠습니다.

|Stage Ratio

기본적으로 resnet의 계층적 구조와 비율은 '경험적' 결과의 산물입니다. 다양한 파라미터와 조건을 테스트해본 결과, 당시 환경과 조건 하에서 가장 높은 성능을 보인 조합인 것입니다. 이로 인해 경험적으로 resnet도 3번째 스테이지가 중요하다는 것을 알고 있었고 모델의 사이즈가 커질수록 3번째 스테이지의 비중이 늘어나게 됩니다. 이러한 3번째 스테이지는 모델의 백본 구조를 이용해 객체 탐지, 시멘틱 세그멘테이션 등 다양한 분야에 세밀하게 적용되는 과정에서 주요한 역할을 합니다. 

convnext에서 stage compute ratio를 조절하는 내용

하지만 resnet의 비율은 swin과는 다소 다릅니다. 기존에 [3:4:6:3]과 같은 식이었던 각 스테이지의 반복 비율을 swin과 같이 [1:1:3s:1]의 비율로 조절하자 성능이 약간 향상되어 채택했다고 합니다. 이를 코드로 구현하려면 다음과 같습니다.

resnet의 HyperParameter

Stage Compute Ratio 조절

|Patchify Stem

여기서 stem은 모델이 최초로 데이터 입력을 받아 임베딩하는 부분을 지칭합니다. 기존 resnet에서는 7x7 kernel과 2 stride, 3 padding을 적용한 convolution과 max pooling을 이용해서 입력 이미지는 1/4로 줄였습니다. 반면 ViT에서는 4사이즈의 작은 패치로 이미지를 분할해 임베딩함으로써 보다 간단하게 1/4로 줄이면서 이후의 convolution에 압축적으로 정보를 넘겨주게 됩니다. 이는 겹치는 정보가 없으면서도 더 큰 커널 크기로 처리하는 것과 같은 역할을 한다고 하며, 실험 결과 성능이 미약하게 높아졌다고 합니다. 이에 대한 코드는 아래와 같습니다.

resnet의 stem(좌)와 convnext의 stem(우)

이렇게 두 개의 Macro Design이 적용된 모델의 파라미터와 성능은 아래와 같습니다.

Macro Design이 적용된 resnet-50

test dataset에 대한 confusion matrix 지표는 소폭 상승했고 학습 및 평가 Loss는 더 낮으며 전체적으로 Loss가 더 안정적으로 수렴하며 하락했습니다. 이에 대한 모델링 코드는 여기, 학습 코드는 여기를 살펴보시면 됩니다.

|STEP 02 : ResNext

resnet에 Macro Design을 이식한 뒤에는 resnet의 후속논문으로 나왔던 ResNext를 이식하는 것입니다. ResNext의 핵심 원칙인 '더 많은 그룹과 더 넓은 구조'에 따라 이 또한 두 개의 구성요소를 가지고 있습니다. 하나는 grouped convolution을 적용하는 것이고, 다른 하나는 모델을 더 wide하게 만드는 것으로 간단히 말해 embed dim을 늘리는 것입니다. 

|Depthwise Convolution

ResNext의 더 많은 그룹(cardinality)은 depthwise convolution으로 구현되며 그 뒤에 따라오는 pointwise convolution까지 고려한다면 이전 CvT 리뷰 포스팅에서 살펴봤던 Depthwise Separable Convolution과 동일한 구조가 적용된 것입니다. 이를 통해 두 번째 블록에서는 공간적 정보를, 세 번째 블록에서는 채널 정보를 중점으로 이해할 수 있도록 합니다. 이를 코드로 구현하면 다음과 같습니다.

depthwise convolution이 적용된 결과

위 그림에서 붉은 펜으로 밑줄 그어진 부분이 기존과 달라진 부분입니다. 또한 이러한 depthwise convolution 연산은 기존 convolution보다 그룹 수에 비례해서 연산량을 줄이는 효과를 얻을 수 있으며, 이는 본 논문에서 depthwise를 적용함으로써 해당 모델이 사용하는 컴퓨팅 자원 지표인 FLOPS를 대폭 줄이는 부분에서 확인할 수 있습니다.

depthwise convolution을 적용함으로써 대폭 낮아진 연산량(컴퓨팅 자원소모)

|Expand Width

다음으로는 모델을 확장하는 것인데요, 이는 더 쉽습니다. Swin 등에서는 모델의 embed dims가 [96,192,384,768]로 resnet의 [64,128,256,512]보다 넓게 돼있는데요, resNext의 '더 넓은 모델'이라는 정신을 계승하고 Depthwise Conv를 통해 아낀 자원을 활용해 모델의 너비를 키우는 코드는 다음과 같습니다.

wider resnet

dims에 들어가는 embed dim을 변경해준 것입니다. 이렇게 ResNext까지 적용된 모델의 파라미터와 성능은 다음과 같습니다.

resnext까지 적용한 100에포크 성능

모델의 성능이 오히려 하락한 것을 볼 수 있는데요, 현대적 기법이 하나씩 적용됨에 따라 모델의 성능을 충분히 학습하기 위해 더 많은 반복수가 필요한 것으로 보입니다. 실제로 본 논문에서 가장 크게 성능이 늘어난 부분 중 하나가 이 파트이기 때문에 학습 에포크를 늘려 200에포크의 성능도 확인해보도록 하겠습니다.

200에포크에서의 성능

이에 대한 모델링 코드는 여기, 학습 코드는 여기를 참조하시면 됩니다.

|STEP 03 : Inverted Bottleneck

기본적으로 기존 convnet 계열의 모델들은 데이터의 너비를 줄였다가 다시 늘림으로써 데이터를 압축, 요약했다가 다시 추상화하는 과정을 가지고 있습니다. 반면 트랜스포머에서는 이와 상반되는 전략을 취하는데요, 바로 데이터의 너비를 오히려 늘렸다가 다시 줄이는 방식입니다. 기존 ConvNet에서 취하던 전략을 Bottleneck이라 하며, Transformer에서 취하던 전략을 Inverted Bottleneck이라고 합니다. 이에 대해서 본 논문에서 제시한 그림을 살펴보면 아래와 같습니다.

bottleneck 구조(a)와 inverted bottleneck 구조(b)

이 전략을 도입하기 위해선 사전에 몇 가지 작업이 필요합니다. 우선 기존 resnet의 구조를 다이어그램으로 요약하면 다음과 같습니다.

|ResNet의 구조

resnet50 structure diagram

왼쪽의 그림에서 볼 수 있다시피, 여태까지 살펴본 resnet은 (Batch, Channel, Width, Height)로 이뤄진 입력 이미지가 Stem을 지나며 Channel이 확장되고, 공간차원은 축소됩니다.

이는 첫 번째 Stage에서는 공간차원은 그대로 두고, Channel만을 늘리지만 이후 두 번째 Stage에서부터는 Channel을 두 배로 늘리며 동시에 공간자원을 가로세로 2배씩 총 4배로 축소하며 진행됩니다.

이때 residual connection을 하는 과정에서 문제가 생길 수 있습니다. 바로 이전 스테이지에서 출력된 값의 형상과 현재 스테이지를 거친 출력 값의 형상이 일치하지 않는 것입니다. 간단히 말해 입력으로 들어온 input x와 출력할 output x의 형상이 다르게 됩니다.

때문에 이러한 이유로 각 스테이지의 처음에는 왼쪽 그림과 같이 residual connection을 적용하기 위한 shortcut conv가 존재하며 각 스테이지의 첫 번째 블록에서 작동함으로써 잔차합이 가능하도록 만듭니다. 

이는 stage 내의 block의 영향을 고려해 만들어진 것인데요, 하지만 inverted bottleneck을 적용할 경우 입력으로 들어온 값과 출력으로 나갈 값의 형상은 동일합니다. 

이는 Stage #1에서는 별도의 shortcut conv가 필요하지 않으며, 이후의 short conv들도 굳이 스테이지의 첫 번째 블록 안에서 적용될 필요가 없음을 의미합니다. 이러한 이유로 inverted bottleneck을 적용한 뒤에는 기존에 잔차합 직전에 적용했던 shortcut conv를 일괄적으로 모든 stage 진행 전에 적용하는 것으로 수정할 수 있습니다.

Separate Downsampling에 대한 설명

참고로 이러한 내용은 본 논문에서 이후의 Micro Design에서 Separate Downsampling이라는 항목으로 제시되는데요, 문맥상 여기서 적용되는 것이 올바른 것 같아 적용했습니다. 이에 대한 내용은 위 그림과 같습니다. 간단히 정리하면 기존에는 bottleneck block에서는 중간의 3x3 Convolution(with 2 stride)을 통해, shortcut connection에서는 1x1 convolution(with 2 stride)을 통해 적용했지만 이를 swin과 같은 방식으로 바꾸기 위해 inverted bottleneck block에서의 stride를 없애고, 그 이전에 2x2 convolution(with 2 stride)만을 통해 공간 차원을 축소한다는 내용입니다.

bottleneck block forward

즉 기존에 bottleneck block 안에서 위 왼쪽 그림과 같이 잔차합 직전에 적용되었던 shortcut conv가 오른쪽 그림처럼 사라지고, 기존의 shortcut conv 중 첫 번째 스테이지를 위해 적용된 부분을 제거하고 기존 stem layer를 붙여 아래 그림처럼 작성할 수 있습니다.

bottleneck 구조(좌)와 inverted bottleneck 구조(우)

그리고 이러한 구조를 받아 inverted bottleneck 구조를 구현하면 다음과 같습니다.

inverted bottleneck이 구현된 block code

위 그림에서 수정된 부분은 입력 채널과 출력 채널의 차이입니다. 기존에는 block1에서 압축되고, block2에서 유지되었다가 block3에서 expansion값만큼 다시 늘어나는 식이었습니다. 하지만 inverted bottleneck에서는 block1에서 expansion값만큼 늘리고, block2에서 유지하고, block3에서 다시 원래 차원만큼 줄이는 식입니다.

이렇게 구현한 모델은 shortcut conv 일부가 제거되는 등의 효과로 인해 아래와 같이 파라미터가 줄어들게 됩니다.

inverted bottleneck이 적용된 모델의 파라미터

해당 모델의 100에포크와 200에포크에서의 성능은 각각 다음과 같습니다.

단 200에포크로 swin 500 에포크 정도의 학습 성능을 보였습니다. 물론 이는 데이터셋이 작아 transformer 기반의 모델에는 불리하고, convolution only 기반의 모델에겐 유리했음을 감안해야 합니다. 이에 대한 모델링 코드는 여기, 학습 코드는 여기를 참조하시면 됩니다.

|STEP 04 : Move Up & Large Kernel

역병목(inverted Bottleneck) 구조를 구현했다면 그 다음은 두 번째에 위치했던 depthwise convolution을 상단으로 이동시키고, 커널의 사이즈를 키우는 것입니다. depthwise convolution을 상단으로 이동시키는 것을 본 논문에서는 Move Up이라고 표현했는데, 이를 그림으로 살펴보면 아래와 같습니다.

move up에 대한 그림

우리가 이전까지 구현한 Inverted Bottleneck Block의 구조가 (b)이고, 이를 개선한 것이 (c)입니다. 단순히 위치 순서를 바꿔준 것에 불과하기 때문에 구현은 어렵지 않습니다. 또한 이때 depthwise convolution의 kernel size를 늘려주는 과정을 추가하게 되는데요, 다양한 사이즈의 커널로 테스트해본 결과 kernel size가 7일 때 가장 적절했다고 합니다. 다만 이에 대해서 직접 실험해본 결과 kernel size가 커질수록 모델의 성능 최적화에 더 많은 학습 에포크가 필요했습니다. 때문에 이번 단계부턴 200에포크까지의 성능을 본 뒤, 500에포크까지 학습을 늘려 진행해보도록 하겠습니다.

우선 위 내용에 대해 구현된 코드는 아래와 같습니다.

move up  & large kernel

이렇게 구현한 모델의 파라미터 수와 200에포크까지의 성능입니다.

모델의 파라미터수

200에포크까지의 학습 중 100에포크과 200에포크 지점에서의 성능

물론 여전히 준수한 성능을 보이기는 합니다만, 500에포크로 늘려서 학습할 경우의 성능도 확인해보도록 하겠습니다. 

500에포크 학습의 결과

작은 데이터셋에 convolution 기반 모델이 더 유리함을 감안해야 하지만, 이제 슬슬 swin의 성능을 따라잡기 시작하는 지점으로 보입니다. 이에 대한 모델링 코드는 여기, 학습 코드는 여기를 참조하시면 됩니다.

|STEP 05 : Micro Design

Micro Design까지 적용되면 본 논문에서 제시한 모델링에 관한 사항을 모두 적용한 것인데요, 여기에 다음 스탭에서 진행할 Regularization(Stochastic Depth, Layer Scale)이 적용되면 본 논문에서 제시한 ConvNext의 형상이 완성됩니다.

|Micro Design #1 : Separate Downsampling Layer

Micro Design에 관한 사항은 총 다섯 가지입니다. 첫 번째로는 우리가 inverted bottleneck을 구현하면서 미리 적용한 downsampling Layer를 분리하는 것입니다. 

separate downsampling layer에 대한 내용

기존에 Bottleneck block 내부에서 잔차합 직전에 적용되던 downsampling을 별도의 레이어로 구분하고, 입출력 형상에 맞는 downsampling layer를 구축합니다. 이때 swin과 같이 2stride가 적용된 2x2 kernel을 통해 공간차원을 줄여나갑니다. 때문에 이 과정에서 기존에 사용된 downsampling layer의 구현은 아래와 같은 방식으로 수정됩니다(resnext 이후 inverted bottleneck을 적용하면서 수정된 내용입니다).

downsampling layer used in Bottleneck block

separate downsampling layer

본 논문에서는 이를 적용함으로써 모델의 성능이 약간 상승했다고 제시하고 있으나, 단순히 위치를 수정한 결과라기보다는 위 코드에서 볼 수 있다시피 downsampling layer의 kernel 사이즈가 증가한 결과로 보입니다.

|Micro Design #2 : activation & normalization

Micro Design의 두 번째는 활성화 함수와 정규화층을 현대적 기법으로 치환하고 사용 개수를 줄이는 것입니다. 구체적으로 기존에 ReLU 활성화 함수를 GELU로 대체하고, Batch Normalization 대신 Layer Normalization을 사용합니다. 또한 그 사용개수 또한 줄이는데요, 이는 아래 그림을 통해 이해하면 쉽습니다.

ResNext 와 ConvNeXt의 Block 비

기존에는 각 블록의 레이어를 지날 때마다 정규화와 활성화 함수가 적용됐던 반면 이제는 첫 번째 depthwise convolution을 지난 후 정규화, pointwise convolution 간 활성화 함수 하나와 같은 방식으로 대폭 그 빈도를 줄였습니다. 이는 swin에서 적용한 것과 같은 방식입니다. 이에 대한 내용을 요약하면 정규화와 활성화 함수를 같이 쓰지 않고, 그 빈도를 최소화하면서 필요한 경우에만 쓴다 정도가 되겠습니다. 이를 코드로 확인하면 아래와 같습니다.

블록 내에서 활성화 함수와 정규화 층이 개선된 코드 비교

또한 블록 외에서도 정규화 층과 활성화 함수가 존재하는데요, 이는 본 논문에서 다운샘플링 레이어를 분리하는 것에 대해 설명하는 부분에서 확인할 수 있습니다.

LN의 위치

위에서 볼 수 있다시피 블록 이외의 영역에 적용되는 정규화 층의 위치는 1)downsampling layer 전, 2)stem 뒤, 3)GAP 뒤입니다. 이를 코드로 구현하면 아래와 같습니다.

downsampling layer 전

본래 batch norm으로 convolution 뒤에 적용됐던 정규화층의 위치를 이전으로 바꿨습니다.

stem 뒤 

기존의 bn을 ln으로 바꿨고, 본 논문의 맥락에 따라 함께 적용되었던 활성화 함수를 제거하였습니다.

gap 이후

기존에 GAP 이후 존재하지 않았던 정규화 층을 추가하였습니다. 이렇게 구현한 모델의 파라미터와 성능은 다음과 같습니다. 참고로 우리가 구현하고 있는 모델은 convnext v1 tiny 모델이며, 이 모델의 파라미터 수는 약 28M입니다.

micro design까지 추가된 convnext 모델

Micro Design이 적용되며 우리가 적용한 작은 데이터셋에 대해서는 오히려 소폭 하락한 성능을 보입니다. 이 과정에 대한 모델링 코드는 여기, 학습 코드는 여기를 참고하시면 됩니다.

|STEP 06 : Regularization(Training Recipe)

본 논문에서 제시한 모델링을 구현하기 위한 마지막 단계입니다. 정식 모델링 단계에서는 언급이 없어 배제했던 Training Recipe와 관련된 사항입니다. 이는 우리가 ViT에서부터 적용했던 DropPath와 LayerScale을 적용하는 것으로 이에 대한 개념은 반복해서 설명하지는 않겠습니다.

DropPath와 LayerScale은 모델의 핵심 블록들을 통과한 값에 대해서 적용되며, 이를 포함해 최종적으로 완성된 invertedBottleneck Block 코드는 다음과 같습니다.

convnext의 block 코드

여기까지 완성된 모델의 파라미터 수와 성능은 다음과 같습니다.

droppath와 layerscale이 적용된 모델의 파라미터 수

모델의 성능이 Micro Design만을 적용했을 때보다 더 떨어졌습니다. 또한 추가적인 변화를 확인할 수 있었는데요, Micro Design이 적용되던 시점부터 모델의 최적화 시점이 조금씩 앞당겨지고 있다는 것입니다. 구체적으로 총 500에포크의 학습 과정에서 Micro Design을 적용했을 때에는 약 400에포크 전후로 성능이 가장 높았고, 정규화까지 적용된 이번 스탭에서는 약 300에포크 전후로 성능이 가장 높았습니다. Loss값 자체는 에포크가 이어질수록 떨어졌지만 그게 성능을 좌우할 정도의 차이는 아니었다는 것으로 해석됩니다.

추가적으로 공식 구현체의 성능을 확인했을 때에도 이와 유사한 패턴을 보였습니다. 제 구현체와 공식 구현체에 다른 부분이 있나 찾아보니 다음과 같이 cosine scheduler with linear warmup 코드가 다소 달랐습니다. 

공식 구현체의 코사인 스케줄러

위 코드에서 볼 수 있다시피 일반적인 스케줄러를 구현하는 공식은 아니지만 핵심 원리는 같습니다. 다만 우리가 적용한 스케줄러와 다른 부분은 위 코드에선 하한선이 있다는 것입니다. 이를 통해 공식 구현체에서는 4e-3까지 warmup이 진행되었다가 1e-6까지 천천히 하락하며 학습이 진행됩니다. 이렇게 될 경우 모델의 학습은 보다 높은 학습율 수준에서 이어지게 됩니다.

다만 이를 바로 테스트해보지는 않을 것인데요, 그 이유는 우리가 아직 본 논문에서 제시한 모든 방법을 적용한 것이 아니기 때문입니다. 아래의 EMA까지 적용하고, 그 결과를 살펴본 뒤 스케줄러를 변경해 성능을 살펴보도록 하겠습니다.

이 부분에 대한 모델링 코드는 여기, 학습 코드는 여기를 참고하시면 됩니다.

|STEP 07 : EMA of Training Recipe

마지막으로 적용할 것은 바로 EMA(Exponential Moving Average)입니다. 이는 일반적으로 시계열 분석 등에 사용되던 방법이었는데요, 이를 학습 과정에서 업데이트되는 파라미터에 적용한 것이 바로 EMA Training입니다. 참고로 이 명칭은 공식적인 것은 아닙니다. EMA를 적용한 학습 방법은 성능 증명에 대한 별도의 논문 등은 찾을 수 없었고, 그저 경험적으로 성능이 좋아 2022년 전후부터 대부분의 논문에서 적용되는 것을 확인할 수 있었습니다. EMA의 공식은 간단한데요, 아래와 같습니다.

EMA 공식[3]

이는 현재 에포크에서 학습한 모델과 이전 에포크까지 학습된 모델의 가중치를 가중평균하는 것으로써 하이퍼 파라미터인 α가 0.99로 지정했다면, 이전 에포크까지 학습된 모델 가중치의 99%와 현재 에포크에서 업데이트된 가중치의 1%를 가중평균해 현재의 최종 가중치를 구현하는 것입니다. 이는 기본적으로 가중치 업데이트 속도를 늦춤으로써 노이즈와 변동성을 제어하고, 일반화 성능을 높이는 기법입니다. 실제 α값에 따른 가중치 업데이트를 가시화하면 아래와 같습니다.

α값에 따른 가중치 변화 추이

위 그림은 50에포크 동안 다른 α값에 따른 가중치 변화를 살펴본 것입니다. α값이 커질수록(1에 가까워질수록) 가중치 정규화 효과가 강하게 발생하는 것을 알 수 있습니다. 이에 대한 실험 코드는 여기를 클릭해 확인할 수 있습니다. 만일 이에 대해서 조금 더 자세히 알고 싶다면 이에 대해 자세히 설명돼 있는 이 블로그 포스팅[3]을 참조하면 좋을 듯 싶습니다. 

EMA를 적용한 모델은 더 빠르게 모델의 성능이 올라가는 결과를 확인할 수 있었는데요, 아래는 각기 100에포크부터 500에포크까지 100에포크 단위로 EMA가 적용되기 전후의 성능을 비교한 것입니다.

|100Epoch

EMA 적용 전(좌)과 후(우)의 성능

EMA를 적용한 결과를 통해 가장 먼저 확인할 수 있는 결과 중 하나는 바로 Convolution 기반 모델들이 가진 초반 loss가 요동치는 것을 줄여주는 것입니다. 위 왼쪽 그림을 보시면 EMA가 적용되기 전에는 Val Loss 값이 큰 폭으로 오가는 반면, 오른쪽 그래프에선 어느정도의 변동은 있으나 비교적 낮은 수준에서 진행되고 있음을 확인할 수 있습니다.

|200Epoch

EMA 적용 전(좌)과 후(우)의 성능

200 에포크 지점에서의 결과입니다. EMA를 적용하지 않은 모델의 경우 보다 빠르게 성능 최적화돼 더 낮은 Val Loss 값을 가지지만, 정작 test dataset에 대한 성능은 떨어지는 것을 확인할 수 있습니다. 이는 모델의 일반화 성능이 EMA를 적용한 쪽이 더 높다고 해석할 수 있습니다.

|300Epoch

EMA 적용 전(좌)과 후(우)의 성능

300에포크 지점에선 EMA를 적용하지 않은 쪽이 오히려 성능이 미세하지만 더 높았습니다만 실질적으로 정확도, 정밀도, 재현율의 값 차이를 살펴보면 아마 한두 개 정도의 테스트 이미지 차이 정도로 보입니다. 참고로 EMA를 적용하지 않았을 경우엔 이 지점이 모델 성능이 가장 높았던 지점입니다.

|400Epoch

EMA 적용 전(좌)과 후(우)의 성능

400에포크 지점에선 STEP 04 단계, 즉 Micro Design이 적용되기 전 가장 높은 성능을 보였던 수준과 유사한 수준에 도달했습니다. 반면 EMA를 적용하지 않은 모델의 경우 300에포크를 기점으로 모델이 과적합되며 성능적으로 하락하기 시작합니다.

|500Epoch

EMA 적용 전(좌)과 후(우)의 성능

마지막에 와서는 EMA를 적용한 모델의 경우에도 성능 하락이 확인되는 모습입니다.

|Cosine Scheduler with Min LR

여기에 이전 스탭에서 언급했었던 하한이 존재하는 Scheduler를 적용하도록 해보겠습니다. 우리가 적용할 scheduler가 그리는 학습율 그래프를 그리면 다음과 같습니다.

하한이 존재하는 scheduler의 lr 그래프

이해가 쉽도록 과하게 높은 min_lr을 지정해 가시화해봤습니다. 위 그림은 우리가 지정한 0.1까지는 선형적으로 증가했다가 그 이후 우리가 지정한 하한까지 cosine curve를 그리며 업데이트되는 것을 의미합니다. 이러한 스케줄러를 이용해 4e-3의 initial lr과 1e-6의 min_lr을 가진 학습율 스케줄러를 구현해 학습에 적용한 결과는 다음과 같습니다.

하한이 존재하는 cosine scheduler를 통해 학습시킨 결과

위 오른쪽 그림의 LR 부분을 보면 학습 종료 시 학습율이 1e-6으로 우리가 지정한 수준으로 종료된 것을 알 수 있습니다. 모델의 최대 성능은 비슷한 수준이지만 모델의 성능이 최적화되는 시점이 보다 뒤로 밀리는 것을 확인할 수 있었습니다. 이러한 스케줄러를 사용한 최근 논문들을 조금 살펴봤는데, 최근에는 보다 낮은 수준의 하한선(1e-7)을 이용하며 더 긴 학습 에포크를 적용하는 것을 확인할 수 있었습니다.

이에 대한 모델링 코드는 여기, 학습 코드는 여기를 참고하시면 됩니다.

|Reference

[1] Zhuang Liu, Hanzi Mao, Chao-Yuan Wu, Christoph Feichtenhofer, Trevor Darrell, Saining Xie. A ConvNet for the 2020s. https://arxiv.org/abs/2201.03545

[2] Kaiming He, Xiangyu Zhang, Shaoqing Ren, Jian Sun. Deep Residual Learning for Image Recognition. https://arxiv.org/abs/1512.03385

[3] Lei Mao. Exponential Moving Average. https://leimao.github.io/blog/Exponential-Moving-Average

[4] Vasilis Vryniotis. How to Train State-Of-The-Art Models Using TorchVision’s Latest Primitives. https://pytorch.org/blog/how-to-train-state-of-the-art-models-using-torchvision-latest-primitives