키미 K3 MoE 구조 해설: 잠재 전문가 혼합과 학습을 지탱한 세 가지 안정화 장치

핵심 메시지

• 키미 K3는 전체 2.8조 파라미터 가운데 토큰마다 1,040억 개 정도만 실제로 사용하는 희소 모델이다.
• 층마다 전문가를 896개까지 늘리고 토큰당 16개만 켜는 대신, 토큰 임베딩을 절반 크기의 잠재 벡터로 압축해 통신 비용을 줄였다.
• 압축과 복원을 거치며 값이 폭주하기 쉬워지자, 전문가 출력에 RMSNorm을 넣고 활성화 함수에 상한을 씌워 학습을 안정화했다.
• 일부 전문가에게만 토큰이 몰리는 문제는 라우터 점수 분포에서 보정값을 직접 계산하는 '분위수 균형' 방식으로 해결했다.

쉽게 이해하기

트랜스포머에서 파라미터가 가장 많이 쌓이는 곳은 각 블록의 피드포워드 신경망이다. 성능을 올리려면 이 층의 폭을 키우면 되지만, 폭을 그대로 늘리면 모든 토큰이 늘어난 가중치를 전부 통과하므로 학습과 추론이 함께 느려지고 메모리도 더 든다. 전문가 혼합(MoE)은 이 큰 신경망을 여러 개의 작은 전문가로 쪼갠 뒤 토큰마다 일부만 켜서, 계산량은 유지한 채 모델의 표현력만 키우는 설계다.

어떤 전문가를 켤지는 라우터라는 작은 신경망이 정한다. 라우터는 전문가마다 점수를 매기고 상위 몇 개를 고르며, 최종 출력은 선택된 전문가들의 결과를 라우터 점수로 가중합한 값이다. 키미 K3는 이 전문가 풀을 896개까지 늘리고 토큰당 16개를 켜는 '세분화된 전문가' 구조를 택했다. 여기에 모든 토큰이 항상 거치는 공유 전문가 두 개를 따로 두어, 공통 지식은 공유 전문가가 맡고 전문화된 지식은 선택된 전문가가 맡도록 역할을 나눴다.

문제는 통신 비용이다. 전문가들이 여러 GPU에 흩어져 있는데 7,168차원짜리 토큰 임베딩을 선택된 전문가마다 통째로 보내야 하기 때문이다. 키미 K3는 하나의 하향 투영 행렬로 토큰을 3,584차원 잠재 벡터로 한 번만 압축한 뒤, 이 작은 벡터를 모든 라우팅 전문가에게 보내고 결과를 다시 원래 차원으로 복원한다. 전문가별로 압축 행렬을 따로 두지 않는다는 점이 핵심이며, 이렇게 아낀 예산을 더 많은 전문가와 더 많은 활성 전문가에 쓴다.

대신 토큰이 하향 투영 → 전문가 내부 신경망 → 상향 투영으로 이어지는 행렬 곱 사슬을 지나면서 값이 단계마다 증폭돼, 낮은 정밀도 수치 형식이 감당할 범위를 넘길 위험이 커진다. 강연은 이를 막는 두 장치를 소개한다. 하나는 전문가 출력을 합친 뒤 상향 투영 직전에 RMSNorm을 넣어 항상 일정한 크기의 입력이 들어가게 하는 것이고, 다른 하나는 전문가 내부 활성화 함수를 손보는 것이다. 흔히 쓰는 SwiGLU는 값이 커질수록 제곱에 가깝게 자라 상한이 없는데, 두 갈래(내용·게이트)의 선형 항에 각각 탄젠트 기반 소프트캡을 씌우면 곱의 크기가 두 상한의 곱을 넘지 못한다. 키미 K3는 게이트에 4, 값에 25를 걸어 활성화가 100을 넘지 않게 했다.

마지막 난관은 부하 균형이다. 학습 초기에는 모든 전문가가 똑같이 서툴기 때문에, 우연히 먼저 선택된 소수가 더 자주 갱신되며 더 잘하게 되고 라우터는 다시 그들에게 몰리는 자기강화가 일어난다. 보조 손실로 균형을 강제하면 언어 모델링 성능과 맞바꾸게 되고, 딥시크 V3처럼 전문가별 편향값을 조금씩 밀어 올리는 방식은 오차의 방향만 알 뿐 크기를 모르므로 전문가가 수백 개로 늘면 수렴이 느리다. 키미 K3의 분위수 균형은 각 토큰의 선택 문턱값을 기준으로 마진을 구한 뒤, 전문가별 목표 부하 위치에 해당하는 마진을 그대로 편향값으로 삼는다. 탐색이나 보폭 조정 없이 한 번에 계산되므로 훨씬 빠르고 안정적으로 균형에 도달한다.

주요 인사이트

• 모델 폭은 토큰당 계산량을 늘리지 않고 파라미터를 더할 수 있는 유일한 축이라, 희소 전문가 구조가 규모 경쟁의 주요 무대가 되고 있다.
• 전문가 수를 늘리는 진짜 비용은 연산이 아니라 GPU 사이를 오가는 통신량이며, 잠재 공간 압축은 그 병목을 겨냥한 해법이다.
• 희소성을 극단으로 밀면 성능보다 수치 안정성이 먼저 무너진다. 정규화 위치와 활성화 함수의 상한 같은 세부가 규모 확장의 전제 조건이 된다.
• 부하 균형을 손실 항으로 강제하면 모델 성능과 경쟁하지만, 라우팅 점수 분포에서 보정값을 직접 계산하면 성능을 깎지 않고 균형을 맞출 수 있다.
• 라우팅 편향은 어떤 전문가가 뽑히는지만 바꾸고 최종 가중합에는 원래 점수를 그대로 쓴다는 점이, 균형과 성능을 분리하는 열쇠다.

자주 묻는 질문

관련 AI 소식