Swin Transformer V2

Paper Review 5 : 모델링, 학습, 튜닝, 3stage

|Intro

이번에 리뷰할 논문은 'Swin Transformer V2: Scaling Up Capacity and Resolution'입니다. 이전에 리뷰한 Swin V1이 통해 다양한 해상도의 이미지에 적용가능하면서도, 단순 이미지 분류뿐 아니라 의미론적 분할(Semantic Segmentation) 등 다양한 태스크에 활용가능한 장점을 가지고 있었지만, 이미지의 사이즈가 일정 이상 커지면서 하드웨어적 한계와 모델의 성능 한계와 같은 문제가 발생했고 이를 해결하기 위한 문제 분석과 그에 대한 방법론이 적용된 V2가 제시되었습니다.

|Paper Summary

본 논문에 적용된 방법론들은 크게 세 개의 결로 나눌 수 있습니다. 각각 모델 구조, 학습 방법, VRAM 최적화입니다. 모델 구조는 기존의 swin v1이 큰 해상도의 이미지를 처리/이해하기 위해 더욱 깊고 큰 레이어를 구성하는 과정에서 발생하는 문제점을 해결하고자 res-post-norm, scaled cosine attention, log spaced continuous RPB(Relative Positional Bias)를 적용한 것입니다. 이를 통해 모델과 이미지 크기 등의 규모를 늘렸을 때의 효율성과 전이 학습에서의 효율성을 개선할 수 있었다고 합니다.

두 번째 학습방법은 기존의 사전지도학습(Supervised Pre-Training) 이후 전이학습(Finetuning, Transfer Learning)을 진행했던 2단계 구조의 제일 앞에 선행 자기주도학습(Self Unsupervised Pre-Training)을 추가해 3단계 구조로 만든 것입니다. 이를 적용하기 위해 SimMIM이라는 기법을 활용했으며, 이를 통해 적은 양의 Labeled Dataset을 통해서도 효과적인 학습이 가능해졌다고 합니다. 또한 두 번째 과정으로 밀려난 사전지도학습의 과정에 Layer-Wise Learning Rate Decay(LLRD)가 적용되며, 자기주도학습의 성능을 함께 끌어올렸습니다.

세 번째 VRAM 최적화는 모델의 구조가 커지고, 이미지의 사이즈가 커지면서 발생하는 문제를 해결하고자 도입된 세 가지 방법론들을 의미합니다. 첫 번째로 복수의 GPU를 사용할 때 모델의 그래디언트, 파라미터, 옵티마이저를 여러 GPU에 복사해 사용하는게 아닌 분할해 사용하는 방식인 Zero-Redundancy Optimizer(ZeRO). 두 번째로 입력 데이터가 모델의 각 레이어를 순차적으로 통과하면서 생성되는 중간결과(활성화-activation)를 추후의 역전파에 활용하기 위해 메모리에 저장해두어야 하는데, 이를 일부만 저장하고 나머지는 역전파가 진행되는 과정에서 실시간으로 다시 계산하는 방식으로 VRAM 사용량을 줄이는 Activation Check-Pointing. 마지막으로 텐서를 이용해 병렬처리되던 Attention 계산을 순차적으로 진행하는 Sequential Self-Attention Computation이 있습니다. 이러한 방법론들은 VRAM의 소모량을 줄일 수 있지만, 그에 따른 연산속도 지연 및 추가 등이 발생해 학습 속도 자체는 느려질 수 있습니다. 또한 추가적으로 이러한 코드가 적용되는 과정에서 모델 구조와 학습 과정이 복잡해질 수 있다는 문제점 또한 지니고 있습니다.

|Modeling

swin v2의 모델 구조는 몇 가지 개선점이 적용되었습니다. 1)기존의 Pre-Norm 대신 Res-Post-Norm이 적용되었고, 2)Scaled dot product attention 대신 Scaled Cosine Attention이 적용되었으며, 3)Parameterized RPB 대신 Log Spaced Continuous RPB가 적용되었습니다.  이를 한 장의 그림으로 살펴보면 아래와 같습니다.

swin v1과 v2의 구조적 차이

또한 본 논문에서는 학습 과정이 크게 3단계로 나뉘어져 있고, 사전 자기주도 학습을 위한 모델 구조가 필요합니다. 때문에 모델 구현에서는 이를 고려해 하나씩 살펴보며 구현토록 해보겠습니다. 이 파트에 대한 전체 코드는 제 github를 참조하시면 됩니다.

|Res-Post-Norm

첫 번째로 살펴볼 것은 정규화를 적용하는 위치의 차이입니다. 이 부분을 집중적으로 살펴보면 아래 그림과 같습니다.

swin v1과 v2의 normalization 위치 차이

특별한 방법이 적용됐다기 보다는 단순히 Normalization의 위치를 바꾼 것입니다. 이는 Residual Connection이 진행되면서 Attention 혹은 MLP 레이어를 통과한 값들끼리 합쳐지는 과정에서 레이어가 깊어질수록 합쳐진 출력값(xl)이 너무 커지는 문제가 발생했는데, 이를 방지하기 위해 잔차합이 적용되기 직전에 정규화를 진행한다는 것입니다.

어느 위치에서 적용하는 정규화가 최적인지에 대해서도 본 논문에서는 테스트를 진행했고, 이에 대한 결과는 아래 이미지와 같습니다.

