Mixture of Experts

Paper Review 12

|Intro

최근 2년 간 봤던 모든 논문에는 MoE라는 방법론을 적용했다는 내용을 종종 확인할 수 있었습니다. 일례로 이전에 리뷰했던 Swin Transformer의 공식 github를 보면 다음과 같이 최초 논문에서는 서술되거나 적용되지는 않았으나, 사후적으로 적용된 기법들이 정리돼 있으며 그 중 하나가 바로 MoE입니다.

swin transformer implementation

MoE는 기존의 예측 및 정리를 위한 FFN(Feed Forward Network)의 성능을 개선하기 위한 게이팅 네트워크(gating network) 기반의 모델 아키텍처입니다. 마치 LSTM에서와 같이 복수의 전문가 레이어와 이들의 가중치를 조절하는 게이팅 네트워크가 결합된 것입니다. 이를 통해 추론 과정에서 일부 전문가 레이어가 비활성화되며 그만큼 컴퓨팅 리소스를 절약하고, 보다 빠르게 실행될 수 있다는 장점을 가지고 있습니다.

MoE를 Transformer 블록에 사용한 예시를 살펴보면 아래와 같습니다.

Transformer Block의 FFN 옵션[2]

위 그림에서는 Transformer Block에 적용되는 FFN을 GLU(Gated Linear Unit)으로 사용하는 방법과 이를 쪼갠 MoE(Mixture of Experts)로 사용하는 방법을 비교하고 있습니다. FFN에 단순 선형층의 결합이 아닌 Gating network를 적용해 가중치를 조절함으로써 선형 변환 과정에서 발생하는 정보 손실을 방지하고, 특정 중요 부분에 대한 가중치를 고려함으로써 유연하면서도 다양한 표현력을 지닐 수 있게 되었습니다. 또한 불필요한 정보 등이 게이팅 네트워크를 통해 줄어들거나 없어지면서, 학습 성능을 높일 수도 있었습니다. 

그리고 이러한 GLU 방식의 FFN를 한 단계 업그레이드시킴으로써 자원의 효율성과 실행속도까지 챙긴 개념이 바로 MoE입니다. 이는 단일 선형층을 복수의 전문가층으로 나눔으로써 각각의 역할을 세분화하고, 이를 게이팅 네트워크를 통해 조율하도록 함으로써 데이터에 맞는 전문가 집합만 선택해 연산량을 줄이고, 여러 전문가가 각각 다른 기능을 학습하도록 유도함으로써, 모델은 더 다양한 패턴을 학습하고 표현하며, 추론 시에는 학습 때와 달리 소수의 전문가층만을 활용해 보다 적은 메모리를 요구하면서도 더욱 빠르게 실행됨으로써 추론 효율성도 높일 수 있게된 것입니다[8].

이전에 리뷰했던 Transformer 포스팅에서 언급했었던 내용 중 하나가 바로 우리가 구축하는 모델에서 가장 많은 파라미터를 가진 부분이 바로 선형층으로 구성된 마지막 전결합층이었는데요, 그 외에도 트랜스포머 내부의 FFN까지 고려하면 MoE를 도입은 거의 필수적인 것이나 다름없습니다. 이를 비교적 최근에 제시된 Mixtral 8x7B 모델[8]의 구조를 통해 살펴보면 다음과 같습니다.

Mixtral 8x7B Structure[6]

위 그림에서 왼쪽 그림은 모델의 구조를 오른쪽 그림은 각 구성 요소들의 파라미터를 의미합니다. 우리가 주목할 부분은 바로 Experts 부분으로, 학습 시에는 총 8개의 전문가 층이 필요해 45,097,156,608개의 파라미터가 필요하지만, 추론 시에는 활성화되는 두 개의 전문가 층만으로도 충분해지게 됩니다. 이 숫자는 11,274,289,152개로 1/4에 해당하는 숫자이며, 전체 파라미터 숫자도 학습 시에는 46.7B에서 12.8B로 대폭 떨어지게 됩니다. (참고로 이때에도 선형층에 해당하는 expert가 전체 파라미터의 대부분을 차지함을 알 수 있습니다.)

이번 스터디를 위해 MoE라는 개념이 제시된 논문들을 찾아보다가 이러한 개념이 제안된지 상당히 오래됐으며, 관련된 논문들이 상당히 많은 것을 알 수 있었습니다. 때문에 본 포스팅에선 MoE에 대한 Survey 논문들[1][2]을 기반으로 MoE의 발전상과 그 구조를 살펴보고, 이를 직접 구현해 학습시키며 테스트하는 식으로 리뷰를 진행하도록 하겠습니다. 

