논문 번역

• 이 글은 제가 읽는 논문을 번역해서 올려놓는 공간입니다.
• 번역은 LLM(현재는 GPT-4)를 통해 번역하고 있습니다.
• 대략의 pipeline은 mathpix를 통해 ocr후 gpt-4로 대략 1000토큰 단위로 나누어 번역하고 있습니다. 논문->ocr을 통한 markdown으로 변경 -> markdown로 작성된 논문을 gpt-4 api로 번역 및 markdown으로 작성
• 현재까지 완성된 파이프라인은 https://github.com/aeolian83/paper_translator 입니다
• 원문 링크는 다음과 같습니다.
  • https://paperswithcode.com/paper/switch-transformers-scaling-to-trillion

스위치 트랜스포머(Switch Transformers): 단순하고 효율적인 희소성(Sparsity)으로 조 단위 매개변수 모델로 확장하기

William Fedus*
LIAMFEDUS@GOOGLE.COM

Barret Zoph*

BARRETZOPH@GOOGLE.COM

Noam Shazeer
NOAM@GOOGLE.COM
Google, Mountain View, CA 94043, USA

편집자: Alexander Clark

초록(Abstract)

딥 러닝(deep learning)에서 모델들은 일반적으로 모든 입력에 대해 동일한 매개변수를 재사용합니다. 전문가 혼합(Mixture of Experts, MoE) 모델들은 이를 거부하고 대신 각 들어오는 예제마다 다른 매개변수를 선택합니다. 그 결과는 상수 계산 비용을 가진 - 엄청난 수의 매개변수를 가진 - 희소 활성화 모델입니다. 그러나 MoE의 여러 주목할만한 성공에도 불구하고, 복잡성, 통신 비용, 그리고 학습 불안정성으로 인해 널리 채택되지 못했습니다. 우리는 이러한 문제들을 스위치 트랜스포머(Switch Transformer)의 도입으로 해결합니다. 우리는 MoE 라우팅 알고리즘을 단순화하고, 통신 및 계산 비용이 줄어든 직관적으로 개선된 모델을 설계합니다. 우리가 제안한 학습 기법은 불안정성을 완화시키고, 우리는 처음으로 낮은 정밀도(bfloat16) 형식으로 큰 희소 모델을 학습할 수 있음을 보여줍니다. 우리는 T5-Base와 T5-Large(Raffel et al., 2019)를 기반으로 한 모델을 설계하여 동일한 계산 자원으로 사전 학습 속도를 최대 $7 \mathrm{x}$까지 향상시킵니다. 이러한 개선은 모든 101개 언어에서 mT5-Base 버전을 능가하는 성과를 측정하는 다국어 설정으로 확장됩니다. 마지막으로, 우리는 "Colossal Clean Crawled Corpus"에서 조 단위 매개변수 모델을 사전 학습하여 T5-XXL 모델보다 4배 빠른 속도를 달성함으로써 현재 언어 모델의 규모를 확장합니다. ${ }^{12}$

키워드(Keywords): 전문가 혼합(mixture-of-experts), 자연어 처리(natural language processing), 희소성(sparsity), 대규모 머신 러닝(large-scale machine learning), 분산 컴퓨팅(distributed computing)

*. 동등한 기여(Equal contribution).

1. 스위치 트랜스포머(Switch Transformer)와 모든 모델 체크포인트는 https://github.com/google-research/t5x 에서 사용할 수 있습니다.
2. 스위치 트랜스포머(Switch Transformer)에 대한 텐서플로우(Tensorflow) 코드는 https://github.com/tensorflow/mesh/blob/master/mesh_tensorflow/transformer/moe.py 에서 사용할 수 있습니다.

목차(Contents)

1. 서론

대규모 학습은 유연하고 강력한 신경 언어 모델(neural language models)을 향한 효과적인 길이었습니다(Radford et al., 2018; Kaplan et al., 2020; Brown et al., 2020). 단순한 구조(simple architectures)는 넉넉한 계산 예산(computational budget), 데이터 세트 크기(data set size), 그리고 매개변수 수(parameter count)를 바탕으로 더 복잡한 알고리즘들을 능가합니다(Sutton, 2019). Radford et al. (2018); Raffel et al. (2019); Brown et al. (2020)에서 따른 접근 방식은 밀집 활성화(densely-activated) 트랜스포머(Transformer)(Vaswani et al., 2017)의 모델 크기를 확장합니다. 이 방법은 효과적이지만, 동시에 매우 계산 집약적입니다(Strubell et al., 2019). 모델 규모의 성공에 영감을 받았지만, 더 큰 계산 효율성을 추구하며, 우리는 대신 희소 활성화(sparse activation) 전문가 모델(expert model)을 제안합니다: 스위치 트랜스포머(Switch Transformer). 우리의 경우 희소성(sparsity)은 각 들어오는 예제(example)에 대해 신경망 가중치(neural network weights)의 부분 집합만을 활성화함으로써 발생합니다.

그림 1: 스위치 트랜스포머(Switch Transformers)의 스케일링과 샘플 효율성. 왼쪽 그래프: 점점 더 희소해지는(더 많은 전문가들을 가진) 스위치 트랜스포머의 스케일링 특성. 오른쪽 그래프: 동일한 계산 예산을 사용하는 T5(Raffel et al., 2019) 모델과 비교한 스위치 트랜스포머의 음의 로그 혼란도(negative log perplexity).

희소 훈련(sparse training)은 활발한 연구 및 공학 분야입니다(Gray et al., 2017; Gale et al., 2020), 하지만 오늘날까지도 머신 러닝(machine learning) 라이브러리와 하드웨어 가속기(hardware accelerators)는 여전히 밀집된 행렬 곱셈(dense matrix multiplications)에 초점을 맞추고 있습니다. 효율적인 희소 알고리즘을 갖기 위해, 우리는 전문가의 혼합(Mixture-of-Expert, MoE) 패러다임(Jacobs et al., 1991; Jordan and Jacobs, 1994; Shazeer et al., 2017)에서 시작하여, 훈련 안정성과 계산적 이점을 얻기 위해 이를 단순화했습니다. MoE 모델은 기계 번역(machine translation)에서 주목할 만한 성공을 거두었습니다(Shazeer et al., 2017, 2018; Lepikhin et al., 2020), 그러나 복잡성, 통신 비용, 그리고 훈련 불안정성으로 인해 널리 채택되지는 못하고 있습니다.

우리는 이러한 문제들을 해결하고, 번역을 넘어서, 이러한 알고리즘 클래스가 자연어(natural language)에서 널리 가치가 있다는 것을 발견합니다. 우리는 자연어 처리(NLP)의 세 가지 체제(pre-training, finetuning, multi-task training)에서 다양한 자연어 작업에 대한 우수한 스케일링을 측정합니다. 이 작업은 스케일에 중점을 두고 있지만, 스위치 트랜스포머 아키텍처가 슈퍼컴퓨터(supercomputers)의 영역에서만 뛰어난 것이 아니라, 몇 개의 계산 코어만 있어도 유익하다는 것을 보여줍니다. 더 나아가, 우리의 큰 희소 모델은 작은 밀집 버전으로 압축(distilled)될 수 있으며, 희소 모델 품질 향상의 $30 \%$를 보존합니다. 우리의 기여는 다음과 같습니다:

• Mixture of Experts를 단순화하고 개선한 Switch Transformer 아키텍처.
• 강력하게 튜닝된 T5 모델 (Raffel et al., 2019)과의 벤치마크를 통한 확장성 특성 및 사전 훈련 속도를 7배 이상 향상시킴. 또한, 최소한 두 개의 전문가만 사용하여 개선 사항이 유지됨을 보여줌.
• 희소 사전 훈련 및 특화된 파인튜닝 모델을 작은 밀집 모델로 압축하는데 성공. 대규모 희소 모델 품질 향상의 $30 \%$를 유지하면서 모델 크기를 최대 $99 \%$까지 축소.
• 사전 훈련 및 파인튜닝 기술의 개선: (1) 낮은 bfloat16 정밀도로 훈련이 가능한 선택적 정밀도 훈련 (2) 더 많은 전문가 수로 확장이 가능한 초기화 체계 (3) 희소 모델 파인튜닝 및 멀티 태스크 훈련을 개선하는 전문가 정규화 증가.
• 101개 언어에 걸쳐 다국어 데이터의 사전 훈련 이점을 측정하여 mT5 기준선 (Xue et al., 2020) 대비 91%의 언어가 mT5 기준선 대비 4배 이상의 속도 향상을 보이는 보편적인 개선을 발견.
• 데이터, 모델 및 전문가 병렬성을 효율적으로 결합하여 1조 개의 매개변수를 갖는 모델을 만들어내는 신경 언어 모델의 규모 확장. 이러한 모델은 강력하게 튜닝된 T5-XXL 기준선의 사전 훈련 속도를 4배 향상시킵니다.

2. 스위치 트랜스포머(Switch Transformer)

스위치 트랜스포머의 설계 원칙은 간단하고 계산적으로 효율적인 방식으로 트랜스포머 모델(Vaswani et al., 2017)의 파라미터 수를 극대화하는 것입니다. 규모의 이점은 Kaplan et al. (2020)에서 철저히 연구되었으며, 모델 크기, 데이터 세트 크기 및 계산 예산과의 멱법칙(powerlaw scaling)을 발견했습니다. 중요한 것은, 이 연구는 상대적으로 작은 양의 데이터에서 큰 모델을 훈련시키는 것이 계산적으로 최적의 접근법이라고 주장합니다.

이러한 결과를 고려하여, 우리는 네 번째 축을 조사합니다: 예제당 부동 소수점 연산(FLOPs)을 일정하게 유지하면서 파라미터 수를 증가시킵니다. 우리의 가설은 총 계산량과는 독립적으로, 파라미터 수는 스케일링해야 할 별도의 중요한 축이라는 것입니다. 우리는 GPU와 TPU와 같이 밀집된 행렬 곱셈에 최적화된 하드웨어를 효율적으로 사용하는 희소 활성화 모델을 설계함으로써 이를 달성합니다. 여기서의 연구는 TPU 아키텍처에 중점을 두고 있지만, 이러한 종류의 모델은 GPU 클러스터에서도 유사하게 훈련될 수 있습니다. 우리의 분산 훈련 설정에서, 희소 활성화 레이어는 다른 장치에서 고유한 가중치를 분할합니다. 따라서 모델의 가중치는 장치의 수와 함께 증가하지만, 각 장치에서 관리 가능한 메모리 및 계산 풋프린트를 유지합니다.

그림 2: 스위치 트랜스포머(Switch Transformer) 인코더 블록의 일러스트레이션입니다. 우리는 트랜스포머(Transformer)에 있는 밀집된 전방 피드 네트워크(dense feed forward network, FFN) 층을 희소한 스위치 FFN 층(light blue)으로 대체했습니다. 이 층은 시퀀스 내의 토큰들에 독립적으로 작동합니다. 우리는 두 개의 토큰($x_{1}=$ "More" 및 $x_{2}=$ "Parameters" 아래)이 네 개의 FFN 전문가들(experts)에 걸쳐 라우팅되는 것을 다이어그램으로 나타냈습니다. 여기서 라우터(router)는 각 토큰을 독립적으로 라우팅합니다. 스위치 FFN 층은 선택된 FFN의 출력을 라우터 게이트 값(router gate value)에 곱한 것을 반환합니다(점선).

2.1 희소 라우팅 단순화하기

전문가 라우팅의 혼합(Mixture of Expert Routing). Shazeer et al. (2017)은 자연어 처리를 위한 전문가들의 혼합(Mixtureof-Experts, MoE) 층을 제안했는데, 이는 토큰 표현 $x$를 입력으로 받아 최상의 상위 $k$ 전문가들에게 이를 라우팅합니다. 이 전문가들은 $N$개의 전문가 집합 $\left\{E_{i}(x)\right\}_{i=1}^{N}$에서 선택됩니다. 라우터 변수 $W_{r}$는 로짓(logits) $h(x)=W_{r} \cdot x$를 생성하며, 이는 그 층에서 사용 가능한 $N$개의 전문가들에 대한 소프트맥스(softmax) 분포를 통해 정규화됩니다. 전문가 $i$에 대한 게이트 값(gate-value)은 다음과 같이 주어집니다.

상위 $k$개의 게이트 값들이 토큰 $x$의 라우팅을 위해 선택됩니다. 만약 $\mathcal{T}$가 선택된 상위 $k$개의 인덱스 집합이라면, 층의 출력 계산은 각 전문가의 토큰에 대한 계산을 게이트 값에 의해 선형적으로 가중한 조합입니다.

