MLP-Mixer: An all-MLP Architecture for Vision

MLP-Mixer: An all-MLP Architecture for Vision

2021. 5. 17. 14:36ㆍ논문 리뷰

728x90

Abstract

Computer Vision에서 CNN은 지배적이었다. 최근에 attention-based의 네트워크들(Vision Transformer 등)이 유명해지기 시작했다. 이 논문에서는 MLP-Mixer라는 Architecture를 제안하며, convolution과 attention이 충분히 좋은 성능을 내지만 필요없다는 것을 보인다. MLP-Mixer는 only MLP로만 만들어졌다.

MLP-Mixer는 두가지 레이어를 가지고 있다.

applied independently to image patches -> mixing per-location features
applied across patches -> mixing spatial information

많은 데이터로 학습을 했을 때, image classification에서 pre-training면이나 inference면에서 sota와 비슷한 성능을 보였다. 이 논문의 저자들은 MLP-mixer가 CNN과 Transformer를 대체할 수 있는 spark가 되길 바란다.

Introduction

Computer Vision 분야에서의 역사가 보여주듯이, 더 큰 dataset의 가용성과 computational capacity의 향상은 종종 패러다임의 변화를 이끌어낸다. 기존의 Vision task들은 CNN이 거의 지배적으로 사용되었고, 최근에서야 ViT와 같이 self-attention layer들로 대체되고 있으며 성능 또한 sota를 달성했다.

이 논문에서는 only MLP만 사용한 MLP-Mixer를 제안한다. Mixer가 only MLP만 사용하지만 우리가 제안한 구조는 spatial location 혹은 feature channels를 아울러 적용되는 MLP기반이다.

위의 그림은 Mixer의 구조를 보여준다. 이는 sequence of linearly projected image patches를 받는다. patches X channels shape이 input의 shape이 되며, 이 dimension을 유지한다. Mixer는 두가지 타입의 MLP layer를 만들어 사용하는데,

channel-mixing MLPs
token-mixing MLPs

이 두가지이다. 먼저 channel-mixing MLPs 같은 경우 다른 채널간의 소통을 가능케 한다. 각 token들 별로 독립적으로 수행하고 각각의 개인적인 row를 input으로 받는다. 다음 section에서 좀 더 자세히 다루도록 하겠다.

두번 째 token-mixing MLPs의 경우 다른 spatial locations간의 communication들을 가능케 한다. 이 역시 다음 section에서 좀 자세히 다루도록 하겠다.

이 두가지 layer를 통해 input dimension의 양쪽으로 interaction이 가능하게 된다.

어떻게 보면 CNN의 특별 케이스 사용으로 볼 수 있다. 1 x 1 convolutions를 사용하여 channel mixing하는 것과 single-channel depth-wise convolution을 사용해서 full receptive field와 token mixing을 위한 parameter sharing을 하는 과정이라고 볼 수 있다. 그러나 반대의 경우 전형적인 CNN의 경우 Mixer의 일환이라고 볼 수는 없다. 게다가 convolution은 MLP보다 훨씬 더 복잡한 matrix multiplication이다.

MLPs의 경우 이 simplicity임에도 불구하고 경쟁할만한 성능을 얻는다. 큰 dataset으로 pre-train 시켰을 때 sota 성능을 달성했다. 하지만 ViT와 마찬가지로 아주 특성화 된 CNN에는 성능이 아직 못 미친다.

Mixture Architecture

Modern deep vision 구조의 layer들은 features를 mix하는 구조로 구성되어 있다.

at a given spatial location
between different spatial locations, or both at once

결국 spatial location 들의 정보 혹은 spatial location 간의 정보들을 어떻게든 섞어 관계를 찾는다는 소리이다.

CNN에서는 spatial location간의 관계를 찾는 것을 N x N convolutions와 pooling으로 구현했다. 더 깊은 위치에 있는 layer가 더 큰 receptive field를 갖는다.

1x1 convolution의 경우 spatial location의 정보를 얻는다. 큰 kernel size를 가질 수록 1, 2 둘 다 수행하게 된다. Vision Transformer 구조의 경우도 self-attention 구조가 spatial location에서의 관계, spatial location들 간의 관계 둘다 수행하게 되고 MLP의 부분에서 spatial location에서의 관계를 보게 된다.

이 논문에서는 spatial location에서의 관계 뿐만 아니라 spatial location들 간의 관계를 파악하는데 only MLP만 사용한다.

