기술 요약: 비전에서 희소 MoE 가 언제 도움이 되는가?

1. 문제 제기

혼합 전문가 (Mixture-of-Experts, MoE) 네트워크는 입력을 전문화된 서브네트워크로 라우팅함으로써 정확도 - 연산량 간의 유리한 트레이드오프를 위한 이론적 경로를 제공합니다. 희소 게이팅이 언어 모델 및 일부 비전 응용 분야에서 성공적으로 확장되었음에도 불구하고, 비전 분류에 대한 실제 배포는 두 가지 지속적인 장애물에 직면해 있습니다:

1. 연산 불균형: 표준 비전 아키텍처에서 공유된 합성곱 백본이 총 연산 비용 (FLOPs) 을 지배합니다. 결과적으로 MoE 가 종종 배치되는 분류 헤드는 총 연산량의 미미한 부분 (종종 0.1% 미만) 을 차지합니다. 완벽한 라우팅이 이루어지더라도 종단 간 비용을 줄일 수 있는 구조적 기회는 최소화되어 라우팅 메커니즘을 개선할 동기를 약화시킵니다.
2. 전문가 붕괴: 라우터는 종종 전문가의 일부 하위 집합으로 붕괴되어 MoE 가 달성하도록 설계된 전문화를 훼손합니다.

본 논문은 MoE 가 보편적으로 유익하다는 가정을 넘어, 하드 용량 제약이 있는 희소 top-$k$ 라우팅이 비전 분류에서 실제로 정확도 향상을 가져오는 구체적인 조건을 조사합니다.

2. 방법론

저자들은 비전에서 MoE 성능을 주도하는 변수들을 분리하기 위해 통제된 실증 연구를 제안합니다.

아키텍처 및 라우팅

• 백본: 본 연구는 표준 합성곱 백본과 깊이별 분리 합성곱 (DW) 을 활용합니다. 깊이별 합성곱은 효율성 때문뿐만 아니라 헤드로부터 백본까지의 FLOPs 비율을 동등한 수준으로 이동시켜 라우팅 가능한 총 연산량의 비율 ($\rho$) 을 증가시키기 위해 사용됩니다.
• 희소 라우터: 입력은 학습된 로짓과 키 유사도 항 ($\cos(Kh, e_i)$) 을 결합한 학습된 라우터를 통해 라우팅됩니다.
• 하드 용량 제약: 전문가 붕괴를 방지하기 위해 저자들은 Switch Transformer 에 따라 배치당 하드 용량 제한을 부과하여 훈련 중 모든 $E$ 개의 전문가에 걸쳐 균형 잡힌 배분을 보장합니다.
• 훈련 목적 함수: 손실 함수는 교차 엔트로피와 부하 균형 및 엔트로피 정규화 항을 결합합니다. 온도 스케줄링 (선형 또는 시그모이드) 은 라우팅 분포를 어닐링하는 데 사용됩니다.

실험 설계

본 연구는 CIFAR-10/100, Tiny-ImageNet, ImageNet-1K 를 포함한 다중 데이터셋 프로토콜을 활용하며, 다음과 같은 통제된 실험을 수행합니다:

1. $\rho$-스윕 (CIFAR-10): 다른 요인들을 고정하면서 총 FLOPs 중 라우팅되는 비율 ($\rho$) 을 변화시키기 위해 표준 및 깊이별 백본에 걸쳐 숨겨진 크기 (hidden-size) 를 스윕합니다.
2. 라우팅 폭 제거 실험 (ImageNet-1K): 아키텍처, 초기화, 그리고 $\rho$를 고정하면서 top-$k$ 선택 ($k=1$ 대 $k=2$) 만을 변화시키는 실험입니다.
3. 배분 진단: 배치 차원을 통해 토큰을 배분하는 표준 Soft MoE 와 이미지당 전문가에 대해 소프트맥스하는 제안된 샘플 단위 변형 간의 비교를 통해 배분 실패 모드를 분리합니다.
4. 백본 통합: 높은-$\rho$ 체제를 테스트하기 위해 MoE 레이어를 ResNet-18 의 레이어 3/4 와 같은 백본 합성곱 내부에 직접 배치합니다.

3. 주요 발견

연산 레버리지 패턴

핵심 관찰은 "연산 레버리지" 패턴입니다: MoE 와 밀집 베이스라인 간의 긍정적인 정확도 차이는 총 FLOPs 의 상당한 비율 $\rho$가 라우팅되어야 발생합니다.

• 낮은 $\rho$ (헤드 전용): 라우팅된 헤드가 총 연산량의 미미한 부분일 때 (예: 표준 백본을 사용한 ImageNet 에서 $\rho < 1\%$), 희소 라우팅은 백본 유형에 관계없이 일관되게 성능을 저하시킵니다.
• 높은 $\rho$ (깊이별/백본): 깊이별 합성곱을 사용하여 $\rho$를 높일 때 (예: CIFAR 에서 약 49%) 또는 MoE 를 백본 내부에 배치할 때 (ImageNet 에서 $\rho \approx 38\%$), 긍정적인 정확도 차이가 나타납니다.

대규모에서의 다중 전문가 라우팅의 필요성

ImageNet 규모에서 높은 $\rho$는 필요하지만 충분하지는 않습니다.

