MoE: 토큰마다 사용할 MLP 선택하기

앞선 글에서는 MLP가 각 토큰의 벡터를 독립적으로 변환하며, 같은 MLP의 가중치가 모든 토큰에 동일하게 적용된다는 점을 살펴봤습니다. 이번 글에서는 이 MLP 자리에 여러 MLP를 두고, 토큰마다 사용할 MLP를 선택하는 MoE(Mixture of Experts)를 살펴보겠습니다. 어떤 MLP를 선택하고, 선택한 MLP의 결과를 어떻게 합치는지 따라가 보겠습니다.

Dense MLP와 MoE

지금까지 살펴본 구조에서는 디코더 블록의 MLP 하나를 모든 토큰에 적용했습니다. 이를 Dense MLP라고 하겠습니다. 반면 이 글에서 다룰 Sparse MoE는 여러 MLP 중 일부만 선택해 각 토큰을 처리합니다. 이때 각각의 MLP를 expert라고 부르고, 사용할 expert를 선택하는 부분을 Router라고 부릅니다.

그림 1은 같은 토큰의 입력 x₁을 두 구조에 넣었을 때의 흐름입니다. 왼쪽에서는 x₁이 하나의 Gated MLP를 거쳐 y₁이 됩니다. 오른쪽에는 E1부터 E4까지 네 expert가 있고, Router가 그중 E1과 E3을 선택합니다. 선택한 두 expert에는 같은 x₁을 입력하지만, 서로 다른 가중치를 사용하므로 각자 다른 변환을 계산합니다. 이 예시의 expert들은 앞서 살펴본 Gated MLP이며, 각각 Gate·Up·Down projection의 가중치를 가집니다.

선택한 expert의 출력에는 Router가 정한 가중치를 곱하고, 같은 위치의 성분끼리 더합니다. 그림에서는 E1의 결과에 약 0.27, E3의 결과에 약 0.73을 곱했습니다. 두 결과 모두 d차원이므로 가중합한 출력 y₁도 d차원입니다. 따라서 MLP 자리를 MoE로 바꾸어도 디코더 블록의 d차원 흐름은 유지되고, 결과를 같은 위치의 Residual stream에 더할 수 있습니다.

MoE의 핵심은 모델이 가진 전체 파라미터 수와 토큰 하나의 계산에 참여하는 파라미터 수를 구분할 수 있다는 점입니다. Expert를 더 두더라도 토큰마다 그중 두 개만 선택하도록 만들 수 있습니다. 선택하지 않은 expert의 파라미터도 저장해야 하므로, 일부 expert만 계산한다고 해서 파라미터를 저장하는 메모리가 줄어들지는 않습니다. 전체 파라미터를 유지하면서 토큰마다 계산에 참여하는 파라미터를 일부로 제한하는 구조입니다.

Router: expert 선택과 가중치 계산

Router는 토큰의 입력 벡터를 받아 어떤 expert를 사용할지, 각 expert의 결과를 어느 비중으로 합칠지 정합니다. 그림 2에서는 d = 8인 토큰 하나와 네 expert를 사용합니다. Router의 가중치 행렬 Wᵣ은 8 × 4 형태입니다. 여기에 입력 x₁을 곱하면 네 개의 점수가 나오고, 각 점수는 expert 하나에 대응합니다.

이 예시의 점수는 E1부터 순서대로 [1, 0, 2, −1]입니다. 이 중 큰 점수 두 개를 고르는 것이 Top-2입니다. 가장 큰 점수 2는 E3에, 다음 점수 1은 E1에 해당하므로 두 expert를 선택합니다. 그림의 Expert 1 (점수 1)에서 앞의 1은 expert 번호이고, 괄호 안의 1은 Router가 계산한 점수입니다.

다음으로 선택한 두 점수에 Softmax를 적용해 결과를 합칠 가중치를 계산합니다. 점수 1과 2의 지수값을 각각 두 지수값의 합으로 나누면 E1은 약 0.269, E3은 약 0.731이 됩니다. 두 가중치의 합은 1이며, 이 예시에서는 점수가 더 높은 E3의 결과를 더 큰 비중으로 반영합니다. 선택하지 않은 E2와 E4는 이 토큰에 대해 계산하지 않습니다.

Router에서 Combine까지

이제 토큰 세 개를 함께 처리하는 흐름을 살펴보겠습니다. 각 토큰은 d = 8인 벡터이고, 세 벡터를 행으로 쌓은 입력 X는 3 × d 형태입니다. 전체 과정은 Router → Dispatch → Expert 계산 → Combine으로 이어집니다. Dispatch는 선택 결과에 따라 입력을 expert별로 모으는 단계이고, Combine은 expert의 출력을 원래 토큰별로 모아 가중합하는 단계입니다.

그림 3의 Router에서는 토큰마다 expert 두 개와 가중치를 정합니다. 토큰 1은 E1과 E3, 토큰 2는 E2와 E3, 토큰 3은 E1과 E3을 선택했습니다. 토큰 1과 토큰 3처럼 같은 expert들을 선택하더라도 입력이 다르면 가중치는 달라질 수 있습니다. 이 예시에서 E1의 비중은 토큰 1에서는 약 0.27, 토큰 3에서는 약 0.88입니다.