스위치 라우팅: 전문가 혼합(Mixture-of-Experts)에 대한 재고찰. Shazeer 외(2017)는 라우팅 함수에 비자명한(non-trivial) 그래디언트를 가지기 위해서는 $k>1$의 전문가에게 라우팅하는 것이 필요하다고 추측했습니다. 저자들은 적어도 두 명의 전문가를 비교할 수 있는 능력 없이는 라우팅 학습이 작동하지 않을 것이라고 직관했습니다. Ramachandran과 Le(2018)는 top-$k$ 결정에 대해 더 깊이 연구하고 모델의 하위 레이어에서 더 높은 $k$ 값이 많은 라우팅 레이어를 가진 모델에 중요하다는 것을 발견했습니다. 이러한 아이디어와는 반대로, 우리는 단 하나의 전문가에게만 라우팅하는 단순화된 전략을 사용합니다. 우리는 이 단순화가 모델 품질을 유지하고, 라우팅 계산을 줄이며, 더 나은 성능을 보인다는 것을 보여줍니다. 이 $k=1$ 라우팅 전략은 나중에 스위치(Switch) 레이어로 언급됩니다. $\mathrm{MoE}$와 스위치 라우팅 모두에서, 방정식 2의 게이트 값 $p_i(x)$는 라우터의 차별화(differentiability)를 허용합니다.

스위치 레이어의 이점은 세 가지입니다: (1) 우리는 토큰을 단 하나의 전문가에게만 라우팅하기 때문에 라우터 계산이 줄어듭니다. (2) 각 토큰이 단 하나의 전문가에게만 라우팅되기 때문에 각 전문가의 배치 크기(전문가 용량)는 적어도 절반으로 줄일 수 있습니다. ${ }^{3}$ (3) 라우팅 구현이 단순화되고 통신 비용이 줄어듭니다. 그림 3은 다른 전문가 용량 인자를 가진 라우팅의 예를 보여줍니다.

그림 3: 토큰 라우팅 동역학의 일러스트레이션. 각 전문가는 용량 인자에 의해 조절되는 고정된 배치 크기의 토큰을 처리합니다. 각 토큰은 가장 높은 라우터 확률을 가진 전문가에게 라우팅되지만, 각 전문가는 (총 토큰 수 / 전문가 수) $\times$ 용량 인자의 고정된 배치 크기를 가집니다. 토큰이 불균등하게 배분되면 특정 전문가들이 오버플로우(점선 빨간 선으로 표시됨)될 수 있으며, 이는 해당 토큰들이 이 레이어에서 처리되지 않는 결과를 초래합니다. 더 큰 용량 인자는 이 오버플로우 문제를 완화시키지만, 계산 및 통신 비용도 증가시킵니다(패딩된 흰색/빈 슬롯으로 표현됨).

2.2 효율적인 희소 라우팅

우리는 Mesh-Tensorflow(MTF)(Shazeer et al., 2018)를 사용하는데, 이는 Tensorflow(Abadi et al., 2016)와 유사한 의미론(semantics)과 API를 가진 라이브러리로, 효율적인 분산 데이터 및 모델 병렬 아키텍처를 용이하게 합니다. 이는 물리적인 코어 집합을 논리적인 프로세서 메쉬(mesh)로 추상화함으로써 이루어집니다. 그 후 텐서(tensor)와 계산은 명명된 차원에 따라 샤딩(sharded)되어 모델을 차원에 걸쳐 쉽게 분할할 수 있습니다. 우리는 정적으로 선언된 크기를 요구하는 TPU를 염두에 두고 모델을 설계했습니다. 아래에서는 우리의 분산 스위치 트랜스포머(Switch Transformer) 구현에 대해 설명합니다.[^0]

분산 스위치 구현. 우리의 모든 텐서 형태는 컴파일 시간에 정적으로 결정되지만, 훈련 및 추론 시 라우팅 결정으로 인해 계산은 동적입니다. 이 때문에 중요한 기술적 고려사항 중 하나는 전문가(expert) 용량을 어떻게 설정할 것인가입니다. 전문가 용량 - 각 전문가가 계산하는 토큰의 수 - 은 배치 내 토큰 수를 전문가 수로 균등하게 나눈 다음 용량 계수(capacity factor)로 추가 확장하여 설정됩니다.

1.0보다 큰 용량 계수는 토큰이 전문가들 사이에 완벽하게 균형을 이루지 않을 때를 대비해 추가 버퍼를 생성합니다. 너무 많은 토큰이 한 전문가에게 라우팅되는 경우(나중에 버려진 토큰(dropped tokens)이라고 함), 계산은 건너뛰어지고 토큰 표현은 잔여 연결(residual connection)을 통해 다음 계층으로 직접 전달됩니다. 그러나 전문가 용량을 늘리는 것은 높은 값으로 인해 계산과 메모리가 낭비될 수 있으므로 단점이 없는 것은 아닙니다. 이러한 트레이드오프(trade-off)는 그림 3에서 설명됩니다. 경험적으로 우리는 희소 전문가 모델(sparse expert-models)의 확장에 있어 버려진 토큰의 낮은 비율을 보장하는 것이 중요하다는 것을 발견했습니다. 우리의 실험에서는 전문가 수에 대한 버려진 토큰 수의 의존성을 발견하지 못했습니다(일반적으로 $<1 \%$). 충분히 높은 계수로 보조 부하 균형 손실(auxiliary load balancing loss)(다음 섹션)을 사용하면 좋은 부하 균형을 보장할 수 있었습니다. 우리는 이러한 설계 결정이 모델 품질과 속도에 미치는 영향을 표 1에서 연구합니다.

차별화 가능한 부하 분산 손실. 전문가들 간에 균형 잡힌 부하를 장려하기 위해 우리는 보조 손실을 추가합니다(Shazeer et al., 2017, 2018; Lepikhin et al., 2020). Shazeer et al. (2018); Lepikhin et al. (2020)에서와 같이, 스위치 트랜스포머(Switch Transformers)는 Shazeer et al. (2017)의 원래 설계를 단순화하여 별도의 부하 균형 및 중요도 가중치 손실을 합쳤습니다. 각 스위치 계층에 대해, 이 보조 손실은 훈련 중 전체 모델 손실에 추가됩니다. $N$개의 전문가가 $i=1$부터 $N$까지 색인되어 있고, $T$개의 토큰을 가진 배치 $\mathcal{B}$가 주어졌을 때, 보조 손실은 벡터 $f$와 $P$ 사이의 스케일이 조정된 내적으로 계산됩니다,

여기서 $f_{i}$는 전문가 $i$로 보낸 토큰의 비율입니다,

그리고 $P_{i}$는 전문가 $i$에 할당된 라우터 확률의 비율입니다,

우리는 $N$개의 전문가에 걸쳐 토큰 배치의 균일한 라우팅을 추구하기 때문에, 두 벡터 모두 $1 / N$의 값을 가지기를 원합니다. 방정식 4의 보조 손실은 균일 분포 하에서 최소화되므로 균일 라우팅을 장려합니다. 목표는 또한 미분 가능한데[^1] $P$-벡터는 미분 가능하지만, $f$-벡터는 미분 가능하지 않습니다. 최종 손실은 전문가 수가 변함에 따라 손실이 일정하게 유지되도록 전문가 수 $N$으로 곱해집니다. 왜냐하면 균일 라우팅 하에서 $\sum_{i=1}^{N}\left(f_{i} \cdot P_{i}\right)=\sum_{i=1}^{N}\left(\frac{1}{N} \cdot \frac{1}{N}\right)=\frac{1}{N}$이기 때문입니다. 마지막으로, 하이퍼파라미터 $\alpha$는 이러한 보조 손실에 대한 곱셈 계수입니다; 이 연구에서는 부하 균형을 보장하는 데 충분히 크면서도 주요 교차 엔트로피 목표를 압도하지 않을 정도로 충분히 작은 $\alpha=10^{-2}$를 사용했습니다. 우리는 $\alpha$의 하이퍼파라미터 범위를 $10^{-1}$부터 $10^{-5}$까지 10의 거듭제곱으로 조사했고, $10^{-2}$가 훈련 손실에 방해를 주지 않으면서 빠르게 부하를 균형있게 만드는 것을 발견했습니다.

2.3 모든 것을 결합하기: 스위치 트랜스포머

우리가 처음으로 스위치 트랜스포머(Switch Transformer)를 테스트한 것은 "Colossal Clean Crawled Corpus" (C4)에서 사전 학습(pre-training)을 시작하면서였습니다. 이는 (Raffel et al., 2019)에서 소개되었습니다. 사전 학습 목표로, 우리는 마스크된 언어 모델링 작업(masked language modeling task) (Taylor, 1953; Fedus et al., 2018; Devlin et al., 2018)을 사용했는데, 여기서 모델은 누락된 토큰을 예측하도록 학습됩니다. 우리의 사전 학습 설정에서는 Raffel et al. (2019)에서 최적으로 결정된 대로, $15\%$의 토큰을 드롭아웃(drop out)하고 마스크된 시퀀스를 단일 센티넬 토큰(sentinel token)으로 대체합니다. 모델들을 비교하기 위해, 우리는 음의 로그 혼란도(negative log perplexity)를 기록합니다. ${ }^{4}$ 논문의 모든 표에서, $\uparrow$는 해당 메트릭(metric)에 대해 높은 값이 더 좋다는 것을 나타내고, 반대의 경우는 $\downarrow$로 표시됩니다. 이 작업에서 연구된 모든 모델의 비교는 표 9에 있습니다.

스위치 트랜스포머와 MoE 트랜스포머(MoE Transformer)의 직접 비교는 표 1에 제시되어 있습니다. 우리의 스위치 트랜스포머 모델은 'T5-Base' (Raffel et al., 2019)와 FLOP-matched(동일한 계산량이 토큰당 적용됨)입니다. 상위-2 라우팅(top-2 routing)을 사용하는 MoE 트랜스포머는 각각의 토큰에 별도의 FFN을 적용하는 두 개의 전문가(experts)를 가지고 있으므로 그 FLOPS는 더 큽니다. 모든 모델은 동일한 하드웨어에서 동일한 스텝 수로 학습되었습니다. 용량 계수(capacity factor)가 2.0에서 1.25로 감소하는 MoE 모델은 위의 실험 설정에서 실제로 느려지는 것(840에서 790)을 볼 수 있는데, 이는 예상치 못한 결과입니다. ${ }^{5}$

표 1에서 세 가지 주요 발견을 강조합니다: (1) 스위치 트랜스포머는 신중하게 조정된 밀집 모델(dense models)과 MoE 트랜스포머보다 속도-품질 기준에서 더 우수한 성능을 보입니다. 고정된 계산량과 실제 시간(wall-clock time)에 대해, 스위치 트랜스포머는 최상의 결과를 달성합니다. (2) 스위치 트랜스포머는 MoE 대응 모델보다 작은 계산적 발자국(computational footprint)을 가집니다. 우리가 그 크기를 MoE 트랜스포머의 학습 속도와 일치시키도록 늘리면, 이는 스텝당 기준으로 모든 MoE 및 밀집 모델보다 더 우수한 성능을 보입니다. (3) 스위치 트랜스포머는 낮은 용량 계수(1.0, 1.25)에서 더 나은 성능을 보입니다. 작은 전문가 용량(expert capacities)은 모델 메모리가 매우 부족하고 용량 계수를 가능한 한 작게 만들고 싶은 대규모 모델 체제(large model regime)에서의 시나리오를 나타냅니다.

2.4 개선된 학습 및 미세 조정 기법(Improved Training and Fine-Tuning Techniques)

희소 전문가 모델(sparse expert models)은 바닐라 트랜스포머(vanilla Transformer)보다 학습에 어려움을 초래할 수 있습니다. 이러한 각 레이어에서의 강력한 스위칭(하드 스위칭, routing) 결정으로 인해 불안정성이 발생할 수 있습니다. 또한, bfloat16과 같은 저정밀도 형식(low precision formats)은 이러한 문제를 악화시킬 수 있습니다 (Wang and Kanwar, 2019).

4. 이 지표에는 로그 밑-e를 사용하여 단위가 nats입니다.
5. 속도 측정은 알고리즘과 구현 세부 사항의 함수입니다. 스위치 트랜스포머(Switch Transformer)는 MoE(Mixture of Experts)에 비해 필요한 계산을 줄이지만(알고리즘), 최종 속도 차이는 저수준 최적화(구현)에 의해 영향을 받습니다.

