(논문리뷰) ConvNeXt V2: Co-designing and Scaling ConvNets with Masked Autoencoders [CVPR 2023]

https://arxiv.org/abs/2301.00808

ConvNeXt V2: Co-designing and Scaling ConvNets with Masked Autoencoders

최근에는 Transformer 계열이 CV로의 확장 가능성과 데이터셋 크기에 대한 확장 가능성으로 인해 ViT계열이 인기를 얻었다. 최근에는 ConvNeXt(CVPR 2022)가 기존 ConvNet을 현대화했으며 순수 conv model도 확장 가능한 아키텍처가 될 수 있음을 보여주었다.

 

이러한 신경망 구조의 발전과정을 따라감에 있어서 밴치마크로 사용되는 것은 ImageNet의 supervised learning의 성능뿐 이라는 한계점이 존재했다. 

 

이와 다른 측면의 연구에서는 Representation Learning을 함에 있어서 pre-text objective를 보면 supervised learning에서 label이 없는 self-supervised pre-training으로 방향이 이동하고 있었고, MAE가 Vision 분야에서 큰 success를 이뤘다. 

하지만 MAE는 Transformer 기반의 enc-dec 구조를 가지고 있고 masking을 하고 남은 token만을 encoder로 받기 때문에 연산량을 줄일 수 있는 이점을 가진 모델이다. 하지만 이미지 전체를 훑으며 지나가는 CNN 방식으로는 적용하기 어려운 단점이 있다. 

 

이를 해결하기 위해 sparse convolution을 사용하여 보이는 패치만 집중적으로 처리하도록 하였다. 이 방식으로 pre-train후  fine-tuning에서는 이 sparse convolution으로 학습된 가중치를 특별한 변환 없이 기존 방식(dense layer)으로 바꿔 사용하였다. 

 

이 방식은 pre-train에서는 이미지의 유용한 특징을 잘 학습할 수 있었지만 fine-tuning 시에는 transformer기반 모델보다는 낮은 성능을 보였다. 이를 보안하기 위해 ConvNeXt의 feature map을 시각화하여 관찰해본 결과 "feature collapse(특징 붕괴)"라는 문제를 발견하였다. 이 문제를 해결하기 위해 Global Response Normalization(GRN)이라는 레이어를 추가하여 모델 내의 채널 간 특징들이 잘 구분되도록 돕도록 하였다.

 

기존의 ImageNet 분류 task에 사용되었던 supervised learning에 맞춘 아키텍처들로는 최적의 성능을 내는 것이 어려웠다. 즉, self-supervised learning 같은 새로운 학습 방식에 맞춰 네트워크 구조를 다시 설계하고 최적화하는 것이 중요해졌고, 본 논문에서는 ConvNeXt와 같은 CNN 아키텍처를 MAE와 GRN등의 기법을 추가로 도입하여 모델을 재설계하는 것을 보여주었다. 

Fully Convolutional Masked Autoencoder

-Masking

Masking strategy로 60%를 masking하고 최소한의 증강기법은 Random crop만을 적용하였다. 그 후 남아있는 나머지를 context를 고려하여 model이 누락된 부분을 예측하고 생성하는 방법을 적용한다. 

- Encoder design

A key observation is that the masked image can be represented as a 2D sparse array of pixels as 3D point clouds. 

(* 3D point cloud: 데이터가 드문드문 분포되어 있어서 전체를 처리하는 것이 아니라 중요한 지점들만 집중적으로 처리하는 방식-> 보이는 부분에만 집중해서 효율적으로 학습)

 

기존의 MAE는 self-supervised learning이기에 라벨링이 필요없어 대규모 데이터로 학습이 가능하다는 장점이 있고, 모델이 더 다양한 특징을 학습할 수 있어서 다양한 downstream task로의 확장이 더 좋아진다. 

 

본 논문에서는 encoder로 ConvNeXt를 사용하였다. 모델링의 중요한 한가지 과제로는 model이 masking 영역에서 정보를 복사하여 붙여넣는 방법을 적용하지 않도록 하는 것이다. 이를 해결하기 위한 가장 simple한 solution은 learnable masked token을 사용하는 것인데 이는 pre-training의 효율성을 떨어뜨리고 train과 test 시간의 일관성을 저해하는 문제를 발생한다고 한다. 이를 해결하기 위해 도입한 방법이 "sparse data perspective"이다. 마스크된 이미지를 2D 희소 배열의 픽셀로 표현할 수 있다는 것이 중요한 통찰이다. 