• 고정된 아키텍처와 $\rho$를 가진 ImageNet-1K 에 대한 통제된 제거 실험은 $k=1$(단일 전문가) 에서 $k=2$(다중 전문가) 로 전환했을 때 다섯 개의 시드 전체에서 정확도 차이를 $-2.08\%$에서 $+1.17\%$로 반전시켰음을 보여주었습니다.
• 이는 대규모에서 단순히 연산량의 큰 비율을 라우팅하는 것만으로는 부족하며, 모델이 이득을 실현하기 위해 동시에 여러 전문가를 활용해야 함을 나타냅니다.

라우팅 안정성 및 실패 모드

• 온도 스케줄링: 온도의 공격적인 선형 어닐링은 후기 단계의 라우팅 붕괴를 초래할 수 있습니다. 안정성을 위해서는 시그모이드 스케줄 또는 더 높은 최소 온도 ($\tau_{min}$) 가 필요합니다.
• 배분 메커니즘: 배치 차원을 통해 토큰을 평균화하는 표준 Soft MoE 는 관련 없는 이미지를 혼합하기 때문에 샘플 단위 CNN 설정에서 치명적으로 붕괴합니다. 샘플 단위 변형 (이미지당 전문가에 대해 소프트맥스) 은 CIFAR-100 에서 성능을 구제하지만 Tiny-ImageNet 과 ImageNet-1K 에서는 밀집 베이스라인을 능가하지 못하여, 규모가 커질 경우 하드 라우팅 없이 가중치 앙상블만으로는 불충분함을 시사합니다.

효율성 대 지연 시간

• FLOPs: 본 연구는 깊이별 백본이 정확도를 유지하거나 향상시키면서 CIFAR-10 에서 최대 22.7% 에 달하는 상당한 FLOPs 감소를 가능하게 함을 확인합니다.
• 지연 시간: 퓨전되지 않은 구현은 전문가별 커널 시작 오버헤드 (Python 루프 및 순차적 CUDA 호출) 로 인해 높은 월클록 지연 시간을 겪습니다. FLOPs 는 낮지만 지연 시간은 밀집 베이스라인보다 현저히 높습니다 (예: CIFAR-100 에서 36ms 대 0.25ms). 저자들은 퓨전된 MoE 커널이 이론적으로 이 격차를 해소할 수 있다고 언급하지만, 현재 구현은 생산 비용의 상한선으로 작용합니다.

4. 기여

1. 실증적 연산 레버리지 관찰: 비전에서 긍정적인 MoE 정확도 차이를 위해서는 상당한 라우팅 비율 $\rho$가 필요함을 확립합니다. 깊이별 백본은 $\rho$를 동등한 수준으로 높여줌으로써 이를 용이하게 합니다.
2. 변수 분리: 두 가지 통제된 실험은 다음을 보여줍니다:
   • 백본 계열 (표준 대 깊이별) 은 활성 변수가 아니며, 오히려 라우팅 비율 $\rho$가 결과를 주도합니다.
   • ImageNet 규모에서 다중 전문가 라우팅 ($k \ge 2$) 은 높은 $\rho$에 추가하여 필수 조건입니다.
3. 다중 데이터셋 검증: CIFAR-10/100, Tiny-ImageNet, ImageNet-1K 전반에 걸쳐 통계적으로 유의미한 결과를 제공하여, 깊이별 공동 최적화가 더 작은 규모에서 그리고 MoE 가 백본에 통합될 때 긍정적인 차이를 생성함을 보여줍니다.
4. 라우팅 안정성 분석: 붕괴를 방지하기 위한 온도 스케줄링의 중요성을 식별하고, 샘플 단위 CNN 에서 Soft MoE 의 지배적인 실패 모드로서 배치 축 배분을 진단합니다.

5. 중요성 및 주장

본 논문은 MoE 가 보편적인 플러그인 개선이 아니라는 관념에 반박하여, 비전에서 희소 MoE 가 언제 도움이 되는지에 대한 엄격하고 통제된 설명을 제공한다고 주장합니다. 저자들은 많은 비전 배포에서 MoE 의 실패는 라우팅 메커니즘 자체 때문이 아니라, 라우팅된 구성 요소가 라우팅 오버헤드나 복잡성을 정당화하기에 너무 작은 구조적 연산 불균형 때문이라고 주장합니다.

본 연구는 비전 분류에 대해 다음과 같이 결론 내립니다:

• 혜택을 보려면 총 FLOPs 의 상당한 비율이 라우팅되어야 합니다.
• 대규모 (ImageNet) 에서 큰 비율을 라우팅하는 것은 필요하지만, 충분하기 위해서는 다중 전문가 선택 ($k \ge 2$) 이 필요합니다.
• 깊이별 분리 합성곱은 FLOP 분포를 이동시킴으로써 이 조건을 위한 실용적인 가능자입니다.

저자들은 절대적 성능에 대해서는 겸손하게 접근하며, 모델이 최첨단 베이스라인과 비교하여 기본 훈련 레시피 (표준 증강, AdamW/CutMix 부재) 를 사용한다고 지적합니다. 그들의 주요 기여는 훈련 레시피 혼란 요인으로부터 라우팅 효과를 분리하여 통제된 조건 하에서 밀집 베이스라인과 일치하는 MoE 의 상대적 개선에 있습니다. 그들은 현재 지연 시간 페널티가 구현에 의존적이며, 실용적인 배포 효율성을 위해서는 퓨전된 커널이 필요함을 인정합니다.