표 1: 스위치 대 MoE 벤치마킹. 스위치 트랜스포머가 MoE 트랜스포머 및 T5 밀집 기준선에 비해 단계별 및 시간별 이점을 측정하는 일대일 비교. 우리는 Neg. Log Perp.=-1.50의 임의로 선택된 품질 임계값에 도달하는 시간과 음의 로그 혼란도로 품질을 측정합니다. 모든 MoE 및 스위치 트랜스포머 모델은 128개의 전문가를 사용하며, 전문가는 매 다른 피드-포워드 층마다 배치됩니다. Switch-Base+의 경우, 모델의 은닉 크기를 768에서 896으로, 헤드 수를 14에서 16으로 늘려 MoE 모델의 속도에 맞춥니다. 모든 모델은 동일한 계산량(32 코어)과 동일한 하드웨어(TPUv3)에서 훈련됩니다. 또한 모든 모델은 -1.50의 우리의 품질 임계값을 달성하기 위해 100k 스텝을 넘는 사전 훈련이 필요했습니다. † T5-Base는 모델이 훈련된 $100 \mathrm{k}$ 스텝에서 이 음의 로그 혼란도를 달성하지 못했습니다.

우리 라우터의 소프트맥스(softmax) 계산에서. 여기서 우리는 훈련의 어려움과 안정적이고 확장 가능한 훈련을 달성하기 위해 사용하는 방법들을 설명합니다.

큰 희소 모델에서의 선택적 정밀도. 모델 불안정성은 효율적인 bfloat16 정밀도를 사용하여 훈련하는 능력을 방해하며, 결과적으로 Lepikhin et al. (2020)은 그들의 MoE Transformer를 float32 정밀도로 훈련합니다. 그러나 우리는 대신 모델의 특정 부분에서만 float32 정밀도로 캐스팅(casting)함으로써 안정성을 달성할 수 있으며, float32 텐서(tensors)의 비싼 통신 비용을 발생시키지 않는다는 것을 보여줍니다. 이 기술은 모델의 특정 부분과 그래디언트 업데이트가 더 높은 정밀도로 이루어지는 현대의 혼합 정밀도 훈련 전략과 일치합니다 Micikevicius et al. (2017). 표 2는 우리의 접근 방식이 bfloat16 훈련과 거의 동일한 속도를 허용하면서 float32의 훈련 안정성을 제공한다는 것을 보여줍니다.

이를 달성하기 위해, 우리는 라우터 입력을 float32 정밀도로 캐스팅합니다. 라우터 함수는 토큰을 입력으로 받아 전문가 계산의 선택과 재결합에 사용되는 디스패치(dispatch)와 결합(combine) 텐서를 생성합니다(자세한 내용은 부록의 코드 블록 15를 참조하세요). 중요한 것은 float32 정밀도가 라우터 함수의 본체 내에서만 사용된다는 것입니다 - 해당 장치에 국한된 계산에서만. 결과적으로 생성된 디스패치와 결합 텐서는 함수의 끝에서 bfloat16 정밀도로 다시 캐스팅되기 때문에, 비싼 float32 텐서

표 2: 선택적 정밀도. 우리는 모델을 안정화시키면서 (불안정한) bfloat16 정밀도 훈련과 거의 동일한 속도를 달성하기 위해 로컬 라우팅 작업을 float32로 캐스팅하면서 다른 곳에서는 bfloat16 정밀도를 유지합니다. 우리는 훈련 초기에 고정된 스텝 수 후 32 전문가 모델의 품질과 속도 성능을 측정합니다. float32로 된 Switch-Base와 선택적 정밀도를 사용한 경우 비슷한 학습 동태를 관찰합니다.

모든 통신 작업을 통해 전체로 전송되지는 않지만, 우리는 여전히 float32의 증가된 안정성의 이점을 누립니다.

안정성을 위한 더 작은 파라미터 초기화. 적절한 초기화는 딥 러닝(deep learning)에서 성공적인 훈련에 매우 중요하며, 특히 스위치 트랜스포머(Switch Transformer)에서 이를 더욱 명확하게 관찰할 수 있습니다. 우리는 가중치 행렬을 초기화할 때, 평균 $\mu=0$과 표준편차 $\sigma=\sqrt{s / n}$을 가지는 절단 정규 분포(truncated normal distribution)에서 요소를 추출하여 초기화합니다. 여기서 $s$는 스케일 하이퍼파라미터(scale hyper-parameter)이고 $n$은 가중치 텐서(weight tensor)의 입력 유닛 수(예: fan-in)입니다. ${ }^{6}$

불안정성을 추가로 해결하기 위해, 기본 트랜스포머(Transformer) 초기화 스케일 $s=1.0$을 10분의 1로 줄이는 것이 좋습니다. 이는 품질을 향상시키고 실험에서 훈련의 불안정성을 줄이는 데 도움이 됩니다. 표 3은 모델 품질의 향상과 훈련 초기 단계에서의 분산 감소를 측정합니다. 우리는 다음과 같이 발견했습니다.

표 3: 초기화 스케일을 줄이면 안정성이 향상됩니다. 초기화 스케일을 줄이면 스위치 트랜스포머(Switch Transformer)의 모델 품질이 더 좋아지고 훈련이 더 안정적입니다. 여기서 우리는 32 전문가 모델의 평균 및 표준편차 모델 품질을 기록했으며, 이는 $3.5 \mathrm{k}$ 단계 후의 음의 로그 혼란도(negative log perplexity)로 측정됩니다(각각 3개의 랜덤 시드).

평균 모델 품질은 Neg. Log Perp.로 측정되며, 실행 간의 분산이 크게 줄어들고 획기적으로 향상되었습니다. 또한, 이와 같은 초기화 방식은 몇 가지 크기의 모델에 대해 광범위하게 효과적입니다. 우리는 $223 \mathrm{M}$ 파라미터 기준선부터 1조 이상의 파라미터를 가진 거대한 모델에 이르기까지 안정적으로 훈련하기 위해 동일한 접근 방식을 사용합니다.

6. 평균에서 두 표준편차 이상 떨어진 값은 재표본됩니다.

큰 희소 모델 정규화하기. 우리 논문은 대규모 코퍼스(corpus)에서 사전 훈련(pre-training)을 한 후 요약이나 질문 응답과 같은 더 작은 하류 작업(downstream tasks)에 미세 조정(fine-tuning)하는 일반적인 자연어 처리(NLP) 접근 방식을 고려합니다. 자연스럽게 발생하는 한 가지 문제는 미세 조정 작업에 예제가 매우 적기 때문에 과적합(overfitting)입니다. 표준 트랜스포머(Transformer)의 미세 조정 중에, Raffel et al. (2019)은 각 계층에서 드롭아웃(dropout) (Srivastava et al., 2014)을 사용하여 과적합을 방지합니다. 우리의 스위치 트랜스포머(Switch Transformers)는 FLOP에 맞춘 밀집 기준선(dense baseline)보다 훨씬 많은 파라미터를 가지고 있어, 이러한 더 작은 하류 작업에서 더 심각한 과적합을 초래할 수 있습니다.

표 4: 미세조정(fine-tuning) 정규화 결과. C4 데이터 세트의 34B 토큰으로 사전 훈련된 Switch Transformer 모델을 미세조정하는 동안 드롭아웃 비율을 조정한 결과입니다(높은 숫자가 더 좋음). 우리는 모든 비전문가(non-expert) 계층에서 표준 드롭아웃 비율을 낮게 유지하고, 전문가(expert) 피드포워드 계층에서 훨씬 높은 드롭아웃 비율을 사용하는 것이 가장 좋은 성능을 낸다는 것을 관찰했습니다.

따라서 우리는 미세조정 중 이 문제를 완화하는 간단한 방법을 제안합니다: 전문가 내부의 드롭아웃을 증가시키는 것으로, 우리는 이를 전문가 드롭아웃(expert dropout)이라고 명명합니다. 미세조정 중에는 각 전문가 계층에서 중간 피드포워드 계산 시에만 드롭아웃 비율을 상당히 증가시킵니다. 표 4는 우리의 전문가 드롭아웃 프로토콜 결과를 보여줍니다. 우리는 모든 계층에서 드롭아웃을 단순히 증가시키는 것이 성능을 떨어뜨린다는 것을 관찰했습니다. 그러나 비전문가 계층에서는 더 작은 드롭아웃 비율(0.1)을 설정하고, 전문가 계층에서는 훨씬 더 큰 드롭아웃 비율(0.4)을 설정하는 것이 네 개의 작은 하류(downstream) 작업에서 성능 향상으로 이어졌습니다.

3. 스케일링 특성

우리는 사전 훈련 중 Switch Transformer 아키텍처의 스케일링 특성에 대한 연구를 제시합니다. Kaplan et al. (2020)에 따르면, 모델이 계산 예산이나 데이터 양에 의해 병목 현상을 겪지 않는 상황을 고려합니다. 데이터 병목을 피하기 위해, 우리는 180B 이상의 대상 토큰을 가진 대규모 C4 코퍼스를 사용하고, 수익이 감소할 때까지 훈련을 진행합니다(Raffel et al., 2019).

전문가의 수는 모델을 확장하는 데 가장 효율적인 차원입니다. 전문가 수를 늘리면 모델이 토큰 당 하나의 전문가만 선택하기 때문에 계산 비용이 대략 고정되어 있습니다. 선택할 전문가의 수에 관계없이 말이죠. 그러나 라우터는 더 많은 전문가에 대한 확률 분포를 계산해야 하지만, 이는 $O\left(d_{\text{model}} \times\right.$ 전문가 수 $)$의 가벼운 계산 비용이 드는 작업입니다. 여기서 $d_{\text{model}}$은 레이어 간에 전달되는 토큰의 임베딩 차원입니다. 이 섹션에서는 고정된 계산 예산으로 단계별 및 시간별로 확장 속성을 고려합니다.

3.1 단계별로 본 확장 결과

그림 4는 고정된 단계 수로 모든 모델을 훈련할 때 전문가 수가 증가함에 따라 일관된 확장 이점을 보여줍니다. 우리는 분명한 추세를 관찰합니다: 토큰 당 FLOPS를 고정하면, 더 많은 매개변수(전문가)가 훈련 속도를 높입니다. 왼쪽 그림은 토큰 당 고정된 FLOPS로 희소 모델 매개변수와 테스트 손실 간의 일관된 확장 속성을 보여줍니다. 이는 희소 모델 매개변수의 추가 축을 따라 확장하는 이점을 드러냅니다. 우리의 오른쪽 그림은 밀집 모델 변형과 네 가지 FLOP-매칭된 희소 변형의 샘플 효율성을 측정합니다. 우리는 전문가 수를 늘리면 더 샘플 효율적인 모델이 된다는 것을 발견했습니다. 우리의 Switch-Base 64 전문가 모델은 T5-Base 모델과 동일한 성능을 단계 $60 \mathrm{k}$에서 단계 $450 \mathrm{k}$로 달성하는데, 이는 단계 시간 측면에서 $7.5 \mathrm{x}$의 가속을 의미합니다. 또한 Kaplan et al. (2020)의 연구 결과와 일치하게, 더 큰 모델들도 더 샘플 효율적이며-고정된 수의 관찰된 토큰에 대해 더 빠르게 학습합니다.

그림 4: 스위치 트랜스포머(Switch Transformer)의 확장 속성. 왼쪽 그림: 우리는 전문가 수를 확장하여 매개변수를 증가시킬 때, 혼란도(perplexity)로 측정한 품질 개선을 측정합니다. 왼쪽 상단 점은 $223 \mathrm{M}$ 매개변수를 가진 T5-Base 모델에 해당합니다. 왼쪽 상단에서 오른쪽 하단으로 이동하면서, 우리는 전문가 수를 2, 4, 8 등으로 두 배씩 늘리고, 오른쪽 하단 점은 14.7B 매개변수를 가진 256 전문가 모델에 이릅니다. 모든 모델이 동일한 계산 예산을 사용함에도 불구하고, 우리는 전문가 수를 확장함으로써 일관된 개선을 관찰합니다. 오른쪽 그림: 전문가 수를 변경하면서 단계별로 음의 로그 혼란도를 측정합니다. 밀집 기준선은 보라색 선으로 표시되며, 우리의 Switch-Base 모델들이 더 나은 샘플 효율성을 보임을 알 수 있습니다.

3.2 시간 기준으로 본 확장 결과

그림 4는 전문가(expert)의 수를 늘릴수록 단계별로 성능이 지속적으로 향상됨을 보여줍니다. 우리의 모델은 기준 모델과 대략적으로 같은 양의 FLOPS(Floating Point Operations Per Second)를 토큰(token) 당 가지고 있지만, 우리의 스위치 트랜스포머(Switch Transformers)는 장치 간 추가적인 통신 비용과 라우팅 메커니즘(routing mechanism)의 추가 계산을 필요로 합니다. 따라서 단계별로 관찰된 샘플 효율성(sample efficiency)의 증가가 반드시 실제 시간(wall-clock)으로 측정된 더 나은 모델 품질로 이어지는 것은 아닙니다. 이는 다음과 같은 질문을 제기합니다:

고정된 훈련 기간과 계산 예산으로, 밀집 모델(dense model)을 훈련시키는 것이 좋을까요, 아니면 희소 모델(sparse model)을 훈련시키는 것이 좋을까요?

그림 5: 스위치 트랜스포머의 속도 이점. 모든 모델은 예제 당 동일한 FLOPS를 가진 32개의 TPUv3 코어에서 훈련되었습니다. 고정된 계산량과 훈련 시간으로, 스위치 트랜스포머는 밀집 트랜스포머(dense Transformer) 기준 모델을 크게 앞집니다. 우리의 64 전문가 스위치-베이스(Switch-Base) 모델은 T5-베이스(T5-Base)의 일곱 분의 일 시간으로 같은 품질을 달성하고 계속해서 개선됩니다.

그림 5와 6은 이 질문에 대해 다룹니다. 그림 5는 시간에 따른 사전 훈련(pre-training) 모델 품질을 측정합니다. 고정된 훈련 기간과 계산 예산으로, 스위치 트랜스포머는 상당한 속도 향상을 제공합니다. 이 설정에서, 우리의 스위치-베이스 64 전문가 모델은 T5-베이스가 비슷한 혼란도(perplexity)를 얻는 데 걸리는 시간의 일곱 분의 일로 훈련됩니다.

3.3 더 큰 밀집 모델과의 스케일링 비교

위의 분석은 계산적으로 매칭된 밀집 모델이 스위치 대응 모델에 뒤처진다는 것을 보여줍니다. 그림 6은 다른 시나리오를 고려합니다: 우리가 대신 더 큰 밀집 모델에 자원을 할당했다면 어떨까요? 이제 우리는 스위치-베이스를 다음 강력한 기준인 T5-라지(T5-Large)와 비교하면서 측정합니다. 그러나 T5-라지가 토큰 당 3.5배 더 많은 FLOPS를 적용함에도 불구하고,

스위치-베이스는 여전히 더 샘플 효율적이며 $2.5 \mathrm{x}$의 속도 향상을 제공합니다. 더욱이, T5-라지에 FLOPS가 매칭된 새롭고 더 큰 희소 버전, 스위치-라지(Switch-Large)를 설계함으로써 더 많은 이득을 얻을 수 있습니다. 우리는 이를 수행하고 다음 섹션에서 우수한 스케일링과 미세 조정(fine-tuning)을 보여줍니다.

그림 6: 스위치 레이어(Switch layers)를 사용한 트랜스포머(Transformer) 모델의 확장 또는 표준 밀집 모델(dense model) 확장. 왼쪽 그래프: Switch-Base는 토큰 당 3.5배 더 많은 FLOPS를 적용하는 T5-Base와 T5-Large 변형보다 샘플 효율성이 더 높습니다. 오른쪽 그래프: 이전과 마찬가지로, 실제 시간(wall-clock) 기준으로 Switch-Base는 여전히 더 빠르며, T5-Large에 비해 $2.5 \mathrm{x}$ 속도 향상을 제공합니다.

4. 하류 결과(Downstream Results)

3장에서는 사전 학습(pre-training) 동안 우수한 확장성을 보여주었지만, 이제 이러한 이득이 하류 작업(downstream tasks)에서 개선된 언어 학습 능력으로 전환되는지를 검증합니다. 우리는 다양한 자연어 처리(NLP) 작업에 미세 조정(fine-tuning)을 시작합니다. 다음으로, 우리는 작고 쉽게 배포할 수 있는 밀집 기준선(dense baselines)으로 압축함으로써 희소 모델(sparse models)의 메모리 사용량을 90% 이상 줄이는 방법을 연구합니다. 마지막으로, 이 장을 다중 작업(multi-task), 다국어(multilingual) 환경에서의 개선 사항을 측정하며 마무리하는데, 여기서 우리는 스위치 트랜스포머(Switch Transformers)가 모든 101개 언어에 걸쳐 다국어 T5-base 모델보다 개선된 강력한 다중 작업 학습자임을 보여줍니다.

4.1 미세 조정(Fine-Tuning)

미세 조정에 사용된 기준선(Baseline) 및 스위치 모델(Switch models). 우리의 기준선은 매우 조정된 223M 매개변수 T5-Base 모델과 739M 매개변수 T5-Large 모델(Raffel et al., 2019)입니다. 두 버전 모두에 대해, 우리는 훨씬 더 많은 매개변수를 가진 FLOP-매칭된 스위치 트랜스포머(Switch Transformer)를 설계했으며, 이는 표 9에서 요약되어 있습니다7. 우리의 기준선은 Raffel et al. (2019)에서와 약간 다릅니다. 왜냐하면 우리는 예제 내 텍스트 중복을 제거하고 따라서 사전 학습 작업으로서의 효과를 증가시키는 개선된 C4 코퍼스(corpus)에서 사전 학습을 하기 때문입니다 Lee et al.

7. FLOPS는 Kaplan et al. (2020)에서 수행된 것처럼 순방향 패스(forward pass)에 대해 계산됩니다.

(2021). 우리의 프로토콜에서는 배치당 $2^{20}(1,048,576)$ 토큰으로 $550k$ 단계까지 사전 학습을 하여 총 $576B$ 토큰을 사용합니다. 그런 다음 우리는 다양한 작업 세트에서 미세 조정을 진행하는데, 스위치 레이어를 제외한 모든 레이어에 대해 0.1의 드롭아웃 비율(dropout rate)을 사용하고, 스위치 레이어는 0.4의 드롭아웃 비율을 사용합니다(표 4 참조). 우리는 배치 크기 $1M$으로 $16k$ 단계 동안 미세 조정을 하고, 각 작업에 대해 200단계마다 모델 품질을 평가하고 검증 세트(validation set)에서 계산된 최고 성능을 보고합니다.

미세 조정 작업 및 데이터 세트. 우리는 질문 응답(question answering), 요약(summarization), 세계에 대한 지식(knowledge about the world)을 포함한 언어 능력을 탐구하는 작업을 선택했습니다. 언어 벤치마크인 GLUE(Wang et al., 2018)와 SuperGLUE(Wang et al., 2019)는 모든 작업이 포함된 복합 혼합물로 처리되며, 각 작업에 존재하는 토큰의 양에 비례하여 혼합됩니다. 이러한 벤치마크는 감정 분석(sentiment analysis)(SST2), 단어 의미 구별(word sense disambiguation)(WIC), 문장 유사도(sentence similarty)(MRPC, STS-B, QQP), 자연어 추론(natural language inference)(MNLI, QNLI, RTE, CB), 질문 응답(question answering)(MultiRC, RECORD, BoolQ), 공동 참조 해결(coreference resolution)(WNLI, WSC) 및 문장 완성(sentence completion)(COPA)과 문장 수용성(sentence acceptability)(CoLA)을 요구하는 작업으로 구성됩니다. CNNDM(Hermann et al., 2015) 및 BBC XSum(Narayan et al., 2018) 데이터 세트는 기사를 요약하는 능력을 측정하는 데 사용됩니다. 질문 응답은 SQuAD 데이터 세트(Rajpurkar et al., 2016)와 ARC 추론 도전(ARC Reasoning Challenge)(Clark et al., 2018)으로 탐구됩니다. 그리고 Roberts et al.(2020)에서와 같이, 우리는 세 개의 클로즈드-북(closed-book) 질문 응답 데이터 세트인 Natural Questions(Kwiatkowski et al., 2019), Web Questions(Berant et al., 2013) 및 Trivia QA(Joshi et al., 2017)에서 미세 조정을 통해 모델의 지식을 평가합니다. 클로즈드-북은 보충 참조 자료나 맥락 자료 없이 제시된 질문을 의미합니다. 모델의 상식 추론(common sense reasoning) 능력을 평가하기 위해 Winogrande Schema Challenge(Sakaguchi et al., 2020)에서 평가합니다. 마지막으로, 우리는 모델의 자연어 추론(natural language inference) 능력을 Adversarial NLI Benchmark(Nie et al., 2019)에서 테스트합니다.

미세 조정 메트릭. 본 논문에서는 다음과 같은 평가 메트릭을 사용합니다: 우리는 GLUE와 SuperGLUE의 모든 하위 작업에 대한 평균 점수를 보고합니다. Rouge-2 메트릭은 CNNDM과 XSum 모두에 사용됩니다. SQuAD와 클로즈드-북 작업(Web, Natural, Trivia Questions)에서는 대상과 정확히 일치하는 답변의 비율을 보고합니다(Roberts et al.(2020)을 참조하여 이 측정의 자세한 내용과 한계를 확인하세요). 마지막으로, ARC Easy, ARC Challenge, ANLI, Winogrande에서는 생성된 응답의 정확도를 보고합니다.

미세조정 결과. 우리는 많은 자연어 처리(natural language) 작업에서 상당한 하류 개선을 관찰합니다. 주목할 만한 개선은 SuperGLUE에서 나타나는데, 여기서 FLOP-matched Switch 변형 모델이 T5-Base와 T5-Large 기준 모델보다 각각 4.4 및 2 퍼센트 포인트 향상되었습니다. 또한 Winogrande, 닫힌 책 형식의 Trivia QA, 그리고 XSum에서 큰 개선을 보였습니다. ${ }^{8}$ 우리의 미세조정 연구에서, 우리가 개선을 관찰하지 못한 유일한 작업은 AI2 추론 챌린지(AI2 Reasoning Challenge, ARC) 데이터 세트로, 여기서 T5-Base는 챌린지 데이터 세트에서 Switch-Base를, T5-Large는 쉬운 데이터 세트에서 Switch-Large를 능가합니다. 전체적으로 볼 때, 우리는 추론과 지식이 많이 필요한 작업 모두에 걸쳐 상당한 개선을 관찰합니다. 이는 우리의 아키텍처가 잘 사전 훈련(pre-train)하는 것뿐만 아니라, 미세조정을 통해 하류 작업에 질적인 개선을 전달할 수 있음을 입증합니다.[^2]

표 5: 파인 튜닝 결과. 다양한 자연어 테스트(검증 세트; 높은 숫자가 더 좋음)에서 T5 기본 모델과 스위치(Switch) 모델의 파인 튜닝 결과입니다. 우리는 FLOP이 일치하는 스위치 모델과 T5-Base, T5-Large 기본 모델을 비교합니다. 고려된 대부분의 작업에서, 스위치 변형 모델의 상당한 개선을 발견했습니다. 우리는 모델 크기와 추론 및 지식이 많이 필요한 언어 작업 모두에서 이득을 관찰합니다.

4.2 증류(Distillation)

수십억 또는 수조 개의 매개변수를 가진 거대한 신경망을 배포하는 것은 불편합니다. 이를 완화하기 위해, 우리는 큰 희소 모델을 작은 밀집 모델로 증류하는 것을 연구합니다. 향후 연구에서는 큰 모델을 더 작은 희소 모델로 증류하는 것을 추가로 연구할 수 있습니다.

증류 기술. 표 6에서는 다양한 증류 기술을 연구합니다. 이 기술들은 Sanh et al. (2019)이 BERT 모델에 대한 증류 방법을 연구한 것을 기반으로 합니다. 우리는 비전문가(non-expert) 가중치로 밀집 모델을 초기화하는 것이 소폭의 개선을 가져온다는 것을 발견했습니다. 이는 모든 모델이 FLOP이 일치하기 때문에 가능하며, 비전문가 계층은 동일한 차원을 가질 것입니다. 전문가 계층은 보통 트랜스포머(Transformer)의 FFN 계층마다 또는 그 사이마다만 추가되기 때문에, 많은 가중치를 훈련된 매개변수로 초기화할 수 있습니다. 또한, 우리는 교사 확률에 0.25, 실제 라벨에 0.75의 혼합을 사용하여 증류 개선을 관찰했습니다. 두 기술을 결합함으로써, 우리는 더 큰 희소 모델로부터의 품질 이득의 $\approx 30 \%$를 보존하면서 매개변수의 $\approx 1 / 20^{th}$만을 사용합니다. 품질 이득은 스위치-베이스(Switch-Base, 교사)와 T5-베이스(T5-Base, 학생) 사이의 품질 차이의 백분율을 나타냅니다. 따라서, $100 \%$의 품질 이득은 학생이 교사의 성능과 동일하다는 것을 의미합니다.

