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

순수 Conv로만 이루어낸 ConvNeXt + MAE

 

arXiv

Github

 

 

Abstract

ConvNet에 masked autoencoder(MAE)와 같은 자체 지도 학습을 도입했을 때, 성능이 낮아졌다.

 

본 논문에서는 ConNeXt 아키텍처에 fully convolutional masked autoencoder framework를 도입하고 채널 간 feature competition을 강화하기 위해 Global Response Normalization (GRN) layer를 추가한다.

 

 

 

Introduction

MAE는 애초에 transformer sequence 처리에 최적화되어 있기 때문에 ConvNet에 적용했을 때 성능이 좋지 않고, 따라서 Conv용 MAE를 설계하는 것이 본 논문의 목표이다. 

 

3D point cloud를 처리하는 데 쓰이는 sparse convolution에 영감을 받아, 마스킹된 입력을 희소 패치 세트로 취급하고 희소 컨볼루션으로 처리한다.

Finetuning 시 특별한 처리 없이 표준 밀도로 변형되고, transformer decoder를 ConvNeXt block으로 대체해 완전한 컨볼루션 설계로 만든다.

 

추가적인 개선으로 마스크된 입력에서 ConvNeXt를 훈련할 때 MLP 계층의 문제를 식별하고, 채널 간 feature competition을 강화하기 위해 GRN 계층을 추가한다.

 

 

 

Related Work

ConvNeXt

MAE

 

 

 

Fully Convolutional Masked Autoencoder

---

Masking

마스킹 비율 0.6

마스크는 가장 낮은 해상도에서 생성되어 업샘플링

증강은 only random resized cropping

---

Encoder Design

Transformer에서는 feature가 이미지 형태로 존재하지 않기 때문에 단순 이미지 복붙을 배우는 것을 막기 쉽지만 conv에서는 그렇지 않다. 마스크 토큰을 학습 가능하게 하는 것은 train-test 불일치를 초래함.

 

해결책은 마스킹된 이미지를 희소 데이터의 관점으로 보는 것이다.

인코더의 컨볼루션 계층을 가시 데이터에만 작동하는 submanifold sparse convolution으로 대체한다.

 

희소 데이터가 아닌 finetuning 단계에서는 정상 컨볼루션처럼 잘 작동한다.

---

Decoder design

디코더는 512차원의 단일 ConvNeXt block을 사용한다.

계층적 구조의 인코더에 비하면 매우 비대칭이지만 복잡한 다른 디코더들 보다 빠르고 우수한 성능을 보여줬다.

---

Reconstruction target

재구성 이미지의 마스킹 되었던 부분과 목표 이미지와의 MSE 계산

---

FCMAE

위의 구성요소들을 조합한 것을 Fully Convolutional Masked AutoEncoder라고 정함.

ConvNeXt-Base에서 pre-train 800ep, fine-tune 100ep 기준으로 ablation study 진행.

 

희소 컨볼루션 좋아요.

 

Fine-tune 지도 학습 300ep에게는 아쉽게 패배...

 

 

 

Global Response Normalization

FCMAE의 문제점을 찾아 강화시켜 보자.

 

Feature collapse

차원 확장 MLP 계층에서 포화된 feature map이 발견됨. (V1 부분)

---

Feature cosine distance analysis

Feature map의 i번째 채널을 X_i라고 할 때, X_i를 HxW 벡터로 바꾸고 각 채널간의 pair-wise cosine distance 계산.

위 식의 값이 클수록 다양한 feature를 나타낸다고 볼 수 있다.

 

확실히 Conv + MAE인 V1 FCMAE는 바닥에 있다.

---

Approach

뇌에서 개별 뉴런의 선택성을 높이고 뉴런 모집단의 반응 다양성을 증진하는 lateral inhibition에서 영감을 받았다고 한다.

 

그렇게 만들어진 GRN(Global Response Normalization)은 세 단계로 구성된다.

• Global feature aggregation
• Feature normalization
• Feature calibration

---

먼저 X_i를 일종의 pooling인 G()로 집계

 

L2 norm이 가장 좋았다고.

---

정규화

 

해당 정규화가 lateral inhibition에 의해 channel competition을 발생시킨다고 한다.

---

적용

 

피라미터 γ, β와 잔차 연결 추가

좋아요.

---

코드 세 줄로 가능

---

ConvNeXt V2

GRN에 정규화가 들어있어서 LayerScale 뺌.

 

좋아요.

---

Relation to feature normalization methods

 

Relation to feature gating methods

 

 

 

ImageNet Experiments