정규화 위치에 따른 모델의 성능

pre-norm은 각 swin v1과 같이 각 레이어(Attention, MLP) 이전에 정규화를 진행하는 것을 의미하며, post-norm은 잔차합이 진행된 이후에, Sandwich는 레이어의 앞뒤에 정규화를 진행하는 것을 의미합니다. 위 결과를 보면 sandwich는 유의미하게 성능이 낮았으며, 이는 정규화 과정에서 너무 많은 표현력(expressiveness)를 잃었기 때문이라고 보고 있습니다.

이를 코드로 적용하면 아래와 같습니다. 우선 기존 v1 transformer block의 forward 코드를 살펴보겠습니다.

[Before - V1]: Pre-Normalization

swin v1 transformer block -forward

위 코드에서 norm1과 norm2의 위치를 찾으면 (row8)각기 attention이 진행되기 전에 norm1이 수행되고, (row31)mlp가 수행되기 전에 norm2가 수행되는 것을 확인할 수 있습니다. 이를 본 논문에서 제시하는 방식으로 변환하면 아래와 같습니다.

[After - V2]: Res-Post-Normalization

swin v1 transformer block -forward

(row29~30)Stochastic Depth를 포함한 각 레이어의 출력 값에 대해 잔차합이 적용되기 전, 정규화를 진행하였습니다.

|Scaled Cosine Attention

두 번째로 적용될 것은 Scaled Cosine Attention입니다. 이는 위에서 적용한 Res-Post-Norm을 보완하기 위한 것입니다. Res-Post-Norm을 적용하면서 스케일업을 진행할 경우, 몇몇 블록과 헤드의 어텐션 맵이 소수의 픽셀에 의해 주도되는(Dominated) 문제가 발생했습니다. 이를 해결하기 위해 본 논문에서는 코사인 유사도를 이용한 어텐션 계산 방식을 제안합니다.

코사인 유사도 공식

위 공식이 코사인 유사도를 계산하기 위한 공식인데요, 두 벡터의 내적을 두 백터의 유클리드 노름(norm)으로 나눈 값입니다. 이는 각 벡터를 길이가 1인 단위벡터로 변환해 비교하는 방식으로, 각 벡터값의 크기가 아닌 방향성만을 고려하게 되며 동시에 일종의 정규화 효과를 얻을 수 있습니다. 추가적으로 두 벡터가 같은 방향일 경우(동일할 경우), 코사인 유사도는 최대값인 1을 나타내며 반대방향일 때는 최소값인 -1을, 그리고 직각일 경우(관련성이 없을 경우)엔 0이 됩니다. 

구체적으로 본 논문에서 적용하는 방식은 아래와 같습니다.

Scaled Cosine Attention 수식

query와 key 사이의 유사도를 측정하고, 이에 대해 학습가능한 파라미터 τ로 나눠줌으로써 스케일링을 진행하며, 스케일링까지 적용된 값에 대해 RPB 값을 더함으로써 attention score를 구하는 식입니다. 이에 대한 코드가 적용되기 전후를 살펴보면 아래와 같습니다.

일반적인 attention score 계산

cosine similarity가 적용된 attention score 계산

본 논문에서는 0.01보다 큰 τ값으로 나눠줬지만, 직관성인 코딩을 위해 torch.clamp를 이용해 최대값을 100(1÷0.01)으로 제한하고 곱하는 식으로 처리했습니다. 또한 값의 범위가 넓을 것을 고려해 scale 인자 등의 계산을 log 함수를 통해 진행하고 최종적으로 attention score에 곱해주기 전에 지수화(exp)하여 곱해주는 식입니다.

위 코드를 통해 구한 attention score에 아래에서 설명할 새로운 방식의 RPB를 더하는 것을 제외하면 이후 softmax를 취한 뒤 value 값에 대해 내적하는 과정은 기존과 동일합니다. 참고로 위 코드에서 t_scale을 정의하는 코드는 아래와 같습니다.

attention score를 조절하기 위한 Scale t

|Log Spaced Continuous RPB

모델링의 마지막 변경점은 바로 Log Spaced Continuous RPB입니다. 이는 전이학습의 과정에서 스케일링을 조정(정확히는 늘릴 때)할 때, 윈도우의 크기 변화에 따른 상대적 위치편향 크기 변화를 보다 자연스럽게 적용하기 위한 것입니다.

이는 크게 두 가지 부분으로 구성돼있습니다. 첫 번째는 Continuous RPB이고, 두 번째는 Log Spaced Coordinates입니다. 이를 각각 살펴보도록 하겠습니다.

|Continuous Relative Position Bias

기존 swin v1에서 상대적 위치 편향은 (2xWindow Size -1)의 제곱 크기의 테이블로 저장되고, 상대적 위치 인덱스는 입력 윈도우 내의 각 위치 쌍에 대해서 계산되었습니다. 이는 고정된 테이블 사이즈에 기반해 최적화되며, 윈도우의 크기나 형태가 변할 경우 정상적으로 작동하지 않습니다. 즉 전이학습이 제대로 이뤄지지 않으며, 유연성 또한 부족하다는 의미입니다. 이를 보완하기 위해 본 논문에서는 Cotinuous RPB를 제시하며, 이에 대한 공식 등은 아래와 같습니다.

continuous RPB에 대한 설명

