5월 8일 공개된 ConvMAE: Masked Convolution Meets MaskedAutoencoders
Abstract
Vision Transformer(ViT)는 다양한 비전 작업을 위해 널리 채택되었다. Feature 사전 훈련 및 multi-scale hybrid convolution-transformer 아키텍처를 위한 Masked auto-encoding은 ViT의 잠재력을 더욱 발휘하여 다양한 작업에서 최첨단 성능을 이끌어낼 수 있다. 본 논문에서 소개할 ConvMAE 프레임워크는 multi-scale hybrid convolution-transformer가 마스크 오토인코딩 체계를 통해 보다 차별적인 표현을 학습할 수 있음을 보여준다. 그러나 원래 마스킹 전략을 직접 사용하면 계산 비용이 많이 들고 pretraining-finetuning 불일치가 발생한다. 이 문제를 해결하기 위해, 마스크된 컨볼루션을 채택하여 컨볼루션 블록의 정보 유출을 방지한다. 계산 효율성을 보장하기 위해 간단한 블록별 마스킹 전략이 제안된다. 또한 multi-scale feature를 활성화하기 위해 인코더의 multi-scale feature를 보다 직접적으로 감독할 것을 제안한다. 사전 훈련된 ConvMAE 모델을 기반으로 한 ConvMAE-Base는 MAE-Base에 비해 ImageNet-1K 미세 조정 정확도를 1.4% 향상한다.
Introduction
Masked Autoencoders(MAE)는 최근 ViT 교육을 위한 대표적인 자체 지도 방법이다. 새로운 비대칭 인코더 및 디코더 아키텍처를 통해 수행되며, 인코더의 마스크된 토큰이 디코더에 의해 재구성된다.
자체 지도 학습의 성공과 함께 ViT의 성능을 향상하기 위해 로컬 유도 편향과 계층적 표현이 탐구된다.
로컬 컨볼루션과 글로벌 트랜스포머 작업의 조합은 다양한 작업에서 상당한 개선을 보여주었다.
MAE와 달리, 로컬 및 글로벌 운영을 기반으로 하는 우수한 성능의 다중 스케일 백본은 주로 지도 학습 접근 방식으로 훈련된다.
본 논문에서는 hybrid convolution-transformer 아키텍처와 마스크된 컨볼루션을 마스크된 오토인코더에 도입하여 ConvMAE 라고 불리는 간단하고 효과적인 자체 지도 학습 프레임워크를 제안한다. 원래 MAE에 대한 수정은 미미하지만 ConvMAE는 다양한 작업의 성능을 높이기 위해 시각적 표현을 사전 훈련하는 데 큰 성공을 보여준다.
ConvMAE의 인코더는 MAE와 달리 입력 이미지를 1,2,3 단계에서 각각 다른 입력 해상도의 다중 스케일 토큰 임베딩으로 점진적으로 추상화하고, 디코더는 다중 스케일 토큰 임베딩을 활용하여 마스크된 토큰에 해당하는 픽셀을 재구성한다. 1단계와 2단계에서의 저수준의 고해상도 토큰 임베딩의 경우 로컬 정보를 인코딩하기 위해 컨볼루션 블록을 채택하고 3단계에서의 고수준 저해상도 토큰 임베딩의 경우 트랜스포머 블록이 사용된다. 따라서 인코더는 서로 다른 단계에서 로컬/글로벌 FOV(field of view)를 가지며 차별적인 다중 스케일 기능을 생성한다.
ConvMAE 인코더로 다양한 hybrid convolution-transformer 백본을 시험했지만 큰 성능 개선이 없었기 때문에 새로운 아키텍처를 설계하는 대신 기본 hybrid convolution-transformer 아키텍처가 마스크 오토인코딩을 위해 작동하도록 하는 데 중점을 두고 다양한 다운스트림 작업에 대한 효과를 입증하기 위해 광범위한 실험을 수행한다.
ConvMAE의 효율적이고 효과적인 훈련은 마스크된 컨볼루션을 사용하는 블록별 마스킹 전략에 의해 가능하다. 기존의 마스크 오토인코딩 프레임워크에 채택된 마스킹 전략은 모든 토큰이 이후 트랜스포머 단계에서 유지되어야 하므로 ConvMAE에 사용될 수 없다. 크고 거대한 모델을 사전 훈련하기 위한 감당할 수 없는 계산 비용이 발생하며, 트랜스포머 인코더에서 마스크된 토큰을 생략하는 MAE의 효율성 이점을 잃게 된다. 또한 convolution-transformer 인코더를 사용하여 직접 사전 훈련을 하면 미세 조정 단계 동안 가시(unmasked) 토큰만 처리되기 때문에 pretraing-finetuning 불일치가 발생한다. 문제를 해결하기 위해 ConvMAE는 블록별 마스킹 전략을 채택했으며 이러한 방식으로 3단계 토큰은 마스크된 토큰과 가시적인 토큰으로 완전히 분리될 수 있고 MAE의 계산 효율성을 상속할 수 있다.
마스크된 컨볼루션은 희소 특징 추출 및 이미지 인페인팅에서 잘 탐구되었으며 그것은 마스킹된 오토인코딩을 가능하게 하기 위해 hybrid convolution-transformer 아키텍처에 자연스럽게 통합될 수 있다.
Approach
A Brief Revisit of MAE
마스크된 오토인코더(MAE)는 가시 패치에서 마스크된 RGB 패치를 재구성하여 ViT를 사전 훈련하기 위한 자체 감독 방법이다. MAE는 단순한 설계를 가지고 있지만, 강력하고 확장 가능한 사전 훈련 프레임워크로 입증되었다. MAE는 트랜스포머 기반 인코더와 디코더로 구성되며, 가시 패치만 인코더로 공급되고 디코더에 의해 마스크 토큰이 처리되어 시각적 표현을 학습한다. 인코더는 가시 토큰만 처리하면 되므로 대규모 비전 모델을 사전 훈련하기 위한 확장성 문제를 완화한다.
ConvMAE
ConvMAE의 목표는 차별적인 다중 스케일 시각적 표현을 학습하고 MAE를 convolution-transformer 네트워크에 적용할 때 pretraing-finetuning 불일치를 방지하는 것이다. 컨볼루션 단계에서 마스크된 컨볼루션과 결합된 계층적 마스킹 전략을 도입하여 소수의 가시 토큰만 트랜스포머 레이어에 입력되도록 한다.
전체 파이프라인은 다음과 같다.
The Hybrid Convolution-transformer Encoder
인코더는 3단계로 구성된다. 처음의 두 컨볼루션 단계는 컨볼루션 블록을 이용해 입력을 토큰 임베딩 E1(H/4 x W/4 x C1), E2(H/8 x W/8 x C2)로 변환한다. 컨볼루션 블록은 self attention을 5x5 depth-wise 컨볼루션으로만 대체함으로써 트랜스포머 블록의 설계 원칙을 따른다.
트랜스포머 단계에서는 토큰 임베딩을 얻기 위해 일반적인 self attention 블록을 사용한다. 1,2 단계에서는 컨볼루션으로 로컬 범위를, 3단계에서는 트랜스포머로 글로벌 범위를 본다.
다른 ViT 모델들과는 달리, 입력에서의 절대 위치 임베딩을 상대 위치 임베딩 또는 컨볼루션으로 대체하고 클래스 토큰 또한 제거한다.
모든 단계 사이에서 stride-2 컨볼루션으로 다운샘플링한다.
Block-wise Masking with Masked Convolutions
기존의 마스킹 전략처럼 입력 토큰을 마스킹하면 단계를 지나면서 빈자리가 채워지기 때문에 3단계에서 모든 토큰이 부분적으로 가시 정보를 가지게 된다. 그래서 대신에 3단계 입력 토큰의 일부를 마스킹하고 마스크를 2배, 4배 업샘플링하여 1, 2단계 토큰에 적용한다. 해당 마스크된 토큰은 인코딩 과정에서 삭제되고 feature 학습을 위해 디코더에 의해서 재구성된다. 이러한 방법으로 MAE의 계산 효율성을 유지한다.
그러나, 처음 두 단계에서 5×5 depth-wise 컨볼루션은 자연스럽게 마스크된 패치보다 큰 수용 필드를 초래하고 마스크된 토큰을 재구성할 때 정보 유출을 일으킨다. 이러한 정보 유출을 방지하기 위해, 처음 두 단계에서 마스크된 컨볼루션을 사용한다. 마스크된 컨볼루션은 마스크된 패치의 가중치를 활성화하지 않음으로써 pretraining-finetuning 불일치를 방지하며 ConvMAE의 성능에 매우 중요한 영향을 끼친다.
The Multi-scale Decoder and Loss
1, 2단계의 출력 E1, E2를 각각 stride-4, stride-2 컨볼루션으로 인코더 출력 E3과 같은 크기로 만들고 linear layer를 통해 융합하여 디코더 입력에 사용할 가시 토큰 Ed를 얻는다.
마스크된 패치를 재구성하는 데 MAE와 동일한 손실이 사용되며 마스크된 패치의 재구성만 고려된다.
(각 패치 t에서의 MSE 계산)
ConvMAE for Object Detection and Semantic Segmentation
그림과 같이 객체 탐지를 위한 ConvMAE를 미세 조정하기 위해서는, 먼저 maxpool로 E4를 얻는다.
3단계에는 과도한 계산 비용이 드는 11개의 global self attention 계층이 있으므로 Benchmarking ViT에 따라 1,4,7,11번째 self attention 계층을 Swin Transformer의 shifted window local self attention으로 교체한다.
그런 다음 E 1,2,3,4를 MaskRCNN으로 공급한다.
의미론적 분할에서는 E 1,2,3을 UperNet에 공급한다.
Conclusion
본 연구는 hybrid local-global 블록의 성능을 높여 차별적인 다중 스케일 feature을 생성할 수 있음을 보여주는 ConvMAE라는 간단한 자체 지도 학습 프레임워크를 제안한다. 기존 MAE의 계산 효율성과 낮은 pretraining-finetuning 차이는 ConvMAE에서 잘 유지될 수 있다. ConvMAE는 다양한 비전 작업에서 상당히 향상된 성능을 보이며 쉽게 구현할 수 있다.
Experments
ImageNet-1K Pretraining and Finetuning
ImageNet1K training 세트에서 ConvMAE를 사전 훈련한다.
마스크 비율 25%.
AdamW optimizer 사용.
대부분의 매개변수의 경우 기존 MAE의 것을 따른다.
유사한 모델 크기의 SOTA 방법들과 분류 성능 비교
(FT = fine tuning, LIN = linear probe accuracy - 모델 가중치 고정한 뒤에 linear layer 붙여서 테스트하는거 말하는듯)
Object Detection
COCO 데이터세트.
헤더로 Mask-RCNN 사용.
대부분의 Benchmarking ViT 설정 따름.
Semantic Segmentation
ADE20K 데이터세트.
헤더로 UperNet 사용.
Ablation Study of ConvMAE
사전 훈련 많이 하면 좋음.
Masked Conv와 Block Masking은 확실히 효과가 있으며,
Conv 크기는 별로 영향이 없어서 비용이 적게 드는 5x5를 기본으로 사용한 듯 하다.
Multi-scale decoder의 영향 검증.(E 1,2,3 다 섞는거 말함)
모델 사이즈 별 성능 비교.
MAE와의 수렴 속도 비교.