표 6: 언어 모델링을 위한 스위치 트랜스포머(Switch Transformers)의 압축. T5-Base를 Switch-Base의 비전문가 가중치로 초기화하고, 교사와 실제 라벨의 혼합으로부터 얻은 손실을 사용하는 것이 가장 좋은 성능을 얻습니다. 우리는 100배 더 많은 매개변수를 가진 큰 희소 모델(sparse model)의 성능 향상을 30% 압축하여 작은 밀집 모델(dense model)로 되돌릴 수 있습니다. 마지막 기준선으로, 전문가 가중치로 초기화된 T5-Base는 압축 없이 정상적으로 훈련했을 때 성능 향상이 없음을 발견했습니다.

달성 가능한 압축 비율. 표 6에 설명된 최고의 압축 기술을 사용하여, 다양한 희소 모델을 밀집 모델로 압축합니다. 우리는 SwitchBase 버전을 압축하며, 전문가의 수를 증가시키는 것을 통해 11억에서 147억 매개변수 사이를 변화시킵니다. 압축을 통해, 우리는 11억 매개변수 모델의 품질 향상의 37%를 유지하면서 82%를 압축할 수 있습니다. 극단적인 경우, 모델을 99% 압축하더라도 여전히 교사 모델의 품질 향상의 28%를 유지할 수 있습니다.

미세 조정된 모델의 압축. 우리는 미세 조정된 희소 모델을 밀집 모델로 압축하는 연구로 이를 마무리합니다. 표 8은 SuperGLUE 작업에 미세 조정된 74억 매개변수 Switch-Base 모델을 2억 2300만 T5-Base로 압축한 결과를 보여줍니다. 사전 훈련 결과와 유사하게, 우리는 FLOP이 일치하는 밀집 변형으로 압축할 때 희소 모델의 이득의 30%를 유지할 수 있음을 발견했습니다. 여기서 고려되지 않은 잠재적인 미래의 방향은, 미세 조정 작업에 사용되는 특정 전문가를 조사하고 이를 추출하여 더 나은 모델 압축을 달성할 수 있을 것입니다.

4.3 다국어 학습

우리의 마지막 하류 실험 세트에서, 우리는 101개의 다른 언어 혼합으로 사전 훈련하는 동안 모델 품질과 속도의 트레이드오프를 측정합니다. 우리는 T5의 다국어 확장인 mT5(Xue et al., 2020)의 최근 작업을 기반으로 구축하고 벤치마크합니다. 우리는 $\mathrm{mT}5$에서 소개된 101개 언어를 포괄하는 다국어 버전의 Common Crawl 데이터 세트 $(\mathrm{mC}4)$에서 사전 훈련을 하지만, 특정 언어 내의 스크립트 변형으로 인해 혼합에는 107개의 작업이 포함됩니다.

그림 7에서 우리는 FLOP이 일치하는 스위치 모델, mSwitch-Base와 T5 베이스 변형, mT5-Base의 모든 언어에 대한 품질 향상을 음의 로그 혼란도(negative log perplexity)로 표시합니다. 이후

표 7: 증류 압축 비율. 우리는 큰 희소 모델을 밀집 기준선으로 증류할 때의 품질을 측정합니다. 우리의 기준선인 T5-Base는 -1.636의 부정 로그 가능도(Neg. Log Perp.) 품질을 가집니다. 오른쪽 열에서는 점점 더 큰 희소 모델을 같은 구조로 증류합니다. 가중치 초기화와 하드 및 소프트 손실의 혼합을 통해, 우리는 $95 \%+$의 희소 선생님을 축소하면서 $30 \%$의 품질 향상을 유지할 수 있습니다. 그러나, 더 나은 성능과 더 큰 사전 훈련된 선생님의 경우, 이러한 압축 비율을 달성하기 위해서는 더 큰 학생 모델이 필요할 것으로 예상됩니다.

표 8: SuperGLUE 모델을 미세 조정한 후 압축하기. 우리는 SuperGLUE 작업에 미세 조정된 Switch-Base 모델을 T5-Base 모델로 압축합니다. 우리는 작은 데이터 세트에서 큰 희소 모델이 압축을 위한 효과적인 교사가 될 수 있다는 것을 관찰합니다. 우리는 교사의 성능의 $30 \%$를 $97 \%$ 압축된 모델에서 다시 달성한다는 것을 발견했습니다.

$1 \mathrm{M}$ 단계로 두 버전을 사전 훈련한 후, 우리는 고려된 모든 101개 언어에서 Switch Transformer가 기준선에 비해 최종 음의 로그 혼란도(negative log perplexity)를 증가시킨다는 것을 발견합니다. 그림 8에서는 다른 관점을 제시하고 이제 Switch Transformer를 사용하여 mT5-Base에 비해 단계별 속도 향상을 히스토그램으로 나타냅니다. ${ }^{9}$ 우리는 mT5-Base에 비해 평균 5배의 속도 향상을 발견하고, $91 \%$의 언어가 적어도 $4 \mathrm{x}$의 속도 향상을 달성한다는 것을 발견했습니다. 이는 Switch Transformers가 효과적인 다중 작업 및 다국어 학습자임을 보여주는 증거입니다.

5. 데이터, 모델, 전문가 병렬성을 가진 모델 설계

임의로 전문가의 수를 증가시키는 것은 수익 감소의 법칙에 따라 제한이 있습니다(그림 4). 여기서 우리는 보완적인 확장 전략을 설명합니다. Transformer를 확장하는 일반적인 방법은 $d_{\text {model }}$ 또는 $d_{ff}$와 같은 차원을 함께 증가시키는 것입니다. 이는 매개변수[^3]를 증가시킵니다.

그림 7: 101개 언어에 대한 다국어 사전 훈련. 101개 언어에 대한 다중 작업 훈련 시 Switch T5 Base 모델이 밀집 기준선에 비해 얼마나 개선되었는지. 우리는 Switch Transformers가 다중 작업 훈련 설정에서 상당히 잘 수행되며 모든 101개 언어에서 개선을 이룬다는 것을 관찰합니다.

그림 8: 101개 언어에 대한 다국어 사전 훈련. 우리는 각 언어에 대해 Switch Transformers가 동일한 품질에 도달하기 위해 FLOP과 일치하는 T5 밀집 기준선에 비해 단계별 속도 향상을 히스토그램으로 나타냅니다. 모든 101개 언어에 대해, 우리는 mT5-Base에 비해 평균 $5 \mathrm{x}$의 단계별 속도 향상을 달성하고, $91 \%$의 언어에서는 최종 혼란도(perplexity)에 도달하기 위해 $4 \mathrm{x}$ 이상의 속도 향상을 기록합니다.

그리고 수행된 계산은 결국 가속기 당 메모리에 의해 제한됩니다. 가속기의 메모리 크기를 초과하면 단일 프로그램 다중 데이터(SPMD) 모델 병렬 처리를 사용할 수 있습니다. 이 섹션에서는 데이터, 모델, 전문가 병렬 처리를 결합하는 것의 트레이드오프에 대해 연구합니다.

피드-포워드 네트워크(Feed-Forward Network, FFN) 레이어 검토하기. 우리는 데이터, 모델 및 전문가 병렬 처리가 Mesh TensorFlow(Shazeer et al., 2018)에서 어떻게 작동하는지 예로 들어 FFN 레이어를 간략히 검토합니다. 배치에서 $B$ 토큰을 가정하며, 각 토큰은 차원 $d_{\text{model}}$을 가집니다. FFN의 입력 $(x)$과 출력 $(y)$은 모두 $\left[B, d_{\text{model}}\right]$ 크기이며, 중간 $(h)$은 $\left[B, d_{ff}\right]$ 크기입니다. 여기서 $d_{ff}$는 일반적으로 $d_{\text{model}}$보다 몇 배 더 큽니다. FFN에서 중간값은 $h=x W_{\text{in}}$이고, 레이어의 출력은 $y=\operatorname{Re}LU(h) W_{\text{out}}$입니다.

따라서 $W_{\text{in}}$과 $W_{\text{out}}$은 각 토큰에 독립적으로 적용되며 크기는 각각 $\left[d_{\text{model}}, d_{ff}\right]$와 $\left[d_{ff}, d_{\text{model}}\right]$입니다.

우리는 파티셔닝의 두 가지 측면을 설명합니다: 가중치와 데이터 배치가 코어에 어떻게 나누어지는지, 그림 9에 나타낸 것처럼. 모든 사용 가능한 코어를 $N$으로 표시하며, Mesh TensorFlow는 이를 논리적인 다차원 프로세서 메쉬로 재매핑할 수 있습니다. 여기서 우리는 두 차원의 논리적 메쉬를 생성하는데, 하나의 차원은 데이터 병렬 샤딩 방법의 수 $(n)$를 나타내고, 다른 하나는 모델 병렬 샤딩 방법의 수 $(m)$를 나타냅니다. 전체 코어는 데이터와 모델 병렬 처리를 걸쳐 샤딩하는 방법의 수와 같아야 합니다. 예를 들어, $N=n \times m$입니다. 레이어를 코어에 걸쳐 샤딩하기 위해, $B$ 토큰의 배치를 포함하는 텐서는 $n$ 데이터 병렬 코어에 걸쳐 샤딩되므로, 각 코어는 $B/n$ 토큰을 포함합니다. $d_{ff}$를 가진 텐서와 변수는 그 다음 $m$ 모델 병렬 코어에 걸쳐 샤딩됩니다. 전문가 레이어가 있는 변형에 대해서는, 최대 $C$ 토큰을 처리할 수 있는 $E$ 전문가를 고려합니다.

5.1 데이터 병렬성(Data Parallelism)

분산 학습(distributed training)의 표준인 데이터 병렬 모델을 학습할 때, 모든 코어는 데이터 병렬 차원에 할당되거나 $n=N, m=1$이 됩니다. 이 방식의 장점은 전체 순전파(forward pass)와 역전파(backward pass)가 끝날 때까지 통신이 필요 없다는 것이며, 그 후에 모든 코어에 걸쳐 그래디언트(gradient)를 집계해야 합니다. 이는 그림 9의 가장 왼쪽 열에 해당합니다.

5.2 모델 병렬성(Model Parallelism)

이제 모든 코어가 모델 병렬 차원에만 전용으로 할당되는 시나리오를 고려해 보겠습니다. 그래서 $n=1, m=N$입니다. 이제 모든 코어는 전체 $B$ 토큰을 유지해야 하며, 각 코어는 가중치의 고유한 조각을 포함하게 됩니다. 각 순전파와 역전파를 위해 이제 통신 비용이 발생합니다. 각 코어는 두 번째 행렬 곱셈 $\operatorname{Re} L U(h) W_{\text {out }}$을 계산하기 위해 $\left[B, d_{\text {model }}\right]$의 텐서를 보냅니다. 왜냐하면 $d_{f f}$ 차원이 분할되어 있고 합산되어야 하기 때문입니다. 일반적인 규칙으로, 코어에 걸쳐 분할된 차원이 합산되어야 할 때마다 순전파와 역전파 모두에 all-reduce 연산이 추가됩니다. 이는 순수 데이터 병렬성(pure data parallelism)과 대조적이며, 여기서 all-reduce는 전체 순전파와 역전파가 끝난 끝에만 발생합니다.

그림 9: 데이터와 가중치 분할 전략. 각 $4 \times 4$ 점선 그리드는 16개의 코어를 나타내고, 그림자가 있는 정사각형은 해당 코어에 포함된 데이터(모델 가중치 또는 토큰 배치)를 나타냅니다. 우리는 각 전략에 대해 모델 가중치와 데이터 텐서가 어떻게 분할되는지를 보여줍니다. 첫 번째 행: 코어에 걸쳐 모델 가중치가 어떻게 분할되는지를 보여줍니다. 이 행의 다양한 크기의 모양은 피드 포워드 네트워크(Feed Forward Network, FFN) 레이어의 더 큰 가중치 행렬을 나타냅니다(예: 더 큰 $d_{ff}$ 크기). 그림자가 있는 정사각형의 각 색깔은 고유한 가중치 행렬을 식별합니다. 코어 당 매개변수의 수는 고정되어 있지만, 더 큰 가중치 행렬은 각 토큰에 더 많은 계산을 적용합니다. 두 번째 행: 코어에 걸쳐 데이터 배치가 어떻게 분할되는지를 보여줍니다. 각 코어는 동일한 수의 토큰을 유지하여 모든 전략에 걸쳐 고정된 메모리 사용량을 유지합니다. 분할 전략은 각 코어가 동일한 토큰을 가지거나 코어에 걸쳐 다른 토큰을 가질 수 있는 다른 속성을 가지고 있으며, 이것이 다른 색깔이 상징하는 것입니다.

5.3 모델 및 데이터 병렬 처리

