NL-342, Switch Transformers: Scaling to Trillion Parameter Models with Simple and Efficient Sparsity, JMLR 2022
◼ Comment
- MoE 초기 논문 이후로, 그 다음 MoE 국밥 논문인듯
- 내가 느끼기엔 근데 방법론 자체는 비슷한데, Transformer 구조에 적용하고 다양한 실험을 해본 것
- 전체 구조는 그림 2와 같고, Transformer에서 FFN layer을 보통 모든 입력에 대해 동일한 layer을 사용하는데
- 여기서는, 입력 토큰별로 다른 FFN layer을 쓴다는 것이다.
- 그럼 어떤 FFN layer을 사용하는지를, Router을 둬서 사용하는것이고, 이 방법은 MoE 방법과 동일하다
- 단, 몇개의 FFN layer을 사용할지 top-k 을 선택해야하는데 기존에는 2개 이상을 쓰는걸로 주로 실험되었지만, 여기서 해보니 top-1을(expert 1개만 사용) 써도 성능 잘나오더라
- 학습 효율
- top-1 expert을 사용하게끔 학습할때, 입력이 하나의 expert로 쏠리는 학습 불안정 현상이 있다는거 같음
- 그래서 이를 어느정도 균등하게 하기 위한 노력이 있는데,
- 1) 식 3처럼, expert당 처리할 수 있는 입력 토큰 개수를 정해두고, 이 제한 입력이 꽉차면, 이 expert로 들어오면 FFN layer의 출력을 0 처리해서 residual 로 바꿔버림.
- 이를 dropped token이라 부름
- 2) 식 4,5,6처럼 penalty loss 개념을 도입한다.
- 이 Loss는 균등하게 Expert로 분배되면 낮아지는 수식이다
- 즉 배치에 구성된 토큰들이 만약 한쪽 expert로 쏠리면, f, P 값이 커지면서 loss가 증가한다
- 여기서 f는 실제 top-1 배분 비율, P는 router가 준 확률의 평균
- 실험 결과는 그냥 스크리닝만 했는데, 기존 dense 스타일인 T5랑 MoE 모델보다 같은 토큰 학습 steps 기준 좋다라고 주장하는거 같음
- 전반적으로 sparse expert는 학습 불안정성이 있지만, 이거를 bfloat16으로 해결하려고 했던거 같고
- 모델 스케일링 실험
- expert을 늘리면 손쉽게 capacity 늘리고 성능도 좋아지지만, 한계가 있음. 너무 많아지면 오히려 안좋아지는듯..?
- 따라서 데이터랑 모델 사이즈도 같이 키우는걸 고려해야한다고 함
- 다운 스트림 실험에서도, 대체적으로 성능이 좋다고 함
Abstract
딥러닝에서 모델은 일반적으로 모든 입력에 대해 동일한 파라미터를 재사용한다. Mixture of Experts(MoE) 모델은 이와 달리 각 입력 예제마다 서로 다른 파라미터를 선택한다. 그 결과, 계산 비용은 일정하게 유지하면서도 엄청나게 많은 파라미터를 가진 희소 활성화 모델이 된다. 그러나 MoE는 여러 주목할 만한 성공에도 불구하고, 복잡성, 통신 비용, 그리고 학습 불안정성 때문에 널리 채택되지 못했다. 우리는 Switch Transformer를 도입하여 이러한 문제를 해결한다. 우리는 MoE 라우팅 알고리즘을 단순화하고, 통신 및 계산 비용을 줄인 직관적인 개선 모델을 설계한다. 또한 제안한 학습 기법은 이러한 불안정성을 완화하며, 대규모 희소 모델을 저정밀도(bfloat16) 형식으로 처음으로 학습할 수 있음을 보인다. 우리는 T5-Base와 T5-Large(Raffel et al., 2019)를 바탕으로 모델을 설계하여, 동일한 계산 자원으로 사전학습 속도를 최대 7배까지 향상시킨다. 이러한 개선은 다국어 설정에도 확장되어, mT5-Base를 기준으로 101개 모든 언어에서 성능 향상을 측정하였다. 마지막으로, “Colossal Clean Crawled Corpus”에서 최대 1조 파라미터 규모의 모델을 사전학습하여 언어 모델의 현재 규모를 한 단계 확장하였고, T5-XXL 모델 대비 4배의 속도 향상을 달성하였다.
1. Introduction
대규모 학습은 유연하고 강력한 신경 언어 모델을 만드는 효과적인 경로였다(Radford et al., 2018; Kaplan et al., 2020; Brown et al., 2020). 충분한 계산 자원, 데이터셋 크기, 파라미터 수를 바탕으로 한 단순한 아키텍처는 더 복잡한 알고리즘을 능가한다(Sutton, 2019). Radford et al. (2018), Raffel et al. (2019), Brown et al. (2020)이 따랐던 접근은 조밀하게 활성화되는 Transformer(Vaswani et al., 2017)의 모델 크기를 확장하는 것이었다. 이러한 방식은 효과적이지만, 계산 비용이 매우 크다(Strubell et al., 2019). 우리는 모델 규모의 성공에서 영감을 받았지만 더 높은 계산 효율을 추구하여, 희소하게 활성화되는 전문가 모델인 Switch Transformer를 제안한다. 이 경우 희소성은 각 입력 예제마다 신경망 가중치의 일부만 활성화한다는 데서 비롯된다.
희소 학습은 활발한 연구 및 공학 분야이지만(Gray et al., 2017; Gale et al., 2020), 현재까지 머신러닝 라이브러리와 하드웨어 가속기는 여전히 조밀한 행렬곱을 중심으로 설계되어 있다. 효율적인 희소 알고리즘을 만들기 위해, 우리는 Mixture-of-Experts(MoE) 패러다임(Jacobs et al., 1991; Jordan and Jacobs, 1994; Shazeer et al., 2017)에서 출발해 이를 단순화함으로써 학습 안정성과 계산상의 이점을 얻는다. MoE 모델은 기계번역에서 주목할 만한 성공을 거두었지만(Shazeer et al., 2017, 2018; Lepikhin et al., 2020), 복잡성, 통신 비용, 학습 불안정성 때문에 널리 채택되지 못했다.
우리는 이러한 문제를 해결하고, 나아가 번역을 넘어 이러한 알고리즘 계열이 자연어 전반에서 폭넓게 유용하다는 것을 보인다. 우리는 다양한 자연어 작업과 NLP의 세 가지 학습 단계—사전학습, 미세조정, 멀티태스크 학습—에서 더 우수한 스케일링을 측정한다. 또한 이 연구는 스케일에 초점을 두지만, Switch Transformer 아키텍처가 슈퍼컴퓨터 환경에서만 뛰어난 것이 아니라, 몇 개의 계산 코어만으로도 유용함을 보인다. 더 나아가, 우리의 큰 희소 모델은 작은 조밀 모델로 증류될 수 있으며, 희소 모델의 성능 향상 중 30%를 유지하면서도 작은 모델로 압축할 수 있다.
우리의 기여는 다음과 같다.
- Mixture of Experts를 단순화하고 개선한 Switch Transformer 아키텍처.
- 강하게 튜닝된 T5 모델(Raffel et al., 2019)과의 비교를 통한 스케일링 특성 및 벤치마크. 동일한 토큰당 FLOPS를 사용하면서도 사전학습 속도를 7배 이상 향상시켰다. 또한 전문가 수가 두 개뿐인 제한된 계산 자원 환경에서도 이러한 개선이 유지됨을 보였다.
- 큰 희소 사전학습 모델과 특화된 미세조정 모델을 작은 조밀 모델로 성공적으로 증류. 모델 크기를 최대 99% 줄이면서도 큰 희소 교사 모델의 품질 향상 중 30%를 보존했다.
- 개선된 사전학습 및 미세조정 기법: (1) 낮은 bfloat16 정밀도에서도 학습을 가능하게 하는 선택적 정밀도 학습, (2) 더 많은 전문가 수로의 확장을 가능하게 하는 초기화 방식, (3) 희소 모델의 미세조정과 멀티태스크 학습을 개선하는 강화된 전문가 정규화.
- 다국어 데이터에서의 사전학습 이점 측정. 모든 101개 언어에서 보편적인 향상을 확인했으며, 그중 91%의 언어에서 mT5 기준선 대비 4배 이상의 속도 향상을 얻었다.
- 데이터, 모델, 전문가 병렬성을 효율적으로 결합하여 최대 1조 개 파라미터를 가진 모델을 만들고, 언어 모델의 규모를 한 단계 끌어올림. 이 모델들은 강하게 튜닝된 T5-XXL 기준선보다 사전학습 속도가 4배 빠르다.
2. Switch Transformer
Switch Transformer의 설계 원리는 Transformer 모델의 파라미터 수를 단순하면서도 계산적으로 효율적인 방식으로 최대화하는 것이다. 스케일의 이점은 Kaplan et al. (2020)에서 철저히 연구되었으며, 그 결과 모델 크기, 데이터셋 크기, 계산 예산에 대해 거듭제곱 법칙 스케일링이 나타난다는 것이 밝혀졌다. 중요한 점은, 이 연구가 계산량 기준으로 볼 때 상대적으로 적은 양의 데이터로 큰 모델을 학습하는 것이 계산적으로 가장 최적이라는 점을 주장한다.
이러한 결과를 바탕으로, 우리는 네 번째 축을 살펴본다. 즉, 예제당 부동소수점 연산(FLOPs)은 일정하게 유지하면서 파라미터 수를 늘리는 것이다. 우리의 가설은 총 계산량과는 독립적으로 파라미터 수 자체가 확장에 있어 별도로 중요한 축이라는 점이다. 우리는 GPU와 TPU처럼 조밀한 행렬곱을 위해 설계된 하드웨어를 효율적으로 사용하는 희소 활성화 모델을 설계함으로써 이를 달성한다. 여기서는 TPU 아키텍처에 초점을 맞추지만, 이러한 종류의 모델은 GPU 클러스터에서도 유사하게 학습될 수 있다. 우리의 분산 학습 설정에서는 희소 활성화 레이어가 고유한 가중치를 서로 다른 디바이스에 분할한다. 따라서 각 디바이스가 감당해야 하는 메모리와 계산 부담을 관리 가능한 수준으로 유지하면서, 모델의 가중치 수는 디바이스 수와 함께 증가한다.
2.1 Sparse Routing 단순화
Mixture of Expert Routing. Shazeer et al. (2017)은 자연어용 Mixture-of-Experts(MoE) 레이어를 제안했는데, 이는 토큰 표현 (x)를 입력으로 받아 (N)개의 전문가 ({E_i(x)}_{i=1}^{N}) 중에서 가장 적합하다고 판단되는 top-(k) 전문가들로 라우팅한다. 라우터 변수 (W_r)는 로짓 (h(x)=W_r \cdot x)를 생성하며, 이는 해당 레이어에서 사용 가능한 (N)개 전문가에 대해 softmax 분포로 정규화된다. 전문가 (i)에 대한 게이트 값은 다음과 같다.
top-(k) 게이트 값이 토큰 (x)의 라우팅에 사용된다. 선택된 top-(k) 인덱스들의 집합을 (T)라고 하면, 레이어의 출력은 각 전문가가 토큰에 대해 계산한 결과를 게이트 값으로 선형 가중합한 값이다.
Switch Routing: Mixture-of-Experts 다시 생각하기. Shazeer et al. (2017)은 라우팅 함수에 대해 비자명한 그래디언트를 얻기 위해서는 (k>1)개의 전문가로 라우팅해야 한다고 추측했다. 저자들은 적어도 두 전문가를 비교할 수 있어야만 라우팅을 학습할 수 있을 것이라고 직관적으로 보았다. Ramachandran and Le (2018)은 더 나아가 top-(k) 결정을 연구했고, 많은 라우팅 레이어를 가진 모델에서는 낮은 레이어에서 더 큰 (k) 값이 중요하다는 사실을 발견했다.
그러나 우리는 이와 반대로, 단 하나의 전문가로만 라우팅하는 단순화된 전략을 사용한다. 우리는 이 단순화가 모델 품질을 유지하고, 라우팅 계산을 줄이며, 더 나은 성능을 낸다는 것을 보인다. 이 (k=1) 라우팅 전략을 이후 Switch layer라고 부른다. MoE Routing과 Switch Routing 모두에서, 식 (2)의 게이트 값 (p_i(x))는 라우터의 미분 가능성을 보장한다.
Switch layer의 장점은 세 가지다.
- (1) 단 하나의 전문가로만 토큰을 라우팅하므로 라우터 계산량이 줄어든다.
- (2) 각 토큰이 하나의 전문가에만 할당되므로 각 전문가의 배치 크기(전문가 용량, expert capacity)를 적어도 절반으로 줄일 수 있다.
- (3) 라우팅 구현이 단순해지고 통신 비용도 감소한다. 그림 3은 서로 다른 expert capacity factor에서의 라우팅 예시를 보여준다.
- 토큰당 하나의 expert만 사용하겠다는 것
- 모델 구조는 Transformer랑 똑같은데 FFN layer을 토큰에 따라 선택되는 방식이지? FFN layer의 후보들이 expert라 보는 것이고
2.2 Efficient Sparse Routing
우리는 Mesh-Tensorflow(MTF)(Shazeer et al., 2018)를 사용한다. MTF는 Tensorflow(Abadi et al., 2016)와 유사한 의미론과 API를 제공하는 라이브러리로, 효율적인 분산 데이터 병렬 및 모델 병렬 아키텍처를 가능하게 한다. 이는 물리적인 코어 집합을 논리적인 프로세서 메쉬로 추상화함으로써 이루어진다. 텐서와 연산은 이름이 붙은 차원별로 분할될 수 있어, 모델을 여러 차원에 걸쳐 쉽게 분할할 수 있다. 우리는 정적으로 크기가 선언되어야 하는 TPU를 염두에 두고 모델을 설계했다. 아래에서는 우리의 분산 Switch Transformer 구현을 설명한다.
Distributed Switch Implementation.
우리의 모든 텐서 shape는 컴파일 시점에 정적으로 결정되지만, 학습 및 추론 시의 라우팅 결정 때문에 연산은 동적이다. 따라서 중요한 기술적 고려사항 중 하나는 expert capacity를 어떻게 설정하느냐이다. expert capacity, 즉 각 expert가 처리하는 토큰 수는 배치의 토큰 수를 expert 수로 균등하게 나눈 뒤, 여기에 capacity factor를 곱해서 정한다.
capacity factor가 1.0보다 크면, 토큰이 expert들 사이에 완벽하게 균형 있게 분배되지 않을 때를 대비한 추가 버퍼가 생긴다. 만약 어떤 expert로 너무 많은 토큰이 라우팅되면(이후 dropped tokens라고 부름), 그 토큰들은 계산이 건너뛰어지고 residual connection을 통해 다음 레이어로 바로 전달된다. 그러나 expert capacity를 늘리는 데는 단점도 있다. 값이 너무 크면 계산과 메모리가 낭비되기 때문이다. 이 트레이드오프는 Figure 3에 설명되어 있다. 경험적으로 우리는 dropped token 비율을 낮게 유지하는 것이 sparse expert model의 스케일링에 중요하다는 것을 발견했다. 실험 전반에서 dropped token 수가 expert 수에 의존하는 경향은 보이지 않았으며(보통 1% 미만), 뒤에서 설명할 auxiliary load balancing loss를 충분히 큰 계수로 사용하면 좋은 load balancing이 가능했다. 이러한 설계 결정들이 모델 품질과 속도에 미치는 영향은 Table 1에서 살펴본다.
가령 한 배치에 1024 토큰이 있고, 모델 설계상 expert 수가 8개라고 하자. 그러면 expert당 기준 토큰 수는 먼저
이 되고, 여기에 capacity factor를 곱해서 각 expert의 최대 수용 토큰 수를 정한다. 그래서
- capacity factor = 1.0이면 expert capacity는 128
- capacity factor = 1.5이면 expert capacity는 192
A Differentiable Load Balancing Loss.
expert들 사이에 균형 잡힌 부하를 유도하기 위해 auxiliary loss를 추가한다(Shazeer et al., 2017, 2018; Lepikhin et al., 2020). Shazeer et al. (2018), Lepikhin et al. (2020)과 마찬가지로, Switch Transformer는 Shazeer et al. (2017)의 원래 설계를 단순화했는데, 원래는 load-balancing loss와 importance-weighting loss가 분리되어 있었다. 각 Switch layer에 대해 이 auxiliary loss는 학습 중 전체 모델 loss에 더해진다.
(i = 1)부터 (N)까지 인덱싱된 (N)개의 expert와, (T)개의 토큰을 가진 배치 (B)가 주어졌을 때, auxiliary loss는 벡터 (f)와 (P)의 scaled dot-product로 계산된다.
여기서 (f_i)는 expert (i)로 dispatch된 토큰의 비율이다.
그리고 (P_i)는 expert (i)에 할당된 router probability의 비율이다.
우리는 (N)개의 expert 전체에 토큰이 균등하게 라우팅되기를 원하므로, 두 벡터 모두 값이 (1/N)이 되길 바란다.
- 식 (4)의 auxiliary loss는 균등 분포에서 최소가 되므로 균등 라우팅을 유도한다.
- 이 목적 함수는 (P)-vector는 미분 가능하지만 (f)-vector는 미분 가능하지 않다는 점에서도 구분된다.
- 최종 loss는 expert 수 (N)을 곱해, expert 수가 변해도 loss의 크기가 일정하게 유지되도록 한다. 균등 라우팅에서는
[
\sum_{i=1}^{N} (f_i \cdot P_i) = \sum_{i=1}^{N} \left(\frac{1}{N}\cdot\frac{1}{N}\right) = \frac{1}{N}
]
이기 때문이다. 마지막으로, (\alpha)는 이러한 auxiliary loss의 곱셈 계수이다. 본 연구에서는 (\alpha = 10^{-2})를 사용했으며, 이는 load balancing을 보장하기에는 충분히 크면서도 주된 cross-entropy objective를 압도하지 않을 만큼 충분히 작은 값이었다. 우리는 (\alpha)를 (10^{-1})부터 (10^{-5})까지 10의 거듭제곱 단위로 탐색했고, (10^{-2})가 학습 loss에 간섭하지 않으면서 빠르게 부하를 균형 있게 만든다는 것을 확인했다.
맞아. 식 (4)(5)(6)은 router가 한 expert로 너무 몰리지 않게 만들기 위한 보조 손실이야. 직관적으로는, “실제로 많이 보내진 expert”와 “router가 많이 보내고 싶어 하는 expert”가 둘 다 한쪽으로 치우치지 않도록 벌점을 주는 방식이라고 보면 된다. 논문도 이 loss가 두 벡터를 uniform distribution, 즉 각 expert가 (1/N)씩 받는 상태로 유도한다고 설명한다.
쉽게 예시로 보자. expert가 4개라고 하고, 한 배치에 토큰이 8개 있다고 하자.
router가 이런 식으로 토큰을 보냈다고 해보자:
맞아. 토큰 8개, expert 4개로 딱 맞춰서 다시 보면 이렇게 돼.
먼저 이 논문의 정의는
[
f_i=\frac{1}{T}\sum_{x\in B}\mathbf{1}{\arg\max p(x)=i},\quad
P_i=\frac{1}{T}\sum_{x\in B}p_i(x)
]
이고, 둘 다 배치 전체 (T)개 토큰에 대해 계산한다. (f)는 실제로 top-1으로 간 비율, (P)는 router가 준 확률의 평균이야.
예를 들어 8개 토큰에 대해 router가 아래처럼 확률을 줬다고 하자.
| token | expert1 | expert2 | expert3 | expert4 | top-1 |
|---|---|---|---|---|---|
| 1 | 0.70 | 0.10 | 0.10 | 0.10 | 1 |
| 2 | 0.60 | 0.20 | 0.10 | 0.10 | 1 |
| 3 | 0.55 | 0.25 | 0.10 | 0.10 | 1 |
| 4 | 0.50 | 0.30 | 0.10 | 0.10 | 1 |
| 5 | 0.20 | 0.50 | 0.20 | 0.10 | 2 |
| 6 | 0.20 | 0.40 | 0.30 | 0.10 | 2 |
| 7 | 0.10 | 0.20 | 0.60 | 0.10 | 3 |
| 8 | 0.10 | 0.15 | 0.10 | 0.65 | 4 |
그러면 실제 배분 비율 (f)는 top-1 개수를 세어서:
[
f=[4/8,;2/8,;1/8,;1/8]=[0.5,;0.25,;0.125,;0.125]
]
가 된다. 이건 실제로 몇 개가 각 expert로 갔는지만 본 값이야.
반면 식 (6)의 (P_i)는 router가 각 expert에 준 확률의 평균이야. 예를 들어 router가 “expert 1이 제일 좋아 보인다”고 계속 높은 확률을 줬다면 (P)도 expert 1 쪽으로 커진다.
이제 (P)는 각 expert의 확률을 8개 토큰에 대해 평균내면 돼:
[
P_1=\frac{0.70+0.60+0.55+0.50+0.20+0.20+0.10+0.10}{8}=0.36875
]
[
P_2=\frac{0.10+0.20+0.25+0.30+0.50+0.40+0.20+0.15}{8}=0.2625
]
[
P_3=\frac{0.10+0.10+0.10+0.10+0.20+0.30+0.60+0.10}{8}=0.2
]
[
P_4=\frac{0.10+0.10+0.10+0.10+0.10+0.10+0.10+0.65}{8}=0.16875
]
즉
[
P=[0.36875,;0.2625,;0.2,;0.16875]
]
이야. 이 값은 **router가 각 expert에 “얼마나 보내고 싶어 했는지”**를 평균낸 거라고 보면 된다.
핵심은 이거야:
(f): 실제 배정 결과
(P): router의 선호도 평균
둘 다 한 expert로 치우치면 식 (4)의 loss가 커져서, router가 더 골고루 보내도록 압력을 받는다. 논문은 이 auxiliary loss를 uniform routing을 유도하는 scaled dot-product로 설명하고, 실제로 (\alpha=10^{-2})를 사용했다고 적고 있다.
식 (4)는 결국
“실제로 많이 간 expert” × “router가 많이 보내고 싶어 한 expert”
가 함께 크면 손해를 주는 형태야. 그래서 expert 1에 토큰이 몰리고 router 확률도 expert 1에 몰리면 loss가 커지고, 반대로 4개 expert에 골고루 퍼지면 loss가 작아진다. 논문에서는 이걸 uniform routing에서 최소화되는 scaled dot-product라고 설명한다.
비유하면 이렇다.
expert 4명이 줄 서 있는데, 한 사람에게만 계속 일감이 몰리면 안 되니까,
식 (5)는 “실제로 누구에게 일이 갔는지”
식 (6)는 “시스템이 누구에게 일을 보내려 했는지”
를 각각 세고,식 (4)는 둘이 한쪽으로 치우칠수록 벌점을 주는 거야.
그래서 이 loss를 학습에 더하면 router가 점점
“한 expert만 고집하지 말고, 비슷한 수준이면 다른 expert도 써라”
쪽으로 유도된다. 논문은 이때 (\alpha = 10^{-2})를 썼고, load balancing에는 충분하면서 메인 loss를 방해하지 않는 값이었다고 적는다.
한 줄로 정리하면:
식 (5)는 “실제 배분”, 식 (6)은 “배분하고 싶은 확률”, 식 (4)는 그 둘이 한 expert로 쏠리면 벌점을 주는 장치야.
2.3 Putting It All Together: The Switch Transformer
Switch Transformer에 대한 우리의 첫 번째 실험은 Raffel et al. (2019)에서 소개된 “Colossal Clean Crawled Corpus”(C4)로 사전학습을 수행하는 것이다. 사전학습 목표로는 masked language modeling task(Taylor, 1953; Fedus et al., 2018; Devlin et al., 2018)를 사용하며, 모델은 누락된 토큰을 예측하도록 학습된다. 사전학습 설정에서는 Raffel et al. (2019)에서 최적이라고 판단한 방식에 따라 전체 토큰의 15%를 dropout하고, 마스킹된 시퀀스를 단일 sentinel token으로 치환한다. 모델들을 비교하기 위해 우리는 negative log perplexity를 기록한다.4 이 논문 전반의 모든 표에서 ↑는 해당 지표가 클수록 좋음을, ↓는 작을수록 좋음을 뜻한다. 이 연구에서 다룬 모든 모델의 비교는 Table 9에 정리되어 있다.
Switch Transformer와 MoE Transformer의 1:1 비교는 Table 1에 제시되어 있다.
우리의 Switch Transformer 모델은 ‘T5-Base’(Raffel et al., 2019)와 FLOP가 맞춰져 있으며(토큰당 동일한 계산량을 사용), MoE Transformer는 top-2 routing을 사용하므로 각 토큰마다 두 개의 expert가 각각 별도의 FFN을 적용하게 되어 FLOPs가 더 크다. 모든 모델은 동일한 하드웨어에서 같은 수의 step 동안 학습되었다. 참고로 MoE 모델은 capacity factor를 2.0에서 1.25로 낮추었을 때, 위 실험 설정에서는 오히려 속도가 느려졌는데(840에서 790으로), 이는 예상 밖의 결과였다.5
Table 1에서 우리는 세 가지 핵심 결과를 강조한다.
- (1) Switch Transformer는 속도-품질 측면에서 세밀하게 튜닝된 dense model과 MoE Transformer 모두를 능가한다. 즉, 고정된 계산량과 wall-clock time에서 Switch Transformer가 가장 좋은 결과를 낸다.
- (2) Switch Transformer는 MoE 대응 모델보다 계산 부담이 더 작다. 만약 Switch Transformer의 크기를 늘려 MoE Transformer와 학습 속도를 맞추면, step 기준으로도 MoE 및 dense model 모두를 능가한다.
- (3) Switch Transformer는 더 낮은 capacity factor(1.0, 1.25)에서 더 좋은 성능을 보인다. 더 작은 expert capacity는 모델 메모리가 매우 부족한 대규모 모델 환경을 나타내며, 이때 capacity factor는 가능한 한 작게 두고 싶어진다.
2.4 Improved Training and Fine-Tuning Techniques
희소한 expert 모델은 일반적인 Transformer보다 학습이 더 어렵다. 각 레이어에서 발생하는 하드 스위칭(hard-switching), 즉 라우팅 결정 때문에 불안정성이 생길 수 있다. 또한 bfloat16(Wang and Kanwar, 2019) 같은 저정밀도 형식은 router의 softmax 계산에서 이러한 문제를 더 악화시킬 수 있다. 우리는 여기서 이러한 학습상의 어려움을 설명하고, 안정적이면서도 확장 가능한 학습을 달성하기 위해 사용한 방법들을 소개한다.
대규모 sparse 모델에서 선택적 정밀도 사용.
모델 불안정성은 효율적인 bfloat16 정밀도로 학습하는 능력을 저해한다. 그 결과 Lepikhin et al. (2020)은 MoE Transformer 전체를 float32 정밀도로 학습했다. 그러나 우리는 모델의 국소적인 일부에서만 선택적으로 float32로 캐스팅하면, float32 텐서를 all-to-all 통신으로 주고받는 큰 비용 없이도 안정성을 얻을 수 있음을 보인다. 이 기법은 모델과 gradient update의 일부를 더 높은 정밀도로 수행하는 현대적 mixed precision 학습 전략과 일치한다. Table 2는 우리의 방법이 bfloat16 학습과 거의 같은 속도를 유지하면서도 float32의 학습 안정성을 제공함을 보여준다.
이를 위해 우리는 router의 입력을 float32 정밀도로 캐스팅한다. router 함수는 토큰을 입력으로 받아 expert 계산을 선택하고 다시 조합하는 데 쓰이는 dispatch 텐서와 combine 텐서를 생성한다. 중요한 점은, float32 정밀도는 오직 router 함수 내부, 즉 해당 디바이스에서의 로컬 계산에만 사용된다는 것이다. 함수가 끝날 때 dispatch 텐서와 combine 텐서는 다시 bfloat16으로 캐스팅되므로, expensive한 float32 텐서들이 all-to-all 통신 연산을 통해 전파되지 않는다. 대신 float32가 주는 안정성의 이점은 그대로 얻을 수 있다.
안정성을 위한 더 작은 초기화.
적절한 초기화는 딥러닝 학습 성공에 매우 중요하며, Switch Transformer에서는 특히 더 그렇다. 우리는 가중치 행렬의 원소를 평균 (\mu=0), 표준편차 (\sigma=\sqrt{s/n})인 truncated normal 분포에서 샘플링하여 초기화한다. 여기서 (s)는 scale 하이퍼파라미터이고, (n)은 가중치 텐서의 입력 유닛 수(fan-in)이다. 추가적인 불안정성 완화책으로, 기본 Transformer 초기화 scale인 (s=1.0)을 10분의 1로 줄일 것을 권장한다. 이는 실험에서 품질을 향상시키고 학습 불안정 가능성을 낮춘다.
Table 3은 초기 학습 초기에 모델 품질의 향상과 분산 감소를 보여준다. 평균 품질(negative log perplexity)은 크게 개선되고, 여러 run 간의 분산도 훨씬 줄어든다. 이 초기화 방식은 여러 자릿수 규모에 걸친 모델에도 폭넓게 유효하며, 우리는 이를 사용해 223M 파라미터급 baseline부터 1조 파라미터를 넘는 거대 모델까지 안정적으로 학습했다.
대규모 sparse 모델의 정규화.
이 논문은 큰 코퍼스로 사전학습한 뒤, 요약이나 질의응답 같은 작은 downstream task에 fine-tuning 하는 일반적인 NLP 절차를 다룬다. 여기서 자연스럽게 생기는 문제는 과적합이다. 왜냐하면 많은 fine-tuning task는 예제가 매우 적기 때문이다. 표준 Transformer의 fine-tuning에서는 Raffel et al. (2019)이 각 레이어에 dropout을 사용해 과적합을 막았다. 그런데 우리의 Switch Transformer는 FLOP가 맞춰진 dense baseline보다 파라미터 수가 훨씬 많아서, 이런 작은 downstream task에서는 과적합이 더 심할 수 있다.
이 문제를 완화하기 위해 우리는 간단한 방법을 제안한다. expert 내부의 dropout을 더 크게 주는 것이다. 이를 expert dropout이라고 부른다. fine-tuning 시에는 각 expert layer의 중간 feed-forward 계산에서만 dropout rate를 크게 올린다.
Table 4는 이 protocol의 결과를 보여준다. 모든 레이어의 dropout을 단순히 올리는 것은 성능을 악화시킨다. 하지만 non-expert layer에는 작은 dropout rate(0.1)를 유지하고, expert layer에는 훨씬 큰 dropout rate(0.4)를 주면 네 개의 작은 downstream task에서 성능 향상을 얻는다.
3. Scaling Properties 결론만 정리하면:
응, 더 정확하게 말하면 이 섹션의 결론은 “Switch Transformer가 같은 계산 예산에서 dense baseline보다 더 효율적이고, MoE보다도 단순하고 효율적인 편이지만, expert 수는 무조건 많을수록 좋은 게 아니라 적당한 지점이 있다”야.
조금 풀어서 보면, 논문은 expert 수를 늘리는 것이 가장 효율적인 스케일링 축이라고 말해. 이유는 토큰당 FLOPs는 거의 고정된 채로, 토큰마다 expert 하나만 고르기 때문이야. router가 더 많은 expert에 대한 확률만 계산하면 되는데, 그 비용은 가볍다고 설명한다. 그래서 같은 compute budget에서 expert 수를 늘리면 test loss가 내려가고 샘플 효율이 좋아진다고 본다.
비교를 더 명확히 하면, 이 논문 실험에서 Switch Transformer는 동일한 계산량/FLOPs 기준으로 dense T5보다 더 빠르고 더 좋은 성능을 보였다. 예를 들어 Switch-Base 64 experts는 T5-Base와 비슷한 성능에 도달하는 데 step 기준 약 7.5배 빠르고, wall-clock 기준으로도 큰 속도 이득을 보였다. 또 T5-Large와 비교해도 Switch-Base가 더 sample-efficient하고 약 2.5배 빠르다고 보고한다.
다만 네가 말한 두 번째 문장도 맞아. expert를 무한히 늘린다고 계속 좋아지는 것은 아니고, diminishing returns가 있다고 논문이 직접 말한다. 그래서 “많을수록 무조건 좋다”가 아니라, 계산/메모리/통신 비용을 고려해 적당한 expert 수를 잡는 게 중요하다는 결론이야. 논문도 작은 규모에서는 2, 4, 8 experts만으로도 dense baseline보다 좋아진다고 보여준다.
정리하면 이렇다.
Switch는 dense보다 효율적: 같은 FLOPs나 같은 시간 기준에서 더 좋은 결과를 내는 실험이 나왔다.
MoE보다도 단순하고 효율적으로 만들었다: top-1 routing으로 계산/통신을 줄였다.
expert 수는 많을수록 무조건 좋지 않다: 성능은 좋아지지만 diminishing returns가 있고, 낮은 expert 수에서도 충분히 이득이 있다.