|Background

MoE는 T5 모델을 시초로 Transformer 기반의 모델들에 중점적으로 적용돼 왔습니다. 다만 그 초기 형태와 최신 논문에 적용된 형태는 다소 차이가 있는데요, 본 포스팅에서 이를 하나하나 업데이트해가며 살펴보기 전에 최초의 형태를 먼저 살펴보도록 하겠습니다.

초기의 Mixture of Experts는 크게 두 가지로 구성돼 있으며, 하나는 복수의 세분화된 작업을 진행할 Expert들과 어떤 전문가에게 일을 배분할지 결정젓는 Router입니다. 모든 MoE는 이 두 개의 역할군을 기반으로 설계되며, 이들을 이용해 구현하는 방법은 다음과 같았습니다.

1. 여러 전문가 레이어를 구축한다                                                           
2. 각 전문가 레이어들의 가중치를 고려할 수 있는 라우팅 레이어를 구축한다  
3. 라우팅 레이어의 가중치를 고려해 각 전문가 레이어를 복합적으로 고려한다

이러한 개념은 직관적이었으며, 기본적인 선형층으로 구성된 FFN보다는 성능이 뛰어났지만 그만큼 계산 복잡도가 높아졌습니다. 이를 코드로 보면 다음과 같습니다.

|Basic MoE

위에서 설명했던 바와 같이 MoE는 기본적으로 전문가 레이어(Expert)와 게이트 네트워크로 구성돼 있습니다. 각각의 전문가들은 일반적으로 단일 선형층만으로 구성되거나 추가적으로 활성화층이 적용되는 식으로 구성됩니다.

그리고 이들을 중개해 가중치를 부여하는 라우팅 레이어는 어떤 전문가를 선택할지를 출력하며, 일반적으로 선형층과 확률 계산을 위한 softmax 함수로 구성됩니다. 이를 코드로 살펴보면 아래와 같습니다. 이를 논문에서는 '입력 (x)를 router의 가중치 행렬 (W)와 곱하고, 각 expert에 대해 SoftMax를 적용해 확률분포 G(x)를 생성한다'와 같이 표현합니다. 

이를 코드로 살펴보면 아래와 같습니다. 이에 대한 전체 코드는 여기에서 확인하실 수 있습니다.

define MoE components

그리고 MoE 레이어에서는 이들을 한데 묶어 여러 전문가 네트워크의 출력을 게이트 네트워크를 통해 가중치를 적용해 활용합니다. 이들을 한데묶는 MoE 레이어 코드는 다음과 같습니다.

define MoE

위 코드를 통해 확인할 수 있듯 코드 자체는 크게 복잡하지 않습니다. 지정된 전문가 수만큼 복수의 전문가 레이어를 정의하고(self.experts), 각 전문가층에 입력값 x를 통과시킨 뒤 이를 게이트 네트워크의 가중치 출력과 가중합산하는 식으로 최종 출력이 결정되는 방식입니다. 

이는 하나의 큰 도메인 내 문제를 해결하기 위해 모델이 구성되고, 그 내부의 FFN 내에서 더 작은 단위로 전문가들을 통해 의견을 제시한 뒤, 어떤 의견을 얼마나 고려할지 라우터가 결정하는 방법입니다. 임의의 데이터를 이용해 출력되는 양상을 살펴보면 다음과 같습니다.

moe's output shape

여기까지가 기본적인 MoE의 구현입니다. 이후에는 이러한 기본적인 MoE를 발전시켜 나가며 본 포스팅을 진행하도록 해보겠습니다.

|Upgrade MoE

위에서 MoE에 대해 '복수의 전문가가 의견을 내고, 이를 얼마나 고려할지 라우터가 결정해 고려한다'와 같이 설명했는데요, 하지만 이 방식은 어찌됐든 모든 전문가들의 의견을 작게든 고려해야 했습니다. 예컨데 회사의 대표가 임원들의 모든 주장을 조금씩이라도 반영해 의사결정해야 한다는 이야기입니다. 이는 다소 번거로울뿐 아니라, 결과적으로 각 전문가들이 독립적인 의견을 내는게 아니라 전문가들끼리 비슷한 의견을 내는 일종의 담합을 유도하거나 특정 전문가들을 아예 배제하는 등의 문제를 야기합니다. 

이는 각각 독립적이고 차별적으로 구성돼야 하는 전문가층이 서로 유사해지거나, 사용되지 않을 수 있다는 것을 의미하며, 간단히 말해 비효율적입니다. 때문에 기존에는 적은 수의 전문가층을 적용하는 식으로 진행했지만, 보다 큰 모델과 보다 큰 데이터 사이즈를 다루기에는 한계가 존재했습니다. 따라서 다음과 같은 방법으로 기존의 MoE를 발전시키게 됩니다.