대규모 모델의 경우 모델과 데이터 병렬 처리를 혼합하는 것이 일반적이며, 이는 가장 큰 T5 모델(Raffel et al., 2019; Xue et al., 2020)과 GPT-3(Brown et al., 2020)에서 수행되었습니다. 총 $N=n \times m$개의 코어를 가진 상태에서 이제 각 코어는 $B / n$개의 토큰과 가중치 및 중간 활성화의 $d_{ff} / m$을 담당하게 됩니다. 순전파(forward)와 역전파(backward) 단계에서 각 코어는 $\left[B / n, d_{\text{model}}\right]$ 크기의 텐서를 all-reduce 연산에서 통신합니다.

5.4 전문가 및 데이터 병렬 처리

다음으로 전문가(expert)와 데이터 병렬성(data parallelism)에 대한 분할 전략을 설명하겠습니다. 스위치 트랜스포머(Switch Transformers)는 모든 코어를 데이터 분할 차원 $n$에 할당할 것이며, 이는 모델 내 전문가의 수와도 일치합니다. 각 토큰(token)에 대해 코어(core)당 라우터(router)는 전문가에게 할당할 작업을 로컬로 계산합니다. 출력은 크기가 $[n, B / n, E, C]$인 이진 행렬(binary matrix)로, 첫 번째 차원을 기준으로 분할되며 전문가 할당을 결정합니다. 이 이진 행렬은 입력 텐서(tensor) $\left[n, B / n, d_{\text {model }}\right]$와 행렬 곱셈을 통해 수집(gather)하는 데 사용됩니다.

결과적으로 최종 텐서는 $\left[n, E, C, d_{\text {model }}\right]$ 형태가 되며, 첫 번째 차원을 기준으로 샤딩(sharded)됩니다. 각 코어가 자신의 전문가를 가지고 있기 때문에, 이제 $n$ 차원 대신 $E$ 차원을 샤딩하기 위해 $\left[E, C, d_{\text {model }}\right]$ 크기의 전체-전체(all-to-all) 통신을 수행합니다. 전방향(forward pass)에서는 다른 코어에 위치한 각 전문가로부터 토큰을 받기 위해 $E \times C \times d_{\text {model }}$ 크기의 bfloat16 텐서에 대한 추가적인 통신 비용이 발생합니다. 전문가 분할 코드에 대한 자세한 분석은 부록 F를 참조하십시오.

5.5 전문가, 모델 및 데이터 병렬성

우리가 최고의 모델을 설계할 때, 토큰 당 FLOPS(Floating-point Operations Per Second)와 파라미터 수를 균형있게 조절하려고 합니다. 전문가의 수를 늘릴 때, 파라미터 수는 증가하지만 토큰 당 FLOPS는 변하지 않습니다. FLOPS를 증가시키기 위해서는 $d_{ff}$ 차원을 늘려야 합니다(이는 파라미터도 증가시키지만, 더 느린 속도로). 이는 트레이드오프(trade-off)를 제시합니다: $d_{ff}$를 증가시키면 코어 당 메모리가 부족해지고, 이는 $m$을 증가시켜야 한다는 것을 의미합니다. 그러나 우리가 고정된 수의 코어 $N$을 가지고 있고, $N=n \times m$이기 때문에, $n$을 줄여야 하며, 이는 더 작은 배치 크기(batch-size)를 사용해야 한다는 것을 의미합니다(코어 당 토큰 수를 일정하게 유지하기 위해).

모델 병렬성(model parallelism)과 전문가 병렬성(expert-parallelism)을 결합할 때, 토큰을 올바른 전문가에게 라우팅하는 데서 발생하는 전체-전체 통신 비용과 모델 병렬성에서 발생하는 내부 all-reduce 통신이 있을 것입니다. FLOPS, 통신 비용, 코어 당 메모리를 균형있게 조절하는 것은 세 가지 방법을 결합할 때 상당히 복잡해지며, 최적의 매핑은 경험적으로 결정됩니다. 전문가 수가 하류 성능에 미치는 영향에 대한 추가 분석은 5.6절에서도 확인할 수 있습니다.

5.6 조(Trillion) 단위 매개변수 모델을 향하여

전문가 병렬 처리(expert parallelism), 모델 병렬 처리(model parallelism), 데이터 병렬 처리(data parallelism)를 결합하여, 각각 3950억 개와 1.6조 개의 매개변수를 가진 두 개의 대규모 스위치 트랜스포머(Switch Transformer) 모델을 설계했습니다. 우리는 이 모델들이 언어 모델로서 상위 스트림 상의 사전 학습(pre-training)과 하위 스트림에서의 미세 조정(fine-tuning) 성능에 어떻게 영향을 미치는지 연구했습니다. 두 가지 다른 모델의 매개변수, 시퀀스 당 FLOPs, 그리고 하이퍼파라미터는 아래 표 9에 나열되어 있습니다. 트랜스포머(Transformer)의 표준 하이퍼파라미터, $d_{\text{model}}, d_{ff}, d_{kv}$, 헤드(heads)의 수, 레이어(layers)의 수가 설명되어 있으며, 덜 일반적인 특징인 $FFN_{GELU}$도 설명되어 있는데, 이는 확장 행렬이 두 세트의 가중치로 대체되고 비선형적으로 결합된 FFN 레이어의 변형을 의미합니다(Shazeer, 2020).

Switch-C 모델은 앞서 5.4절에서 설명한 바와 같이, 전문가 병렬 처리만을 사용하여 설계되었고, 모델 병렬 처리는 사용하지 않았습니다. 결과적으로, 폭(width),

표 9: 스위치 모델 디자인과 사전 학습 성능. T5 모델들과 우리의 스위치 트랜스포머(Switch Transformer) 변형 모델들의 하이퍼파라미터와 사전 학습 성능을 비교합니다. 마지막 두 열은 각각 250k 단계와 500k 단계 후 $\mathrm{C}4$ 데이터 세트에서의 사전 학습 모델 품질을 기록합니다. 우리는 스위치C 트랜스포머(SwitchC Transformer) 변형 모델이 같은 계산 예산으로 고정된 혼란도(perplexity)에 도달하는 데 T5-XXL 모델보다 $4 \mathrm{x}$ 빠르며, 학습이 진행됨에 따라 그 격차가 더 커지는 것을 관찰합니다.

깊이(depth), 헤드(heads)의 수 등은 T5-XXL 모델보다 훨씬 작습니다. 반면에, 스위치-XXL(Switch-XXL)은 T5-XXL 모델과 FLOP이 일치하도록 설계되어 하이퍼파라미터의 더 큰 차원을 허용하지만, 모델 병렬성(model-parallelism)에 의해 유발되는 추가적인 통신 비용을 초래합니다(자세한 내용은 5.5절을 참조하세요).

T5-XXL 대비 샘플 효율성. 표 9의 마지막 두 열에서는 각각 250k 단계와 500k 단계 후 $\mathrm{C}4$ 코퍼스(corpus)에서의 음의 로그 혼란도(negative log perplexity)를 기록합니다. 250k 단계 후, 우리는 두 스위치 트랜스포머(Switch Transformer) 변형 모델이 T5-XXL 버전의 음의 로그 혼란도를 0.061 이상 개선하는 것을 발견했습니다. 0.061의 격차가 얼마나 중요한지를 이해하기 위해, T5-XXL 모델이 0.052를 증가시키기 위해 추가적으로 250k 단계를 학습해야 했다는 점을 지적합니다. 추가 학습이 진행됨에 따라 격차는 계속 증가하며, 500k 단계에서 스위치-XXL(Switch-XXL) 모델이 T5-XXL을 0.087로 능가합니다.

학습 불안정성. 그러나 서론에서 설명한 바와 같이, 큰 희소 모델은 불안정할 수 있으며, 규모를 늘리면서 간헐적인 문제에 직면하기도 합니다. 우리는 1.6T 파라미터와 2048 전문가(experts)를 가진 더 큰 스위치-C(Switch-C) 모델이 전혀 학습 불안정성을 보이지 않는다는 것을 발견했습니다. 대신에, 시퀀스 당 거의 10배 더 큰 FLOPs를 가진 스위치 XXL(Switch XXL) 버전은 때때로 불안정합니다. 결과적으로, 이 모델이 단계 기준으로 우리의 더 나은 모델이지만, T5(Raffel et al., 2019)의 최종 보고된 결과와 일치하도록 전체 $1 \mathrm{M}$ 단계까지 사전 학습을 진행하지 않습니다.

10. 이 보고된 품질 차이는 하한값이며, 실제로는 더 클 수 있습니다. T5-XXL은 중복된 스니펫(snippets)이 포함되어 있어 쉽게 복사할 수 있는 더 쉬운 $\mathrm{C}4$ 데이터 세트에서 사전 학습되었습니다.

추론 세부 조정 성능. 모델 품질의 예비 평가로, 우리는 T5-XXL 모델이 사용한 텍스트의 대략 절반인 503B 토큰에 대해 부분적으로 사전 학습된 Switch-XXL 모델을 사용합니다. 이 체크포인트를 사용하여, 모든 작업을 효율적으로 동시에 학습하는 멀티태스크 학습을 수행하는데, 이는 개별적으로 세부 조정하는 것이 아닙니다. 우리는 SQuAD 정확도가 검증 세트에서 89.7로 증가했으며, 이는 최신 기술(state-of-the-art)인 91.3에 비해 증가한 것을 발견합니다. 다음으로, SuperGLUE 테스트 평균 점수는 최신 기술인 90.0(Wang et al., 2019)에 비해 T5 버전이 89.3을 얻은 것에 비해 87.5로 기록됩니다. ANLI(Nie et al., 2019)에서는 Switch XXL이 이전 최고 기술인 49.4(Yang et al., 2020)에 비해 65.7 정확도로 개선되었습니다. 우리는 SwitchXXL이 상류 사전 학습 과제에서 최신 기술인 Neg. Log Perp.를 가지고 있지만, 그 이득이 아직 최신 기술의 하류 성능으로 완전히 전환되지 않았음을 지적합니다. 이 문제에 대해 부록 E에서 더 자세히 연구합니다.

지식 기반 세부 조정 성능. 마지막으로, 우리는 추가적인 사전 학습 없이 Salient Span Masking(Guu et al., 2020)을 사용하여 세 가지 폐쇄형 지식 기반 작업인 Natural Questions, WebQuestions, TriviaQA에서 모델의 지식을 초기 검토합니다. 세 경우 모두에서, 우리는 이전 최고 기술인 T5-XXL 모델(SSM 없이)보다 개선된 것을 관찰합니다. Natural Questions 정확한 일치는 이전 최고인 32.8에 비해 34.4로 증가하고, Web Questions은 37.2에서 41.0으로 증가하며, TriviaQA는 42.9에 비해 47.5로 증가합니다.

요약하자면, 다른 모델들이 사용한 데이터의 절반 미만으로 학습에도 불구하고, 우리는 이미 비교 가능하고 때로는 최신 기술의 모델 품질을 발견합니다. 현재 Switch Transformer는 지식 기반 작업에 대해 상류 이득을 추론 작업보다 더 잘 전환합니다(부록 E 참조). 대규모 전문가 모델에서 더 강력한 세부 조정 성능을 추출하는 것은 활발한 연구 질문이며, 사전 학습 perplexity는 미래의 개선이 가능해야 함을 나타냅니다.

6. 관련 연구

신경망에서 규모의 중요성은 널리 인식되어 왔으며, 여러 접근 방식이 제안되었습니다. 최근의 연구들은 모델 병렬성(model parallelism)을 사용하여 수십억 개의 매개변수를 가진 모델을 확장했습니다(예: 여러 코어에 가중치와 텐서를 분할) (Shazeer et al., 2018; Rajbhandari et al., 2019; Raffel et al., 2019; Brown et al., 2020; Shoeybi et al., 2019). 대안으로, Harlap et al. (2018); Huang et al. (2019)은 파이프라인 기반 모델 병렬성(pipeline based model parallelism)을 제안했는데, 여기서 다른 레이어들이 장치에 분할되고 마이크로 배치(micro-batches)가 다른 레이어로 파이프라인됩니다. 마지막으로, Product Key 네트워크(Lample et al., 2019)는 주어진 레이어로 들어오는 토큰 표현에 기반하여 학습 가능한 임베딩을 조회함으로써 신경망의 용량을 확장하는 것을 제안했습니다.