위 공식은 기존에 단순 연산되던 편향값들로 이뤄진 테이블에 작은 메타 네트워크를 적용함으로써 유연성을 확보한다는 것인데요, 여기서 작은 메타 네트워크는 2개의 선형층과 한 개의 ReLU 엑티베이션 층으로 구성된 레이어로 정의하고 있습니다. 이를 코드로 구현하면 아래와 같습니다.

small meta network code

첫 번째 선형층은 가로세로 2차원 좌표 입력을 받기 때문에 입력 차원의 개수는 2입니다. 두 번째 활성화 층인 ReLU는 입력값이 0보다 클 경우 그대로 출력하고, 0보다 작을 경우 0을 출력하는 방식으로 이전 입력값을 저장하는 것에 대한 의미가 없기 때문에 조금이라도 VRAM을 아끼기 위해 inplace 옵션을 True로 지정했습니다. 참고로 이러한 inplace 옵션은 ReLU와 같이 이전 값을 기억할 필요가 없는 단순한 활성화 층이 아니라면, 디버깅 포인트를 찾기 어렵거나 필요 이상으로 모델의 연산속도가 떨어질 우려가 있으니 조심히 사용해야 합니다. 마지막 선형층은 head의 개수만큼을 출력합니다. 이때 출력하는 값은 위치 바이어스의 역할을 수행하기 때문에 추가적인 바이어스를 적용할 경우 불필요한 연산이 추가되거나, 원하는 역할을 적절하게 수행하지 못할 우려가 있어 마지막 선형층의 bias는 False로 지정합니다.

|Log Spaced Coordinates

파인튜닝의 결과, 특히 보다 큰 이미지셋으로의 파인튜닝이 이뤄질 때에는 일반적으로 윈도우의 사이즈도 함께 증가하는 것이 효율적입니다. 하지만 이는 모델 구조가 확장된 상대적 위치의 범위를 이해해야 한다는 것을 의미하며, 자연스레 늘어난 윈도우 사이즈의 제곱만큼의 복잡성이 적용돼 비효율성이 발생하게 됩니다. 이를 해결하고자 본 논문에서는 아래 그림과 같 상대적 위치 좌표를 log를 통해 처리하는 방식을 제안하고 있습니다.

log spaced coordinates에 대한 설명

이를 통해 윈도우 크기가 증가해도 증가하는 좌표의 범위가 줄어들어 효율적인 파인튜닝이 가능하다는 것입니다. 예컨데 8x8 윈도우를 16x16으로 전환할 때, swin v1의 선형좌표 기준에서는 좌표 범위가 [-7,7]에서 [-15,15]로 확장되어야 합니다. 반면 swin v2에서는 [-2.079,2.079]에서 [-2.773,2.773]으로 확장되어야 합니다. 각 버전에서 확장된 비율은 1.14와 0.33으로 log space를 취함으로써 약 4배의 효율성을 얻을 수 있게 되는 셈입니다. 이는 본 논문에서 제시한 위치 편향별 성능 지표를 도식화한 아래 그림을 통해서도 확인할 수 있습니다.

위치편향 방법론에 따른 성능지표

이제 코드로 살펴보겠습니다. 먼저 v1에서의 코드는 아래와 같았습니다.

swin v1 - window attention - rpb part

v1에서 rpb는 일정한 기준을 가지고 상대적 위치 좌표 인덱스를 구한 후 이를 적용한 학습가능한 파라미터인 편향 테이블을 활용하는 방식으로 진행됐었습니다. 하지만 v2는 다소 다른데요, 우선 v2의 편향 테이블은 학습가능한 파라미터가 아닙니다. 학습가능한 부분은 이전에 진행되는 small metal network뿐이며, 그 외에는 윈도우 사이즈에 따라 동적으로 조정되는 매트릭스일뿐입니다. 

swin v2 - window attention - rpb part

위 v2의 코드를 살펴보면 달라진 점은 크게 두 가지가 있습니다. 하나는 small meta network인 crpb_mlp가 적용됨으로써 상대적 위치 정보를 담은 테이블은 더이상 학습가능한 파라미터가 아니게 되었다는 것입니다. 하지만 이로 인해 상대적 위치테이블 정보를 저장된 모델에 포함할 수 없기 때문에 매번 새롭게 연산시키기 보다는 buffer로 등록함으로써 모델이 저장될 때 포함되도록 했습니다. 그렇지 않을 경우, 모델이 로드될 때마다 연산하도록 할 순 있지만 시간적으로 비효율적이며 연산이 이뤄지는 디바이스나 환경이 달라질 경우 추가적인 문제가 발생할 수 있습니다.

두 번째 차이점은 사전학습된 데이터셋의 윈도우 사이즈에 따라 이를 조정하는 매커니즘이 추가되었다는 것입니다. 이를 통해 윈도우 사이즈가 동적으로 계산되며, 사전학습된 정보가 포함된 상태로 늘어난 부분에 대한 정보만을 추정함으로써 파인튜닝 시 모델의 성능을 유지할 수 있습니다.

|Train

학습과 관련된 파라미터는 아래와 같습니다. 대부분의 파라미터나 세팅은 swin v1과 같으며, 아래 그림의 붉은 색으로 강조한 gradient clipping의 max norm 수치만 다릅니다. 이 파트에 대한 전체 코드는 제 github를 참조하시면 됩니다.

swin tiny, small, base training args

이러한 파라미터를 적용해 학습한 결과는 아래와 같습니다.