|Upgrade #1 : Dense → Sparse

MoE 구조에서 첫 번째로 개선된 부분은 기존의 모든 Experts들을 Router에 의해 조정된 가중치를 적용해 적용하던 방식에서 소수의 전문가들만 사용하는 방식으로 변경한 것입니다. 모든 전문가들의 의견을 조금씩이라도 고려하는 것 대신, 가장 좋은 의견을 제시한 전문가 한두 명의 의견만을 고려하는 식입니다. 

이를 그림으로 먼저 살펴보면 아래와 같습니다.

Dense → Sparse (from [6])

이러한 개선을 통해 얻을 수 있는 이득은 크게 두 가지로, 하나는 모델이 추론 시에 일부 전문가 층을 비활성화함으로써 그만큼 연산속도가 빨라지고 컴퓨팅 리소스를 효율적으로 사용할 수 있다는 것입니다. 그리고 다른 하나는 일부 우선순위가 낮은 전문가들의 의견을 배제함으로써 학습 과정에서 중요성이 낮은 노이즈로 인한 영향력을 최소화할 수 있다는 것이었습니다.

즉 헛소리하는 전문가들의 의견을 배제하고, 굳이 매번 모든 문제에 대해 모든 전문가들의 의견을 듣는게 아니라 해당 문제를 잘 풀었던 소수의 전문가들의 의견만을 통해 빠르게 결정내릴 수 있다는 것입니다. (다만 이는 코드로 바로 구현하긴 어렵습니다. 그 이유는 모델을 로딩할 때, 기본적으로 모든 전문가층을 포함한 일체를 로드해야 하기 때문입니다. 또한 라우터에 의해 선택된 모델만을 선택해서 하는 경우도 병렬성을 해치기 때문에 오히려 속도 저하가 있을 수 있습니다. 다만 복수의 GPU를 통해 각 GPU별로 독립적인 전문가층을 배치하는 식으로 해결가능하긴 합니다[4]. 그 구체적인 방법이 궁금하다면 Tutel-MoE[4] 논문을 참조하시길 바랍니다.)

이를 기존의 코드에 적용하려면 다음과 같습니다. 우선 가장 뛰어난 의견을 제시한 상위 1명의 전문가의 의견만을 수렴하는 케이스입니다. 이해가 쉽도록 간단히 구현했습니다. 코드는 여기를 통해 확인할 수 있습니다.

to Sparse 업데이트 ① topk 1 only & ModuleList

위 코드는 expert를 굳이 별도의 층으로 두지 않고 ModuleList를 이용해 MoE 레이어 내에서 구현했습니다. 간단히 설명하자면 (14) gate_outputs의 출력에 대해 상위 1개의 전문가 층의 값과 인덱스를 가져와 (5)ModuleList 형태로 관리되고 있는 experts들 중에서 (20)인덱스에 해당하는 전문가만 골라 가져와 (20)하게 되는 식입니다. 참고로 본문에서 제시하는 괄호 내의 숫자는 위 코드 상의 라인 넘버를 의미합니다.

다만 위 구현은 오로지 단 하나의 전문가 의견만을 들을 수 있는 구조인데요. 만일 2 이상의 전문가 의견을 고려하고 싶다면 반복문을 이용해 다음과 같이 구현해야 합니다. 전체 코드는 여기를 통해 확인할 수 있습니다.

toSparse 업데이트 ② topk

위 코드는 (17) 라우터에 의해 선택된 k개의 전문가들의 인덱스를 고려해, (17)행렬 연산을 위해 전체가 0으로 구성된 출력 형태의 텐서로 초기화하고, (20~26) 라우터가 제시한 전문가의 인덱스에 해당하는 전문가의 의견만을 라우터에서 제시한 전문가의 가중치를 곱해 가중합산해 출력하는 방식입니다.

하지만 이러한 ModuleList에 의한 반복 연산은 모델의 병렬성을 저해합니다. 위에서 언급했던 바와 같이 많일 복수의 gpu를 통해 MoE를 구현함으로써 각 gpu 별로 독립적인 expert들을 배치하는 경우를 제외한다면, 최대한 병렬성을 유지할 수 있도록 하는 전략이 필요합니다. 