우리의 연구는 입력에 기반하여 동적으로 계산 결정을 내리는 조건부 계산(conditional computation)을 수행하는 방법들의 클래스에서 특정 모델을 연구합니다. Cho와 Bengio (2014)는 모델의 은닉 상태(hidden-states)에서 발생하는 특정 비트 패턴에 기반하여 가중치를 적응적으로 선택하는 것을 제안했습니다. Eigen et al. (2013)은 밀집 행렬 곱셈과 ReLU 활성화를 가진 스택된 전문가 레이어(stacked expert layers)를 구축하고 지터링된 MNIST와 단조로운 음성에서 유망한 결과를 보였습니다. 컴퓨터 비전 분야에서 Puigcerver et al. (2020)은 상류 사전 훈련(upstream pre-training) 동안 의미적 클래스에 기반하여 수동으로 토큰을 라우팅하고, 그 다음 하류 작업(downstream task)에 따라 관련 전문가를 선택합니다.

전문가의 혼합(Mixture of Experts, MoE)은 현대 딥 러닝 아키텍처의 맥락에서 Shazeer et al. (2017)에 의해 효과적임이 입증되었습니다. 그 연구에서는 LSTM (Hochreiter and Schmidhuber, 1997) 레이어 사이에 MoE 레이어를 추가했고, 토큰들은 별도로 전문가들의 조합으로 라우팅되었습니다. 이것은 언어 모델링과 기계 번역 벤치마크에서 최첨단 결과를 가져왔습니다. MoE 레이어는 Mesh Tensorflow 라이브러리(Shazeer et al., 2018)에 의해 Transformer 아키텍처로 재도입되었는데, 여기서 MoE 레이어는 FFN 레이어의 대체물로 도입되었지만, NLP 결과는 동반되지 않았습니다. 더 최근에는, 기계 학습 인프라의 발전을 통해, GShard (Lepikhin et al., 2020)는 XLA 컴파일러를 확장하여 100개 언어에 걸친 기계 번역을 대폭 개선하는 데 MoE Transformer를 사용했습니다. 마지막으로 Fan et al. (2021)은 모델 매개변수를 언어 그룹별로 중복되지 않게 분할하는 다른 결정론적 MoE 전략을 선택했습니다.

시퀀스 길이 차원 $(L)$에서 Transformer 주의 패턴의 희소성(sparsity)은 $O\left(L^{2}\right)$의 주의 복잡성을 줄이는 데 성공한 기술입니다(Child et al., 2019; Correia et al., 2019; Sukhbaatar et al., 2019; Kitaev et al., 2020; Zaheer et al., 2020; Beltagy et al., 2020). 이는 이전보다 더 긴 시퀀스를 학습할 수 있게 해주었습니다. 이 버전의 Switch Transformer는 주의 희소성(attention sparsity)을 사용하지 않지만, 이 기술들은 서로 보완적이며, 향후 작업에서 이를 결합하여 긴 맥락을 필요로 하는 작업에서 학습을 개선할 수 있을 것입니다.

7. 토론

우리는 Switch Transformer와 희소 전문가 모델(sparse expert models)에 대해 질문을 던지고 토론합니다. 여기서 희소성(sparsity)은 주의 패턴이 아닌 가중치에 관한 것입니다.

Switch Transformer가 단순히 매개변수 수가 많아서 더 나은 것이 아닌가요? 네, 그리고 의도된 바입니다! 매개변수는 총 FLOPs와 독립적으로, 신경 언어 모델을 확장하는 유용한 축입니다. 큰 모델은 더 나은 성능을 보인다는 것이 충분히 입증되었습니다(Kaplan et al., 2020). 하지만 이 경우, 우리의 모델은 같은 계산 자원을 사용하면서 더 샘플 효율적이고 빠릅니다.

슈퍼컴퓨터에 접근할 수 없는데, 이게 여전히 유용한가요? 이 작업은 매우 큰 모델에 초점을 맞추고 있지만, 단 두 개의 전문가(experts)를 가진 모델도 성능을 향상시키면서 일반적으로 사용 가능한 GPU나 TPU의 메모리 제약 내에서 쉽게 맞출 수 있다는 것을 발견했습니다(자세한 내용은 부록 D에 있습니다). 따라서 우리는 우리의 기술이 소규모 환경에서도 유용하다고 믿습니다.

희소 모델이 속도-정확도 파레토 곡선(speed-accuracy Pareto curve)에서 밀집 모델을 능가하나요? 네. 다양한 모델 크기에 걸쳐, 희소 모델은 단계별로 밀집 모델을 능가하며 실제 시간에 있어서도 그렇습니다. 우리의 통제된 실험은 고정된 계산량과 시간에 대해, 희소 모델이 밀집 모델을 능가한다는 것을 보여줍니다.

1조 개의 매개변수 모델을 배포할 수 없는데 — 이 모델들을 축소할 수 있나요? 모델 품질을 완전히 보존할 수는 없지만, 우리의 희소 모델을 밀집 모델로 압축하여 전문가 모델의 품질 향상의 $\approx 30 \%$를 달성하면서 10에서 $100 \mathrm{x}$의 압축률을 달성할 수 있습니다.

모델-병렬 밀집 모델 대신 Switch Transformer를 사용하는 이유는 무엇인가요? 시간 기준으로, Switch Transformers는 샤드된 매개변수를 가진 밀집 모델보다 훨씬 효율적일 수 있습니다(그림 6). 또한, 이 결정은 상호 배타적이지 않다는 점을 지적합니다 - 우리는 Switch Transformers에서 모델-병렬성을 사용할 수 있고, 실제로 사용합니다. 이는 토큰당 FLOPs를 증가시키지만, 전통적인 모델-병렬성의 느려짐을 초래합니다.

왜 희소 모델(sparse models)이 이미 널리 사용되지 않았을까요? 희소 모델을 시도하려는 동기는 밀집 모델(dense models)의 대규모 성공에 의해 저해되었습니다(이 성공은 부분적으로 Hooker(2020)에서 주장한 바와 같이 딥 러닝 하드웨어와의 공진화에 의해 주도되었습니다). 더욱이, 희소 모델은 (1) 모델 복잡성, (2) 학습의 어려움, 그리고 (3) 통신 비용과 같은 여러 문제에 직면해 있습니다. 스위치 트랜스포머(Switch Transformer)는 이러한 문제들을 완화하기 위해 진전을 이루었습니다.

8. 미래의 연구

이 논문은 간소화된 구조, 개선된 학습 절차, 그리고 희소 모델이 어떻게 확장되는지에 대한 연구를 제시합니다. 그러나 여전히 많은 열린 미래 방향들이 있으며, 여기서 간략하게 설명합니다:

1. 가장 큰 모델들의 학습 안정성을 더욱 향상시키는 것은 상당한 도전입니다. 우리의 안정성 기술은 Switch-Base, Switch-Large, Switch-C 모델들에 대해서는 효과적이었습니다(불안정성 관찰되지 않음), 하지만 Switch-XXL에 대해서는 충분하지 않았습니다. 우리는 이러한 모델들을 안정화하는 초기 단계를 밟고 있으며, 이는 큰 모델들에 대해 일반적으로 유용할 수 있다고 생각합니다. 여기에는 안정성을 향상시키기 위한 정규화 기법과 그래디언트 클리핑(gradient clipping)의 적응형 형태 사용이 포함되지만, 이 문제는 여전히 해결되지 않았습니다.
2. 일반적으로 우리는 개선된 사전 학습 품질이 더 나은 하류 결과를 가져온다는 것을 발견합니다(부록 E 참조), 비록 때때로 놀라운 이상 현상을 마주하기도 합니다. 예를 들어, $\mathrm{C}4$ 데이터 세트를 모델링하는 데 비슷한 혼란도(perplexities)에도 불구하고, 1.6T 파라미터의 Switch-C는 SQuAD에서 87.7의 정확한 일치 점수만을 달성하는데, 이는 더 작은 Switch-XXL 모델의 89.6에 비해 불리하게 비교됩니다. 주목할 만한 차이점은 SwitchXXL 모델이 Switch-C 모델보다 토큰당 $\approx 10 \mathrm{x}$ 만큼의 FLOPS를 적용한다는 것입니다, 비록 그것이 $\approx 4 \mathrm{x}$ 더 적은 고유 파라미터를 가지고 있음에도 불구하고 ( $395 \mathrm{~B}$ 대 $1.6 \mathrm{~T}$ ). 이는 미세 조정 품질, 토큰당 FLOPS, 그리고 파라미터 수 사이의 잘 이해되지 않은 의존성을 시사합니다.
3. 데이터, 모델, 전문가 병렬 처리를 혼합하는 아키텍처 설계를 안내하는 스케일링 관계에 대한 종합적인 연구를 수행합니다. 이상적으로, 하드웨어 구성(계산, 메모리, 통신)의 사양을 감안할 때, 더 빠르게 최적의 모델을 설계할 수 있을 것입니다. 그리고 반대로, 이는 미래 하드웨어의 설계에도 도움이 될 수 있습니다.
4. 우리의 작업은 적응형 계산 알고리즘의 범주에 속합니다. 우리의 접근 방식은 항상 동일하고 균질한 전문가들을 사용했지만, 미래의 설계(더 유연한 인프라에 의해 용이해짐)는 이질적인 전문가들을 지원할 수 있습니다. 이는 더 많은 계산이 원하는 경우—아마도 더 어려운 예시들에 대해—더 큰 전문가들로 라우팅함으로써 더 유연한 적응을 가능하게 할 것입니다.
5. 트랜스포머(Transformer)의 FFN 층 외부에서 전문가 층을 조사합니다. 우리는 이것이 유사하게 모델 품질을 향상시킬 수 있다는 초기 증거를 발견했습니다. 부록 A에서, 우리는 Self-Attention 층 내부에 이들을 추가하여 품질 개선을 보고합니다. 여기서 우리의 층은 Q, K, V를 생성하는 가중치 행렬을 대체합니다. 그러나, bfloat16 포맷으로 인한 학습 불안정성 때문에, 우리는 대신 이를 미래의 연구 분야로 남겨둡니다.
6. 새로운 모달리티와 다른 모달리티에서 스위치 트랜스포머(Switch Transformer)를 검토합니다. 우리는 지금까지 언어에 대해서만 고려했지만, 모델 희소성이 새로운 모달리티에서도 장점을 제공할 수 있으며, 멀티 모달 네트워크에서도 마찬가지일 것이라고 믿습니다.

이 목록은 쉽게 확장될 수 있지만, 우리가 생각하고 있는 도전 유형과 앞으로 유망한 방향에 대한 맛을 제공하기를 바랍니다.

9. 결론

스위치 트랜스포머(Switch Transformers)는 확장 가능하고 자연어 학습에 효과적인 모델입니다. 우리는 전문가의 혼합(Mixture of Experts)을 단순화하여 이해하기 쉽고, 훈련하기 안정적이며, 동등한 크기의 밀집 모델(dense models)보다 훨씬 더 샘플 효율이 높은 구조를 만들었습니다. 이 모델들이 다양한 자연어 처리 작업에서 뛰어난 성능을 보이며, 사전 훈련(pre-training), 미세 조정(fine-tuning), 다중 작업 훈련(multi-task training) 등 다양한 훈련 체제에서도 탁월함을 발견했습니다. 이러한 발전은 수백억에서 조 단위의 매개변수를 가진 모델을 훈련할 수 있게 하며, 밀집 T5 기준 모델(dense T5 baselines)에 비해 상당한 속도 향상을 달성할 수 있게 합니다. 우리는 이 작업이 희소 모델(sparse models)을 효과적인 구조로서 동기를 부여하고, 연구자와 실무자들이 자연어 작업뿐만 아니라 그 이상의 분야에서 이러한 유연한 모델을 고려하도록 장려하기를 바랍니다.

감사의 말

저자들은 알고리즘 개선에 대한 핵심적인 통찰과 경험적 연구를 위한 제안을 몇 달 동안 제공한 Margaret Li에게 감사를 표합니다. 초안에 대한 현명한 조언과 명확한 코멘트를 해준 Hugo Larochelle, 상세한 코멘트와 신중한 수정을 해준 Irwan Bello, 신경 언어 모델(neural language models)과 T5 코드 베이스에 대한 적시적인 조언을 해준 Colin Raffel과 Adam Roberts, 적응형 계산(adaptive computation) 연구에 대한 조언과 격려를 해준 Yoshua Bengio, 새로운 대규모 모델의 안정화와 논문 수정을 위한 흥미로운 새 방향을 제시한 Jascha Sohl-dickstein, 그리고 논문에 대한 유용한 토론을 해준 Google Brain Team에게 감사합니다. 모델의 훈련 성능을 프로파일링하고 개선하는 데 귀중한 도움을 준 Blake Hechtman에게도 감사합니다.

A. 주의(Attention)를 위한 스위치