직접 구현체(좌)와 timm 구현체(우)로 학습한 결과

이 수치는 swin v1보다 오히려 소폭 하락한 수치입니다. 이는 Swin v2 Tiny 구현체의 파라미터 수는 약 2,764만 개 Swin v1 Tiny 구현체의 파라미터 수는 약 2,827만 개로 약간 적은 것에서 오는 결과입니다. swin v2는 모델과 데이터셋의 사이즈가 커지는 것에 초점을 맞춰진 버전이며, 우리가 테스트한 모델이 Tiny이기 때문에 큰 사이즈의 모델을 사용해 그 성능을 비교할 경우 swin v1의 큰 모델에 비해 파라미터 효율적인 성능을 보일 수 있습니다.

참고로 위 학습 과정에서는 본 논문에서 제시된 SimMIM을 이용한 사전 자기주도학습 과정이 배제돼있지만, SimMIM과 같은 Self Supervised Learning 방식은 모델의 성능 자체를 높여주진 않습니다. 이러한 방식들은 모델의 효율성을 높여 최적화에 이르는 비용을 낮춰주는 효과가 있습니다. 

다만 예외적으로 보다 높은 성능을 위해 적용되는 경우, 즉 SimMIM과 같은 SSL 방식이 요구되는 경우가 있습니다. 예컨데 레이블이 있는 데이터는 소수이고, 대다수의 데이터가 레이블이 없는 상황이거나 여기에 모델의 사이즈와 데이터셋의 사이즈 자체도 큰 경우가 그렇습니다. 간단히 말해 더 적은 라벨 데이터, 더 적은 epoch로 적용하는 모델 구조가 가진 성능의 최대치에 도달할 수 있는 것입니다. 개인적으로 이러한 사전 자기주도학습은 딥러닝 구조를 이용한 효율적인 파라미터 초기화에 가깝다고 보고 있습니다.

이러한 사전 자기주도학습(Self-Supervised Learning)을 줄여서 SSL이라고 부르며, 이러한 방법론들을 모아놓고 사용하기 쉽도록 정리한 라이브러리인 Lightly SSL을 참고하면 쉽게 다양한 SSL 방법들을 적용해볼 수 있습니다. 하지만 해당 라이브러리는 사용하기는 쉽지만 SimMIM의 자세한 구조를 구현하면서 이해하기는 다소 어려운 부분이 있기 때문에 여기서는 SimMIM의 공식 github에서 제시한 코드를 참조해 SimMIM을 적용하고 테스트해보도록 하겠습니다.

|Stage-1 : SimMIM

본 논문에서는 SimMIM을 적용했다고 언급하고 넘어갈뿐 자세하게 설명돼있지는 않습니다. 때문에 이를 이해하기 위해선 SimMIM 논문을 먼저 살펴봐야 하는데요, SimMIM 논문에서 제시된 구조는 아래와 같습니다.

SimMIM의 주요 구조에 대한 그림

SimMIM은 본래 이미지의 일부에 마스크를 적용하고 이를 주변 이미지를 통해 추론해 다시 그리도록 함으로써 학습하는 데이터셋의 주요 특징을 학습하는 것입니다. 이때 사전 최적화를 적용할 모델 구조를 Encoder로 사용하고, 아주 간단한 구조의 Decoder(위 그림의 Prediction Head)를 붙여 복구한 이미지를 본래 이미지와 비교합니다. 이때 보다 효율적인 연산을 위해 마스킹을 칠한 부분에 대해서만 loss 연산을 진행합니다. 이를 코드로 구현하기 위해서는 이미지에 적절한 변형(transform)을 적용하고 마스크를 생성해 같이 전달하는 데이터 로더, 사전학습시키고자 하는 모델을 상속받은 모델과 로스 등을 구현해야 합니다. 

이에 대한 개념과 코드를 하나씩 살펴보겠습니다. 참고로 여기 제시된 코드는 SimMIM github의 공식코드를 거의 그대로 가져와 사용가능하도록 몇몇의 수정만을 가해 작성되었습니다. (본 포스팅이 SimMIM 논문 리뷰가 아니기 때문입니다.) 이 파트에 대한 코드는 여기를 참조하시면 됩니다.

|SimMIM Modeling

SimMIM 코드 구현

위 코드는 인코더로 기존에 구현한 swin 모델을 전달받고, 1x1 Convolution 이후 nn.PixelShuffle이라는 간단한 레이어를 통해 decoder의 영향력을 최소화한 구조입니다. 이를 통해 SimMIM 모델은 Encoder의 성능을 높이는 것에 초점을 맞추게 됩니다. 이때 Encoder를 우리가 구현한 모델 그대로 사용하기보다는 mask를 이용한 연산 등을 효율적으로 적용하기 위해, 기존 모델을 상속한 새로운 모델을 구현하면 아래와 같습니다.

SimMIM에 Encoder로 사용할 코드

위 코드를 통해 입력 이미지와 마스크를 전달받고, 마스크 영역을 학습가능한 파라미터인 mask token으로 대체합니다. 이를 통해 새로 정의한 모델이 mask 영역을 재구성하도록 하며, 그 외의 부분은 기존의 구조를 동일하게 적용하다가 마지막 부분에 본래의 이미지 구조로 복구시키게 됩니다. 이렇게 정의한 모델이 Encoder로 입력받게 되며, 추후에 학습이 완료된 이후 mask_token을 제외한 파라미터가 전달되게 됩니다. 