때문에 본래 순차적인 연산 등을 위해 필요했던 ModuleList 대신 하나의 큰 nn.Linear로 전체 전문가들의 의견을 생성하고, 라우터에 의해 선택된 영역에 해당하는 부분만 가져와 사용하는 방식으로 사용하겠습니다. 코드는 아래와 같습니다. 해당 코드는 여기를 통해 확인하실 수 있습니다.

toSparse 업데이트 ③ Parallel Structure

위 코드에서 변경점을 살펴보면 우선 (9)nn.ModuleList 형태로 관리되던 복수의 전문가 레이어를 하나의 큰 단일 선형층으로 구현했습니다. (12~13)이때 임의로 레이어가 초기화되는 과정에서 대칭적 초기화 등으로 인해 일부 전문가가 유사해질 가능성을 고려해, xavier uniform 초기화를 통해 균등하게 펼쳐지도록 함으로써 전문가들 간의 다양성을 확보하고자 했습니다.

이후에는 (29~33) torch tensor에서 행렬 단위로 슬라이싱해 가져올 수 있는 torch.gather를 활용해 전체 전문가 출력 결과에 대해 라우터가 제시한 전문가의 인덱스에 대해 가져옵니다. (36) 그리고 이렇게 가져온 [batch_size, topk, output_dim] 형태의 결과물을 라우터에서 제시한 전문가별 확률(topk_probs)과 가중합산해 최종 출력 결과를 만드는 식입니다.

이렇게 변환한 결과를 간단히 테스트해보면 다음과 같은 실행 속도 차이를 확인할 수 있습니다.

병렬화 전후의 결과

위 그림들에서 두 번째 모델(moe2 & Model B)가 병렬화가 적용된 Sparse MoE입니다. 여태까지 구현한 Sparse MoE 레이어에 대한 개념을 그림으로 살펴보면 아래와 같습니다. 아래 그림은 topk=1로 지정된, 그러니까 4명의 전문가 중 가장 두드러진 전문가 1명을 선발하는 방식을 예를 들고 있습니다.

Sparse MoE[6]

위 그림에서는 단일 전문가가 선택됐음에도 전문가 확률이 곱해지는 식으로 진행됐는데, 실제로는 topk=1일 경우 라우터에 의해 결정된 확률을 곱하지 않고 그대로 쓰는 것이 논리적으로 합당해보입니다. 이에 따라 기존의 가중 합산 코드를 다음과 같이 수정해야 합니다. 전체 코드는 여기를 통해 확인할 수 있습니다.

topk를 온전히 반영하기 위한 가중합산

또한 코드를 여기까지 작성했을 때, 제가 한 가지 놓치고 있다는 것을 떠올렸는데요. 논리적으로 생각해보면 선택된 전문가들끼리만 다시 모아서 softmax를 한 뒤에 고려해야 하는 것 아닌가 하는 것이었습니다. 만약 그렇다면 라우터에서 softmax를 취하는 건 중복이 될 수 있어, 라우터는 단일 선형층만으로 구성할 수 있으며, 굳이 그렇다면 별도의 클래스로 두지 말고 MoE 내부로 가져오는 식으로 전체 코드를 구성할 수 있었습니다.

이를 코드로 확인하면 다음과 같습니다. 이전에는 Gating Network 등을 별도로 구현된 것을 사용했고, 이를 활용함을 가정했지만 아래 코드는 MoE 및 그 구성요소들을 포함한 전체입니다. 아래 코드는 여기를 통해 확인할 수 있습니다.

단일 클래스로 정의된 MoE

이렇게 구현된 모델은 그 직전의 버전보다 아주 미약하게나마 빠릅니다. 하지만 이렇게 구현된 모델에도 여전히 부족한 부분이 존재하는데요, 그것은 바로 여전히 전문가들끼리 비슷한 의견을 제시하거나 특정 전문가가 반복적으로 선택되지 않는 등의 비효율성이 여전히 남아있다는 것입니다. 이를 해결하기 위해 다음의 로드 밸런서가 필요해지게 됩니다.

|Upgrade #2 : Load Balancer

그 다음 개선 사항은 바로 최대한 다양한 전문가가 선택되도록 함으로써, 모델이 하나의 문제를 세분화하고 각각의 분야에 대해 전문가에 이식하는 것입니다. MoE를 적용하다보면 소수의 or 특정 전문가 층만을 채택하는 현상이 확인되곤 합니다. 이는 특히 가중합을 하는 게 아닌 특정 전문가만을 고르는 Sparse MoE를 구축하는 전략을 채택했을 때, 불필요하게 모델의 사이즈를 키우고 다른 전문가 층을 활용하지 못하는 등의 문제를 야기합니다. 