즉, 1) sparse convolution을 먼저 적용하여 가려지지 않은 픽셀에만 연산을 하는 방식이다. 이렇게 되면 연산량이 줄어들어 효율적으로 모델을 학습할 수 있게된다. 그리고 이렇게 pre-training을 완료하면 다시 standard conv를 적용하여 학습한 가중치를 그대로 일반적인 방식을 적용하는 것이다. 

2) binary masking을 conv 연산 전후에 적용하여 가려진 부분을 제거하고 보이는 부분만 계산하는 방법이다. 이는 sparse conv와 비슷한 효과를 보여주지만 계산량이 더 많을수도 있는 문제가 있다. 하지만 TPU와 같이 sparse한 연산보다 dense한 연산에 최적화되어 있는 하드웨어의 경우 성능 면에서는 2) 기법이 더 유리할 수 있다. 

-Decoder design.

위 표를 보면 Decoder design을 할 수 있는데 MAE에서 제안했듯이 비대칭적인 enc-dec를 설계하고 dec는 비교적 가볍게 설계하는 것이 유리한 것을 알 수 있다. 

본 논문에서는 ConvNeXt block을 사용했고, 1개의 block만을 256 또는 512 차원으로 설정하였다. 

 

-Reconstruction target

이렇게 설계된 enc-dec 구조로 input img와 output img간의 MSE를 최소화하는 방향으로 enc-dec을 reconstruct하는 방법을 사용한다. 

 

-FCMAE

결국 FCMAE는 Swin-Transformer에서 masking기법을 사용하여 self-supervised learning을 가능하게 했던 점에 주목하여 이를 conv network에서도 차용한 것인데, Transformer에서 random하게 token을 masking을 했다면 CNN에서는 sparse convolution를 적용하여 이미지의 중요한 부분만을 선택적으로 학습할 수 있도록 설계한 것이다. 

FCMAE의 성능을 평가하기 위해 ImageNet-1K을 사용하여 800 epochs동안 사전 학습을 진행한 후, 100 epochs 동안 fine-tuning을 수행했다. 모든 실험에서는 224x224 크기의 중심 이미지를 사용해 top-1 IN-1K를 평가했다. 

위 실험 결과를 통해 마스크된 부분에서 정보가 유출되면 모델이 올바른 방식으로 학습되지 않기 때문에, 이를 방지하는 것을 알게 되었다. 

FCMAE로 학습된 모델은 랜덤 초기화한 모델에 비해 더 나은 초기 성능을 제공했다(82.7 → 83.7)

그러나 여전히 supervised 학습에서 300 에포크 동안 학습한 ConvNeXt 모델의 성능에는 미치지 못하는 결과를 보여주면 서 Transformer기반 모델처럼 self-supervised학습에서 뛰어난 성능을 보여주는데는 아직 어려움을 겪고 있다는 것을 보여준다. 

앞으로 ConvNeXt 모델의 구조적 특성을 더 깊이 이해하고, 이를 반영한 최적의 masked autoencoder 설계를 통해 FCMAE의 성능을 더욱 끌어올릴 필요가 있다는 것을 시사한다.

-GRN (Global Response Normalization)

총 레이어 수가 아키텍처마다 다르기 때문에, 정규화된 레이어 인덱스에 대해 거리 값을 계산한다. ConvNeXt V1 FCMAE 사전 학습된 모델은 심각한 feature collapse를 보인다. 이말은 channel을 확장하는 MLP계층에서 dead or saturated feature map이 있는 것이다. 

이를 보완하기 위해서 GRN을 제안하였고, 아래와 같은 3단계로 구성된다. 

• Global feature aggregation : GAP와 유사한 것으로 L2 norm을 적용하여 feature map을 scalar로 표현하도록 함. 
• Feature normalization : simple divisive normalization을 적용하여 정규화하고 
• Feature calibration : 원본 입력값을 피라미터 γ, β를 추가하여 residual connection으로 출력과 연결함. 

- ConvNeXt V2

 

ConvNeXt V2에서는 차원 확장 MLP 레이어 후에 GRN 레이어를 추가 했기에 불필요해진 LayerScale을 제거한다.