|SimMIM's MaskGenerator

그 다음으로 정의할 것은 SimMIM의 개념이 적용된 데이터 로더입니다. SimMIM의 데이터 로더는 기본적인 증강이 적용된 이미지 데이터와 마스크 데이터의 쌍으로 출력이 되어야 합니다. 이를 위해 Mask를 생성하는 클래스를 정의해 mask를 정의하고, 일반적인 이미지 변형을 가해 변형된 이미지를 출력하는 transform 함수를 구현하면 아래와 같습니다. 우선 마스크 생성과 관련된 클래스입니다.

Mask 생성 클래스 코드

위 코드는 입력된 이미지에서 이미지의 사이즈와 모델의 패치 사이즈, 마스크의 사이즈와 마스크 비율을 입력으로 받아 마스크를 생성하는 코드입니다. 모델의 패치 사이즈를 고려해 축소한 형태의 매트릭스를 출력하며, 마스크 패치 사이즈를 고려해 마스크 비율에 근소한 비율의 마스크 결과를 만들어내게 됩니다. 이해를 위해 위 클래스를 실행해보도록 하겠습니다.

[Mask Generator Test]

SimMIM을 위한 모델, 로더 등을 불러와 SimMIM을 학습시키고, 파인튜닝하는 과정에 대한 코드는 여기를 참조하시면 됩니다.

Mask Generator의 실행결과

이해가 용이하도록 이미지 패치 사이즈를 줄여 생성된 마스크의 크기를 8x8 매트릭스로 조정하였습니다. 생성된 마스크 매트릭스의 전체 크기는 이미지 입력 사이즈를 모델의 패치 사이즈로 나눈 값(224÷28=8)이 됩니다. 이렇게 생성한 매트릭스는 0으로만 구성돼 있고, 해당 매트릭스에 1로 구성된 마스크를 작성하는 식입니다. 위 코드에서는 마스크 패치 사이즈가 모델 패치 사이즈의 2배이기 때문에 가로세로 2 사이즈의 마스크를 랜덤으로 겹치지 않게 생성하게 됩니다. 만약 마스크 패치 사이즈를 28로 줄이면 다음과 같이 보다 60%에 근사한 비율의 마스크가 생성되게 됩니다.

mask size가 1인 경우

|SimMIM's Transforms

Mask Generator가 생성되었다면 간단한 변형을 가해 이미지를 증강하는 코드와 합쳐 이미지와 마스크 쌍으로 전달하는 Transformer(모델 구조를 의미하는 게 아닙니다)를 생성해야 합니다. 모델 구조를 의미하는 transformer와 동일한 이름이기 때문에 구분하고자 transforms라고 칭하도록 하겠습니다. 본 논문에서는 아래와 같이 SimMIM Transforms를 위한 방법도 제시되어 있습니다.

Stage 1 self-supervised pre training에 대한 설명

위 설명을 보면 이미지에 0.67~1 범위의 스케일과 3/4~4/3 범위의 aspect ratio를 적용한 Random Cropping and Resize 이후 Random Flip과 Color Normalization을 적용했다고 제시돼 있습니다. 이러한 사항들과 Mask Generator를 포함해 Image와 Mask를 전달하는 SimMIMTransform 클래스는 다음과 같습니다.

SimMIM Transform code

|SimMIM's Hyper Parameters

다음으로는 본격적인 학습을 위한 하이퍼 파라미터 설정입니다. 이에 대한 본 논문에서 제시한 사항들은 아래와 같습니다.

모델의 Hyper parameter 등

위 하이라이트된 부분을 살펴보면 이미지는 본래(224)보다 작은 192 사이즈로 구성하였고, AdamW Optimizer와 Cosine Scheduler with Warmup을 적용하였습니다. Initial learning rate는 0.0014이고, weight decay는 0.1, max norm은 100입니다. 이러한 본 논문의 사항과 다소 다르게 적용한 부분이 있는데요, 바로 weight decay와 max norm입니다. 본 논문에서 제시된 부분은 SwinV2-G에 대한 세팅을 설명한 것으로 저희가 구현하는 Tiny의 파라미터가 27백만개인 것과 달리 무려 30억 개의 파라미터로 구성된 11,000배 큰 대형 모델이며, 데이터셋의 규모도 훨씬 작기 때문입니다. 그만큼 깊고, 큰 모델이기 때문에 저만큼의 weight decay와 max norm이 적용되어야 하지만, 우리가 적용할 base 모델과 sports 데이터셋에서는 실험결과 오히려 성능이 떨어졌습니다. 때문에 우리는 다소 낮은 0.05의 weight decay와 5의 max norm을 사용할 것입니다. 

SimMIM 학습결과

위 SimMIM 학습 결과를 보면 총 100epoch를 학습시켰고, 이미지 사이즈를 줄이고 더 큰 배치 사이즈를 활용했으며, 학습 구조도 간단해서 epoch당 학습속도가 20초도 걸리지 않습니다. 100epoch를 학습시키는데, 약 1,900초 정도가 걸렸습니다.