이러한 문제를 해결하기 위해 기본적으로 제시하는 방법은 바로 전문가들 간의 균형을 맞춰주기 위해 단순히 상위 k개의 전문가를 고르는 것을 넘어 추가적인 규제를 추가하는 KeepTopK[7] 전략을 적용하고, 별도의 손실(auxiliary loss)을 정의해 이를 역전파 과정에서 고려하는 방식입니다. 하나씩 살펴보겠습니다.

|KeepTopK

우선 KeepTopK 전략에 대해 살펴봅시다. 이러한 전략은 router를 통해 구현되는데요, 기존에 입력 x에 가중치 행렬  W를 곱해 H(x)를 출력하던 결과부터 시작됩니다. 이 과정에서 반복적으로 동일한 전문가가 선택되는 문제를 해결하기 위해 간단한 규제를 적용할 것인데요, 바로 약간의 노이즈를 추가하는 것입니다.

이를 구현하기 위해 일반적으로 학습가능한 노이즈 매트릭스를 활용하는데요, 이를 그림으로 살펴보면 아래와 같습니다. 참고로 아래 그림의 데이터 형상은 x(Batch_size, input_dim) @ W(input_dim, num_experts) + n(Batch_size, num_experts) = H(x) (Batch_size, num_experts) 입니다.

add noise to router[6]

이후에는 상위 k개의 전문가를 제외한 전문가들의 가중치를 -∞로 설정해 일종의 마스크를 씌움으로써 이후의 SoftMax를 적용할 때 0이 나오도록 해줍니다. 이를 그림으로 살펴보면 아래와 같습니다.

Mask to H(x)[6]

노이즈를 추가함으로써 보다 다양한 전문가들이 선택될 여지가 생겼지만, 아직 충분하지 않은데요. 이를 해결하기 위해 router의 출력값인 각 전문가들 간의 공분산을 추가적인 손실로 고려하는 전략을 추가적으로 적용합니다. 구체적으로 Importance Score 혹은 auxiliary loss라고 부르는 토큰별 각 전문가들의 확률에 대해 공분산을 계산하고 이를 낮추는 방향으로 학습함으로써 특정 전문가가 소외되거나 배제되지 않도록 유도하는 것입니다.

참고로 여기까지의 구현 중에서 노이즈를 추가하는 부분을 제외한 다른 부분은 위에서 구현한 것과 논리적인 흐름은 같습니다. 다만 위 그림 및 논문들에서와 달리 저는 선택되지 않은 전문가 행렬을 -∞로 마스킹하고 이를 softmax를 취하라고 했지만, 이러한 연산을 구현하는 것은 우선 코드의 가독성이 떨어지고 연산적으로 비효율적이기 때문에 위에서 구현한 코드를 그대로 유지하도록 하겠습니다.

기존 코드에 KeepTopK 전략을 취해 노이즈를 추가한 코드는 아래와 같습니다. 이 부분에 대한 코드는 여기서 확인할 수 있습니다.

add noise code

위 코드에서 (2) noise_std라는 파라미터를 추가함으로써 생성하는 가우시안 노이즈의 수준을 입력으로 줄 수 있도록 했고, (23~24) 이 값이 0보다 크면 라우터의 출력에 가우시안 노이즈를 생성해 더하는 식으로 구현했습니다.

이러한 노이즈의 영향으로 비슷한 수준의 라우터 확률을 가지는 전문가들 간의 교체가 이뤄지게 됩니다. 하지만 이는 선택되는 전문가들 간의 다양성을 어느정도 보장해줄뿐, 아예 낮은 확률로 선택에서 배제돼 버리는 전문가가 생기는 문제를 해결할 순 없었기에 여기에 새로운 방법이 추가됩니다.

전문가 확률 분포에 따른 공분산 비교[6]

바로 라우터가 제시하는 전문가 확률 혹은 가중치 간의 공분산을 최소화하도록 유도함으로써 전체적으로 고르게 전문가들을 선택하도록 유도하는 것입니다. 이를 통해 지속적으로 배제되던 전문가들에게도 충분히 기회를 주며, 동시에 각 전문가층을 고루 활용하도록 모델을 훈련시킬 수 있게 됩니다. 

이를 코드로 구현하면 아래와 같습니다. 여타 부분은 동일하기 때문에 공분산 개념이 구현된 MoE 클래스의 forward 부분입니다. 공분산을 통한 auxiliary loss가 적용된 전체 코드는 여기에서 확인할 수 있습니다.

로드밸런서가 구현된 MoE의 Forward

