[CoIn] 논문 리뷰 | Sparse Upcycling: Training Mixture-of-Experts from Dense Checkpoints (Komatsuzaki et al., 2022)

https://arxiv.org/abs/2212.05055

Sparse Upcycling: Training Mixture-of-Experts from Dense Checkpoints

 

Abstract.

 

앞선 글들에서 말했듯, 대규모의 깊은 신경망들을 수렴할 때까지 훈련시키는 것은 계산 비용이 매우 비쌉니다. 따라서 MoE, 즉 희소하게 활성화되는(sparsely activated) 모델들이 dense 모델에 대한 대안이 되고 있습니다. MoE가 계산 비용 측면에서 효율적이지만, sparse 모델들은 여전히 데이터 부족에 시달리며(data-hungry), 대규모의 모델에서 처음부터 학습시키기에는 비용이 여전히 많이 든다는 문제가 있습니다.

 

Note: sparse 모델들이 데이터 부족에 시달린다는 것은 절대적인 데이터 양이 부족하다는 것이 아닌, dense 모델의 경우 하나의 토큰이 모든 FFN들을 통과하는 반면, sparse 모델의 경우는 하나의 토큰이 몇 개(소수)의 FFN만을 통과하기 때문에, 전체 FFN들을 학습하는데 필요한 데이터의 양이 더 많아야 한다는 의미입니다.

 

이러한 문제를 해결하기 위해 해당 연구에서는 sparse upcycling을 제안합니다. 이는 dense model의 FFN checkpoints로부터 sparse model의 FFN(experts)들을 초기화함으로써 훈련 비용을 재사용하는 방법입니다. 해당 방법은 dense 모델의 훈련비용의 ~50%만을 사용하면서도 성능을 뛰어넘었고, 처음부터 학습된(해당 논문에서 제안한 방법이 아닌 방법으로 초기화하여 학습된)sparse 모델들 대비 성능이 좋았습니다.

 

Introduction.

 

scale은 딥러닝에서 더 나은 성능을 위한 중요한 동력입니다. NLP 분야에서 BERT, GPT-3, PaLM을 비롯하여, 비전 분야에서 AlexNet, ViT-G에 이르기까지 성능의 발전은 더 큰 하드웨어와 데이터셋 그리고 더 큰 아키텍처로부터 얻어졌습니다. 하지만 대부분의 신경망들은 처음부터(from-scratch) 훈련됩니다. 즉, 무작위로 초기화된 가중치에서 시작한다는건데, 이는 네트워크를 훈련하는 비용을 크게 증가시킵니다(BERT-Large는 0.5ZFLOPS, GPT-3는 314ZFLOPS, PaLM은 2527ZFLOPS). 이는 새로운 LLM에 대한 연구를 많은 자원을 사용할 수 있는 소수의 팀으로 제한시킵니다.

 

해당 논문에서는 model upcycling을 탐구합니다. 이는 상대적으로 적은 추가 계산 예산으로 기존 모델을 업그레이드하는 방법인데, 특히 dense 모델을 더 큰 MoE 모델로 upcycling하는 것에 대해 탐구합니다. 논문의 저자들은 Transformer checkpoints를 활용하여, 원래 모델을 훈련하는 데 사용하는 것 보다 훨씬 적은 비용으로, 원래 모델보타 더 고성능인 MoE 모델을 만들었습니다. 이는 다음의 두 가지 시나리오에서 유의미합니다:

• 사전 훈련된 Transformer에 접근할 수 있고, 제한된 계산 예산으로 이를 개선하고자 하는 경우
• MoE는 고성능이지만 훈련하기가 더 힘듭니다. 따라서 우선 dense 모델을 훈련한 후, MoE 모델로 upcycling하는 방법을 사용할 수 있습니다.

 

Background.

 

Dense 모델은 모든 입력에 대해 모든 파라미터를 적용합니다. 따라서, 모델 용량을 키우는대로 계산 비용 또한 커지게 됩니다. 그에 반해, Sparse 모델은 각 입력에 대해 파라미터의 부분집합만을 활성화함으로써 이 문제를 해결합니다. 전형적으로 MoE 모델들은 표준이 되는 dense transformer 블록과 MoE 블록을 번갈아 배치합니다. MoE 블록은 서로 다른 학습 가능한 파라미터를 갖는 다수의 experts(MLP)와 각 개별 토큰에 어떤 experts가 적용될지 결정하는 작은 신경망인 router가 존재합니다. 

 