SimMIM을 조금 더 효율적으로 쓰려면 SimMIM을 다시 여러 단계로 나눠서 첫 단계에서는 작은 마스크 패치 더 높은 마스크 비율, 그 다음 단계에선 더 큰 마스크 패치와 더 낮은 마스크 비율, 그 다음 단계에선 더더 큰 마스크 패치와 더더 낮은 마스크 비율과 같은 식으로 단계적으로 적용하면 보다 효율적으로 SimMIM을 학습시킬 수 있습니다. 이런 방법을 적용했을 때에는 동일한 시간과 총 100epoch를 학습시키더라도 약 0.2~0.3 수준의 Loss까지 떨어뜨릴 수 있습니다. 다만 이러한 단계적 SimMIM은 말 그대로 데이터가 충분히 많을 때 사용해야 하며, 자칫 과적합에 빠질 우려가 있음을 인지해야 합니다.

SimMIM 모델의 하이퍼 파라미터에 대한 세팅은 model/config 폴더 이하의 pretrain.yaml을 참조하면 자세히 알 수 있습니다. SimMIM의 성능을 확인하기 위해 SimMIM을 적용하지 않고 150epoch를 학습시킨 경우와 SimMIM 100epoch 이후 100epoch를 학습하는 경우를 비교해보도록 하겠습니다. 우선 SimMIM 없이 150epoch를 학습한 경우입니다.

150epoch를 학습한 경우

Validation Dataset에 대한 Loss는 0.7625 수준, Test Dataset에 대한 F1-Score는 0.8467이었으며 총 소요시간은 약 135분이었습니다. 다음으로는 SimMIM을 적용한 경우입니다. 자세한 코드는 여기를 참고하시면 됩니다.

SimMIM을 통해 본래 사이즈보다 작은 192 이미지 사이즈에 6 윈도우 사이즈로 100epoch 학습시킨 후에 해당 가중치를 224 이미지 사이즈에 7 윈도우 사이즈로 조정하기 위해서는 다음 코드를 통해 사용하지 않을 가중치는 배제하고, 사용못하는 가중치는 초기화하거나 대체하고, 사용가능한 가중치만을 조정합니다. 사용가능한 가중치를 조정할 때에는 bicubic interpolation을 사용하는데요, 이는 swin v1 paper에서부터 제시된 방법입니다.

swin v1에서 제시된 interpolation에 대한 설명

또한 모든 부분을 재조정해서 사용하는 것은 아닌데요, 구체적으로 Relative Position Bias Table과 Absolution Position Embedding을 조정하고, Relative Position Index 및 Relative Coords Table, 그리고 Attention Mask는 초기화해 사용합니다. 

어떤 가중치는 조정해서 사용하고, 어떤 가중치는 초기화하는 이유를 간단히 설명하자면, Relative Bias Table과 Absolute Positional Embedding은 픽셀 간의 상대적/절대적 위치 정보를 이해하기 쉽도록 하는 것으로서 이미지의 특정 부분이 다른 부분에 비해 얼마나 멀고 가까운지를 의미하는 편차가 됩니다. 이는 이미지의 사이즈가 바뀌더라도 대체로 유사하게 사용할 수 있어, 기존에 학습한 가중치로 대체한 후 거기서 학습/조정하는 것이 효율적입니다.

그 외에 Relative Position Index 및 Relative Coords Table은 모델이 어떻게 특정 위치에 주의를 기울일지 결정하는데 사용되며, 입력 이미지의 크기가 바뀌면 픽셀들 간의 절대적인 위치 관계가 변경되어 무용해지게 됩니다. Attention Mask도 마찬가지입니다. 테스트해봤을 때에는 굳이 보간해서 가져와도 성능차이가 존재하지 않거나 아주 미약했습니다. 또한 이러한 요소들은 가중치 형태가 아닌 buffer로 등록되서 사용되는 논리적으로 계산되는 구조이기 때문에 모델이 정의되며 새롭게 정의하고, 이후의 meta network를 제대로 불러와 통과시키도록 하는 게 효율적입니다.

이를 적용한 코드는 아래와 같습니다.

다른 이미지 사이즈와 윈도우 사이즈를 가진 모델의 가중치를 적용하기 위한 코드

위 코드는 크게 다섯 부분으로 이뤄져있습니다. 첫 번째로는 SwinTransformerForSimMIM에서 사용할 가중치인 encoder부분만을 가져오고, 'encoder.'이 붙은 prefix를 제거하는 것입니다. 두 번째로는 사용하지 않고 초기화해서 사용할 Relative Position Index와 Relative Coords Table, Attention Mask를 제거합니다. 세 번째로는 Relative Position Bias Table을 bicubic method를 이용해 보간하고, 네 번째로는 Absolute Positional Embedding도 동일하게 bicubic method를 이용해 보간합니다. 이때 만약 Absolute Positional Embedding을 추가적용하지 않은 모델구조라면 이 부분은 제외해도 됩니다. 마지막으로는 Classifier를 조정합니다. 

이렇게 다섯 개의 사전작업을 처리한 뒤에는 load_state_dict의 strict 파라미터 옵션은 False로 하면 state dict의 key와 shape가 동일한 가중치를 불러오게 됩니다. 이에 대한 실행결과는 다음과 같으며, 우리가 제거한 Relative Position Index와 Relative Coords Table, Attention Mask가 없다는 메세지입니다.

simmim 가중치를 불러왔을 때의 출력결과

추가로 위 코드가 정상적으로 작동했다는 것을 확인하기 위해 swin v2를 로드하자마자의 가중치 결과와 로드된 뒤의 가중치 결과가 다르다는 것을 살펴보도록 하겠습니다.