위 코드는 기존 코드의 forward 부분에서 (13~24)에 해당하는 공분산을 구하는 부분이 추가된 코드입니다. 구체적으로 각 샘플별로 라우터에 의해 제시된 전문가별 확률을 평균하고, 이들 간의 공분산을 구해 제시하는 방식입니다. (47) 이렇게 산출된 로드 밸런스 손실은 output과 함께 반환되며, 모델의 학습 과정에서 추가적인 손실로 함께 고려되게 됩니다.

|Experiments - replace auxiliary loss with GRN

하지만 이러한 방법은 로드 밸런서 손실을 얼마나 고려할 것인지에 대해 하이퍼 파라미터인 알파(α)값을 통해 결정해야 하고, 최적의 알파값을 찾기 위해 추가적인 학습 리소스가 요구될 수도 있습니다. 또한 각 레이어를 고르게 활용해야 한다는 개념은 우리가 이전 ConvNeXt V2 포스팅에서 구현했던 GRN의 개념과도 어울립니다. 

다만 CV를 이용한 방식에서는 배치 사이즈가 충분히 크고 데이터 로더에 의해 입력되는 배치 데이터가 충분히 다양하다면 설령 하나의 샘플에서 소수의 전문가만 선택되더라도 다른 유형의 문제를 통해 보간될 수 있습니다. 하지만 GRN은 모든 샘플에 대해서 모든 전문가층이 최대한 활성화되도록 유도하기 때문에 소수의 전문가를 각각의 케이스에 특화시킨다라는 방향에는 어울리지 않을 수도 있습니다.

하지만 일단 해보면 재밌을 것 같기 때문에 cv를 고려하는 것 대신 GRN을 이용해 MoE를 구현한 후, 이들을 간단히 테스트해보도록 하겠습니다. 우선 CV 대신 GRN을 적용한 코드를 구현하기 위해 우리가 convnext v2에서 구현했던 GRN을 다시 정의해보겠습니다.

GRN code

위 코드는 convnext v2 포스팅 때와 달라진 것이 없습니다. 이에 대해 잘 모르신다면 해당 포스팅을 참조하는 것을 권장드립니다. 이렇게 구현한 grn을 이용해 공분산을 대체하는 식으로 MoE를 구현하면 다음과 같습니다. 전체 코드는 여기를 통해 확인할 수 있습니다.

cv -> grn

위 코드는 공분산을 적용하기 전 MoE에서 (32) 가우시안 노이즈가 추가된 gate_logits에 대해 grn을 적용한 결과입니다. 이제 이렇게 정의한 코드들의 성능을 확인하는 결과만이 남았는데요. tutel 라이브러리를 통해 MoE 클래스를 구현하고, 이에 걸맞게 우리가 구현한 MoE 클래스들을 조금 조정한 뒤 간단하게 cifar10 데이터셋에 학습시켜 그 성능을 확인해보겠습니다.

|Performance Test ①

우선 성능 측정을 하기 전에 크게 두 개의 조정을 먼저 해야 합니다. 첫 번째로 tutel을 통해 공식적으로 구현된 MoE 레이어의 경우, 전문가 레이어가 단일 선형층으로 구성된 것이 아닌 '(선형) 레이어 → 활성화 → (선형) 레이어'와 같은 식으로 구현돼 있습니다. 이로 인한 모델의 규모에서 차이가 날 수 있기 때문에 우리가 구현한 모델도 이와 비슷하게 맞춰주도록 하겠습니다. 

또한 우리가 테스트할 데이터셋이 rgb 3개 색 차원을 가진 이미지 데이터셋이기 때문에 입력 데이터의 형상을 조절하고, 최종 분류기 층을 하나 더 추가해주도록 하겠습니다. 이러한 과정은 우리가 구현한 original moe와 grn moe, 그리고 tutel moe 모두 동일하게 진행됩니다.

|Model Update - 2 layers with activation

구현된 tutel 코드는 아래와 같습니다. 기본적으로 활성화 함수를 사이에 낀 두 개의 선형층으로 구성돼 있으며, 최종 예측을 위해 분류기 역할을 하는 선형층 하나를 추가한 구조입니다. 아래 코드는 여기에서 확인할 수 있습니다.

tutel이 적용된 MoE 모델

이와 맞춰주기 위해 구현된 GRN MOE는 아래와 같습니다. 전체 코드는 여기에서 확인할 수 있습니다.

수정된 GRN 코드

위 코드를 통해 확인할 수 있듯 수정된 사항은 크게 세가지 입니다. (14) 하나는 전문가 층을 두 개로 늘렸고, (20) 두 번째로 최종 분류를 위한 분류층을 추가했습니다. 그리고 (34) 마지막으로 입력 데이터가 3개의 색 채널을 가진 데이터임을 고려해 데이터의 형상을 변환해주는 코드입니다.