위의 구조 그림에서 나와있다 싶이, S sequence의 image patch들을 input으로 받는다. 각각은 hidden dimension C로 projected 된다.

X ∈ R^(S x C) -> input dimension이 S(sequence length), C(hidden dimension)

만약 original input image 가 (H, W)의 사이즈를 가지고 있고 각 patch가 (P, P)로 나뉘게 된다면, S = HW / P^2 가 된다.

모든 patch들은 같은 projection matrix에 의해 projected된다. Mixer는 multiple layer들로 구성되어 있고, 각 layer는 두개의 MLP block으로 구성되어진다.

첫번째는 token-mixing MLP block, 이 block은 X를 transpose 하여, column 들간의 MLP를 돌리게 되는데 이렇게 할 경우 모든 sequence에 대해 MLP를 하게 된다.

두번째는 channel-mixing MLP block, 이 block은 다시 Transpose를 시켜, 각 sequence별 관계를 파악하게 된다.

위의 과정을 수식으로 표현하면 이렇게 된다. 하지만 위의 과정 그림을 보는게 더 이해가 빠르다.(그렇게 어려운 개념이 아니므로..)

여기서 activate function은 GELU를 사용했다. D_S의 경우 token-mixing시 hidden width를 나타내고 D_C의 경우 channel-mixing의 경우 hidden width를 나타낸다. D_S의 경우 input patch들의 갯수에 독립적으로 선택되고 D_C의 경우 patch size에 독립적으로 설정 된다. 이를 통해 Mixer의 복잡도는 ViT에 비해 훨씬 가볍게 된다.

MLP layer들과 별개로 Mixer는 skip-connection과 Layer Normalization 구조 또한 활용한다. 게다가 ViT와는 달리, Mixer는 position embedding을 사용하지 않는다.

class PreNormResidual(nn.Module):
    def __init__(self, dim, fn):
        super().__init__()
        self.fn = fn
        self.norm = nn.LayerNorm(dim)

    def forward(self, x):
        return self.fn(self.norm(x)) + x

def FeedForward(dim, expansion_factor = 4, dropout = 0., dense = nn.Linear):
    return nn.Sequential(
        dense(dim, dim * expansion_factor),
        nn.GELU(),
        nn.Dropout(dropout),
        dense(dim * expansion_factor, dim),
        nn.Dropout(dropout)
    )

def MLPMixer(*, image_size, patch_size, dim, depth, num_classes, expansion_factor = 4, dropout = 0.):
    assert (image_size % patch_size) == 0, 'image must be divisible by patch size'
    num_patches = (image_size // patch_size) ** 2
    chan_first, chan_last = partial(nn.Conv1d, kernel_size = 1), nn.Linear

    return nn.Sequential(
        Rearrange('b c (h p1) (w p2) -> b (h w) (p1 p2 c)', p1 = patch_size, p2 = patch_size),
        nn.Linear((patch_size ** 2) * 3, dim),
        *[nn.Sequential(
            PreNormResidual(dim, FeedForward(num_patches, expansion_factor, dropout, chan_first)),
            PreNormResidual(dim, FeedForward(dim, expansion_factor, dropout, chan_last))
        ) for _ in range(depth)],
        nn.LayerNorm(dim),
        Reduce('b n c -> b c', 'mean'),
        nn.Linear(dim, num_classes)
    )

MLP Mixer를 코드로 나타내 본 것이다. 그렇게 복잡한 형태는 아니다. 단순한 MLP를 아이디어로 탈바꿈한 것이라고 보여진다.

Experiments

이 논문에서 MLP-Mixer model에 대한 성능을 평가한다. 작은 모델부터 큰 모델까지의 pre-trained 시킨 것, 작은 것과 중간 사이즈의 downstream classification task에 대한 성능 평가를 진행한다. 흥미롭게 볼 것은 downstream task에서의 정확도, total computational cost of pre-training, inference 진행 시간 이 세가지이다. 이 논문의 목표는 SOTA를 증명하는데 있는 것이 아니라 간단한 MLP-based model 만으로도 충분히 오늘날의 좋은 성과를 내고 있는 convolution이나 attentino-based model과 경쟁할만한 모델이 나온다는 것이다.

Specifications of the Mixer architectures used in this paper

다양한 실험을 통해 성능을 증명했다.

실험 내용들은 읽어보도록하자..