[논문 리뷰] Sparse-MLP: A Fully-MLP Architecture with Conditional Computation.

Yuxuan Lou, Fuzhao Xue외 2명|arXiv (Cornell University)|2021. 09. 05.

Domain Adaptation and Few-Shot Learning참고 문헌 41인용 수 5

한 줄 요약

Sparse-MLP는 완전히 MLP 기반 아키텍처로, MLP-Mixer에 희소 Mixture-of-Experts (MoE) 메커니즘을 통합하여 일부 밀집 MLP 블록을 두 단계 MoE 레이어로 대체함으로써 차원 간 특징 혼합을 수행한다. Re-represent 레이어를 도입하여 전문가의 능력과 라우팅 효율성을 향상시킴으로써, 파라미터 수가 적고 계산 비용이 낮은데도 불구하고 밀집 MLP 기준 모델보다 ImageNet Top-1 정확도가 2.5% 높아졌다.

ABSTRACT

Mixture-of-Experts (MoE) with sparse conditional computation has been proved an effective architecture for scaling attention-based models to more parameters with comparable computation cost. In this paper, we propose Sparse-MLP, scaling the recent MLP-Mixer model with sparse MoE layers, to achieve a more computation-efficient architecture. We replace a subset of dense MLP blocks in the MLP-Mixer model with Sparse blocks. In each Sparse block, we apply two stages of MoE layers: one with MLP experts mixing information within channels along image patch dimension, one with MLP experts mixing information within patches along the channel dimension. Besides, to reduce computational cost in routing and improve expert capacity, we design Re-represent layers in each Sparse block. These layers are to re-scale image representations by two simple but effective linear transformations. When pre-training on ImageNet-1k with MoCo v3 algorithm, our models can outperform dense MLP models by 2.5\% on ImageNet Top-1 accuracy with fewer parameters and computational cost. On small-scale downstream image classification tasks, i.e. Cifar10 and Cifar100, our Sparse-MLP can still achieve better performance than baselines.

연구 동기 및 목표

• Mixture-of-Experts (MoE)를 통한 희소 조건부 계산을 도입하여 MLP-Mixer의 효율성과 확장성을 향상시키는 것.
• 비전 트랜스포머 및 MLP-Mixer 아키텍처에서 밀집 MLP 블록의 계산 비효율성과 제한된 능력 문제를 해결하는 것.
• 희소 블록 내에서 두 단계 MoE 레이어를 통해 채널 차원과 패치 차원 모두에서 더 효과적인 특징 혼합을 가능하게 하는 것.
• 대표성 재척화를 위한 Re-represent 레이어를 도입하여 라우팅 오버헤드를 줄이고 전문가 활용도를 향상시키는 것.
• 밀집 MLP 모델보다 파라미터 수와 FLOPs가 적은데도 이미지 분류 벤치마크에서 최신 기술 수준의 성능을 달성하는 것.

제안 방법

• MLP-Mixer의 일부 밀집 MLP 블록을 두 단계 MoE를 구현하는 희소 블록으로 대체함: 하나는 패치 간 채널 방향 혼합을 위해, 다른 하나는 채널 간 패치 방향 혼합을 위해.
• MoE 라우팅 이전과 이후에 적용되는 두 개의 단순 선형 변환인 Re-represent 레이어를 도입하여 특징 표현을 재척화하고 전문가의 능력 및 라우팅 효율성을 향상시킴.
• 희소 MoE 환경에서 강력한 자기지도 학습을 보장하기 위해 ImageNet-1k에서 MoCo v3를 사용하여 사전학습을 수행함.
• 각 MoE 레이어에서 토큰당 상위-k 전문가를 동적으로 선택하는 라우터 메커니즘을 적용하여 조건부 계산을 구현하고 FLOPs를 감소시킴.
• 희소 활성화된 전문가를 통해 고용량 특징 학습이 가능하도록 하면서도 계산 효율성을 유지하는 희소 블록을 설계함.
• 입력 표현을 안정화하고 정교화하여 전문가 선택 이전에 Re-represent 레이어를 적용함으로써 라우팅 과정을 최적화함.

실험 결과

연구 질문

• RQ1 MLP-Mixer와 같은 완전히 MLP 기반 아키텍처에 희소 MoE 메커니즘이 효과적으로 통합될 수 있는가?
• RQ2 채널과 패치 차원을 별도로 처리하는 두 단계 MoE 설계 방식이 특징 표현과 모델 정확도에 어떤 영향을 미치는가?
• RQ3 Re-represent 레이어가 희소 MLP 아키텍처에서 전문가 활용도를 향상시키고 라우팅 오버헤드를 얼마나 줄이는가?
• RQ4 희소 MLP 아키텍처가 ImageNet-1k 및 소규모 벤치마크에서 정확도, 파라미터 수, FLOPs 측면에서 밀집 MLP 기준 모델을 능가할 수 있는가?
• RQ5 제안된 Sparse-MLP는 최소한의 아키텍처 변경으로도 CIFAR-10 및 CIFAR-100과 같은 후행 작업에서 강력한 일반화 성능을 유지하는가?

주요 결과

• Sparse-MLP는 MoCo v3를 사용해 ImageNet-1k에서 사전학습한 후, 밀집 MLP-Mixer 기준 모델보다 ImageNet Top-1 정확도가 2.5% 높다.
• 이 성능 향상은 파라미터 수가 적고 계산 비용이 낮은데도 달성되었다.
• CIFAR-10 및 CIFAR-100에서는 이들 데이터셋의 크기가 작음에도 불구하고 기존 기준 모델을 능가하는 성능을 보였다.
• Re-represent 레이어는 MoE 메커니즘에서 전문가의 능력 향상과 라우팅 관련 계산 오버헤드 감소에 효과적으로 기여했다.
• 두 단계 MoE 설계는 FLOPs를 크게 증가시키지 않으면서도 효율적인 차원 간 특징 혼합을 가능하게 하여 표현 학습을 향상시켰다.
• 희소 아키텍처는 다양한 벤치마크에서 강력한 일반화 성능을 유지하며 확장성과 효율성의 우수함을 입증했다.