[논문 리뷰] Cross-token Modeling with Conditional Computation

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

Advanced Neural Network Applications인용 수 4

한 줄 요약

이 논문은 두 단계의 믹스처 오브 응용(MoE)을 사용하는 모든-MLP 비전 모델인 Sparse-MLP를 제안한다. 이는 토큰 간 상호작용을 모델링하기 위해 공간적 맥락을 위한 토큰 믹싱 MoE와 채널 간 정보 흐름을 위한 채널 믹싱 MoE를 사용한다. ViT-L/16보다 62.8% 적은 파라미터와 85.5% 적은 사전학습 비용을 사용하여 ImageNet Top-1 정확도 79.2%를 달성하며, 계산 비용이 더 낮은 상황에서 MLP-Mixer를 2.5% 뛰어넘는 성능을 보였다.

ABSTRACT

Mixture-of-Experts (MoE), a conditional computation architecture, achieved promising performance by scaling local module (i.e. feed-forward network) of transformer. However, scaling the cross-token module (i.e. self-attention) is challenging due to the unstable training. This work proposes Sparse-MLP, an all-MLP model which applies sparsely-activated MLPs to cross-token modeling. Specifically, in each Sparse block of our all-MLP model, we apply two stages of MoE layers: one with MLP experts mixing information within channels along image patch dimension, the other with MLP experts mixing information within patches along the channel dimension. In addition, by proposing importance-score routing strategy for MoE and redesigning the image representation shape, we further improve our model's computational efficiency. Experimentally, we are more computation-efficient than Vision Transformers with comparable accuracy. Also, our models can outperform MLP-Mixer by 2.5\% on ImageNet Top-1 accuracy with fewer parameters and computational cost. On downstream tasks, i.e. Cifar10 and Cifar100, our models can still achieve better performance than baselines.

연구 동기 및 목표

• 자기주의 어텐션 기반의 비전 트랜스포머에서 MoE 학습의 불안정성을 해결하기 위해, 토큰 간 어텐션을 조건부 계산으로 대체한다.
• 스케일러블하고 효율적인 모든-MLP 아키텍처를 설계하여, 공간(토큰) 및 채널 차원에서 MoE를 적용함으로써 모델링 능력을 향상시킨다.
• 중요도 점수 라우팅 및 표현 형태 재설계를 통해 MoE의 효율성을 향상시켜 라우팅 계산을 감소시킨다.
• ViT와 MLP-Mixer보다 더 낮은 파라미터 및 계산 비용으로 ImageNet 및 후속 비전 벤치마크에서 최신 기술 수준의 성능을 달성한다.

제안 방법

• 각 채널 내에서 패치 차원에 걸쳐 희소 전문가를 적용하는 토큰 믹싱 MoE(MoE_S)를 도입하여 공간적 맥락을 모델링한다.
• 각 패치 내에서 채널에 걸쳐 희소 전문가를 적용하는 채널 믹싱 MoE(MoE_C)를 도입하여 채널 간 상호작용을 모델링한다.
• 라우팅 계산을 줄이기 위해 중요도 점수로 토큰 또는 채널을 순위 매기는 중요도 점수 라우팅 전략을 제안한다.
• Sparse 블록 내에서 이미지 표현 형태를 재설계하여 토큰 믹싱 MoE 게이팅 네트워크의 효율성을 향상시킨다.
• 계산 비용을 감소시키고 학습 안정성을 향상시키기 위해 Sparse 블록에 리스케일 서브레이어를 적용한다.
• 다운스트림 성능 평가를 위해 ImageNet-1k에서 MoCo v3를 사용하여 자기지도 학습을 수행한다.

실험 결과

연구 질문

• RQ1 비전 트랜스포머에서 자기주의 어텐션을 대체하기 위해 믹스처 오브 응용(MoE)을 효과적이고 안정적으로 토큰 간 상호작용 모델링에 적용할 수 있는가?
• RQ2 공간적 및 채널 차원 양쪽에 MoE를 적용함으로써, 모든-MLP 아키텍처에서 모델링 능력과 효율성이 향상되는가?
• RQ3 중요도 점수 라우팅이 모델 성능을 훼손시키지 않으면서 MoE의 라우팅 계산을 줄일 수 있는가?
• RQ4 이미지 표현 형태의 재설계가 토큰 믹싱 MoE의 효율성을 향상시키는가?
• RQ5 두 단계의 MoE를 갖춘 모든-MLP 모델이 정확도와 효율성 면에서 ViT와 MLP-Mixer를 능가할 수 있는가?

주요 결과

• Sparse-B는 77.9%의 ImageNet Top-1 정확도를 달성하였으며, 계산 비용이 유사한 MLP-Mixer-B/16보다 2.0% 높고, 계산 비용이 더 낮은 ViT-B/16보다 1.2% 높다.
• Sparse-L은 79.2%의 Top-1 정확도를 기록하여, MLP-Mixer-L/16를 2.5% 뛰어넘었고, 파라미터 수는 62.8% 줄였으며, 사전학습 비용은 85.5% 감소시켰다.
• Sparse-L은 ViT-L/16보다 1.8% 높은 성능을 보였고, 사전학습 비용은 반 이하로 감소시켰다.
• 제거 실험 결과, 리스케일 서브레이어가 사전학습 비용을 30.5% 감소시키고 정확도를 76.9%에서 77.9%로 향상시켰다.
• 중요도 점수 라우팅은 모델 용량을 유지하면서 MoE의 라우팅 계산을 줄여 효율성을 향상시켰다.
• 공간적 및 채널 기반 조건부 계산을 통한 두 단계의 MoE 설계는 효과적이었으며, 이는 모든-MLP 모델이 ViT와 MLP-Mixer를 성능과 효율성 면에서 능가할 수 있음을 입증했다.