1. 서론
MoE 모델은 입력 토큰을 여러 전문가(expert)에게 분배하여 처리하는 방식으로, 대규모 언어모델의 계산 비용을 효과적으로 관리할 수 있는 방법을 제공한다. 본 논문에서는 MoE 모델의 전문가 수를 확장할 때의 성능 개선을 조사하고, 이를 통해 최적의 전문가 수를 도출하고자 한다.
2. 배경
2.1 MoE 모델의 메커니즘
MoE는 각 입력 토큰을 K명의 전문가 중 일부에게만 전달하는 라우터와 전문가 집합으로 구성된다. 이 구조는 Transformer 모델의 FFN 레이어를 MoE 레이어로 대체하여 구현된다. 전문가의 아키텍처는 대체된 FFN과 동일하다.
2.2 스케일링 법칙
기존의 연구는 Transformer 모델의 파라미터 수와 훈련 토큰 수가 모델의 성능에 미치는 영향을 멱법칙(Power Law)으로 설명하고 있다. 이런 연구를 바탕으로, MoE 모델에 대한 새로운 스케일링 법칙을 개발하였다.
3. 방법
3.1 실험 설정
다양한 전문가 수(4, 8, 16, 32)를 갖는 모델을 대규모 파라미터로 훈련시켜 MoE 모델의 스케일링 행동을 관찰한다. training dataset은 SlimPajama에서 균등하게 샘플링한다.
3.2 MoE의 스케일링 법칙 공식화
MoE 모델의 성능 개선을 관찰하고 다음과 같은 새로운 스케일링 법칙을 제안한다.
MoE (Mixture of Experts) 모델의 성능 개선을 설명하는 새로운 스케일링 법칙
\[\log L(N, D, E) = \log\left(A N^\alpha (E-1) + B \hat{E}^\beta + C D^\gamma\right) + F\]본 논문의 MoE 스케일링 법칙에 따르면
MoE 스케일링 법칙 Scaling Law 색인마킹
4. 인퍼런스 비용 추정
MoE 모델의 높은 성능은 인퍼런스 단계에서의 비용 증가를 수반한다. 인퍼런스 비용을 다음과 같이 정의하고 모델링한다.
\[C_{\text{Model}, G} = \frac{G \cdot C_0}{T_{\text{Model}}(G)}\]$G$는 사용된 GPU 수, $C_0$는 GPU 운영 비용, $T_{\text{Model}}(G)$는 모델의 처리량이다.
5. 결과
모델의 훈련 비용과 인퍼런스 비용 간의 트레이드오프를 고려하여, 손실 최적화된 설정에서 벗어나 과학습된 설정을 사용하여 더 낮은 인퍼런스 비용으로 높은 성능을 달성한다. 이 접근 방식은 인퍼런스 비용을 상당히 절감하면서도 비교적 높은 성능을 유지한다.
6. 관련 작업 이전 연구는 Transformer 모델의 스케일링 행동을 광범위하게 탐구하였으나, MoE 모델의 특성을 고려한 연구는 제한적이다. 본 연구는 MoE 모델의 스케일링 법칙을 체계적으로 분석하고, 훈련과 인퍼런스 단계에서의 비용을 종합적으로 고려한 새로운 접근 방식을 제안한다.
Appendix
A. 밀집 모델을 위한 최적 할당
밀집 모델에 대한 최적 할당 문제는 라그랑지안 프레임워크를 사용하여 훈련 예산을 최적화하는 문제로 정의됩니다. 이 과정에서 파라미터의 수 \(N\), 훈련 토큰 \(D\), 그리고 훈련 비용 \(C\) 사이의 관계를 고려합니다. 최적화 문제는 다음과 같이 표현됩니다.
\[L(N, D, \lambda) = L_0 + A N^\alpha + B D^\beta + \lambda (6ND - C)\]\(\lambda\)는 라그랑지 승수로, 주어진 제약조건 \(6ND = C\)를 만족하는 \(N\)과 \(D\)의 값을 찾는 데 사용됩니다. 이 함수를 \(N\)과 \(D\)에 대해 미분하여 다음과 같은 조건을 도출합니다.
\(\frac{\partial L}{\partial N} = -\alpha A N^{\alpha-1} + 6\lambda D = 0\) \(\frac{\partial L}{\partial D} = -\beta B D^{\beta-1} + 6\lambda N = 0\)
이 두 미분을 0으로 설정하고, 주어진 제약 \(6ND = C\)을 사용하여 최적의 \(N\)과 \(D\)를 계산합니다. 이를 통해 학습 비용과 성능 간의 최적 균형을 찾을 수 있습니다.
[[Check]]
라그랑지안 멀티플라이어 (Lagrange Multiplier) 상세 [참고자료 1]
라그랑지안 멀티플라이어 (Lagrange Multiplier)
라그랑지안 멀티플라이어는 제약 조건이 있는 최적화 문제에서 사용되는 기법입니다. 이 방법은 제약 조건을 만족하면서 목적함수(objective function)의 극값(최대값 또는 최소값)을 찾을 때 사용됩니다.
예를 들어, 훈련 비용 \(C\)를 상수로 고정하고, 주어진 비용 내에서 모델 성능(e.g., 검증 손실)을 최소화하는 파라미터 \(N\)과 훈련 토큰 \(D\)를 찾고자 할 때 이 기법이 유용합니다.
최적화 문제의 설정
최적화 문제는 다음과 같이 설정됩니다.
\[L(N, D, \lambda) = L_0 + A N^\alpha + B D^\beta + \lambda (6ND - C)\]수식의 미분
목적함수 \(L\)을 \(N\)과 \(D\)에 대해 각각 편미분하여 최소값을 찾습니다.
이 식에서 \(N\)에 대한 최적값을 찾기 위해 위 식을 \(N\)에 대해 풀어야 합니다.
마찬가지로, \(D\)에 대한 최적값을 찾기 위해 이 식을 \(D\)에 대해 풀어야 합니다.
제약 조건의 활용
제약 조건 \(6ND = C\)을 사용하여 \(N\) 또는 \(D\) 중 하나를 다른 변수로 표현할 수 있습니다. 예를 들어, \(D\)를 \(N\)과 \(C\)를 사용하여 표현할 수 있습니다.
\[D = \frac{C}{6N}\]이를 통해 위의 미분 식들에 대입하여 \(N\)만의 함수로 문제를 간소화하고, 이를 통해 \(N\)에 대한 최적값을 계산할 수 있습니다. 이후 \(D\)의 최적값도 찾을 수 있습니다.
이렇게 계산된 \(N\)과 \(D\)는 주어진 비용 \(C\) 내에서 모델 성능을 최적화하는 파라미터입니다.
B. 훈련 세부 사항
MoE 모델의 아키텍처는 여러 전문가 네트워크(FFN)로 구성되어 있으며, 각 입력 토큰은 최대 \(K\)개의 전문가에게 라우팅됩니다. training dataset로는 고품질의 SlimPajama 데이터셋을 사용하며, 이 데이터셋은 인터넷 콘텐츠, 코드, 위키백과, 책, Arxiv, StackExchange에서 수집된 자료를 포함하고 있습니다. 모델은 Megatron-Deepspeed 프레임워크를 사용하여 32개의 GPU에서 텐서 병렬화와 데이터 병렬화를 활용하여 훈련됩니다.
C. 인퍼런스 비용 추정
MoE 모델의 크기는 해당 밀집 모델의 크기를 참조하며, 모델의 총 파라미터 수는 대략 \(N \times (1 + (E-1)c)\)로 표현될 수 있습니다. \(c\)는 모델 아키텍처에 따라 달라집니다. 모델은 각 토큰을 처리할 때 필요한 KV-캐시 메모리를 고려하여 최대 배치 크기를 계산하고, 이를 바탕으로 인퍼런스 비용을 추정합니다.
이런 분석은 MoE 모델의 훈련과 인퍼런스 단계에서 비용과 성능 간의 최적의 균형을 찾는 데 중요한 역할을 합니다. 효율적인 모델 설계를 위해서는 이런 이론적 배경과 계산 과정이 필수적입니다.
[참고자료 1] 라그랑지안
문제 정의
훈련 예산 \(C\)를 최적화하기 위해 다음과 같은 목적함수를 고려합니다.
\[L(N, D, \lambda) = L_0 + A N^\alpha + B D^\beta + \lambda (6ND - C)\]라그랑지안 함수 설정
라그랑지안 함수는 목적함수와 제약 조건을 결합한 형태로, 다음과 같이 설정됩니다.
\[L(N, D, \lambda) = L_0 + A N^\alpha + B D^\beta + \lambda (6ND - C)\]미분을 통한 최적화
라그랑지안 함수를 각 변수에 대해 편미분합니다.
\(N\)에 대한 편미분
\(\frac{\partial L}{\partial N} = \alpha A N^{\alpha-1} + 6\lambda D = 0\) \(\alpha A N^{\alpha-1} = -6\lambda D\) \(\lambda = -\frac{\alpha A N^{\alpha-1}}{6D}\)
\(D\)에 대한 편미분
\(\frac{\partial L}{\partial D} = \beta B D^{\beta-1} + 6\lambda N = 0\) \(\beta B D^{\beta-1} = -6\lambda N\) \(\lambda = -\frac{\beta B D^{\beta-1}}{6N}\)
제약 조건 활용
제약 조건 \(6ND = C\)를 사용하여 \(D\)를 \(N\)의 함수로 표현합니다.
\[D = \frac{C}{6N}\]이 값을 \(\lambda\)의 표현식에 대입하여 \(\lambda\)를 \(N\)의 함수로 표현합니다.
두 식의 비교
두 식을 비교하여 최적의 \(N\)과 \(D\)를 구합니다.
\[-\frac{\alpha A N^{\alpha-1}}{6D} = -\frac{\beta B D^{\beta-1}}{6N}\]\(D = \frac{C}{6N}\)를 대입하여 다음과 같이 정리합니다.
\(\frac{\alpha A N^{\alpha-1}}{D} = \frac{\beta B D^{\beta-1}}{N}\) \(\alpha A N^{\alpha-1} \cdot N = \beta B \cdot \left(\text{C}{6N}\right)^{\beta-1} \cdot D\) \(\alpha A N^\alpha = \beta B \left(\text{C}{6N}\right)^{\beta-1} \cdot \frac{C}{6N}\) \(\alpha A N^\alpha = \beta B \left(\text{C}{6N}\right)^\beta\) \(\alpha A N^\alpha = \beta B \left(\text{C^\beta}{(6N)^\beta}\right)\) \(\alpha A N^\alpha = \beta B \frac{C^\beta}{6^\beta N^\beta}\) \(\alpha A N^\alpha = \beta B \frac{C^\beta}{6^\beta N^\beta}\) \(\alpha A = \beta B \frac{C^\beta}{6^\beta N^{\beta + \alpha}}\)
최적의 $N$과 $D$ 값 계산
이제 최적의 \(N\)을 구할 수 있습니다.
\(N^{\alpha + \beta} = \frac{\beta B C^\beta}{\alpha A 6^\beta}\) \(N = \left( \frac{\beta B C^\beta}{\alpha A 6^\beta} \right)^{\frac{1}{\alpha + \beta}}\)
최적의 \(N\) 값을 찾은 후, 이를 \(D = \frac{C}{6N}\)에 대입하여 최적의 \(D\) 값을 계산합니다.
\(D = \frac{C}{6N}\) \(D = \frac{C}{6 \left( \frac{\beta B C^\beta}{\alpha A 6^\beta} \right)^{\frac{1}{\alpha + \beta}}}\)
이 과정을 통해 \(N\)과 \(D\)의 최적 값을 찾을 수 있습니다. 이 값들은 주어진 훈련 예산 \(C\) 내에서 모델 성능을 최적화하는 파라미터입니다.
[참고자료 2] 라그랑지안 멀티플라이어 (Lagrange Multiplier)
기본 이론
라그랑지안 멀티플라이어 방법은 제약 조건 하에 목적함수를 최대화하거나 최소화할 때 사용됩니다. 이 방법은 제약 조건을 활용하여 목적함수에 추가적인 항(라그랑지안 항)을 도입하고, 이 새로운 함수의 극값을 찾음으로써 원래 문제의 해를 구합니다.
아래 예시는 간단한 선형방정식이지만, 다차원의 방정식의 경우 라그랑지안을 사용하는 것이 더 체계적이며 접근가능한 방법일 수 있습니다.
[예제 2-1] 제한된 예산 하에서 이익 최대화
문제 정의
한 기업이 두 가지 유형의 제품 \(A\)와 \(B\)를 생산하고 있고, 각 제품의 이익(단위당)은 각각 100원, 150원입니다.
기업은 제품 \(A\)를 만드는 데 2만 원, 제품 \(B\)를 만드는 데 3만 원의 비용이 든다고 가정하고, 기업의 총 예산이 60만 원이라면 기업이 최대 이익을 얻기 위해 각 제품을 몇 개씩 생산해야 하나?
수식으로 표현
라그랑지안 함수 설정
라그랑지안 함수 \(L\)을 다음과 같이 설정합니다.
\[L(A, B, \lambda) = 100A + 150B - \lambda(2A + 3B - 60)\]\(\lambda\)는 라그랑지안 멀티플라이어입니다.
미분을 통한 최적화
라그랑지안 함수를 각 변수에 대해 편미분합니다.
해 찾기
위의 첫 두 미분을 이용해 \(\lambda\)를 구합니다.
\(\lambda\) 값을 제약 조건 식에 대입하여 \(A\)와 \(B\)를 구합니다.
[예제 2-2] 제한조건 하에서의 최적화
한 농부가 두 가지 작물, 즉 옥수수와 콩을 재배하려고 합니다. 옥수수는 단위당 5달러의 이익을 내고, 콩은 단위당 3달러의 이익을 냅니다. 농부는 100 제곱미터의 경작지를 소유하고 있고, 옥수수는 제곱미터당 1일, 콩은 제곱미터당 2일의 작업일이 필요하며, 농부는 최대 150일의 작업할 수 있습니다. 농부가 최대 이익을 얻기 위해 각 작물을 몇 제곱미터씩 재배해야 하는지 구하시오.
문제 수식화
라그랑지안 함수 설정
라그랑지안 함수를 설정하면 다음과 같습니다.
\[L(x, y, \lambda) = 5x + 3y - \lambda_1(x + y - 100) - \lambda_2(x + 2y - 150)\]\(\lambda_1\)와 \(\lambda_2\)는 라그랑지안 멀티플라이어입니다.
라그랑지안 함수를 각 변수에 대해 편미분하면,
두 식을 연립하여 풀면: \(\lambda_2 = 2\) \(\lambda_1 = 3\)
이제 제약 조건을 고려하여 \(x\)와 \(y\)의 값을 구합니다.
위 두 식을 연립하면 다음과 같고, \(y = 50\) \(x = 50\)
따라서 농부는 옥수수와 콩을 각각 50 제곱미터씩 재배해야 최대 이익을 얻을 수 있습니다. 이 경우 총 이익은 다음과 같습니다.
\[P = 5(50) + 3(50) = 250 + 150 = 400\]라그랑지안 멀티플라이어 \(\lambda_1\)와 \(\lambda_2\)는 제약 조건이 목적함수에 미치는 한계 이익 예를 들어, \(\lambda_1 = 3\)는 총 경작지 제약 조건이 1단위 증가할 때마다 총 이익이 3만큼 증가한다는 것을 의미하고, \(\lambda_2 = 2\)는 최대 작업일 제약 조건이 1단위 증가할 때마다 총 이익이 2만큼 증가한다는 것을 의미합니다.
라그랑지안 멀티플라이어 방법을 사용할 때, 특히 선형 제약 조건과 선형 목적함수를 가진 문제에서는 여러 최적 해가 존재할 수 있습니다. 이는 제약 조건을 만족하는 동안 목적함수가 동일한 값을 가지는 여러 포인트가 존재할 수 있기 때문입니다. 라그랑지안 멀티플라이어를 사용하여 이 문제의 모든 해를 찾기 위해서는 추가적인 해석이나 다른 방식으로 모든 가능한 해를 검토해야할 수 있습니다.