Shazeer et al. (2018); Lepikhin et al. (2020)은 트랜스포머(Transformer)의 밀집 피드포워드 네트워크(dense feedforward network, FFN) 계산에 MoE 계층을 추가함으로써 MoE 트랜스포머를 설계했습니다(Shazeer et al., 2017). 마찬가지로, 우리의 작업도 트랜스포머의 FFN 계층을 대체했지만, 여기서는 대안적인 설계를 간략히 탐구합니다. 우리는 트랜스포머 자기 주의(self-attention) 계층에 스위치 계층을 추가합니다. 이를 위해 쿼리(queries), 키(keys), 값(values)을 생성하는 훈련 가능한 가중치 행렬을 그림 10에서 볼 수 있듯이 스위치 계층으로 대체합니다.

표 10은 고정된 수의 단계 후의 품질과 여러 변형에 대한 훈련 시간을 기록합니다. 개선 사항을 발견했지만, bfloat16 정밀도를 사용할 때 이 계층들이 더 불안정하다는 것도 발견했으므로 최종 변형에는 포함하지 않았습니다.

그림 10: 주의(attention)에서 스위치 레이어. 자기 주의(self-attention) 트랜스포머 블록에 스위치 레이어를 통합하는 방법을 다이어그램으로 보여줍니다. 각 토큰(여기서는 두 토큰, $x_{1}=$ "더(More)"와 $x_{2}=$ "매개변수(Parameters)")에 대해, 한 세트의 가중치가 쿼리를 생성하고 다른 고유한 가중치 세트가 공유 키(keys)와 값(values)을 생성합니다. 우리는 각 전문가(expert)가 선형 연산이 되는 것과 이 작업 전반에 걸쳐 그랬던 것처럼 FFN이 되는 것을 실험했습니다. 이를 사용하여 품질 개선을 발견했지만, 저정밀도 숫자 형식을 사용할 때 더 불안정하다는 것을 발견했으므로, 이는 미래의 작업으로 남겨두었습니다.

그러나 이러한 레이어가 안정적으로 훈련될 때, 우리는 예비적인 긍정적 결과가 약속하는 미래의 방향을 제시한다고 믿습니다.

표 10: 스위치 주의 레이어 결과. 모든 모델은 32개의 전문가를 가지고 있으며, 배치당 524k 토큰으로 훈련합니다. 전문가 FF는 트랜스포머에서 FFN을 전문가가 대체하는 경우로, 이 논문 전반에 걸쳐 우리의 표준 설정입니다. 전문가 FF + 주의는 전문가가 FFN과 자기 주의 레이어 모두를 대체할 때입니다. bfloat16 정밀도로 훈련할 때 전문가 주의를 가진 모델은 발산합니다.

B. No-Token-Left-Behind로 토큰 드롭 방지

TPU 가속기의 소프트웨어 제약으로 인해 우리의 텐서(Tensors)는 정적으로 크기가 정해져 있어야 합니다. 그 결과, 각 전문가(expert)는 토큰 표현을 처리할 수 있는 유한하고 고정된 용량을 가지게 됩니다. 그러나 이것은 실행 시간에 토큰을 동적으로 라우팅하는 우리 모델에 문제를 제기합니다. 이로 인해 전문가들에게 토큰이 불균등하게 분배될 수 있습니다. 전문가에게 보내진 토큰의 수가 전문가의 용량보다 적으면, 계산은 단순히 패딩될 수 있습니다 - 하드웨어 사용에 비효율적이지만, 수학적으로는 정확합니다. 그러나 전문가에게 보내진 토큰의 수가 그 용량을 초과하는 경우(전문가 오버플로우), 이를 처리할 프로토콜이 필요합니다. Lepikhin et al. (2020)은 Mixture-of-Expert 모델을 적용하고 잔여 연결을 통해 처리하지 않고 다음 레이어로 표현을 전달함으로써 전문가 오버플로우를 처리하는 방법을 제시했으며, 우리도 이를 따릅니다.

우리는 토큰에 아무런 계산도 적용되지 않는 것이 특히 한 전문가에 오버플로우가 발생하면 다른 전문가는 여분의 용량을 가질 것이기 때문에 매우 낭비적일 수 있다고 의심했습니다. 이러한 직관을 바탕으로 우리는 No-Token-Left-Behind를 만들었으며, 이는 처음에 오버플로우가 발생하는 전문가에게 라우팅된 토큰을 반복적으로 재라우팅합니다. 그림 11은 이 방법의 그래픽 설명을 보여주며, 이를 통해 훈련 및 추론 중 거의 토큰이 누락되지 않도록 거의 보장할 수 있습니다. 우리는 이것이 성능을 향상시키고 훈련을 더 안정화시킬 수 있다고 가설을 세웠지만, 실제로는 이점을 발견하지 못했습니다. 우리는 네트워크가 다른 토큰과 전문가 간의 연관성을 학습한 후, 이 연관성이 변경되면(예: 토큰을 두 번째로 높은 전문가에게 보내는 경우) 성능이 저하될 수 있다고 의심합니다.

C. 전문가들 간의 탐색 장려

각 전문가 레이어에서, 라우터는 토큰을 어느 전문가에게 보낼지 결정합니다. 이것은 토큰의 표현에 대한 정보를 조건으로 하는 사용 가능한 전문가들에 대한 이산적 결정입니다. 들어오는 토큰 표현을 바탕으로, 라우터는 최적의 전문가를 결정하지만, 다른 전문가를 선택했을 때 얼마나 잘 할 수 있었을지에 대한 반증적 정보는 받지 못합니다. 강화 학습(reinforcement learning)에서처럼, 클래식한 탐색-이용(exploration-exploitation) 딜레마가 발생합니다(Sutton and Barto, 2018). 이러한 문제는 Rosenbaum et al. (2017)에 의해 비슷하게 지적되고 다르게 해결되었으며, 다중 작업 학습에서 성공을 입증했습니다. 이 특정 설정은 문맥적 밴딧(contextual bandit) (Robbins, 1952)과 가장 밀접하게 일치합니다. 항상 최고의 전문가를 결정적으로 선택하는 것은 이용적 전략에 해당합니다 - 우리는 더 나은 전문가 할당을 찾기 위해 탐색을 균형있게 고려합니다.

탐색을 도입하기 위해, 우리는 여러 접근 방식을 고려합니다: 1) 결정론적 또는 argmax 2) 소프트맥스 분포에서 샘플링 3) 들어오는 표현에 대한 입력 드롭아웃 4) 들어오는 표현에 대한 곱셈적 지터 노이즈. 모델 품질에 미치는 영향은 표 11에 보고되어 있습니다. 이 작업을 통틀어, 우리는 입력 지터를 사용하여 노이즈를 주입하는 것이 경험적으로 가장 잘 수행되는 것으로 발견했습니다.

D. 낮은 컴퓨트 환경에서의 스위치 트랜스포머(Switch Transformers)

스위치 트랜스포머(Switch Transformer)는 수천 개의 코어(core)와 수조 개의 매개변수(parameter)를 가진 환경뿐만 아니라 작은 규모에서도 효과적인 구조입니다. 우리의 이전 실험들은 대부분

그림 11: No-Token-Left-Behind 라우팅 다이어그램. 1단계는 스위치 라우팅(Switch routing)과 동일하며, 토큰(token)들이 라우터(router)로부터 가장 높은 확률을 가진 전문가(expert)에게 라우팅됩니다. 2단계에서는 오버플로된 모든 토큰들을 살펴보고 두 번째로 높은 확률을 가진 전문가에게 라우팅합니다. 두 번째로 높은 전문가에게도 너무 많은 토큰이 있으면 여전히 오버플로될 수 있지만, 대부분의 토큰들이 라우팅될 수 있도록 합니다. 이 과정은 거의 모든 토큰이 떨어지지 않도록 반복될 수 있습니다.

표 11: 라우터 탐색 전략. 스위치 트랜스포머의 품질은 전문가를 선택하는 데 사용되는 무작위성 전략에 따라 다르며, 음의 로그 혼란도(negative log perplexity)로 측정됩니다(낮을수록 좋음). 변형 간에는 속도 성능 차이가 거의 없습니다.

$10 \mathrm{~B}+$ 매개변수 모델의 규모에서도 실험을 했지만, 그림 12에서 볼 수 있듯이 단 2개의 전문가만으로도 FLOP-matched 상대방보다 뛰어난 성과를 보여줍니다. 슈퍼 컴퓨터(super computer)가 즉시 사용 가능하지 않더라도, 우리가 일반적으로 권장하는 대로 코어 당 한 명의 전문가를 사용하여 2, 4, 또는 8개의 전문가로 스위치 트랜스포머를 훈련시키면 T5 밀집 기준선(dense baselines)에 비해 탄탄한 개선을 이룰 수 있습니다.

그림 12: 소수의 전문가를 가진 스위치 트랜스포머. 매우 적은 수의 전문가를 가지고도 스위치 트랜스포머는 기준선을 뛰어넘는 성능을 보여줍니다. 여기서는 매우 작은 규모에서의 확장성을 보여주며, 2, 4, 8개의 전문가를 사용하여 T5-Base 모델을 뛰어넘는 성능을 보여줍니다.

E. 상류 모델과 하류 모델 성능의 관계

모델의 사전 훈련(pre-training) 목표에 대한 품질이 하류 작업(downstream task) 결과로 이어진다는 보장은 없습니다. 그림 13은 $\mathrm{C}4$ 사전 훈련 작업에서의 상류 모델 품질과 두 가지 하류 작업 척도인 SuperGLUE 평균 성능과 TriviaQA 점수 간의 상관관계를 보여줍니다. 우리는 이 두 작업을 선택했는데, 하나는 모델의 추론 능력을, 다른 하나는 사실 지식을 평가하기 때문입니다.

그림 13: 상류 사전 훈련 품질과 하류 모델 품질. 우리는 상류 성능과 SuperGLUE 및 TriviaQA(SOTA는 SSM 없이 기록됨)의 하류 품질을 상관시켜, 각각 추론 및 지식 중심 벤치마크(검증 세트)에 대해 연관성을 찾습니다. 우리는 기준선과 마찬가지로 Switch 모델이 상류 사전 훈련 작업에서의 개선과 함께 확장된다는 것을 발견했습니다. SuperGLUE의 경우, 음의 로그 혼란도(negative log perplexity)와 평균 SuperGLUE 점수 사이에 느슨한 선형 관계를 찾았습니다. 그러나, 고정된 혼란도에 대해, 특히 대규모 체제에서는 밀집 모델(dense model)이 종종 더 나은 성능을 보입니다. 반대로, 지식 중심 작업인 TriviaQA에서는 Switch Transformer가 개선된 확장 관계를 따를 수 있음을 발견했습니다 - 주어진 상류 혼란도에 대해, 밀집 대응 모델보다 더 나은 성능을 보입니다. 이러한 관찰을 확인하기 위해서는 추가 통계(수집 비용이 많이 들고 향후 작업으로 남겨진)가 필요합니다.

우리는 일관된 상관관계를 발견했으며, 이는 기준선과 Switch 모델 모두에서 개선된 사전 훈련이 더 나은 하류 결과로 이어진다는 것을 나타냅니다. 또한, 고정된 상류 혼란도에서는 Switch와 밀집 모델이 소형에서 중형 모델 크기 체제에서 비슷한 성능을 보인다는 것을 발견했습니다. 그러나, 가장 큰 모델 체제(T5-11B/T5-XXL)에서는, 5.6절에서 언급한 바와 같이, 우리의 가장 큰 Switch 모델은 항상 상류 혼란도를 SuperGLUE 작업에 대한 하류 미세 조정(fine-tuning)으로 잘 전환하지 못합니다. 이는 희소 모델(sparse models)의 잠재력을 완전히 실현하기 위해 향후 조사와 연구가 필요함을 시사합니다. 전문가 모델(expert-models)과 함께 미세 조정 동역학을 이해하는 것은 매우 복잡하며, 정규화(regularization), 부하 균형(load-balancing), 미세 조정 하이퍼파라미터(fine-tuning hyper-parameters)에 의존합니다.

F. 스위치 트랜스포머를 위한 의사 코드(Pseudo Code)

Mesh Tensorflow(Shazeer et al., 2018)에서 스위치 트랜스포머를 위한 의사 코드입니다. 아래 코드에는 모델 병렬 처리(model parallelism)가 사용되지 않습니다(자세한 내용은 5.4절 참조).

그림 14: 메쉬 텐서플로우(Mesh Tensorflow)에서 스위치 트랜스포머(Switch Transformers)를 위한 로드 밸런스 손실(load balance loss)의 의사 코드.

그림 15: 메쉬 텐서플로우(Mesh Tensorflow)에서 스위치 트랜스포머(Switch Transformers)를 위한 라우터(router)의 의사 코드.

그림 16: 메쉬 텐서플로우(Mesh Tensorflow)에서 스위치 트랜스포머(Switch Transformer) 레이어의 의사 코드.