MoE(Mixture of Experts) by hugging face

https://huggingface.co/blog/moe

Mixture of Experts Explained

위 글에서 알고자 하는 건 크게 2가지이다. 이건 내가 알고 싶은 거고 위 글쓴이가 알려주고자 하는 건 TL;DR에 있다.

1. Experts가 어떻게 학습되는 지(GPU 자원을 어찌 활용하는지)
2. Inference 할 때 발생하는 Tradeoff가 뭔지

what is a Mixture of Experts?

• MoE 모델은 기본적으로 2가지 구성으로 이뤄져있는데

1. Sparse MoE layers - 기존 Transformer 모델에서 Feed Forward Network를 그림처럼 MoE Layer로 수정하고 MoE Layer 내부에 Experts가 존재하는데 얘들이 FFN과 동일한 역할을 한다.
2. A gate network or router - 그림에서 볼 수 있는 것처럼 router나 gate network를 통해 각각의 토큰들이 적절한 expert로 전송된다. router나 gate network는 학습 가능한 파라미터로 전체 신경망이 학습될 때 같이 학습된다.

MoE는 기본 Transformer(dense models)에 비해서 효율적으로 pretraining이 가능하고 더 빠른 추론이 가능하지만 아래와 같은 문제가 있다.

• Training : routing이 제대로 되지 않을 경우 기본 Dense Model에 비해 적은 파라미터를 가지는 Experts들이 overfitting 될 수 있다.
• Inference : Training과 마찬가지로 추론 동안에도 사용되는 Experts만 사용될 수 있다. 이러면 추론은 빨라질 수 있지만 RAM에 모든 파라미터가 load 되야 하기 때문에 메모리 요구사항이 높아 질 수 있다. <추가적으로 8x7B이지만 필요한 VRAM은 56B가 아니라 47B인 이유는 Experts들만 다른 파라미터이고 다른 파라미터는 공유하기 때문이다.>

A Brief History of MoEs

MoE는 "Adaptive Mixture of Local Experts" (1991)에서 제안된 방법을 기초로 발전했다.

이해하기 쉽게 이야기 하면 input space에 우리가 해석하기 어려운 형태로 특정 영역이 존재하고 이를 각각의 sub-networks들이 학습하도록 하고 최종적으로 모든 network를 결합하여 복잡한 문제를 해결하려고 하는 접근법이라고 한다. 어떤 식으로 network를 선택하는지는 알 수 없지만 Supervised Learning으로 하면 gating을 포함해서 잘 된다고 한다.