이러한 방식으로 우리가 구현한 MoE도 수정해주면 다음과 같습니다. 코드는 여기에서 확인할 수 있습니다.

수정된 MoE 코드

이렇게 구현된 코드들에 더해 간단하게 별도의 선형층으로만 구성된 모델을 구현할 것인데요, 우리가 구현한 모델들의 성능을 비교하기 위한 대조군의 역할입니다. 그에 대한 코드는 아래와 같습니다. 아래 코드 및 모델 학습에 대한 코드는 여기에서 확인할 수 있습니다.

대조군 역할을 하는 선형 모델

대조군 역할을 하는 간단한 모델은 기존의 MoE를 다시 선형 레이어로만 구성한 것으로써 파라미터 수는 동일한 수준으로 맞춰지도록 구현했습니다. 실제 각 모델들의 파라미터 수를 살펴보면 다음과 같습니다.

|Model Info

구현한 모델들의 파라미터 수

왼쪽 위가 일반 선형 레이어로 구성된 모델이며, 오른쪽 위가 직접 구현한 MoE가 적용된 모델, 왼쪽  아래가 MoE에서 cv를 이용한 로드 밸런서 손실 대신 grn을 적용한 모델, 마지막으로 오른쪽 아래가 공식적으로 MoE를 지원하는 tutel 라이브러리를 통해 구현한 모델의 파라미터입니다.

Tutel에서 제시한 문서를 기반으로 모델을 구현했고, 모델 구조 자체를 뜯어봤을 땐 구조상 큰 차이는 없는 것으로 봤는데 파라미터 수가 10,000,000개 정도가 차이가 납니다. 아마 Tutel에서 구현된 FusedExpertNetwork에 그 이유가 있는 것 같습니다(추측이며, 추후에 확인된다면 다시 수정토록 하겠습니다).

|Model Performance

위 모델들을 간단히 cifar10 데이터셋에 학습시켜 그 성능을 1차적으로 확인해보도록 하겠습니다.

각 모델들의 성능

우선 우리가 구현한 모델은 convolution과 같이 2차원 데이터를 이해하는 데 적합한 모델도 아니며, 이미지를 학습하고 이해할 수 있을만큼 모델의 scalability가 높지도 않습니다. 또 그렇다고 scheduler와 같은 별다른 학습 기법을 추가적으로 적용시키지도 않았으며, 학습 epoch도 30으로 아주 적습니다. 하지만 그럼에도 우리는 각 모델들이 가진 성능을 대략적으로 확인할 수는 있었습니다.

|Performance Test ①

위에서 간략히 구현한 결과를 살펴보면 tutel로 구현한 모델이 파라미터 때문인지 가장 성능이 낮았고, 이후로 선형층으로 구성된 모델, MoE가 적용된 모델, MoE에 GRN이 적용된 모델의 순으로 성능 차이가 났습니다. 그렇다면 일반적인 모델에 적용된 결과는 어떨까요?

제대로 된 모델과의 결합을 테스트해보고, 추가적으로 충분한 수준의 학습을 진행해보도록 하겠습니다.

|Model Update - ConvNeXt V2 + α

이전 ConvNeXt V2 포스팅을 통해 리뷰하고 구현했던 ConvNeXt V2 구조를 가져오겠습니다. 기본적으로 MoE는 Transformer Block 구조 내의 FFN을 대체하기 위한 모듈입니다. 하지만 여타 블로그나 공식 코드를 통해 NLP와 CV 분야 모두에 걸쳐 고루 성능이 올라간다고 실험돼 있는만큼, 그렇다면 transformer 구조가 아닌 convnext 구조에 적용하고 최종 출력 전의 선형층만을 변경하면 어떨까 하는 생각에서 진행한 실험입니다.

우선 이러한 실험 전에 ConvNeXt에 존재하는 두 개의 PointWise Convolution을 MoE로 대체해 구현하고, 이를 다양한 epoch나 학습 조건 하에서 테스트했으나 성능이 유사하거나 오히려 떨어지는 모습을 보였습니다. 때문에 마지막으로 최종 출력을 위한 선형층 전에 MoE를 추가하는 식으로 코드를 업데이트하였고, 코드는 아래와 같은 식입니다. 이에 대한 코드는 여기에서 확인할 수 있습니다.

ConvNeXt V2 + MoE

