A.X K2는 한국 정부의 소버린 AI(Sovereign AI) 프로젝트의 일환으로 개발된 대규모 Mixture-of-Experts(MoE) 기반의 에이전트형 파운데이션 모델입니다. 총 688B 파라미터 중 33B 파라미터가 활성화되는 구조를 가지며, 복잡한 추론을 위한 'Thinking mode'와 저지연 응답을 위한 'Non-thinking mode'를 하나의 모델에서 제어하는 Think-Fusion 기법을 도입했습니다.

핵심 방법론 및 기술적 특징

1. 아키텍처 및 MoE 구조:
   • 256개의 라우팅 전문가(Routed Experts)와 1개의 공유 전문가(Shared Expert)로 구성되며, 토큰당 8개의 전문가를 활성화합니다.
   • Multi-head Latent Attention (MLA)를 기반으로 하며, 학습 안정성과 효율성을 위해 Gated Transformer Block 설계를 채택했습니다.

2. Sparse Gated Attention (SGA):
   • 긴 문맥(Long-context) 추론 효율을 높이기 위해 Sparse Attention과 Gated Attention을 결합했습니다.
   • 경량 인덱서(Indexer)가 top- $k$ ( $k=2048$ )개의 토큰을 선택하여 어텐션 연산량을 대폭 절감합니다.
   • 헤드별 출력 게이트(Head-specific output gate)를 통해 어텐션 싱크(Attention sinks)를 억제하여 모델이 문맥상 중요한 토큰에 집중하도록 최적화했습니다.

3. Gated Norm (GN):
   • RMSNorm 직후에 입력 의존적 게이트를 배치하여 이상치(Outliers)를 억제합니다. 이는 대규모 학습의 안정성을 높이고, FP8이나 NVFP4와 같은 저정밀도 추론 환경에서 양자화 오차를 줄이는 데 기여합니다.

4. Think-Fusion 및 학습 전략:
   • 추론 깊이를 동적으로 제어하기 위해 paired SFT(Supervised Fine-Tuning) 데이터셋을 활용하여 하나의 모델이 <think>...</think> 태그를 통한 단계적 추론과 직접적인 응답을 모두 수행할 수 있도록 설계되었습니다.
   • 8.2조 개의 토큰을 활용한 3단계 커리큘럼 학습을 거쳤으며, 기본적으로 MXFP8 (E4M3) 형식을 사용하여 학습 및 체크포인트 저장을 수행했습니다.

5. 장문 문맥 처리:
   • 128K 네이티브 학습 후, YaRN 스케일링을 통해 최대 256K까지의 문맥 길이를 지원합니다. RULER 및 Needle-in-a-Haystack(NIAH) 테스트에서 높은 성능을 입증했습니다.

데이터 및 활용

• 학습 데이터는 영어(주력), 한국어(약 15%), 코드, 기타 다국어로 구성되었으며, 고품질 데이터 선별 및 PII 마스킹 파이프라인을 거쳤습니다.
• Hugging Face Transformers 및 vLLM 프레임워크를 통해 별도의 추가 양자화 과정 없이 배포 가능하며, Apache-2.0 라이선스로 제공되어 연구 및 상업적 목적 모두 활용 가능합니다.