모델 로드 직후 초기화된 가중치(로드할 때마다 변화)

simmim을 조정해 로드했을 때의 가중치(로드할 때마다 동일)

위 그림은 모델을 로드할 때 초기화된 가중치이며, 그렇기 때문에 매번 로드할 때마다 값이 조금씩 변화하게 됩니다. 반면 아래의 그림은 simmim으로 학습시킨 결과를 불러온 것이기 때문에 늘 일정한 값이 로드되는 것을 확인할 수 있습니다.

|Stage-2 : Supervised Pre-Training

SimMIM을 이용해 Self Superivsed pre-training하는 단계를 본 논문에서는 Stage-1이라 정의하며, 그 다음으로 레이블 데이터로 학습시키는 과정을 Stage-2로 정의합니다. 이 단계에서도 기존과 다른 몇 가지가 추가로 제시되는데요, 이에 대한 내용은 다음과 같습니다.

본 논문에서 제시한 Stage-2에 대한 설명

위 그림에서 하이라이트되지 않은 부분은 이전의 세팅과 동일한 부분이며, 노란색 하이라이트로 칠한 부분은 간단한 파라미터 값 수정을 통해 조정할 수 있는 부분, 붉은색 하이라이트로 칠한 부분은 새롭게 정의해 적용해야 하는 부분입니다. 여기서 새롭게 제시된 내용이 바로 Layer-Wise Learning Rate Decay입니다.

|Layer-Wise Learning Rate Decay(LLRD)

Layer-Wise Learning Rate Decay는 간단히 말해 모델의 뒷부분을 중점으로 학습하기 위한 방법입니다. 이를 위해 모델의 앞부분으로 갈수록 학습율을 감소시키게 되는데요, 이에 대한 자세한 설명은 여기[4]를 살펴보시는 것을 추천드립니다. 이에 대한 개념을 그림으로 살펴보면 다음과 같습니다.

layer-wise learning rate decay에 대한 그림(출처:[4])

위 그림은 만약 모델의 initial learning rate가 0.01이고 layer가 앞단의 레이어일수록 이전 레이어 학습율의 1/10이 되는 구조입니다. 이를 통해 이전에 모델이 학습하는 과정(Stage-1)에서 학습된 가중치 중 모델의 앞단 구조는 크게 변화시키지 않으면서 뒷부분만을 조정해 모델의 성능을 튜닝할 수 있습니다. 이러한 방법은 파인튜닝과 같이 여러 단계로 구성된 학습 구조로 모델을 학습시킬 때, 특히 모델이 크고 복잡할수록 모델의 성능을 빠르게 최적화시킬 수 있습니다.

이를 코드로 구현하면 다음과 같습니다.

layer-wise learning rate decay code

위 코드를 통해 모델의 각 파라미터의 이름을 추출하고, 뒷 레이어부터 시작해서 학습율을 낮추기 위해 추출한 이름 리스트를 뒤집습니다. 뒤집은 이름 리스트를  이용해 각 파라미터의 레이어 그룹명이 같다면 동일한 학습율을 적용하며, 레이어 그룹명이 다르다면 그 다음 레이어에 곱할 가중치를 적용한 학습율을 연산해 적용합니다. 본 논문에서는 초기 학습율 0.0014에 0.87의 가중치를 적용하라고 하여 그대로 적용했고, 그 외의 weight decay 0.1의 경우만 0.01로 낮춰서 적용했습니다. (데이터셋과 모델의 사이즈가 작아 weight decay를 0.1로 해 적용해보니 최적화 속도가 오히려 느려졌습니다)

이렇게 적용한 코드의 결과 일부를 살펴보면 다음과 같습니다.

layer-wise learning rate decay의 결과

위 결과 이미지를 통해 알 수 있듯 모델의 끝부분, 즉 classifier 레이어는 초기 학습율인 0.0014가 되고 그 이후부터 가중치인 0.84씩 곱해진 값이 각 레이어 구성요소들의 학습율이 된 것을 확인할 수 있습니다.

참고로 Layer-Wise Learning Rate Decay(LLRD)의 효과는 꽤 커서 SimMIM만으로는 충분히 줄어들지 않던 최적화 속도를 가속화해주는 효과가 있습니다. 예컨데 SimMIM 없이 바로 학습시킨 모델이 100epoch 정도에 달성한 성능을 SimMIM만을 적용했을 때에는 90~95epoch 정도, SimMIM과 LLRD를 적용했을 때에는 70epoch 정도에 달성했습니다. 물론 그렇다고 LLRD만 적용하면 최적화 속도가 기존보다도 더 느려지거나 과적합에 빠지고 맙니다.

|학습결과 비교

이렇게 가중치를 불러온 모델을 완전히 동일한 환경과 세팅에서 학습시켜볼 것인데요, 본래 SimMIM을 적용하지 않았을 때에는 약 150epoch를 학습시켜 0.7625의 val_loss와 0.8467의 F1-Score를 얻었습니다. 여기서는 SimMIM을 적용한 후 각각 10, 50, 100epoch를 학습시킨 뒤의 성능을 평가해보도록 하겠습니다.

10 Epoch

SimMIM+LLRD 적용 전(좌)과 후(우)의 10epoch 성능