기존의 ConvNeXt V2 구조에서 최종 분류를 위한 선형층을 통과시키기 전에 MoE를 적용하는 층을 추가해주었습니다. 그리고 파라미터의 숫자를 맞춰주기 위해 대조군으로 구현한 ConvNeXt V2 모델에서는 MoE 대신 선형 레이어를 추가해주었습니다. 그리고 이전에 구현했던 공분산을 이용한 auxiliary loss 대신 grn을 적용한 모델까지 총 3개의 모델을 테스트해보도록 하겠습니다.

|Model Info

각 모델의 파라미터 수를 살펴보면 다음과 같습니다. 이 파트 이하의 코드들은 여기에서 확인할 수 있습니다.

ConvNeXt original(좌) vs MoE 적용 vs MoE&GRN 적용

파라미터 수는 위에서 볼 수 있다시피 대동소이 합니다. 또한 조금이라도 파라미터 수가 큼으로 인해 모델의 성능이 좋아지는 것을 방지하기 위해 대조군 모델(일반 ConvNeXt v2 구조)이 조금이라도 많은 파라미터를 지니고 있습니다.

|Model Performance

8개의 전문가들과 그 중 2개의 모델들의 예측을 통해 구현한 모델의 결과는 다음과 같습니다. ConvNeXt V2 구조를 적용함에 따라 단 10에포크의 학습만으로도 아래와 같은 성능을 보입니다.

10에포크 학습 결과

이처럼 적은 수의 epoch를 구현했을 때에는 일반 선형 모델을 적용했을 때와 MoE를 적용했을 때가 비슷한 수준의 성능을 보였고, GRN을 적용했을 땐 조금 더 빠르게 성능이 높아지는 결과를 확인할 수 있었습니다. 그렇다면 100epoch의 학습에서는 어땠을까요? 참고로 실험 결과 1,000 에포크까지 학습시킨 결과와 100에포크 학습 결과는 대동소이했습니다.

100에포크 학습 결과

이번에는 100에포크 학습 결과입니다. 이번에도 여전히 MoE를 적용한 쪽이 성능이 훨씬 좋았습니다. 그리고 이때 확인됐던 점은 적은 수의 학습이 진행될 경우에는 2개의 전문가층을 활용하는 topk2가 더 효율적이었으며, 4개 전문가 중 1개 전문가의 의견을 채택하는 것보다 8개 전문가 중 2개 전문가 의견을 채택하는 케이스가 성능이 훨씬 뛰어났습니다.

반대로 충분한 수의 에포크가 보장된다면 topk1이 유의미하게 성능이 높았습니다. 마지막으로 1,000에포크 학습한 결과를 살펴보고 이번 포스팅을 마치도록 하겠습니다.

1,000 에포크 학습 결과

1,000 에포크의 학습을 진행하니 MoE에 GRN을 추가한 모델의 성능은 미약하게 올라갔으나 큰 차이가 없던 반면, 선형 모델을 적용한 경우나 일반적인 MoE를 적용한 경우엔 꽤 유의미한 수준으로 성능이 올라간 것을 확인할 수 있었습니다. 이는 아무래도 GRN의 적용이 초기에 지나치게 빠른 성능 수렴을 유도하는게 아닐까하는 생각이 들기도 합니다.

참고로 위 수행결과들은 이 코드를 통해 argparse로 모델 타입과 학습율, 그리고 학습 에포크 등에 대한 정보를 주어 tmux 환경에서 수행되었습니다.

|Reference

[1] Weilin Cai, Juyong Jiang, et al. "A Survey on Mixture of Experts." 

https://arxiv.org/pdf/2407.06204. 

[2] William Fedus, Jeff Dean, Barret Zoph. "A REVIEW OF SPARSE EXPERT MODELS IN DEEP LEARNING." https://arxiv.org/pdf/2209.01667.

[3]  Changho Hwang, et al. "Tutel: Adaptive Mixture-of-Experts at Scale." 

https://arxiv.org/abs/2206.03382.

[4]  Changho Hwang, et al. "Switch Tranformer." https://arxiv.org/abs/2206.03382.

[5] Brian Chu, et al. "Training MoEs at Scale with Pytorch." 

https://pytorch.org/blog/training-moes/.

[6] Maarten Grootendorst. "A Visual Guide to Mixture of Experts (MoE)." 

https://newsletter.maartengrootendorst.com/p/a-visual-guide-to-mixture-of-experts.

[7] Noam Shazeer1, et al. "OUTRAGEOUSLY LARGE NEURAL NETWORKS: THE SPARSELY-GATED MIXTURE-OF-EXPERTS LAYER." https://arxiv.org/pdf/1701.06538

[8] Jiang, Albert Q., et al. "Mixtral of experts." arXiv preprint arXiv:2401.04088.