입력 X에서 Router와 Dispatch로 화살표가 각각 내려오는 이유는 같은 입력을 두 용도로 사용하기 때문입니다. Router는 X를 사용해 선택과 가중치를 계산합니다. Dispatch는 그 선택 정보를 받아, X에 들어 있는 원래 토큰 벡터를 해당 expert에 전달합니다. Expert에 들어가는 것은 Router가 출력한 점수가 아니라 토큰의 입력 벡터입니다.

Dispatch에서는 같은 expert가 처리할 토큰을 묶습니다. E1에는 x₁과 x₃, E2에는 x₂, E3에는 x₁·x₂·x₃가 모입니다. E4는 어느 토큰도 선택하지 않아 이 예시에서는 입력이 없습니다. 토큰 하나가 expert 두 개를 선택했으므로, 같은 입력 벡터가 두 expert의 입력 묶음에 각각 들어갑니다. 색은 expert가 아니라 원래 토큰을 구분합니다.

Expert 계산에서는 모인 토큰마다 해당 expert의 MLP를 적용합니다. 그림의 E₁(x₁)은 E1에 x₁을 입력한 결과를 뜻합니다. 같은 묶음에 있는 x₁과 x₃에는 E1의 가중치를 동일하게 적용하지만, 두 토큰의 값을 서로 섞지는 않습니다. 앞서 여러 토큰을 행으로 쌓아 MLP를 계산했을 때와 같이, 각 출력은 해당 토큰의 입력으로 계산합니다.

마지막 Combine에서는 각 expert에서 나온 결과를 원래 토큰별로 모읍니다. 토큰 1의 출력 y₁은 E₁(x₁)에 약 0.27을 곱하고, E₃(x₁)에 약 0.73을 곱해 더한 결과입니다. 그림 하단의 점선은 Router가 계산한 가중치가 이 가중합에 사용된다는 뜻입니다. 두 d차원 벡터를 성분별로 더하므로 출력도 d차원이고, y₁·y₂·y₃를 원래 토큰 순서로 놓으면 전체 출력 Y는 다시 3 × d 형태가 됩니다.

이 과정에서 expert별로 여러 토큰을 묶지만, 서로 다른 토큰의 정보를 결합하는 것은 아닙니다. 이 예시의 Router는 각 토큰의 입력으로 선택과 가중치를 계산하고, expert도 각 토큰을 독립적으로 변환합니다. Combine 역시 같은 토큰에서 나온 결과끼리만 합칩니다. 다른 토큰의 정보를 가져와 반영하는 Attention과 달리, 여기서는 한 토큰에 어떤 변환을 적용할지가 달라집니다.

Shared expert: 모든 토큰에 적용하는 변환

지금까지 살펴본 expert들은 Router가 선택한 토큰만 처리합니다. 이런 expert를 Routed expert라고 부릅니다. 일부 MoE 모델은 여기에 Router의 선택과 관계없이 모든 토큰에 적용하는 Shared expert를 함께 둡니다. 그림 4에서는 기존 routed expert 네 개와 Top-2 선택을 유지하고, shared expert 하나를 추가했습니다.

왼쪽의 Shared expert S에는 x₁·x₂·x₃가 모두 들어갑니다. 같은 MLP를 각 토큰에 독립적으로 적용하므로, 모든 토큰을 처리한다는 점에서는 앞서 살펴본 Dense MLP와 같습니다. 오른쪽은 그림 3의 과정을 간략하게 표시한 것입니다. Router가 expert를 선택하고, 선택된 expert의 결과를 Combine에서 가중합합니다. 이 routed 경로의 결과를 R(x)라고 하겠습니다.

두 경로는 같은 입력을 사용하고, 마지막에는 같은 토큰의 결과끼리 성분별로 더합니다. 토큰 1의 출력 y₁은 S(x₁) + R(x₁)입니다. 두 결과가 모두 d차원이므로 합한 출력도 d차원이며, 토큰 2와 3에도 같은 방식이 적용됩니다. 이 예시에서 토큰 하나는 shared expert 한 개와 선택된 routed expert 두 개, 총 세 expert의 계산을 거칩니다. Shared expert는 routed expert의 Top-2 선택에 포함되지 않습니다.

Shared expert를 두는 이유는 여러 입력에 공통으로 필요한 변환을 별도로 학습하도록 하기 위해서입니다. DeepSeekMoE 논문은 공통 지식을 shared expert가 학습하도록 하여 routed expert들이 비슷한 내용을 중복해서 학습하는 것을 줄이려는 설계라고 설명합니다. 이는 각 expert가 담당하는 내용을 미리 정해 준다는 뜻이 아니라, 공통으로 사용하는 경로와 선택해서 사용하는 경로를 나누는 설계 의도입니다.

MoE는 디코더 블록의 MLP 자리에 여러 expert를 두고, 토큰마다 일부를 선택해 계산하는 구조입니다. 여기에 모든 토큰에 적용하는 shared expert를 함께 둘 수도 있습니다. 어느 경우든 같은 토큰의 expert 결과를 합치며, 출력은 원래 토큰의 위치와 d차원을 유지합니다. 다음 글에서는 지금까지 살펴본 구성 요소를 다시 연결해 모델 전체 흐름을 정리하겠습니다.