4. Downstream Results 핵심
이 섹션의 결론은, Switch Transformer의 사전학습 이득이 실제 downstream 작업에도 상당 부분 이어진다는 거야. 다만 항상 100% 깔끔하게 전이되지는 않고, 특히 가장 큰 모델 구간이나 일부 추론형 과제에서는 예외가 보인다.
먼저 fine-tuning에서는, FLOP을 맞춘 T5 대비 Switch 변형이 대부분의 과제에서 더 좋았다. 특히 SuperGLUE에서 T5-Base보다 4.4점, T5-Large보다 2점 정도 개선이 있었고, Winogrande, closed-book TriviaQA, XSum에서도 큰 향상이 있었다. 반대로 ARC 계열은 예외가 있어서, 여기서는 T5가 더 나은 경우가 있었다. 즉, reasoning/knowledge-heavy 과제 전반에서 개선되지만, 전 과제가 다 좋아지는 것은 아니다라는 결론이야.
둘째, distillation에서는 아주 큰 sparse 모델을 작은 dense 모델로 꽤 잘 압축할 수 있었다. 기본적으로 모델 크기를 10~100배 줄이는 압축이 가능했고, 최적 방법을 쓰면 teacher의 성능 향상분 중 약 30%를 유지할 수 있었다. 예를 들어 Switch-Base를 T5-Base로 distill했을 때, 파라미터는 크게 줄이면서도 성능 일부를 보존했다. fine-tuned SuperGLUE 모델도 distillation이 가능했고, 97% 압축된 모델에서도 teacher 향상의 약 30%를 남겼다.
셋째, multilingual learning에서도 Switch가 강했다. 101개 언어 전부에서 dense mT5-Base보다 좋았고, 평균적으로 step 기준 5배 속도 향상, 그리고 91%의 언어에서 4배 이상 speedup을 보였다. 즉, 이 모델은 단일 언어뿐 아니라 멀티태스크·다국어 환경에서도 효율적이라는 걸 보여줬다.
마지막으로 이 섹션을 한 줄로 묶으면:
Switch Transformer는 pretraining에서 얻은 효율을 downstream fine-tuning, 압축, 다국어 학습으로 꽤 잘 이어갔고, 특히 SuperGLUE·Winogrande·TriviaQA·XSum·다국어에서 강했지만, ARC 같은 일부 과제와 아주 큰 스케일에서는 전이가 덜 안정적이었다.
5. Designing Models with Data, Model, and Expert-Parallelism
이 섹션의 주제는 “스위치 모델을 더 크게 만들려면, 단순히 expert 수만 늘리는 것 말고도 data parallelism, model parallelism, expert parallelism을 어떻게 조합할지 설계해야 한다”는 거야. 그리고 결론은 expert 수를 늘리는 건 가장 효율적인 축이지만, 어느 순간부터는 수익이 줄어들어서 다른 병렬화 전략과 같이 써야 한다는 거다.
먼저 data parallelism은 가장 단순한 방식이야. batch를 여러 core로 나눠서 각 core가 같은 모델을 복제해 계산하고, 마지막에 gradient만 모은다. 통신이 거의 없어서 단순하지만, 모델 크기를 키우는 데는 한계가 있다.
model parallelism은 반대로 가중치 자체를 core들에 나눠서 저장하는 방식이야. 메모리는 절약되지만, forward/backward마다 core 간 통신이 들어간다. 즉, 모델을 키울 수는 있지만 통신 비용이 늘어난다.
그다음 model + data parallelism은 이 둘을 섞는 방식이다. 큰 T5나 GPT-3 같은 모델이 이 방법을 썼고, 각 core는 batch의 일부와 가중치의 일부를 함께 담당한다. 큰 모델을 학습하려면 사실상 이런 조합이 필요하다는 뜻이야.
Switch Transformer에서 중요한 건 expert + data parallelism이야. 이 논문에서는 각 core가 expert 하나씩 맡는 형태로 분할하고, router가 각 토큰을 로컬에서 expert에 배정한다. 그다음 expert가 다른 core에 흩어져 있으므로 all-to-all communication이 일어난다. 핵심은 expert 수를 늘려도 토큰당 FLOPs는 거의 안 늘고, 파라미터 수만 늘릴 수 있다는 점이야.
하지만 expert 수를 늘린다고 무한정 좋은 것은 아니고, 이제는 FLOPs를 늘리려면 (d_{ff}) 같은 모델 폭도 키워야 하고, 그러면 메모리가 부족해져서 model parallelism까지 같이 써야 한다. 이때는 통신 비용, 메모리, FLOPs의 균형을 맞춰야 해서 최적의 배치는 이론보다 실험적으로 결정된다고 설명한다.
마지막으로 이 섹션의 하이라이트는 trillion-parameter 모델이야. 이들은 395B 파라미터의 Switch-XXL과 1.571T 파라미터의 Switch-C를 만들었고, 둘 다 C4 사전학습에서 T5-XXL보다 더 좋은 perplexity를 보였다. 특히 Switch-C는 expert-parallelism만 사용한 설계이고, Switch-XXL은 T5-XXL과 FLOP를 맞춘 버전이야. 또 Switch-XXL은 250k/500k step에서 T5-XXL보다 더 좋은 성능을 보였지만, 큰 sparse 모델은 여전히 학습 불안정성이 있을 수 있다고 적고 있다.
한 줄로 요약하면:
이 섹션은 “Switch는 expert 수를 늘려 파라미터를 크게 키우되, 실제로는 data/model/expert parallelism을 조합해서 메모리와 통신을 맞춰야 한다”는 설계 원리를 보여주고, 그 결과 395B~1.6T급 모델까지 확장 가능함을 보인 부분이야.
6. Related Work
신경망에서 스케일의 중요성은 널리 인식되어 왔고, 이를 위해 여러 접근이 제안되었다. 최근 연구들은 모델 병렬화(예: 가중치와 텐서를 여러 코어에 분할)를 이용해 수십억 파라미터 규모까지 모델을 확장했다(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)은 파이프라인 기반 모델 병렬화를 제안했는데, 여기서는 서로 다른 레이어를 디바이스에 나누고 마이크로 배치를 각 레이어로 파이프라인처럼 흘려보낸다. 마지막으로 Product Key networks (Lample et al., 2019)는 입력 토큰 표현을 바탕으로 학습 가능한 임베딩을 조회하는 방식으로 신경망의 용량을 확장하도록 제안되었다.
우리의 연구는 조건부 계산(conditional computation)을 수행하는 한 종류의 방법을 다룬다. 여기서 계산 결정은 입력에 따라 동적으로 이루어진다. Cho and Bengio (2014)는 모델의 은닉 상태에 나타나는 특정 비트 패턴에 따라 가중치를 적응적으로 선택하는 방법을 제안했다. Eigen et al. (2013)은 밀집 행렬곱과 ReLU 활성화를 사용하는 expert 층을 여러 층 쌓아, jittered MNIST와 monotone speech에서 유망한 결과를 보였다. 컴퓨터 비전 분야에서는 Puigcerver et al. (2020)이 상위 단계 사전학습 동안 의미 클래스에 따라 토큰을 수동으로 라우팅한 뒤, downstream task에 맞는 expert를 선택하는 방식을 사용했다.
현대 딥러닝 아키텍처 맥락에서 Mixture of Experts(MoE)는 Shazeer et al. (2017)에서 효과적인 것으로 입증되었다. 그 연구는 LSTM(Hochreiter and Schmidhuber, 1997) 층 사이에 MoE 층을 추가했고, 토큰을 여러 expert 조합으로 개별 라우팅했다. 그 결과 언어모델링과 기계번역 벤치마크에서 state-of-the-art 성능을 달성했다. 이후 Mesh TensorFlow 라이브러리(Shazeer et al., 2018)에 의해 MoE 층이 Transformer의 FFN 층 대체물로 다시 도입되었지만, 당시에는 그에 상응하는 NLP 결과는 없었다. 더 최근에는 머신러닝 인프라의 발전을 바탕으로 GShard(Lepikhin et al., 2020)가 XLA 컴파일러를 확장하고 MoE Transformer를 사용해 100개 언어 기계번역 성능을 크게 향상시켰다. 또한 Fan et al. (2021)은 모델 파라미터를 겹치지 않는 언어 그룹으로 나누는 다른 결정적 MoE 전략을 선택했다.
Transformer의 attention 패턴에서 sequence length 차원 (L)에 대한 sparsity는 attention 복잡도를 (O(L^2))에서 줄이는 데 성공적인 기법이었다(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. Discussion
우리는 Switch Transformer와 일반적인 sparse expert model에 대해 몇 가지 질문을 던지고 논의한다. 여기서 sparsity는 attention 패턴이 아니라 가중치(weights) 에 적용되는 것을 뜻한다.
Switch Transformer는 단지 파라미터 수가 많아서 더 좋은 것 아닌가?
그렇다, 그리고 그것은 의도된 설계다. 총 FLOPs와는 독립적으로 파라미터 수 자체는 신경 언어 모델을 확장하는 데 유용한 축이다. 큰 모델이 더 잘 작동한다는 것은 이미 충분히 입증되었다(Kaplan et al., 2020). 하지만 이 경우 우리의 모델은 같은 계산 자원을 사용하면서도 더 sample efficient하고 더 빠르다.
슈퍼컴퓨터가 없는데도 이게 유용한가?
이 연구는 매우 큰 모델에 초점을 맞추었지만, expert가 2개만 있어도 성능이 향상되며, 일반적인 GPU나 TPU의 메모리 제약 안에도 쉽게 들어간다는 점을 확인했다(자세한 내용은 Appendix D). 따라서 이 기법은 소규모 설정에서도 유용하다고 본다.
Sparse model이 speed-accuracy Pareto curve에서 dense model보다 낫나?
그렇다. 다양한 모델 크기에서, sparse model은 step 기준으로도, wall-clock 기준으로도 dense model보다 더 좋은 성능을 보였다. 통제된 실험에서, 같은 계산량과 같은 시간을 기준으로 sparse model이 dense model보다 우수했다.
1조 파라미터 모델을 배포할 수 없다면, 이 모델을 줄일 수 있나?
모델 품질을 완전히 그대로 보존할 수는 없지만, sparse model을 dense model로 distillation 하면 10배에서 100배까지 압축할 수 있고, 동시에 expert model이 얻은 품질 향상분의 **약 30%**를 유지할 수 있다.
모델 병렬 dense model 대신 Switch Transformer를 쓰는 이유는?
시간 기준으로 보면, Switch Transformer는 파라미터를 분할한 dense model보다 훨씬 효율적일 수 있다(Figure 6). 또한 이 둘은 배타적인 선택이 아니다. Switch Transformer 안에서도 model-parallelism을 사용할 수 있으며, 그 경우 token당 FLOPs는 늘지만 전통적인 model parallelism의 속도 저하도 함께 따른다.
왜 sparse model이 이미 널리 쓰이지 않았나?
dense model을 확장하는 방식이 매우 성공적이었기 때문에 sparse model을 시도하려는 동기가 약해졌다. 이 성공은 부분적으로 딥러닝 하드웨어와의 상호 적응에 의해 강화되었다(Hooker, 2020). 또한 sparse model은 (1) 모델 복잡성, (2) 학습 난이도, (3) 통신 비용 같은 문제를 안고 있었다. Switch Transformer는 이 문제들을 완화하는 데 진전을 보인다.
한 문장으로 요약하면, 이 섹션은 Switch Transformer가 단순히 파라미터만 큰 모델이 아니라, 같은 계산 자원에서 dense/model-parallel baseline보다 효율적이고, 작은 expert 수에서도 쓸모가 있으며, sparse model의 복잡도·학습 안정성·통신 비용 문제를 상당 부분 줄였다는 점을 정리한다.
8. Future Work
이 논문은 단순화된 아키텍처, 개선된 학습 절차, 그리고 sparse model이 어떻게 스케일하는지에 대한 연구를 제시했다. 그러나 앞으로 풀어야 할 방향은 여전히 많으며, 여기서는 그중 일부를 간단히 설명한다.
가장 큰 모델들에서의 학습 안정성을 더 개선하는 것이 중요한 과제다. 우리의 안정화 기법은 Switch-Base, Switch-Large, Switch-C에서는 효과적이었고 불안정성이 관찰되지 않았지만, Switch-XXL에는 충분하지 않았다. 우리는 정규화 기법을 사용해 안정성을 높이거나, gradient clipping을 변형하는 등 초기적인 시도를 했지만, 이 문제는 아직 해결되지 않았다.
일반적으로는 사전학습 품질이 좋아질수록 downstream 결과도 좋아졌지만, 때때로 흥미로운 이상 현상도 나타났다. 예를 들어 C4 데이터셋에 대해 비슷한 perplexity를 보였음에도, 1.6T 파라미터의 Switch-C는 SQuAD exact match에서 87.7에 그쳐, 더 작은 Switch-XXL의 89.6보다 낮았다. 눈에 띄는 차이 중 하나는 Switch-XXL이 Switch-C보다 토큰당 FLOPs를 약 10배 더 사용하지만, 고유 파라미터 수는 약 4배 적다는 점이다(395B 대 1.6T). 이는 fine-tuning 품질, 토큰당 FLOPs, 파라미터 수 사이의 관계가 아직 충분히 이해되지 않았음을 시사한다.
데이터, 모델, expert parallelism을 함께 섞는 아키텍처를 설계할 수 있도록 스케일링 관계를 더 체계적으로 연구할 필요가 있다. 이상적으로는 하드웨어 구성의 사양(연산, 메모리, 통신)을 보고 더 빠르게 최적 모델을 설계할 수 있어야 하고, 반대로 이런 연구가 미래 하드웨어 설계에도 도움을 줄 수 있어야 한다.
우리의 접근은 adaptive computation 알고리즘 계열에 속한다. 현재는 항상 동일하고 균일한 전문가(homogeneous experts)만 사용했지만, 앞으로는 더 유연한 인프라를 통해 이종 전문가(heterogeneous experts) 를 지원하는 설계가 가능할 수 있다. 이렇게 되면 더 많은 계산이 필요한, 예를 들어 더 어려운 예제에 대해서는 더 큰 expert로 라우팅하는 식의 유연한 적응이 가능해진다.
Transformer의 FFN layer 외부에 expert layer를 두는 것도 탐구할 가치가 있다. 초기 결과로는 이것도 모델 품질을 개선할 수 있다는 신호가 있었다. Appendix A에서는 Self-Attention layer 안에 이를 넣는 경우의 품질 향상을 보고했는데, 여기서는 해당 layer가 Q, K, V를 생성하는 가중치 행렬을 대체한다. 다만 bfloat16에서 학습 불안정성이 있었기 때문에, 이 방향은 future work로 남겨두었다.
언어뿐 아니라 새로운 modality에서도 Switch Transformer를 검토할 수 있다. 지금까지는 언어만 고려했지만, 모델 sparsity는 다른 modality나 멀티모달 네트워크에서도 이점을 제공할 수 있다고 본다.
이 목록은 더 확장될 수 있지만, 우리가 현재 생각하고 있는 문제들과 유망한 미래 방향의 감을 보여주기에는 충분하다고 본다.
9. Conclusion
Switch Transformers는 확장 가능하고 효과적인 자연어 학습 모델이다. 우리는 Mixture of Experts를 단순화하여, 이해하기 쉽고 학습이 안정적이며, 동일한 크기의 조밀한 모델보다 훨씬 더 샘플 효율적인 아키텍처를 제안했다. 우리는 이 모델이 사전학습, 미세조정, 멀티태스크 학습을 포함한 다양한 자연어 작업과 서로 다른 학습 단계에서 모두 우수함을 확인했다. 이러한 진전은 수천억에서 1조 개 파라미터 규모의 모델을 학습할 수 있게 해주며, 조밀한 T5 기준선보다 상당한 속도 향상을 달성하게 한다. 우리는 우리의 연구가 sparse model을 효과적인 아키텍처로서 널리 활용하는 계기가 되기를 바라며, 연구자와 실무자들이 자연어 작업은 물론 그 밖의 영역에서도 이러한 유연한 모델을 고려하길 기대한다.
Reference














댓글
댓글 쓰기