10epoch 지점에서의 결과입니다. 왼쪽이 바로 sports dataset에 학습시켰을 때, 오른쪽이 SimMIM 이후 LLRD를 적용해 학습시키는 경우입니다. 3Epochs의 평균 Val Loss 기준 약 15% 수준의 성능차이가 나고 있음을 알 수 있습니다.

50 Epoch

SimMIM+LLRD 적용 전(좌)과 후(우)의 50epoch 성능

50epoch 지점에서의 결과입니다. SimMIM+LLRD의 효율이 더 커져가며 3Epochs의 평균 Val Loss 기준 약 20% 이의 성능 차이가 나고 있음을 알 수 있습니다. 참고로 바로 데이터셋에 학습시켰을 경우, 위의 성능을 내기 위해선 약 70epoch에 이르러야 했습니다.

100 Epoch

SimMIM+LLRD 적용 전(좌)과 후(우)의 100epoch 성능

100epoch 지점에서의 결과입니다. 이미 SimMIM과 LLRD가 적용된 성능은 기존 모델이 150epoch를 학습했을 때보다도 높아진 상태입니다. 참고로 이때의 confusion matrix 지표를 측정해 비교해보면 다음과 같습니다.

바로 학습의 150epoch에서의 지표와 SimMIM+LLRD의 100epoch 지표

Validation Dataset에 대한 Loss는 150epoch보다도 낮았지만 F1-Score는 아직 소폭 낮은 수준입니다. 이는 다른 지표들은 소폭 높은 수준이지만 Precision이 조금 더 낮은 결과입니다. 참고로 여기까지 진행했을 때, 학습에 사용된 리소스를 학습시간으로 환산해 비교하면 다음과 같습니다.

각 모델 학습에 사용된 리소스 계산 with A100 80GB 운용시간(초)

완전히 동일한 환경과 장비를 통해 실험했기 때문에 유사한 수준의 성능에 이르기까지 영향을 미친 것은 모델의 구조 혹은 학습 방법밖에 없으며, 모델의 구조 또한 직접 구현한 swin v2로 동일하기 때문에 성능의 차이는 학습 방법에서 기인한 것입니다. 이는 약 12.22%의 리소스를 절약했음에도 본래 이상의 성능을 낼 수 있다는 점을 시사하는 것입니다. 또한 이렇게 학습시킨 모델을 여타 다양한 태스크에 다시 파인튜닝함으로써 우리가 원하는 성능을 낼 수 있으며, 모델의 구조가 크고 데이터셋의 규모가 클 경우(+UnLabeled Data가 존재할 경우) 긍정적인 효과를 얻을 수 있습니다.

110epoch(좌)와 120epoch(우)에서의 성능

10에포크를 더 학습시킬 경우, 전체 학습에 소요한 컴퓨팅 리소스는 8,500(초)로 150epoch를 한 번에 학습한 경우보다 낮지만 성능은 위의 왼쪽 그림과 같이 더 높아지게 됩니다. 여기에 10에포크를 더 학습시켜 120epoch를 학습시키게 되면 전체 학습에 소요한 컴퓨팅 리소스는 9,100초로 기존 방법에 비해 1.1% 초과하지만 성능은 F1-Score 기준 2.7% 향상되며, Accuracy 기준 2.6% 향상됩니다. 이는 투입한 자원 대비 효율적인 결과입니다.

|Reference

[1] Ze Liu, Han Hu, Yutong Lin, Zhuliang Yao, Zhenda Xie, Yixuan Wei, Jia Ning, Yue Cao, Zheng Zhang, Li Dong, Furu Wei, Baining Guo. Swin Transformer V2: Scaling Up Capacity and Resolution. https://arxiv.org/abs/2111.09883

[2] Zhenda Xie, Zheng Zhang, Yue Cao, Yutong Lin, Jianmin Bao, Zhuliang Yao, Qi Dai, Han Hu. SimMIM: A Simple Framework for Masked Image Modeling. https://arxiv.org/abs/2111.09886

[3] Ze Liu, Yutong Lin, Yue Cao, Han Hu, Yixuan Wei, Zheng Zhang, Stephen Lin, Baining Guo. Swin Transformer: Hierarchical Vision Transformer using Shifted Windows. https://arxiv.org/pdf/2103.14030.pdf

[4] Nikita Kozodoi. Layer-Wise Learning Rate in PyTorch. https://kozodoi.me/blog/20220329/discriminative-lr

|Log

2023.12.31 | 논문 읽기 - swin v2

2024.01.02 | 논문 읽기 - swin v2

2024.01.03 | 논문 읽기 - swin v2

2024.01.04 | 논문 읽기 - swin v2

2024.01.05 | 논문 읽기 - swin v2

2024.01.06 | 논문 읽기 - SimMIM

2024.01.08 | 논문 읽기 - SimMIM

2024.01.09 | 논문 읽기 - SimMIM

2024.01.10 | 논문 읽기 - SimMIM

2024.01.11 | Swin v2 Modeling ①

2024.01.12 | Swin v2 Modeling ② → Training

2024.01.13 | SimMIM 구현 ①

2024.01.15 | SimMIM 구현 ② → Training

2024.01.16 | SimMIM 결과 비교, 초고 작성

2024.01.17 | 코드 정리, 코드 자료 추가 → 다른 스케줄러, 학습율, weight decay 테스트

2024.01.18 | 테스트 결과 확인 및 발행