[논문 리뷰] Spatial Mixture-of-Experts

Nikoli Dryden, Torsten Hoefler|arXiv (Cornell University)|2022. 11. 24.

Human Mobility and Location-Based Analysis인용 수 5

한 줄 요약

이 논문은 공간적 인도적 편향을 통해 데이터의 미세한 공간적 구조를 학습하는 데 중점을 두고, 텐서 라우팅 함수를 사용해 픽셀 수준에서 전문가를 라우팅하는 희박하게 게이팅된 신경망 구성요소인 공간 혼합 전문가(SMoE) 레이어를 소개한다. 자기지도 라우팅 손실과 전문가 오차 감쇠를 통해 공간적 인도적 편향을 통합함으로써 중거리 기상 예측 및 앙상블 후처리 과제에서 최신 기술 수준(SOTA) 성능을 달성하였으며, 밀도 높은 회귀 과제에서 컨볼루션 및 표준 MoE 모델보다 뛰어난 성능을 보였다.

ABSTRACT

Many data have an underlying dependence on spatial location; it may be weather on the Earth, a simulation on a mesh, or a registered image. Yet this feature is rarely taken advantage of, and violates common assumptions made by many neural network layers, such as translation equivariance. Further, many works that do incorporate locality fail to capture fine-grained structure. To address this, we introduce the Spatial Mixture-of-Experts (SMoE) layer, a sparsely-gated layer that learns spatial structure in the input domain and routes experts at a fine-grained level to utilize it. We also develop new techniques to train SMoEs, including a self-supervised routing loss and damping expert errors. Finally, we show strong results for SMoEs on numerous tasks, and set new state-of-the-art results for medium-range weather prediction and post-processing ensemble weather forecasts.

연구 동기 및 목표

• 위치 기반의 구조를 가진 데이터에서 미세한 공간적 상관관계를 포착하지 못하는 표준 신경망의 한계를 해결하기 위해.
• 밀도 높은 회귀 과제에서 위치에 따라 달라지는 패턴을 학습하지 못하는 컨볼루션 및 표준 혼합 전문가 모델의 실패를 극복하기 위해.
• 모든 공간적 위치에서 전문가가 특화되도록 하되 계산 효율성을 유지하는 라우팅 메커니즘을 개발하기 위해.
• 새로운 훈련 기법을 통해 공간적 인도적 편향을 통합함으로써 기상 예측 과제에서 모델 성능을 향상시키기 위해.

제안 방법

• 각 픽셀 수준에서 입력을 공유된 전문가 집합으로 희박하게 라우팅하는 학습된 게이팅 함수를 사용하는 공간 혼합 전문가(SMoE) 레이어를 제안한다.
• 입력 도메인의 공간적 구조를 명시적으로 모델링하는 가벼우면서 효과적인 라우팅 함수인 텐서 라우팅을 도입한다.
• 라우팅 오류를 예측함으로써 게이트에 더 강한 학습 신호를 제공하기 위해 자기지도 라우팅 분류 손실을 적용한다.
• 잘못 라우팅된 샘플로부터의 기울기 업데이트를 억제함으로써 견고성과 성능을 향상시키기 위해 전문가 오차 감쇠를 적용한다.
• 유사한 공간적 치수를 가진 레이어들 간에 라우팅 게이트를 공유함으로써 파rameter를 줄이고 일반화 성능을 향상시킨다.
• U-Net 및 ResNet 아키텍처 내부에 SMoE 레이어를 훈련시키며, 표준 컨볼루션을 대체함으로써 공간 특화를 가능하게 한다.

실험 결과

연구 질문

• RQ1 혼합 전문가 아키텍처는 위치 기반의 구조를 가진 데이터에서 미세한 공간적 상관관계를 효과적으로 학습할 수 있는가?
• RQ2 MoE에서 미세한 픽셀 수준의 라우팅은 굵은 픽셀 수준의 라우팅이나 표준 컨볼루션보다 밀도 높은 회귀 과제에서 더 높은 성능을 낼 수 있는가?
• RQ3 엔드 투 엔드 라우팅이 공간적 구조를 포착하지 못할 경우, 자기지도 라우팅 손실이 게이트 훈련을 향상시킬 수 있는가?
• RQ4 라우팅 오류가 존재하는 상황에서 전문가 오차 감쇠는 모델 수렴과 성능에 어떤 영향을 미치는가?
• RQ5 SMoE 레이어는 더 적은 파라미터로 기상 예측 및 이미지 분류를 포함한 다양한 과제에 일반화 가능한가?

주요 결과

• SMoE는 중거리 기상 예측을 위한 WeatherBench 벤치마크에서 최신 기술 수준(SOTA) 성능을 달성하였다.
• 앙성 기상 예측 후처리를 위한 ENS-10 데이터셋에서, 모든 리드 타임과 앙성 크기에서 CRPS 및 EECRPS 점수가 모두 낮아 새로운 SOTA를 수립하였다.
• 위치 기반의 확산 계수를 가진 제어된 열 확산 과제에서, 표준 CNN 및 MoE는 평균 확산 계수로 수렴하는 반면, SMoE는 공간적 구조를 성공적으로 학습하였다.
• ImageNet-1k에서, 모든 3×3 컨볼루션을 대체할 때 SMoE는 파라미터를 56% 줄이며 경쟁 가능한 정확도를 유지하면서 LRLCN보다 뛰어난 성능을 보였다.
• 텐서 라우팅, 자기지도 라우팅 손실, 전문가 오차 감쇠의 조합은 특히 미세한 공간 모델링이 필요한 과제에서 뚜렷한 성능 향상을 이끌었다.
• SMoE는 글로벌 기상 데이터와 같은 큰 공간 영역으로도 효과적으로 확장되며 성능이나 효율성의 손실 없이 작동한다.