전문가 혼합(MoE) 토큰 라우팅 원리: 라우터·희소성·용량 초과 토큰 드롭까지

핵심 메시지

• 허깅페이스 허브에는 이미 2,800개가 넘는 MoE 모델이 있을 만큼 이 구조가 널리 쓰인다.
• MoE 층은 하나의 MLP를 여러 개로 복제해 '전문가'로 만든 것으로, 각 토큰은 모든 전문가가 아니라 일부(상위 k개)에게만 라우팅된다. 이것이 희소성의 핵심이다.
• 라우터는 토큰을 선형 투영하는 층이며, 그 결과인 라우터 로짓은 각 토큰이 어느 전문가로 갈 가능성이 높은지를 담는다.
• 상위 k 선택 후 나머지에 음의 무한대를 넣고 소프트맥스를 적용하면 비활성 전문가는 0이 되어 라우터 확률(가중치)이 만들어진다.
• 전문가마다 처리 용량(슬롯)이 정해져 있어, 용량을 초과해 몰린 토큰은 효율을 위해 드롭된다.

쉽게 이해하기

허깅페이스의 오리트라가 진행하는 이 영상은 전문가 혼합(Mixture of Experts, MoE)에서 가장 핵심이자 까다로운 부분인 '토큰 라우팅'을 슬라이드와 코랩 노트북 코드로 함께 뜯어본다. 발표자는 허깅페이스 허브에 이미 2,813개의 MoE 모델이 올라와 있을 만큼 지금이 이 구조를 이야기할 적기라고 말한다. 학습이나 추론 방법이 아니라 라우팅 알고리즘 하나에 집중한다.

먼저 MoE 층이 무엇인지 짚는다. 트랜스포머 층에서 원래 MLP(피드포워드)가 있던 자리에, 그 MLP를 여러 개로 복제해 각각을 '전문가'라고 부르면 그것이 MoE 층이다. 중요한 건 각 토큰이 모든 전문가로 가지 않고 하나·둘·셋 정도로만 라우팅된다는 점이다. 일부만 켜고 나머지는 끄는 희소성으로 학습과 추론의 계산 예산을 아낀다.

예제는 토큰 4개, 전문가 4개, 토큰당 상위 2개 전문가 선택, 전문가당 슬롯(용량) 2개로 설정한다. 라우터는 토큰을 선형 투영해 '라우터 로짓'을 만든다. 이 행렬은 각 토큰(행)이 각 전문가(열)로 갈 가능성을 나타낸다. 여기서 상위 k개만 남기고 나머지 자리에는 음의 무한대를 넣은 뒤 소프트맥스를 적용하면, 비활성 전문가는 0이 되고 활성 전문가만 가중치를 갖는 '라우터 확률'이 된다. 0은 곧 그 토큰이 그 전문가로 갈 확률이 전혀 없다는 뜻이며, 이것이 MoE 희소성의 본질이다.

다음은 슬롯 선택과 토큰 드롭이다. 각 토큰의 전문가 선택을 원핫 인코딩하고 우선순위(1순위·2순위)를 분리한 뒤, 누적합에서 1을 빼는 방식으로 각 토큰이 전문가의 몇 번 슬롯에 들어가는지 계산한다. 이때 어떤 전문가가 용량(슬롯 2개)을 넘겨 과다 배정되면, 초과한 토큰은 효율을 위해 드롭된다. 예제에서는 토큰 2와 토큰 3이 초과로 드롭되고, 이들의 가중치도 0으로 처리해 전파되지 않게 한다. 마지막으로 전문가·슬롯·토큰·가중치 정보를 한데 담은 최종 가중치 행렬을 만들어, 실제로 어떤 토큰이 어느 전문가의 어느 슬롯에 자리 잡는지를 완성한다.

주요 인사이트

• MoE의 마법은 특별한 새 부품이 아니라, MLP를 여러 개 복제한 뒤 '전부가 아니라 일부에게만' 토큰을 보내는 라우팅에 있다.
• 라우터는 거창한 것이 아니라 토큰을 전문가 점수로 바꾸는 선형 층이며, 음의 무한대 + 소프트맥스 조합이 비활성 전문가를 0으로 만들어 희소성을 구현한다.
• 전문가마다 용량(슬롯)이 있어 무한정 받지 못한다. 용량을 넘겨 몰린 토큰을 드롭하는 것은 버그가 아니라 효율을 위한 설계다.
• 2차원 행렬 하나가 순열과 누적합 같은 연산을 거치며 토큰·전문가·슬롯·가중치라는 여러 정보를 동시에 표현한다는 점이 라우팅 구현의 묘미다.

자주 묻는 질문

관련 AI 소식