Kimi K3 구조 해설: 잠재 MoE와 분위수 균형으로 2.8조 파라미터를 안정 학습한 방법

핵심 메시지

• Kimi K3는 총 2.8조 파라미터 규모지만 토큰 하나가 실제로 통과하는 전문가는 896개 중 16개뿐이라, 모델 크기와 계산량을 분리한 전형적인 전문가 혼합(MoE) 설계다.
• 전문가 수를 늘리면 토큰 임베딩을 여러 GPU로 실어 나르는 통신 비용이 급증하는데, Kimi K3는 임베딩을 7,168차원에서 3,584차원으로 한 번만 압축해 모든 전문가에 보내는 '잠재 MoE'로 이 비용을 줄인다.
• 압축·전문가·복원으로 이어지는 행렬 곱 사슬은 값이 증폭되기 쉬워, 전문가 출력 합에 RMS 정규화를 걸고 전문가 내부 활성화에는 상한을 씌워 저정밀도 연산에서의 오버플로를 막았다.
• 특정 전문가만 계속 선택되는 쏠림을 막기 위해, 편향값을 조금씩 더듬어 찾는 대신 라우터 점수 분포에서 목표 부하에 해당하는 분위수를 직접 계산해 한 번에 맞추는 '분위수 균형'을 쓴다.
• 정규화·활성화 상한·분위수 균형 세 가지가 함께 작동해야 896개 전문가가 모두 실제로 학습되며, 이 조합이 조 단위 확장을 가능하게 한 핵심이라는 것이 영상의 결론이다.

쉽게 이해하기

이 영상은 최근 공개된 대형 언어모델 Kimi K3의 구조 가운데 '모델 폭'을 키우는 부분, 즉 전문가 혼합(Mixture of Experts) 설계를 집중적으로 뜯어본다. 발표자는 Kimi K3가 코딩과 지식, 추론 벤치마크에서 최상위 비공개 모델과 견줄 만한 성능을 보인다는 점을 출발점으로 삼는다. 시퀀스 방향으로는 멀티헤드 잠재 어텐션과 Kimi 델타 어텐션을 1:3으로 섞고, 깊이 방향으로는 일반적인 잔차 연결 대신 어텐션 잔차를 써서 앞선 블록의 표현에 직접 접근하게 했다고 정리한다.

트랜스포머 블록의 파라미터 대부분은 피드포워드 층에 몰려 있다. 은닉층 차원을 그냥 키우면 능력은 오르지만 모든 토큰이 그 가중치를 전부 쓰게 되어 학습·추론 속도와 메모리가 함께 나빠진다. 전문가 혼합은 이 큰 피드포워드 층을 여러 개의 작은 전문가로 쪼갠 뒤, 라우터가 매긴 점수 상위 몇 개만 토큰마다 활성화한다. Kimi K3는 층마다 896개의 세분화된 전문가를 두고 토큰당 16개만 켜며, 여기에 모든 토큰이 항상 거치는 공유 전문가 2개를 더해 공통 지식과 전문 지식을 나눠 맡긴다.

문제는 활성 전문가 수를 늘릴수록 통신 비용이 커진다는 점이다. 7,168차원짜리 토큰 임베딩을 선택된 전문가마다 GPU 간에 보내야 하기 때문이다. 잠재 MoE는 전문가별로 따로 압축하는 것이 아니라 하나의 하향 투영 행렬로 토큰을 절반 크기인 3,584차원으로 한 번만 압축하고, 그 작은 벡터를 선택된 전문가 모두에게 보낸다. 전문가 출력을 라우터 확률로 가중합한 뒤 상향 투영으로 원래 차원에 되돌리고, 공유 전문가의 출력과 더해 최종 결과를 만든다.

이렇게 아낀 자원으로 전문가를 더 많이, 더 자주 켤 수 있게 되지만 대신 학습이 불안정해진다. 하향 투영·전문가 내부 연산·상향 투영으로 이어지는 곱셈 사슬에서 값이 단계마다 증폭되면 저정밀도 형식이 담을 수 있는 범위를 넘어서기 때문이다. Kimi K3는 두 곳에서 이를 막는다. 먼저 전문가 출력을 합친 직후 상향 투영 전에 RMS 정규화를 넣어, 라우팅에 따라 들쭉날쭉한 합의 크기를 일정하게 맞춘다. 발표자는 이 단순한 변경이 안정성뿐 아니라 검증 손실과 다운스트림 성능까지 개선했다고 설명한다.

두 번째는 전문가 내부의 활성화 함수다. 게이트와 콘텐츠를 곱하는 게이트 선형 유닛 계열은 두 항이 모두 무한정 커질 수 있어 입력이 커지면 출력이 제곱처럼 튄다. Kimi K3는 두 갈래 각각에 소프트캡(tanh 계열 상한)을 씌워 게이트에는 4, 값에는 25라는 상한을 두고, 결과 활성화가 100을 넘지 못하게 한다. 원점 근처에서는 기존 활성화와 거의 같게 동작하되 큰 값에서만 눌리는 방식이라 표현력을 크게 해치지 않는다는 설명이다.

주요 인사이트

• MoE의 진짜 병목은 계산량이 아니라 전문가 사이를 오가는 통신량이며, 토큰을 한 번만 압축해 공유하는 잠재 MoE는 '아낀 통신을 전문가 수로 바꿔 쓰는' 교환에 가깝다.
• 라우터에 보조 손실을 붙여 균형을 맞추는 기존 방식은 언어모델 손실과 같은 가중치를 잡아당겨 성능과 균형을 맞바꾸게 되는데, 편향값 조정 방식은 최종 가중합에는 원래 점수를 그대로 쓰기 때문에 이 상충을 피한다.
• 편향을 조금씩 더하고 빼는 방식은 오차의 부호만 알 뿐 크기를 모르기 때문에, 크게 틀어진 전문가나 거의 맞은 전문가나 같은 속도로 움직여 수렴이 느리고 하이퍼파라미터에 민감하다.
• 분위수 균형은 토큰마다 상위 K 진입선을 기준으로 마진을 계산해 토큰별 점수 편차를 지운 뒤, 목표 부하에 해당하는 순위의 마진을 바로 편향으로 삼는다. 탐색도 학습률도 없이 한 번에 값을 정하는 셈이다.
• 결국 세 장치는 각각 다른 실패 모드를 막는다. 정규화는 층 사이 폭주를, 활성화 상한은 전문가 내부 폭주를, 분위수 균형은 전문가 사멸을 막으며 이 중 하나만 빠져도 조 단위 확장은 어렵다.

자주 묻는 질문

관련 AI 소식