한국어 고성능 MoE 언어모델 A.X K1: 생각·비생각 전환이 가능한 519B 파라미터 모델

초록

A.X K1은 519 억 파라미터(활성 33 억) 규모의 Mixture‑of‑Experts(LMoE) 모델로, 10 조 토큰을 75일간 1 024대 H200 GPU에서 사전학습하였다. 스케일링 법칙을 기반으로 160 K 토큰 vocab을 설계하고, Think‑Fusion 훈련 레시피를 도입해 “생각 모드”(깊은 추론)와 “비생각 모드”(빠른 응답)를 하나의 모델에서 사용자 제어가 가능하도록 했다. 한국어 벤치마크에서 기존 오픈소스 모델을 능가하며, 효율적인 추론과 고품질 데이터 파이프라인을 강조한다.

상세 분석

A.X K1은 현재 공개된 가장 큰 MoE 기반 한국어 모델 중 하나로, 총 파라미터 519 억 중 실제 활성화되는 파라미터는 33 억에 불과하도록 설계했다. 이는 Tian et al. (2025)의 MoE 스케일링 법칙을 그대로 적용한 결과이며, 고정된 FLOPs 예산(≈2.55 × 10²⁴ FLOPs) 하에서 “파라미터‑데이터 비율”을 최적화했다는 점이 핵심이다. 특히, 전문가(Expert) 용량을 d model = 7 168, d expert = 2 048으로 설정해 granularity G = 7을 선택했는데, 이는 이상적인 8‑12 범위보다 약간 낮지만, 부하 불균형 상황에서 안정성을 높이는 전략으로 해석된다.

아키텍처는 Multi‑head Latent Attention(MLA)를 채택해 KV‑cache 효율을 개선하고, RMSNorm을 MoE 블록 입·출력에만 적용하는 듀얼 정규화 방식을 도입했다. 이는 Gemma 팀(2025)의 두 단계 정규화 아이디어를 간소화한 형태로, 초기 학습 단계에서 손실 급등을 현저히 억제한다. 또한, 64개의 어텐션 헤드와 공유된 dense expert를 사용해 메모리 오버헤드를 감소시키면서도 지식 공유 효과를 얻었다.

토크나이저는 BBPE 기반 160 K vocab을 사용했으며, 이는 Tao et al. (2024)의 파생 기반 최적화 결과를 25 % 확대한 선택이다. vocab 크기를 128의 배수로 맞춰 하드웨어 정렬을 최적화하고, 다국어(영어·한국어·중·일·스페인어) 커버리지를 확보했다.

데이터 파이프라인은 원시 웹·PDF·코드·STEM 등 다양한 소스를 수집한 뒤, 품질·도메인·난이도 분류기 3단계 필터링을 거쳐 10 조 토큰을 구축했다. 특히, 자체 VL 모델을 활용한 PDF 레이아웃 파싱과, 경량 LLM 기반 시드‑코퍼스 재구성·주제‑기반 합성 데이터 생성으로 고품질·고난이도 학습 데이터를 대량 확보했다.

Think‑Fusion 레시피는 모델 머징과 SFT(Instruction‑tuned) 단계에서 “Thinking”과 “Non‑Thinking” 라벨을 혼합 학습함으로써, 추론 시 토큰 생성 길이와 계산량을 동적으로 조절할 수 있게 한다. 이는 복잡한 문제에서는 깊은 체인‑오브‑사고를, 간단한 질의에서는 최소 토큰으로 빠른 응답을 제공한다는 실용적 장점을 만든다.

평가 결과, A.X K1은 영어·한국어 전반의 벤치마크(지식, 지시, 수학, 코딩)에서 최신 오픈소스 모델과 동등하거나 우수한 성능을 보였으며, 특히 한국어 MMLU·KoBench 등에서 현저히 앞섰다. 추론 효율성 측면에서도 활성 파라미터 33 억 기반의 높은 토큰당 처리량을 유지하면서, Think‑Fusion 모드 전환에 따른 지연 시간 차이를 2‑3배 수준으로 제어했다.

전반적으로 A.X K1은 제한된 컴퓨팅 예산 하에서 MoE 스케일링 법칙을 실증적으로 적용한 사례이며, 한국어 특화 고성능 모델을 오픈소스로 제공함으로써 주권 AI 생태계 구축에 기여한다는 전략적 의미도 크다.