Expert Choice Routing(token이 experts를 선택하는 Top-K routing과 달리, ECR은 experts가 token을 선택하는 것)은 다음과 같이 작동합니다. E를 MoE layer의 experts의 개수라고 하고, n을 토큰을 개수라고 할 때, router는 routing 확률을 가진 행렬 r을 생성합니다(r≥0, ∑r = 1). 그런 다음, 모든 expert e는 독립적으로 e에 대해 가장 높은 확률을 가진 T개의 토큰을 선택하고, 그것들을 처리합니다. T는 C(n/E)로 계산하며, C는 expert당 더 많거나 적은 토큰을 선택하기 위해 제어하는 용량 인자(capacity factor)입니다. C=1이면, 각 expert는 정확히 n/E개의 토큰을 처리합니다(이때 어떤 토큰은 여러 전문가에 의해 처리될 수 있는 반면, 다른 토큰들은 아무에게도 처리되지 않고 residual connection을 통해 다음 layer로 넘어갈 수도 있습니다). C>1로 설정하는 것은 더 높은 계산 비용을 허용함으로써 더 높은 성능을 달성할 수 있게 합니다(하지만 communication costs 등 여러 가지 overheads가 존재합니다).

 

(... 아키텍처에 대한 설명 ...)

 

The Upcycling Algorithm.

 

Figure 1: The upcycling initialization process. All parameters, and optionally their optimizer state, are copied from the original checkpoint, except those corresponding to the MoE router, which does not exist in the original architecture. In particular, the experts in the new MoE layer are identical copies of the original MLP layer that is replaced.

 

upcycling algorithm은 그림1. 에 설명되어 있습니다. 모델을 upcycling하기 위해, dense 모델의 파라미터(checkpoints)들을 이용하여, router를 제외(router는 무작위로 초기화)하고는 모두 그대로 복사합니다. 새 모델이 로드되고 초기화 된 후에는 이용 가능한 예산과 자원들에 따라 계속 훈련시키면 됩니다.

 

upcycling된 모델의 성능은 MoE layer들의 구성에 의해 심하게 영향을 받습니다. 일반적으로 upcycling된 layer의 수 또는 전문가의 수, 전문가 용량을 늘림으로써 모델 용량을 늘리는 것은 더 높은 성능의 모델로 이어집니다. 하지만 이는 layer들의 급격한 재구성으로 인해, 계산 비용을 증가시키며, 더 큰 초기 성능 하락을 초래합니다.

 

router는 capacity factor가 2인 expert choice routing을 사용합니다. 또한 훈련 시간(전체 배치 teacher forcing)과 추론 시간(단일 토큰 auto-regressive decoding)간의 불일치를 피하기 위해, 언어 디코더에서는 Top-K(K=2) routing을 사용합니다.

 

Note: Top-K는 앞서 설명했듯, 토큰이 expert를 고르는 방식입니다(예: 나는 1번, 2번 전문하가한테 갈래!). 반면 ECR은 expert가 토큰을 고릅니다(예: expert A가 "자, 여기 있는 100개 토큰 중에 내가 제일 잘 아는 애 5명 와!"하고 데려감). 이때 LLM의 training time에는 정답 문장이 "I love AI"라면, 모델은 "I", "love", "AI"를 한꺼번(batch)에 봅니다. 이때는 ECR이 가능한데, 토큰들이 마낳이 모여 있으니, expert들이 토큰들을 나눠갖을 수 있습니다. 하지만 inference time에는 모델이 "I"를 예측하고 "love"를 예측하고자 할 때, ECR은 불가능합니다. 토큰들 중에서 선택하고 싶지만, 고를 수 있는 토큰이 "love"밖에 없기 때문입니다(현대의 LLM은 KV Cache를 수행하기 때문에, auto-regressive 단계에서 모델 안에 존재하는 feature tensor의 sequence lenght 차원은 1입니다. 이전에 생성된 토큰들은 cached되어서 attention score를 계산하는데만 사용됩니다). 따라서 decoding 단계에서는 Top-K routing을 사용하는 것입니다.

 

더 많은 MoE layer들을 추가하는 것은 모델의 용량을 극적으로 증가시킵니다. 하지만, 이는 원래 dense 모델과 비교하여 급격한 모델의 재구성으로 인해 upcycled model의 초기 성능을 떨어뜨립니다. 따라서 해당 논문의 저자는 기존 모델 내에서의 MLP layer들의 절반만을 MoE layer로 대체합니다.

 