여러 방향으로 MoE는 연구되었고 137B LSTM에 MoE가 적용된 연구가 진행되었었다(Shazeer et al. (2017, with “et al.” including Geoffrey Hinton and Jeff Dean, Google’s Chuck Norris). 하지만 아무래도 파라미터가 많으니 높은 communication cost가 발생했고 학습이 불안정하게 이뤄지는 문제가 발생했다.

What is Sparsity?

희소성은 결국 조건적인 계산으로 특정 조건에 맞는 Experts만 사용하겠다는 것으로 생각될 수 있다. 그래서 Sparsity는 계산량 증가 없이 모델의 크기를 키울 수 있게 하지만 여러 문제들을 불러왔다고 한다. 예를 들어 배치 크기가 증가하면 일반적으로 모델의 성능이 증가하지만 MoE 같은 경우 배치화된 input의 token이 각각 다른 Experts에 전달 될 수 있기 때문에 각 Experts가 처리하는 데이터의 양은 전체 배치 크기보다 작아지기 때문이다.

 

이러한 문제를 해결하는 방법은 데이터가 Expert들에게 골고루 분배되도록 routing mechanism을 개선하는 방법이 있다. Shazeer의 방법만 보도록 하자. 전통적으로 Gating은 아래와 같은 수식으로 Experts들을 선택했다고 할 수 있다.

 

하지만 Shazeer는 이걸 Noizy Top-k Gating방법으로 수정한다.

Shape은 간단하게 생각해서 G(x) = d_model x experts num으로 될텐데 이때 k개의 expert를 선택하는 거니 d_model x k 아니면 그냥 k 일 수도 있다. 이렇게 해서 k가 1 또는 2 같이 작은 숫자가 된다면 기존보다 더 빠르게 학습이나 추론을 진행 할 수 있게 되는 것이다. 그런데 왜 Noize를 추가해주는 것일까? 그냥 Top k만 선택하면 되는 거 아닌가? 하는 생각이 들 수 있다. 이를 Load balancing tokens 때문이라고 설명한다.

Load balancing tokens for MoEs

기본 MoE 학습에서는 Gating network가 소수의 expert를 활성화하도록 하는 경향이 존재하며 이게 계속 선택되면 결국 overfitting이 되는 문제가 발생한다. 따라서 이를 방지하기 위해 보조 loss를 추가하여 모든 experts가 동일한 중요성을 가지도록 유도한다고 한다. 이 loss가 모든 expets가 동일한 학습 예시를 받을 수 있도록 보장한다고 한다. 이걸 위해서 Balancing loss를 추가해주는데 이게 과도하게 적용되면 모델이 학습 초기 단계에서는 routing 과정이 급격하게 변해서 학습이 불안정해질 수 있다고 한다.

MoEs and Transformers

이건 그냥 google에서 MoE 를 엄청 크게 늘리고 GPU간에 파라미터를 공유함으로써 계산 효율성을 늘리는 방식이기 때문에 너무 깊게 생각하지 않고 넘어가도록 하겠다.

Switch Transformers

 

이게 그 Switch Transformer인데 MoE 가 학습과 fine-tuning에서 불안정성이 있다는 문제를 해결하고자 제안된 방법이다.

요 방법은 한번 Switching FFN layer로 input을 넣을 때 2개의 input을 동시에 넣는 것이다. 이렇게 하면 router 계산 감소, 각 Expert의 batch size를 최소한 절반으로 만들 수 있고(이게 잘못 해석한 거 같은데 내 생각엔 불균형했던 batch size를 최소한 절반으로 감소한다는 뜻인듯하다. ) 여기서 배치 크기라는 표현이 좀 어색하지만 결국 할당되는 token의 개수라고 생각하면 좋을 듯 하다. 이외에도 성능 유지나 communication cost 감소라는 효과가 있다고 한다.

 

그리고

이걸 통해서 각 expert들이 받을 수 있는 token의 정도를 제한할 수 있는데 1보다 크면 완벽하게 balanced되지 않을 때 buffer를 제공할 수 있다고 한다. capacity가 커질수록 inter-device communication이 증가하기 때문에 trade-off가 생긴다고 한다. 흠.. 근데 저걸 늘리면 좋은게 뭐지?

특정 전문가가 더 많은 토큰을 처리해야 하는 경우, capacity factor가 높으면 이를 수용할 수 있어서 특정 전문가가 과부하로 인해 성능이 저하되는 것을 방지하는 효과가 있다고 한다. 복잡한 데이터에서 특정 전문가가 집중적으로 학습해야 할 가능성이 높아지는 경우, 더 많은 데이터를 처리할 수 있는 여유를 주어 모델 성능을 유지하거나 개선할 수 있다고 한다.

 

이외에도 router를 학습할 때는 정확한 성능을 위해 full precision을 사용했다고 한다.

How does scaling the number of experts impact pretraining?

전문가가 많을수록 샘플 효율성이 향상되고 속도가 빨라진다. 하지만 이러한 이점은 256이나 512개의 전문가를 뒀을 때 이후에 감소하고 많은 VRAM이 필요하게 된다. dimension은 얼마나 하는게 제일 좋을까? / 그리고 많은 expert를 둔 실험과 적은 수의 expert를 둔 실험에서 동일한 특성이 나왔다고 한다.

Parallelism

 

이건 구현이 너무 어려우므로 pass하도록 하겠다.

그림에서 볼 수 있는 것처럼 위의 파란색은 모델 weight, 아래는 data를 의미한다. 색깔이 다른 건 다른 experts를 의미. 이걸 보면 이해하는데 크게 어렵지는 않다.

Fine-Tuning MoEs (Mixture of Experts)

희소 모델(Sparse Models)은 조밀 모델(Dense Models)보다 오버피팅에 더 취약.

전문가(Experts) 내에서 더 높은 드롭아웃(Dropout) 등 강한 정규화 기법을 활용 가능.

보조 손실(Auxiliary Loss) 사용 여부가 중요하지 않을 수도 있으며, **토큰 드롭(Token Dropping)**이 오버피팅 방지에 기여할 수 있음.

희소 모델은 추론(reasoning) 기반 작업(SuperGLUE)에서 성능이 떨어지지만, 지식(knowledge) 기반 작업(TriviaQA)에서는 더 우수한 성능을 보임.

MoE 레이어를 고정하고 다른 파라미터만 업데이트하면 학습 속도가 빨라지고 메모리 사용량이 줄어듬.

반대로 MoE 레이어만 업데이트하면 성능 저하.

이는 MoE 구조의 대부분의 파라미터(80%)가 전문가 레이어에 집중되어 있지만, 각 토큰이 적은 전문가만 참조하기 때문. (MoE 레이어만 업데이트하면, 모델의 성능이 저하됩니다. 이는 각 토큰이 소수의 전문가(Top-2)만 참조하기 때문입니다. 즉, 일부 전문가의 파라미터만 계속 업데이트되며, 나머지 전문가들은 학습되지 않아 일부 전문가에 의존하는 편향된 학습이 발생할 수 있습니다. by GPT)

**MoEs Meets Instruction Tuning (2023)**에서 다양한 작업을 통해 희소 모델을 세밀하게 조정.

Flan-T5 기반 희소 모델은 T5 기반 모델보다 훨씬 더 좋은 성능을 보임.

• Instruction Tuning은 희소 모델에서 조밀 모델보다 더 큰 성능 향상을 가져옴.

Serving Techniques (MoE 배포 기법)

• 추후에도 정리하지 않을 것 같다.

Efficient Training Techniques (효율적 학습 기법)

• 추후에도 정리하지 않을 것 같다.

이렇게 해서 MoE에 대해서 알아보는 시간을 가졌다. 맨 처음 가졌던 질문을 보면

1. Experts가 어떻게 학습되는 지(GPU 자원을 어찌 활용하는지) - supervised로 학습이 되고 이 과정 중에 load balancing과 capacity 보정이 들어간다.
2. Inference 할 때 발생하는 Tradeoff가 뭔지 - capacity scale factor를 키울 수록 비용이 많이 든다고 한다. 얘를 키우면 각 expert들이 처리할 수 있는 token의 수가 증가하기 때문에 특정 전문가가 집중적으로 학습해야 할 가능성이 높아지는 경우, 더 많은 데이터를 처리할 수 있는 여유를 주어 모델 성능을 유지하거나 개선할 수 있어 trade off가 발생한다고 한다.