새로운 experts는 추가하는 것은 모델의 용량을 확장하는 좋은 방법입니다. expert capacity(expert당 처리하는 토큰의 수)는 expert의 수 E에 반비례하기 때문에, 더 많은 experts들을 추가하는 것은 모델의 FLOPS나 실행 시간에 크게 영향을 미치지 않습니다. 하지만, experts의 수가 매우 많아진다면, upcycling된 모델은 baseline dense 모델 대비 더 큰 초기 성능 하락을 초래합니다(아무래도 구조가 많이 달라지며, 랜덤 초기화되는 router에 부담이 커지기 때문에). 이는 충분한 upcycling computing이 주어진다면 해결할 수 있으며, 해당 논문의 저자는 초기 baseling dense 모델의 계산 비용의 +20%에서 +100%를 사용할 때, 32명의 experts가 적당하다는 것을 발견했습니다.

 

expert capacity를 조절함으로써 평균적으로 각 토큰을 처리하는 experts들의 수를 조절할 수 있습니다. 더 큰 expert capacity는 일반적으로 더 나은 성능을 보이지만, FLOPS와 실행 시간을 증가시킵니다. 해당 논문의 저자들은 C=2에서 일반적으로 계산 시간 기준 좋은 품질을 달성할 수 있다는 사실을 발견했습니다.

 

비전 모델의 경우 원래의 dense checkpoints로부터 optimizer의 상태를 resume하는 것이 성능을 높여줍니다(LLM에서는 optimizer의 상태를 resume하는 것이 성능에 도움을 주지는 않았습니다). 또한 비전 모델의 경우는 upcycling 모델을 만들 때, 성능 하락을 줄이기 위해, 각 토큰의 router 결합 가중치들을 1로 정규화합니다. 이는 이전 dense 모델에서 오직 하나의 expert에 의해서만 추론이 되었다는 사실에서 기반합니다(router 결합 가중치를 1로 정규화하는 것은 LLM에서 사용하지 않습니다).

 

Experiments.

 

Figure 2: Pretraining performance achieved by the dense continuation and upcycling methods, for different Transformer variants. The left plot shows the performance on the vision task and the right one on the text task. The x-axis shows the extra pretraining time (TPU-core-days), with respect to the total time needed to train the original dense checkpoints, for each size. The horizontal lines indicate the quality (y-axis) of the original dense checkpoints.

 

그림2. 는 upcycling된 모델들과 dense 모델들의 성능에 대한 지표입니다. 상당한 양의 추가 연산을 적용하면, upcycling된 모델이 dense 모델에 비해 더 강한 이득을 보여주는 패턴이 드러납니다. 

 

Figure 4: Pretraining performance achieved by the upcycling method and a MoE model trained from scratch, for B/16 (left plot, vision task) and Base (right plot, text task). The x-axis shows the extra pretraining time (TPU-core-days), with respect to the total time needed to train the original dense checkpoints. The MoE model trained from scratch only catches up to the language upcycled model after about 120% of the original dense checkpoint computation budget (second last orange and green dots from the right).

 

그림 4.는 처음부터 학습된 MoE(sparse) 모델과 upcycling된 모델을 비교한 것입니다. 처음부터 훈련하는 것은 dense model을 학습하는데 사용했던 계산 비용을 재사용하지 않기 때문에 upcycling된 모델을 따라잡는데 더 오래걸립니다. 처음부터 훈련된 언어 MoE 모델은 upcycling된 모델을 따라잡기 위해서 원래 dence checkpoints 계산 예산의 약 120%를 필요로 합니다.

 

비전에 대해서 Batch Prioritized Routing (BPR)을 갖춘 Top-K routing이 스텝 단위 기준으로는 Expert choice routing과 성능이 비슷해도, 속도는 더 느립니다. 한편, capacity factor C를 늘리는 것이 스텝 단위 기준으로는 좋은 성능을 기록하지만, 해당 논문의 저자는 C=2에서 계산 시간 기준으로 제일 좋은 성능을 기록하는 것을 확인했습니다.

 

또한 모델을 upcycling할 때는 절반의 layer만을 MoE로 사용하는 것이 가장 좋았으며, 더 많은 expert 다양성을 촉진하기 위해 가우시안 노이즈를 각각의 experts들에게 추가했지만, 노이즈가 너무 강하면 성능이 떨어졌으며, 노이즈가 작으면 효과가 거의 없습니다. 더 많은 experts들을 추가하는 것은 앞서 말했듯, 모델의 구조가 급격히 바뀌기 대문에, 더 큰 초기 성능 하락을 초래합니다. 따라서 32개의 experts가 적당하다는 것을 해당 논문의 저자들은 발견합니다.