TL;DR

이 연구를 시작하게 된 배경과 동기는 무엇입니까?

대규모 언어 모델(LLM)의 발전은 인공지능 분야에서 혁명적인 변화를 예고하고 있습니다. 그러나 기존 모델들은 정적인 데이터 모방에 그치며, 실제 환경에서 자율적으로 추론하고 행동하는 능력에 한계를 보였습니다. 특히 도구 사용, 소프트웨어 개발, 복잡한 다단계 추론과 같은 에이전틱 인텔리전스 영역에서 기존 모델들의 성능은 매우 제한적이었습니다. 연구팀은 이러한 한계를 극복하고, 모델이 단순한 응답 생성을 넘어 실제 환경과 상호작용하며 학습하고 적응할 수 있는 새로운 접근법의 필요성을 인식했습니다.

대규모 언어 모델의 훈련과 정렬에는 여러 근본적인 도전과제가 존재합니다. 고품질 데이터의 제한된 가용성, 훈련 불안정성, 복잡한 에이전틱 능력의 학습 어려움 등이 주요 장애물로 작용해왔습니다. 특히 토큰 효율성과 모델의 일반화 능력을 동시에 확보하는 것은 대규모 언어 모델 연구에서 가장 중요한 과제 중 하나로 인식되어 왔습니다. 이러한 배경에서 연구팀은 에이전틱 인텔리전스의 새로운 패러다임을 제시하고자 했습니다.

이 연구에서 제시하는 새로운 해결 방법은 무엇입니까?

Kimi K2 연구는 에이전틱 인텔리전스를 위한 혁신적인 세 가지 핵심 접근법을 제시합니다. 첫째, MuonClip 옵티마이저는 대규모 언어 모델의 훈련 불안정성 문제를 해결하기 위해 개발되었습니다. 이 옵티마이저는 가중치 클리핑 메커니즘을 통해 어텐션 로짓의 폭발적 증가를 제어하면서, 동시에 모델의 토큰 효율성을 극대화합니다. 둘째, 대규모 에이전틱 데이터 합성 파이프라인은 실제 환경과 유사한 도구 사용 시나리오를 대규모로 생성할 수 있는 혁신적인 방법론을 제시합니다. 이 파이프라인은 다양한 도구, 에이전트, 작업을 체계적으로 시뮬레이션하여 고품질의 에이전틱 훈련 데이터를 생성합니다.

셋째, 통합 강화학습 프레임워크는 검증 가능한 보상과 자기 비판 메커니즘을 결합하여 모델의 학습 능력을 획기적으로 향상시킵니다. 이 프레임워크는 객관적 성능 지표와 주관적 평가 기준을 동시에 고려하여, 모델이 다양한 작업 도메인에서 일관되고 유연한 성능을 발휘할 수 있도록 설계되었습니다. 특히 자기 비판 루브릭 보상 메커니즘은 모델이 스스로의 출력을 평가하고 개선할 수 있는 혁신적인 접근법을 제시합니다.

제안된 방법은 어떻게 구현되었습니까?

Kimi K2의 구현은 복잡하고 정교한 다단계 접근법을 따릅니다. 사전 훈련 단계에서는 320억 개의 활성화된 파라미터와 1조 개의 총 파라미터를 가진 Mixture-of-Experts(MoE) 모델 아키텍처를 채택했습니다. MuonClip 옵티마이저를 사용하여 15.5조 개의 토큰으로 모델을 훈련했으며, 이 과정에서 단 한 번의 손실 스파이크도 발생하지 않았습니다. 데이터 처리에서는 웹 텍스트, 코드, 수학, 지식 영역의 고품질 데이터를 큐레이션하고, 데이터 재구성 기법을 통해 토큰 유용성을 극대화했습니다.

후훈련 과정은 지도 학습 파인튜닝과 강화학습의 하이브리드 접근법을 채택했습니다. 대규모 에이전틱 데이터 합성 파이프라인을 통해 도구 사용 시나리오를 생성하고, 검증 가능한 보상 시스템과 자기 비판 메커니즘을 결합하여 모델의 학습을 진행했습니다. 특히 강화학습 인프라는 코로케이션 아키텍처, 효율적인 엔진 전환 메커니즘, 에이전틱 롤아웃 지원 등 혁신적인 시스템 설계를 포함합니다.

이 연구의 결과가 가지는 의미는 무엇입니까?

Kimi K2는 에이전틱 인텔리전스 분야에서 새로운 최첨단 성능을 달성했습니다. 다양한 벤치마크에서 오픈소스 및 클로즈드소스 모델들을 능가하는 성과를 보여주었습니다. 특히 SWE-bench Verified에서 65.8%, ACEBench에서 76.5%, AIME 2025에서 49.5%, GPQA-Diamond에서 75.1%의 뛰어난 성적을 기록했습니다. LMSYS Arena 리더보드에서는 3,000명 이상의 사용자 투표를 바탕으로 1위 오픈소스 모델이자 전체 5위를 차지했습니다.

이 연구의 의의는 단순한 성능 향상을 넘어서는 것입니다. 오픈소스로 공개된 기본 모델과 후훈련된 체크포인트는 연구 커뮤니티가 에이전틱 인텔리전스를 대규모로 탐색하고 개선할 수 있는 중요한 기반을 제공합니다. MuonClip 옵티마이저, 에이전틱 데이터 합성 파이프라인, 통합 강화학습 프레임워크 등의 혁신적인 접근법은 향후 대규모 언어 모델 연구의 새로운 방향을 제시합니다. 이 연구는 AI가 단순한 데이터 모방을 넘어 실제 환경에서 자율적으로 학습하고 추론할 수 있는 가능성을 보여주는 중요한 이정표로 평가됩니다.

Kimi K2: 오픈 에이전틱 인텔리전스

초록

Kimi K2는 320억 개의 활성화된 파라미터와 1조 개의 총 파라미터를 가진 Mixture-of-Experts(MoE) 대규모 언어 모델입니다. 이 모델은 혁신적인 MuonClip 옵티마이저를 도입하여 기존 Muon 알고리즘의 뛰어난 토큰 효율성을 유지하면서도 새로운 QK-clip 기법을 통해 훈련 불안정성 문제를 해결했습니다.

MuonClip을 기반으로 K2는 15.5조 개의 토큰으로 사전 훈련되었으며, 단 한 번의 손실 스파이크도 발생하지 않았습니다. 후훈련 과정에서 K2는 대규모 에이전틱 데이터 합성 파이프라인과 결합된 강화학습 단계를 포함한 다단계 후훈련 프로세스를 거쳤습니다. 이 과정에서 모델은 실제 환경과 합성 환경에서의 상호작용을 통해 능력을 향상시켰습니다.

Kimi K2는 오픈소스 비사고(non-thinking) 모델 중에서 최첨단 성능을 달성했으며, 특히 에이전틱 능력에서 강점을 보입니다. 주목할 만한 성과로는 Tau2-Bench에서 66.1점, ACEBench(En)에서 76.5점, SWE-Bench Verified에서 65.8점, SWE-Bench Multilingual에서 47.3점을 기록하여 비사고 설정에서 대부분의 오픈소스 및 클로즈드소스 기준선을 능가했습니다.

또한 K2는 코딩, 수학, 추론 작업에서도 강력한 능력을 보여주었습니다. LiveCodeBench v6에서 53.7점, AIME 2025에서 49.5점, GPQA-Diamond에서 75.1점, OJBench에서 27.1점을 기록했으며, 이 모든 결과는 확장된 사고 과정 없이 달성되었습니다. 이러한 결과는 Kimi K2를 현재까지 가장 유능한 오픈소스 대규모 언어 모델 중 하나로 위치시키며, 특히 소프트웨어 엔지니어링과 에이전틱 작업에서 뛰어난 성능을 보입니다.

서론

대규모 언어 모델(LLM)의 발전은 에이전틱 인텔리전스(Agentic Intelligence)라는 근본적인 패러다임 전환을 겪고 있습니다. 에이전틱 인텔리전스는 모델이 복잡하고 동적인 환경에서 자율적으로 인지하고, 계획하며, 추론하고, 행동할 수 있는 능력을 의미합니다. 이러한 전환은 정적인 모방 학습에서 벗어나 상호작용을 통해 적극적으로 학습하고, 훈련 분포를 넘어서는 새로운 기술을 습득하며, 경험을 통해 행동을 적응시키는 모델로의 변화를 나타냅니다.

이러한 접근법은 AI 에이전트가 정적인 인간 생성 데이터의 한계를 넘어서서 자체적인 탐색과 활용을 통해 초인간적 능력을 획득할 수 있게 한다고 여겨집니다. 따라서 에이전틱 인텔리전스는 도구 사용, 소프트웨어 개발, 실세계 자율성 등 광범위한 영역에 걸쳐 차세대 파운데이션 모델의 핵심 능력으로 빠르게 부상하고 있습니다.

에이전틱 인텔리전스를 달성하는 것은 사전 훈련과 후훈련 모두에서 도전과제를 제시합니다. 사전 훈련은 제한된 고품질 데이터라는 제약 하에서 모델에 광범위한 범용 사전 지식을 부여해야 하므로, 토큰당 학습 신호인 토큰 효율성이 중요한 스케일링 계수로 부각됩니다. 후훈련은 이러한 사전 지식을 실행 가능한 행동으로 변환해야 하지만, 다단계 추론, 장기 계획, 도구 사용과 같은 에이전틱 능력은 자연 데이터에서 드물고 확장하기에 비용이 많이 듭니다. 구조화된 고품질 에이전틱 궤적의 확장 가능한 합성과 선호도 및 자기 비판을 통합하는 일반적인 강화학습 기법이 이러한 격차를 해소하는 데 필수적입니다.

본 연구에서는 이러한 핵심 도전과제를 해결하고 에이전틱 능력의 경계를 확장하기 위해 의도적으로 설계된 1.04조 파라미터 MoE LLM인 Kimi K2를 소개합니다. 우리의 기여는 사전 훈련과 후훈련 영역 모두에 걸쳐 있습니다.

MuonClip 옵티마이저: 토큰 효율적인 Muon 알고리즘과 QK-Clip이라는 안정성 향상 메커니즘을 통합한 새로운 옵티마이저를 제시합니다. MuonClip을 사용하여 15.5조 토큰에서 단 한 번의 손실 스파이크 없이 Kimi K2를 성공적으로 사전 훈련시켰습니다.

대규모 에이전틱 데이터 합성 파이프라인: 시뮬레이션된 환경과 실제 환경을 통해 도구 사용 시연을 체계적으로 생성하는 대규모 에이전틱 데이터 합성 파이프라인을 도입합니다. 이 시스템은 다양한 도구, 에이전트, 작업, 궤적을 구성하여 대규모로 고충실도의 검증 가능하게 정확한 에이전틱 상호작용을 생성합니다.

일반적인 강화학습 프레임워크: 검증 가능한 보상(RLVR)과 자기 비판 루브릭 보상 메커니즘을 결합한 일반적인 강화학습 프레임워크를 설계합니다. 모델은 외부적으로 정의된 작업뿐만 아니라 자신의 출력을 평가하는 것으로부터도 학습하여, 정적 영역에서 개방형 영역으로 정렬을 확장합니다.

Kimi K2는 광범위한 에이전틱 및 최첨단 벤치마크에서 강력한 성능을 보여줍니다. 비사고 평가 설정에서 대부분의 오픈소스 및 클로즈드 가중치 기준선을 능가하며, Claude 4 Opus 및 Sonnet과의 격차를 줄였습니다. 코딩, 수학, 더 넓은 STEM 영역에서도 뛰어난 성과를 거두어 일반적인 작업에서의 능력을 더욱 부각시킵니다.

LMSYS Arena 리더보드(2025년 7월 17일)에서 Kimi K2는 3,000명 이상의 사용자 투표를 바탕으로 1위 오픈소스 모델이자 전체 5위를 기록했습니다. 에이전틱 인텔리전스의 추가적인 발전을 촉진하기 위해, 커뮤니티가 에이전틱 인텔리전스를 대규모로 탐색, 개선, 배포할 수 있도록 기본 모델과 후훈련된 체크포인트를 오픈소스로 공개합니다.

사전 훈련

Kimi K2의 기본 모델은 15.5조 개의 고품질 토큰으로 사전 훈련된 1조 파라미터 규모의 Mixture-of-Experts(MoE) 트랜스포머 모델입니다. 고품질 인간 데이터의 가용성이 점점 제한되는 상황에서, 토큰 효율성이 대규모 언어 모델 스케일링의 핵심 계수로 부상하고 있습니다. 이를 해결하기 위해 토큰 효율성을 극대화하도록 명시적으로 설계된 사전 훈련 기법들이 도입되었습니다.

구체적으로는 토큰 효율적인 Muon 옵티마이저를 채택하고, QK-Clip 도입을 통해 훈련 불안정성을 완화했습니다. 또한 사용 가능한 고품질 토큰에서 지능을 더욱 추출하기 위해 합성 데이터 생성을 통합했습니다. 모델 아키텍처는 경험적 스케일링 법칙 분석에서 도출된 DeepSeek-V3와 유사한 Multi-head Latent Attention(MLA)을 가진 초희소 MoE를 따릅니다.

MuonClip: 가중치 클리핑을 통한 안정적 훈련

Kimi K2는 가중치 감쇠와 일관된 업데이트 RMS 스케일링을 통합한 토큰 효율적인 Muon 옵티마이저를 사용하여 훈련됩니다. 이전 연구인 Moonlight에서의 실험들은 동일한 계산 예산과 모델 크기, 즉 동일한 양의 훈련 데이터 하에서 Muon이 AdamW를 상당히 능가함을 보여주어, 대규모 언어 모델 훈련에서 토큰 효율성을 개선하는 효과적인 선택임을 입증했습니다.

Muon 스케일링 시 훈련 불안정성

효율성에도 불구하고, Muon 훈련을 확장하면 폭발하는 어텐션 로짓으로 인한 훈련 불안정성이라는 도전과제가 드러납니다. 이는 실험에서 Muon에서는 더 자주 발생하지만 AdamW에서는 덜 발생하는 문제입니다. 기존 완화 전략들은 불충분합니다. 예를 들어, 로짓 소프트 캡은 어텐션 로짓을 직접 클리핑하지만, 쿼리와 키 간의 내적은 캡핑이 적용되기 전에 여전히 과도하게 증가할 수 있습니다. 반면 Query-Key Normalization(QK-Norm)은 추론 중에 Key 행렬이 완전히 구체화되지 않기 때문에 Multi-head Latent Attention(MLA)에 적용할 수 없습니다.

QK-Clip으로 Muon 길들이기

이 문제를 해결하기 위해 어텐션 로짓을 명시적으로 제약하는 새로운 가중치 클리핑 메커니즘인 QK-Clip이 제안되었습니다. QK-Clip은 업데이트 후 쿼리와 키 투영 가중치를 재스케일링하여 어텐션 로짓의 증가를 제한하는 방식으로 작동합니다.

트랜스포머 레이어의 입력 표현을 $\mathbf{X}$라 하면, 각 어텐션 헤드 $h$에 대해 쿼리, 키, 값 투영은 다음과 같이 계산됩니다.

여기서 $\mathbf{W}{q}, \mathbf{W}{k}, \mathbf{W}_{v}$는 모델 파라미터입니다. 어텐션 출력은 다음과 같습니다.

배치 $B$에서 헤드별 스칼라인 최대 로짓을 softmax에 대한 최대 입력으로 정의합니다.

여기서 $i, j$는 훈련 샘플 $\mathbf{X}$에서 서로 다른 토큰의 인덱스입니다.

QK-Clip의 핵심 아이디어는 $S_{\max}^{h}$가 목표 임계값 $\tau$를 초과할 때마다 $\mathbf{W}{k}, \mathbf{W}{q}$를 재스케일링하는 것입니다. 중요한 점은 이 연산이 현재 단계의 순방향/역방향 계산을 변경하지 않는다는 것입니다. 단지 최대 로짓을 가중치 증가를 제어하는 강도를 결정하는 안내 신호로 사용할 뿐입니다.

단순한 구현은 모든 헤드를 동시에 클리핑합니다.

여기서 $\gamma = \min(1, \tau/S_{\max})$이고 $S_{\max} = \max_{h}S_{\max}^{h}$이며, $\alpha$는 일반적으로 0.5로 설정되어 쿼리와 키에 동등한 스케일링을 적용하는 균형 파라미터입니다.

그러나 실제로는 소수의 헤드만이 폭발하는 로짓을 보입니다. 모델 훈련에 대한 개입을 최소화하기 위해, 헤드별 스케일링 팩터 $\gamma_{h} = \min(1, \tau/S_{\max}^{h})$를 결정하고 헤드별 QK-Clip을 적용하기로 선택했습니다.

이러한 클리핑은 일반적인 Multi-head Attention(MHA)에서는 직관적입니다. MLA의 경우, 공유되지 않은 어텐션 헤드 구성요소에만 클리핑을 적용합니다.

• $\textbf{q}^{C}$와 $\textbf{k}^{C}$ (헤드별 구성요소): 각각 $\sqrt{\gamma_{h}}$로 스케일링
• $\textbf{q}^{R}$ (헤드별 회전): $\gamma_{h}$로 스케일링
• $\textbf{k}^{R}$ (공유 회전): 헤드 간 영향을 피하기 위해 그대로 유지

위 그림들은 QK-Clip의 효과를 보여줍니다. 왼쪽 그림에서는 중간 규모 훈련 실행 중 어텐션 로짓이 빠르게 1000을 초과하여 잠재적인 수치적 불안정성과 훈련 발산을 초래할 수 있음을 보여줍니다. 반면 오른쪽 그림에서는 MuonClip과 $\tau = 100$을 사용한 Kimi K2의 최대 로짓이 빠르게 100의 캡핑된 값까지 증가하고, 훈련 단계의 약 30% 후에만 안정적인 범위로 감소하여 QK-Clip의 효과적인 조절 효과를 보여줍니다.

MuonClip: 새로운 옵티마이저

가중치 감쇠, 일관된 RMS 매칭, QK-Clip을 단일 옵티마이저로 통합한 것을 MuonClip이라고 합니다. 여러 스케일링 실험을 통해 MuonClip의 효과를 입증했습니다.

먼저 바닐라 Muon을 사용하여 90억 개의 활성화된 파라미터와 530억 개의 총 파라미터를 가진 중간 규모 MoE 모델을 훈련했습니다. 앞서 보여준 그림에서 확인할 수 있듯이, 최대 어텐션 로짓이 빠르게 1000의 크기를 초과하여 이 규모에서 Muon 훈련에서 어텐션 로짓 폭발이 이미 명백함을 관찰했습니다. 이 수준의 최대 로짓은 일반적으로 상당한 손실 스파이크와 간헐적인 발산을 포함한 훈련 중 불안정성을 초래합니다.

다음으로 QK-Clip이 모델 성능을 저하시키지 않음을 보여주고 MuonClip 옵티마이저가 손실 궤적에 악영향을 주지 않으면서 Muon의 최적화 특성을 보존함을 확인했습니다.

마지막으로 $\tau = 100$으로 MuonClip을 사용하여 대규모 MoE 모델인 Kimi K2를 훈련하고 전체 훈련 실행에서 최대 어텐션 로짓을 모니터링했습니다. 처음에는 QK-Clip으로 인해 로짓이 100에서 캡핑됩니다. 훈련 과정에서 최대 로짓은 $\tau$에 대한 조정 없이 점진적으로 일반적인 작동 범위로 감소합니다. 중요한 것은 훈련 손실이 부드럽고 안정적으로 유지되며 관찰 가능한 스파이크가 없다는 점입니다.

위 그림은 Kimi K2의 단계별 훈련 손실 곡선을 보여주며, 전체 훈련 과정에서 스파이크가 없음을 확인할 수 있습니다. 이는 MuonClip이 대규모 언어 모델 훈련에서 어텐션 역학에 대한 견고하고 확장 가능한 제어를 제공함을 검증합니다.

사전 훈련 데이터: 재구성을 통한 토큰 유용성 개선

사전 훈련에서 토큰 효율성은 훈련 중 소비되는 각 토큰에 대해 달성되는 성능 개선의 정도를 의미합니다. 토큰 유용성, 즉 각 토큰이 기여하는 효과적인 학습 신호를 증가시키는 것은 모델 업데이트에 대한 토큰당 영향을 직접적으로 개선하여 토큰 효율성을 향상시킵니다. 이는 고품질 토큰의 공급이 제한되어 최대한 활용되어야 할 때 특히 중요합니다.

토큰 유용성을 증가시키는 단순한 접근법은 동일한 토큰에 반복적으로 노출시키는 것이지만, 이는 과적합과 일반화 감소로 이어질 수 있습니다. Kimi K2의 사전 훈련 데이터에서 Kimi K1.5에 비한 주요 발전은 토큰 유용성을 증가시키기 위한 합성 데이터 생성 전략의 도입입니다. 구체적으로, 상당한 과적합을 유발하지 않으면서 고품질 토큰의 양을 증폭시키기 위해 신중하게 설계된 재구성 파이프라인이 사용됩니다.

지식 데이터 재구성

지식 집약적인 자연 텍스트에 대한 사전 훈련은 트레이드오프를 제시합니다. 단일 에포크는 포괄적인 지식 흡수에 불충분하지만, 다중 에포크 반복은 수익 감소와 과적합 위험 증가를 가져옵니다. 고품질 지식 토큰의 토큰 유용성을 개선하기 위해 다음 핵심 구성요소로 이루어진 합성 재구성 프레임워크가 제안됩니다.

스타일 및 관점 다양성 프롬프팅: 사실적 무결성을 유지하면서 언어적 다양성을 향상시키기 위해 신중하게 설계된 다양한 프롬프트를 적용합니다. 이러한 프롬프트는 대규모 언어 모델이 다양한 스타일과 서로 다른 관점에서 원본 텍스트의 충실한 재구성을 생성하도록 안내합니다.

청크별 자기회귀 생성: 긴 문서에서 전역 일관성을 보존하고 정보 손실을 피하기 위해 청크 기반 자기회귀 재작성 전략을 채택합니다. 텍스트를 세그먼트로 나누고, 개별적으로 재구성한 다음, 다시 연결하여 완전한 구절을 형성합니다. 이 방법은 LLM에 일반적으로 존재하는 암시적 출력 길이 제한을 완화합니다.

위 그림은 긴 입력 발췌문에 대한 자기회귀 청크별 재구성 파이프라인을 보여줍니다. 입력이 컨텍스트가 보존된 작은 청크로 분할되고, 순차적으로 재작성된 후 완전한 재작성된 구절로 연결됩니다.

충실성 검증: 원본과 재작성된 내용 간의 일관성을 보장하기 위해 각 재구성된 구절과 그 소스의 의미적 정렬을 비교하는 충실성 검사를 수행합니다. 이는 훈련 전 초기 품질 제어 단계 역할을 합니다.

데이터 재구성과 다중 에포크 반복을 비교하기 위해 SimpleQA에서 해당 정확도를 테스트했습니다. K2의 초기 체크포인트로 실험하여 세 가지 훈련 전략을 평가했습니다. (1) 원본 데이터셋을 10 에포크 반복, (2) 데이터를 한 번 재구성하고 10 에포크 반복, (3) 데이터를 10번 재구성하고 단일 훈련 패스.

위 표에서 보듯이 이러한 전략들에서 정확도가 일관되게 향상되어 재구성 기반 증강의 효과를 입증합니다. 이 방법을 다른 대규모 지식 코퍼스로 확장했으며 유사하게 고무적인 결과를 관찰했고, 각 코퍼스는 최대 두 번 재구성됩니다.

수학 데이터 재구성

수학적 추론 능력을 향상시키기 위해 SwallowMath에서 도입된 방법론에 따라 고품질 수학 문서를 "학습 노트" 스타일로 재작성합니다. 또한 다른 언어의 고품질 수학 자료를 영어로 번역하여 데이터 다양성을 증가시켰습니다.

데이터셋의 재구성된 하위 집합을 사용한 초기 실험은 유망한 결과를 보여주지만, 지속적인 스케일링을 위한 전략으로서 합성 데이터의 사용은 여전히 활발한 조사 영역입니다. 주요 도전과제에는 사실적 정확성을 손상시키지 않으면서 다양한 소스 도메인에 접근법을 일반화하는 것, 환각과 의도하지 않은 독성을 최소화하는 것, 대규모 데이터셋에 대한 확장성을 보장하는 것이 포함됩니다.

전체 사전 훈련 데이터

Kimi K2 사전 훈련 코퍼스는 웹 텍스트, 코드, 수학, 지식의 네 가지 주요 도메인에 걸친 15.5조 개의 큐레이션된 고품질 토큰으로 구성됩니다. 대부분의 데이터 처리 파이프라인은 Kimi K1.5에서 설명된 방법론을 따릅니다. 각 도메인에 대해 엄격한 정확성과 품질 검증을 수행하고 큐레이션된 데이터셋이 높은 다양성과 효과를 모두 달성하도록 목표 데이터 실험을 설계했습니다.

모델 아키텍처

Kimi K2는 320억 개의 활성화된 파라미터와 1.04조 개의 총 파라미터를 가진 Mixture-of-Experts(MoE) 트랜스포머 모델입니다. 이 아키텍처는 DeepSeek-V3와 유사한 설계를 따르며, Multi-head Latent Attention(MLA)을 어텐션 메커니즘으로 채택하고 있습니다. 모델의 은닉 차원은 7168이고 MoE 전문가 은닉 차원은 2048입니다.

스케일링 법칙 분석을 통해 희소성의 지속적인 증가가 상당한 성능 향상을 가져다준다는 것을 발견했으며, 이는 전문가 수를 DeepSeek-V3의 256개에서 384개로 증가시키는 동기가 되었습니다. 추론 중 계산 오버헤드를 줄이기 위해 어텐션 헤드 수를 DeepSeek-V3의 128개에서 64개로 줄였습니다.

다음 표는 Kimi K2와 DeepSeek-V3 간의 아키텍처 파라미터를 상세히 비교한 것입니다.

희소성 스케일링 법칙

Muon을 사용하여 Mixture-of-Experts(MoE) 모델 패밀리에 맞춘 희소성 스케일링 법칙을 개발했습니다. 희소성은 총 전문가 수와 활성화된 전문가 수의 비율로 정의됩니다. 신중하게 제어된 소규모 실험을 통해 고정된 활성화된 파라미터 수(즉, 일정한 FLOP) 하에서 총 전문가 수를 증가시키는 것(즉, 희소성 증가)이 훈련 및 검증 손실을 일관되게 낮추어 전체 모델 성능을 향상시킨다는 것을 관찰했습니다.

위 그림은 희소성 스케일링 법칙을 보여줍니다. 활성화된 전문가 수를 8개로, 공유 전문가 수를 1개로 고정하고 총 전문가 수를 변화시켜 서로 다른 희소성 수준을 가진 모델들을 만들었습니다. 그래프에서 확인할 수 있듯이 희소성이 증가할수록 모델 성능이 향상됩니다.

구체적으로, 계산 최적 희소성 스케일링 법칙 하에서 동일한 검증 손실 1.5를 달성하기 위해 희소성 48은 희소성 수준 8, 16, 32에 비해 각각 1.69배, 1.39배, 1.15배의 FLOP을 감소시킵니다. 희소성 증가가 더 나은 성능으로 이어지지만, 이러한 이득은 인프라 복잡성 증가를 수반합니다. 모델 성능과 비용의 균형을 맞추기 위해 Kimi K2에서는 희소성 48을 채택하여 순방향 패스당 384개 전문가 중 8개를 활성화합니다.

어텐션 헤드 수

DeepSeek-V3는 메모리 대역폭을 더 잘 활용하고 계산 효율성을 향상시키기 위해 어텐션 헤드 수를 모델 레이어 수의 대략 두 배로 설정합니다. 그러나 컨텍스트 길이가 증가함에 따라 어텐션 헤드 수를 두 배로 늘리는 것은 상당한 추론 오버헤드를 초래하여 더 긴 시퀀스 길이에서 효율성을 감소시킵니다. 이는 효율적인 긴 컨텍스트 처리가 필수적인 에이전틱 애플리케이션에서 주요 제한사항이 됩니다.

예를 들어, 시퀀스 길이가 128k일 때 총 전문가 수를 384개로 고정한 상태에서 어텐션 헤드 수를 64개에서 128개로 증가시키면 추론 FLOP이 83% 증가합니다. 이러한 설계의 영향을 평가하기 위해 어텐션 헤드 수가 레이어 수와 같은 구성과 헤드 수가 두 배인 구성을 다양한 훈련 FLOP 하에서 비교하는 제어된 실험을 수행했습니다.

위 그림은 어텐션 헤드 수가 레이어 수와 같은 모델과 헤드 수가 두 배인 모델의 스케일링 곡선을 보여줍니다. 동일 토큰 훈련 조건 하에서 어텐션 헤드를 두 배로 늘리는 것이 서로 다른 계산 예산에서 검증 손실에서 0.5%에서 1.2% 범위의 미미한 개선만을 가져다준다는 것을 관찰했습니다.

희소성 48이 이미 강력한 성능을 제공하는 상황에서, 어텐션 헤드를 두 배로 늘리는 것으로부터의 한계 이득은 추론 비용을 정당화하지 못합니다. 따라서 64개의 어텐션 헤드를 선택했습니다.

훈련 인프라

컴퓨팅 클러스터

Kimi K2는 NVIDIA H800 GPU가 장착된 클러스터에서 훈련되었습니다. H800 클러스터의 각 노드는 2TB RAM과 노드 내에서 NVLink와 NVSwitch로 연결된 8개의 GPU를 포함합니다. 서로 다른 노드 간에는 8×400 Gbps RoCE 인터커넥트를 활용하여 통신을 촉진합니다.

모델 스케일링을 위한 병렬화

대규모 언어 모델의 훈련은 종종 동적 자원 가용성 하에서 진행됩니다. 특정 자원량 하에서만 적용 가능한 하나의 병렬화 전략을 최적화하는 대신, Kimi K2가 32의 배수인 임의의 노드 수에서 훈련될 수 있도록 하는 유연한 전략을 추구합니다.

이 전략은 가상 스테이지를 가진 16-way Pipeline Parallelism(PP), 16-way Expert Parallelism(EP), ZeRO-1 Data Parallelism의 조합을 활용합니다. 이 설정 하에서 BF16으로 모델 파라미터를 저장하고 FP32로 그래디언트 누적 버퍼를 저장하는 것은 256개 GPU의 모델 병렬 그룹에 분산되어 약 6TB의 GPU 메모리를 필요로 합니다.

옵티마이저 상태의 배치는 훈련 구성에 따라 달라집니다. 총 훈련 노드 수가 클 때는 옵티마이저 상태가 분산되어 디바이스당 메모리 사용량을 무시할 수 있는 수준으로 줄입니다. 총 훈련 노드 수가 적을 때(예: 32개)는 일부 옵티마이저 상태를 CPU로 오프로드할 수 있습니다.

이러한 접근법을 통해 소규모 및 대규모 실험 모두에 동일한 병렬화 구성을 재사용할 수 있으며, 각 GPU가 모든 상태에 대해 약 30GB의 GPU 메모리를 보유하도록 합니다. 나머지 GPU 메모리는 활성화에 사용됩니다. 이러한 일관된 설계는 시스템을 단순화하고 실험 반복을 상당히 가속화하므로 연구 효율성에 중요합니다.

인터리브된 1F1B를 통한 EP 통신 오버랩: 웜업 마이크로 배치 수를 증가시킴으로써 표준 인터리브된 1F1B 스케줄 하에서 EP all-to-all 통신을 계산과 오버랩할 수 있습니다. 이에 비해 DualPipe는 파라미터와 그래디언트에 필요한 메모리를 두 배로 늘려 이를 보상하기 위한 병렬화 증가를 필요로 합니다. PP 증가는 더 많은 버블을 도입하고, EP 증가는 더 높은 오버헤드를 발생시킵니다. 1조 파라미터가 넘는 대규모 모델 훈련에는 추가 비용이 너무 높아 DualPipe를 사용하지 않기로 했습니다.

그러나 인터리브된 1F1B는 모델을 더 많은 스테이지로 분할하여 무시할 수 없는 PP 통신 오버헤드를 도입합니다. 이 비용을 완화하기 위해 각 마이크로 배치의 역방향 패스에서 가중치-그래디언트 계산을 분리하고 해당 PP 통신과 병렬로 실행합니다. 결과적으로 웜업 단계를 제외하고 모든 PP 통신을 효과적으로 오버랩할 수 있습니다.

더 작은 EP 크기: 1F1B 단계에서 완전한 계산-통신 오버랩을 보장하기 위해 K2의 감소된 어텐션 계산 시간(DeepSeek-V3의 128개 헤드에 비해 64개 어텐션 헤드)은 EP 연산 시간을 최소화할 필요가 있습니다. 이는 가장 작은 실행 가능한 EP 병렬화 전략, 구체적으로 EP = 16을 채택함으로써 달성됩니다. 더 작은 EP 그룹을 활용하는 것은 또한 전문가 균형 제약을 완화하여 추가 튜닝 없이 거의 최적의 속도를 달성할 수 있게 합니다.

활성화 감소

파라미터, 그래디언트 버퍼, 옵티마이저 상태를 위한 공간을 예약한 후, 각 디바이스의 나머지 GPU 메모리는 전체 MoE 활성화를 보유하기에 불충분합니다. 특히 1F1B 웜업 단계에서 가장 큰 활성화를 누적하는 초기 파이프라인 스테이지에서 활성화 메모리가 제약 내에 맞도록 하기 위해 다음 기법들이 사용됩니다.

선택적 재계산: LayerNorm, SwiGLU, MLA 업-프로젝션을 포함한 저비용, 고메모리 사용량 스테이지에 재계산이 적용됩니다. 또한 훈련 중 활성화 메모리를 더욱 줄이기 위해 MoE 다운-프로젝션이 재계산됩니다. 선택사항이지만, 이 재계산은 적절한 GPU 메모리를 유지하여 초기 훈련 단계의 전문가 불균형으로 인한 충돌을 방지합니다.

민감하지 않은 활성화를 위한 FP8 저장: MoE 업-프로젝션과 SwiGLU의 입력은 FP32 스케일을 가진 1×128 타일에서 FP8-E4M3로 압축됩니다. 소규모 실험에서는 측정 가능한 손실 증가가 없음을 보여줍니다. 예비 연구에서 관찰된 성능 저하의 잠재적 위험으로 인해 계산에서는 FP8을 적용하지 않습니다.

활성화 CPU 오프로드: 나머지 모든 활성화는 CPU RAM으로 오프로드됩니다. 복사 엔진이 오프로드와 온로드 스트리밍을 담당하며, 계산 및 통신 커널과 오버랩됩니다. 1F1B 단계에서는 이전 마이크로 배치의 순방향 활성화를 오프로드하면서 다음 마이크로 배치의 역방향 활성화를 프리페치합니다. 웜업과 쿨다운 단계도 유사하게 처리됩니다.

위 그림은 서로 다른 PP 단계에서 계산, 통신, 오프로딩이 오버랩되는 패턴을 보여줍니다. 오프로딩이 PCIe 트래픽 혼잡으로 인해 EP 트래픽에 약간 영향을 줄 수 있지만, 테스트 결과 EP 통신은 여전히 완전히 오버랩됨을 보여줍니다.

훈련 레시피

모델은 MuonClip 옵티마이저와 WSD 학습률 스케줄을 사용하여 4,096 토큰 컨텍스트 윈도우로 사전 훈련되었으며, 총 15.5조 토큰을 처리했습니다. 처음 10조 토큰은 500단계 웜업 후 2e-4의 일정한 학습률로 훈련되었고, 이어서 5.5조 토큰은 2e-4에서 2e-5로의 코사인 감쇠로 훈련되었습니다. 가중치 감쇠는 전체적으로 0.1로 설정되었고, 전역 배치 크기는 6700만 토큰으로 유지되었습니다.

사전 훈련 말기에는 어닐링 단계와 긴 컨텍스트 활성화 단계를 수행했습니다. 배치 크기는 6700만 토큰으로 일정하게 유지되었고, 학습률은 2e-5에서 7e-6으로 감쇠되었습니다. 이 단계에서 모델은 4k 시퀀스 길이로 4000억 토큰에 대해 훈련되었고, 이어서 32k 시퀀스 길이로 추가 600억 토큰에 대해 훈련되었습니다. 컨텍스트 윈도우를 128k로 확장하기 위해 YaRN 방법을 사용했습니다.

후훈련

Kimi K2의 후훈련 과정은 사전 훈련된 기본 모델을 실제 사용 가능한 에이전틱 인텔리전스로 변환하는 핵심 단계입니다. 이 과정은 지도 학습 기반 파인튜닝과 강화학습을 통한 정렬을 포함하는 다단계 접근법을 채택하여, 모델이 복잡한 도구 사용과 다단계 추론 작업을 수행할 수 있도록 합니다.

지도 학습 기반 파인튜닝

후훈련에서는 이전 연구에서 Muon으로 사전 훈련된 체크포인트가 Muon 파인튜닝에서 최고의 성능을 보인다는 결론에 따라 Muon 옵티마이저를 사용합니다. 프롬프트 다양성 극대화와 높은 응답 품질 보장이라는 두 가지 핵심 원칙에 따라 다양한 도메인에 걸친 대규모 지시 튜닝 데이터셋을 구축했습니다.

이를 위해 서로 다른 작업 도메인에 맞춤화된 데이터 생성 파이프라인 모음을 개발했으며, 각각은 인간 주석, 프롬프트 엔지니어링, 검증 프로세스의 조합을 활용합니다. K1.5와 기타 사내 도메인 특화 전문가 모델을 채택하여 다양한 작업에 대한 후보 응답을 생성하고, 이어서 LLM 또는 인간 기반 판정자가 자동화된 품질 평가와 필터링을 수행합니다.

에이전틱 데이터의 경우, 다단계 상호작용 추론을 통해 모델에게 도구 사용 능력을 가르치기 위한 데이터 합성 파이프라인을 구축했습니다.

도구 사용 학습을 위한 대규모 에이전틱 데이터 합성

현대 LLM 에이전트의 핵심 능력은 익숙하지 않은 도구를 자율적으로 사용하고, 외부 환경과 상호작용하며, 추론, 실행, 오류 수정을 통해 반복적으로 행동을 개선하는 것입니다. 에이전틱 도구 사용 능력은 실세계 시스템과의 동적 상호작용이 필요한 복잡하고 다단계적인 작업을 해결하는 데 필수적입니다.

ACEBench와 $\tau$-bench와 같은 최근 벤치마크들은 포괄적인 도구 사용 평가의 중요성을 강조했으며, ToolLLM과 ACEBench 같은 프레임워크들은 수천 개의 도구를 효과적으로 사용하도록 모델을 가르치는 잠재력을 보여주었습니다. 그러나 이러한 능력을 대규모로 훈련하는 것은 상당한 도전과제를 제시합니다. 실세계 환경은 풍부하고 진정성 있는 상호작용 신호를 제공하지만, 비용, 복잡성, 프라이버시, 접근성 제약으로 인해 대규모로 구축하기 어려운 경우가 많습니다.

합성 데이터 생성에 관한 최근 연구들(AgentInstruct, Self-Instruct, StableToolBench, ZeroSearch)은 실세계 상호작용에 의존하지 않고 대규모 데이터를 생성하는 데 유망한 결과를 보여주었습니다. 이러한 발전을 바탕으로 하고 ACEBench의 포괄적인 데이터 합성 프레임워크에서 영감을 받아, 실세계 도구 사용 시나리오를 대규모로 시뮬레이션하는 파이프라인을 개발하여 수만 개의 다양하고 고품질의 훈련 예제를 생성할 수 있게 했습니다.

위 그림들은 도구 사용을 위한 데이터 합성 파이프라인을 보여줍니다. 첫 번째 그림에서는 실세계 도구와 LLM에서 도구 사양을 합성하고, 도구 저장소에서 에이전트와 작업을 생성하는 과정을 보여줍니다. 두 번째 그림에서는 사용자 에이전트가 작업과 상호작용하고, 에이전트가 사용자 에이전트의 행동과 호출을 관찰하여 궤적을 생성하며, 판정자 에이전트가 데이터를 필터링하여 고품질 데이터를 생성하는 다중 에이전트 파이프라인을 보여줍니다.

위 그림들은 도구 임베딩의 t-SNE 시각화를 보여줍니다. 첫 번째 그림에서는 실세계 MCP 도구들이 원래 소스 카테고리에 따라 자연스러운 클러스터링을 보이며, 두 번째 그림에서는 합성 도구들이 사전 정의된 도메인 카테고리로 체계적으로 조직되어 도구 공간의 포괄적인 표현을 보장함을 확인할 수 있습니다.

데이터 합성 파이프라인은 세 단계로 구성됩니다.

도구 사양 생성: 먼저 실세계 도구와 LLM 합성 도구 모두에서 도구 사양의 대규모 저장소를 구축합니다.

에이전트 및 작업 생성: 도구 저장소에서 샘플링된 각 도구 세트에 대해 도구 세트를 사용할 에이전트와 해당 작업들을 생성합니다.

궤적 생성: 각 에이전트와 작업에 대해 에이전트가 도구를 호출하여 작업을 완료하는 궤적을 생성합니다.

도메인 진화와 도구 생성

두 가지 상호 보완적인 접근법을 통해 포괄적인 도구 저장소를 구축합니다. 첫째, GitHub 저장소에서 3000개 이상의 실제 MCP(Model Context Protocol) 도구를 직접 가져와 기존의 고품질 도구 사양을 활용합니다. 둘째, 계층적 도메인 생성 프로세스를 통해 합성 도구를 체계적으로 진화시킵니다. 핵심 카테고리(예: 금융 거래, 소프트웨어 애플리케이션, 로봇 제어)부터 시작하여 각 카테고리 내에서 여러 특정 애플리케이션 도메인을 진화시킵니다. 그 다음 각 도메인에 대해 명확한 인터페이스, 설명, 운영 의미론을 가진 특화된 도구들을 합성합니다. 이 진화 프로세스는 20,000개 이상의 합성 도구를 생성합니다.

앞서 보여준 t-SNE 임베딩을 통한 도구 컬렉션의 다양성 시각화는 MCP와 합성 도구 모두가 도구 공간의 상호 보완적인 영역을 커버함을 보여줍니다.

에이전트 다양화

도구 저장소에서 다양한 도구 조합으로 다양한 시스템 프롬프트를 합성하여 수천 개의 서로 다른 에이전트를 생성합니다. 이는 다양한 능력, 전문 분야, 행동 패턴을 가진 에이전트들의 다양한 집단을 만들어 잠재적 사용 사례의 광범위한 커버리지를 보장합니다.

루브릭 기반 작업 생성

각 에이전트 구성에 대해 단순한 것부터 복잡한 작업까지 다양한 작업을 생성합니다. 각 작업은 성공 기준, 예상되는 도구 사용 패턴, 평가 체크포인트를 명시하는 명시적 루브릭과 쌍을 이룹니다. 이러한 루브릭 기반 접근법은 에이전트 성능의 일관되고 객관적인 평가를 보장합니다.

다중 턴 궤적 생성

여러 구성요소를 통해 현실적인 도구 사용 시나리오를 시뮬레이션합니다.

사용자 시뮬레이션: 서로 다른 의사소통 스타일과 선호도를 가진 LLM 생성 사용자 페르소나가 에이전트와 다중 턴 대화에 참여하여 자연스러운 상호작용 패턴을 만듭니다.

도구 실행 환경: 정교한 도구 시뮬레이터(기능적으로 세계 모델과 동등)가 도구 호출을 실행하고 현실적인 피드백을 제공합니다. 시뮬레이터는 각 도구 실행 후 상태를 유지하고 업데이트하여 지속적인 효과를 가진 복잡한 다단계 상호작용을 가능하게 합니다. 성공, 부분적 실패, 엣지 케이스를 포함한 다양한 결과를 생성하기 위해 제어된 확률성을 도입합니다.

품질 평가 및 필터링

LLM 기반 판정자가 각 궤적을 작업 루브릭에 대해 평가합니다. 성공 기준을 충족하는 궤적만이 훈련을 위해 보존되어 작업 완료 전략의 자연스러운 변화를 허용하면서 고품질 데이터를 보장합니다.

실제 실행 환경과의 하이브리드 접근법

시뮬레이션이 확장성을 제공하지만, 시뮬레이션 충실도의 고유한 한계를 인정합니다. 이를 해결하기 위해 진정성이 중요한 시나리오, 특히 코딩과 소프트웨어 엔지니어링 작업에서 시뮬레이션된 환경을 실제 실행 샌드박스로 보완합니다. 이러한 실제 샌드박스는 실제 코드를 실행하고, 진정한 개발 환경과 상호작용하며, 테스트 스위트 통과율과 같은 객관적 메트릭을 통해 실측 피드백을 제공합니다.

이러한 확장 가능한 시뮬레이션과 목표화된 실세계 실행을 결합한 하이브리드 파이프라인을 활용하여, 커버리지와 진정성의 균형을 맞춘 다양하고 고품질의 도구 사용 시연을 생성합니다. 합성 데이터 생성의 규모와 자동화는 실제 실행 환경에서 제공되는 근거와 결합되어 품질 필터링 프로세스를 통한 대규모 거부 샘플링을 효과적으로 구현합니다.

이러한 고품질 합성 데이터를 지도 학습 파인튜닝에 사용할 때, 광범위한 실세계 애플리케이션에서 모델의 도구 사용 능력에 상당한 개선을 보여주었습니다.

강화학습

Kimi K2의 후훈련 과정에서 강화학습은 지도 학습 파인튜닝을 통해 구축된 기본 능력을 더욱 발전시키는 핵심 단계입니다. 이 과정은 검증 가능한 보상과 자기 비판 메커니즘을 결합한 포괄적인 강화학습 프레임워크를 통해 모델이 다양한 시나리오에서 학습할 수 있도록 설계되었습니다.

강화학습 프레임워크는 확장 가능한 Gym과 유사한 구조를 채택하여 다양한 작업 도메인에서 일관된 학습 경험을 제공합니다. 이 시스템은 객관적으로 평가 가능한 작업에 대해서는 검증 가능한 보상을 사용하고, 주관적 선호도가 중요한 작업에 대해서는 자기 비판 보상 메커니즘을 활용합니다. 이러한 이중 접근법을 통해 작업 다양성과 훈련 FLOP 모두를 확장할 수 있습니다.

검증 가능한 보상

검증 가능한 보상 시스템은 객관적으로 평가할 수 있는 다양한 작업 도메인을 포괄합니다. 이 시스템은 수학, STEM, 논리적 추론, 복잡한 지시 따르기, 충실성, 코딩, 소프트웨어 엔지니어링, 안전성 작업을 다루며, 각각에 대해 자동화된 검증 시스템을 구축했습니다.

수학 및 STEM 작업에서는 하이브리드 규칙 검증 시스템을 구현했습니다. 이 시스템은 수치적 답안에 대해서는 정확한 매칭을 수행하고, 대수적 표현에 대해서는 기호적 동등성을 검증합니다. 복잡한 수학 문제의 경우 단계별 추론 과정도 평가하여 최종 답안뿐만 아니라 해결 과정의 논리적 일관성도 확인합니다.

논리적 추론 작업에서는 전제와 결론 간의 논리적 연결성을 자동으로 검증하는 시스템을 개발했습니다. 이는 형식 논리학의 원칙을 바탕으로 추론 체인의 각 단계가 논리적으로 타당한지 확인하며, 오류가 있는 추론 단계를 정확히 식별할 수 있습니다.

복잡한 지시 따르기에서는 다단계 지시사항의 각 구성요소가 올바르게 수행되었는지 체계적으로 평가합니다. 이는 지시사항을 구조화된 체크리스트로 분해하고, 각 항목의 완료 여부를 자동으로 검증하는 방식으로 구현됩니다.

충실성 평가에서는 문장 수준의 충실성 판정자를 구현했습니다. 이 시스템은 생성된 텍스트의 각 문장이 주어진 소스 자료와 얼마나 일치하는지 세밀하게 분석하며, 사실적 오류나 왜곡을 정확히 탐지할 수 있습니다.

코딩 및 소프트웨어 엔지니어링에서는 가장 정교한 검증 인프라를 구축했습니다. Kubernetes 기반의 샌드박스 인프라를 통해 10,000개 이상의 동시 인스턴스를 지원하며, 각 코드 솔루션을 안전하고 격리된 환경에서 실행하여 정확성을 검증합니다. 이 시스템은 단순한 테스트 케이스 통과뿐만 아니라 코드 품질, 효율성, 보안성까지 종합적으로 평가합니다.

안전성 작업에서는 다층적 안전성 검증 시스템을 구현했습니다. 이는 유해한 콘텐츠 생성 방지, 편향성 검사, 윤리적 가이드라인 준수 등을 자동으로 평가하며, 각 응답이 안전성 기준을 충족하는지 실시간으로 모니터링합니다.

검증을 넘어서: 자기 비판 루브릭 보상

검증 가능한 보상만으로는 다룰 수 없는 복잡하고 주관적인 작업들을 위해 자기 비판 루브릭 보상 프레임워크를 도입했습니다. 이 시스템은 K2가 자신의 출력을 핵심 루브릭, 규범적 루브릭, 인간 주석 루브릭에 대해 평가하는 일반적인 강화학습 프레임워크입니다.

자기 비판 시스템의 핵심은 폐쇄 루프 비판자 개선 과정입니다. 이 과정에서 K2는 먼저 자신의 응답을 생성하고, 이를 다양한 기준에 따라 자체 평가합니다. 그 다음 이 평가 결과를 바탕으로 응답을 개선하고, 개선된 응답을 다시 평가하는 반복적 과정을 거칩니다. 이러한 반복적 자기 개선 과정은 검증 가능한 신호에 근거하여 이루어지므로, 복잡한 비검증 가능 목표와의 정렬을 가능하게 합니다.

핵심 루브릭은 명확성과 관련성, 대화적 유창성과 참여도, 객관적이고 근거 있는 상호작용이라는 세 가지 기본 평가 기준을 포함합니다. 명확성과 관련성 기준은 응답이 사용자의 질문에 직접적이고 간결하게 답하는지 평가하며, 불필요한 정보나 모호한 표현을 피하도록 유도합니다. 대화적 유창성과 참여도 기준은 자연스러운 대화 흐름을 유지하면서도 사용자의 관심을 끌 수 있는 응답을 생성하도록 합니다. 객관적이고 근거 있는 상호작용 기준은 메타 코멘터리나 과도한 칭찬 없이 핵심에 집중하는 응답을 장려합니다.

규범적 루브릭은 보상 해킹을 방지하기 위한 금지된 행동들을 명시합니다. 여기에는 사용자나 질문에 대한 초기 칭찬과 응답 품질을 설명하기보다는 내용을 기술하는 명시적 정당화 문장들이 포함됩니다. 이러한 규범적 제약은 모델이 표면적인 개선보다는 실질적인 가치 제공에 집중하도록 유도합니다.

인간 주석 루브릭은 복잡한 주관적 판단이 필요한 작업에 대해 인간 평가자들이 제공한 기준을 포함합니다. 이는 창의성, 공감능력, 문화적 민감성 등 자동화된 평가로는 측정하기 어려운 측면들을 다룹니다.

자기 비판 시스템의 한계도 인정해야 합니다. 모호한 상황에서 자기 자격 부여를 피하고 명확성을 선호하는 경향으로 인해 인식론적 겸손이 더 적절한 애매한 상황에서 과도한 확신을 보일 수 있는 잠재적 편향이 있습니다.

강화학습 알고리즘

K2의 강화학습에서는 K1.5에서 사용된 정책 최적화 알고리즘을 기반으로 하되, 여러 개선사항을 도입했습니다. 주요 개선사항으로는 응답 길이를 관리하기 위한 예산 제어, 고품질 데이터의 망각을 방지하기 위한 PTX 손실, 훈련 중 탐색에서 활용으로 전환하는 온도 감쇠 스케줄이 있습니다.

예산 제어 메커니즘은 모델이 과도하게 긴 응답을 생성하는 것을 방지합니다. 이는 응답 길이에 대한 소프트 제약을 도입하여, 정보의 밀도와 간결성 사이의 균형을 맞춥니다. 제약은 다음과 같이 수식화됩니다.

여기서 $\lambda_{budget}$는 예산 제어 강도를 조절하는 하이퍼파라미터이고, $target_{length}$는 목표 응답 길이입니다.

PTX(Pre-Training miXture) 손실은 강화학습 과정에서 사전 훈련된 지식이 손실되는 것을 방지합니다. 이는 원래 사전 훈련 데이터의 일부를 강화학습 과정에 지속적으로 포함시켜 모델이 기본적인 언어 이해 능력을 유지하도록 합니다.

여기서 $L_{RL}$은 강화학습 손실, $L_{PTX}$는 사전 훈련 손실, $\alpha_{PTX}$는 두 손실 간의 균형을 조절하는 가중치입니다.

온도 감쇠 스케줄은 훈련 초기에는 높은 온도로 탐색을 장려하고, 훈련이 진행됨에 따라 온도를 점진적으로 낮춰 활용에 집중하도록 합니다.

여기서 $T(t)$는 시간 $t$에서의 온도, $T_{init}$는 초기 온도, $\gamma$는 감쇠율입니다.

이러한 알고리즘적 개선사항들은 모델이 다양한 작업에서 안정적이고 효과적으로 학습할 수 있도록 하며, 특히 에이전틱 작업에서 요구되는 복잡한 다단계 추론 능력을 향상시키는 데 기여합니다.

강화학습 인프라

Kimi K2의 강화학습 훈련을 위해서는 대규모 모델의 정책 최적화와 환경 상호작용을 효율적으로 처리할 수 있는 특별한 인프라가 필요합니다. 이를 위해 동기화된 강화학습 훈련을 지원하는 코로케이션 아키텍처를 개발했으며, 효율적인 엔진 전환, 시스템 시작 최적화, 장기간 다중 턴 작업을 위한 에이전틱 롤아웃 지원을 포함합니다.

코로케이션 아키텍처

강화학습 훈련의 핵심 도전과제 중 하나는 정책 업데이트와 환경 상호작용 간의 효율적인 조정입니다. 전통적인 분리된 아키텍처에서는 훈련 클러스터와 추론 클러스터가 별도로 운영되어 통신 오버헤드와 자원 활용 비효율성이 발생합니다. 이를 해결하기 위해 훈련 엔진과 추론 엔진이 워커를 공유하는 하이브리드 코로케이션 아키텍처를 구현했습니다.

이 아키텍처에서는 중앙화된 컨트롤러가 데이터 생성과 파라미터 업데이트 간의 조정을 담당합니다. 컨트롤러는 다음과 같은 핵심 기능을 수행합니다.

데이터 생성 단계: 추론 엔진이 활성화되어 현재 정책을 사용해 환경과 상호작용하며 경험 데이터를 수집합니다. 이 과정에서 에이전트는 다양한 작업을 수행하고, 보상 신호를 받으며, 궤적 데이터를 생성합니다.

정책 업데이트 단계: 수집된 데이터를 바탕으로 훈련 엔진이 활성화되어 정책 그래디언트 계산과 파라미터 업데이트를 수행합니다. 이 단계에서는 PPO(Proximal Policy Optimization) 알고리즘을 기반으로 한 정책 최적화가 이루어집니다.

코로케이션 아키텍처의 주요 장점은 자원 활용 효율성입니다. 동일한 GPU 클러스터가 추론과 훈련을 번갈아 수행함으로써 하드웨어 자원을 최대한 활용할 수 있으며, 네트워크 통신 오버헤드를 최소화할 수 있습니다.

효율적인 엔진 전환

코로케이션 아키텍처에서 가장 중요한 기술적 도전과제는 추론 엔진과 훈련 엔진 간의 빠르고 안정적인 전환입니다. Kimi K2와 같은 1조 파라미터 규모의 모델에서는 파라미터 업데이트와 동기화가 상당한 시간을 소요할 수 있어, 이를 효율적으로 처리하는 것이 전체 훈련 효율성에 결정적입니다.

이를 위해 분산 체크포인트 엔진을 개발했습니다. 이 시스템은 파이프라인된 파라미터별 전송과 클러스터 전체의 브로드캐스트 연산을 사용하여 Kimi K2의 전체 파라미터 업데이트를 30초 이내에 완료할 수 있습니다.

분산 체크포인트 엔진의 작동 원리는 다음과 같습니다.

파라미터 분할: 모델의 전체 파라미터를 여러 청크로 분할하여 병렬 처리가 가능하도록 합니다. 각 청크는 독립적으로 전송되고 업데이트될 수 있어 전체 처리 시간을 단축시킵니다.

파이프라인 전송: 파라미터 업데이트를 파이프라인 방식으로 처리하여 한 청크가 전송되는 동안 다른 청크의 계산이 동시에 진행될 수 있도록 합니다. 이는 전체 대기 시간을 크게 줄입니다.

브로드캐스트 최적화: 업데이트된 파라미터를 클러스터의 모든 노드에 효율적으로 전파하기 위해 계층적 브로드캐스트 알고리즘을 사용합니다. 이는 네트워크 대역폭을 최적으로 활용하면서 모든 노드가 일관된 파라미터를 유지하도록 보장합니다.

효율적인 시스템 시작

대규모 분산 시스템에서 시스템 시작 시간은 전체 훈련 효율성에 중요한 영향을 미칩니다. 특히 강화학습에서는 정책 업데이트 후 새로운 정책으로 환경 상호작용을 재시작해야 하므로, 빠른 시스템 재시작이 필수적입니다.

시작 시간을 최적화하기 위해 다음과 같은 전략을 구현했습니다.

선택적 체크포인트 읽기: 훈련 워커들이 체크포인트의 필요한 부분만 선택적으로 읽고, 이를 동료 노드들과 브로드캐스트하는 방식을 채택했습니다. 이는 전체 체크포인트를 모든 노드에서 읽는 것보다 훨씬 효율적입니다.

집단적 디스크 읽기: 체크포인트 엔진이 디스크에서 집단적으로 읽기를 수행하고 추론 엔진을 업데이트하는 방식을 통해 단일 장애점에 대한 견고성을 보장합니다. 이는 특정 노드의 실패가 전체 시스템에 미치는 영향을 최소화합니다.

점진적 로딩: 모델 파라미터를 점진적으로 로딩하여 시스템이 부분적으로라도 빠르게 작동을 시작할 수 있도록 합니다. 이는 전체 로딩이 완료되기 전에도 일부 작업을 수행할 수 있게 하여 전체 대기 시간을 줄입니다.

에이전틱 롤아웃

에이전틱 작업의 특성상 장기간에 걸친 다중 턴 상호작용이 필요하며, 이는 전통적인 강화학습 환경보다 훨씬 복잡한 요구사항을 제시합니다. 이를 지원하기 위해 특별히 설계된 에이전틱 롤아웃 시스템을 구현했습니다.

중량급 환경 서비스: 복잡한 도구 사용과 외부 시스템 상호작용을 지원하기 위해 중량급 환경 서비스를 구축했습니다. 이러한 서비스는 실제 소프트웨어 개발 환경, 데이터베이스 시스템, API 서비스 등과의 상호작용을 시뮬레이션하거나 실제로 연결할 수 있습니다.

동시 롤아웃: GPU 활용률을 최대화하기 위해 여러 에이전트가 동시에 서로 다른 환경에서 롤아웃을 수행할 수 있도록 합니다. 이는 단일 에이전트의 순차적 상호작용으로 인한 자원 낭비를 방지합니다.

부분 롤아웃 기법: 장기간 궤적에서 발생할 수 있는 롱테일 문제를 해결하기 위해 부분 롤아웃 기법을 도입했습니다. 이는 매우 긴 상호작용 시퀀스를 적절한 지점에서 분할하여 처리함으로써 메모리 사용량을 제어하고 훈련 안정성을 향상시킵니다.

통합 OpenAI Gym 인터페이스: 다양한 환경과의 통합을 위해 OpenAI Gym에서 영감을 받은 통합 인터페이스를 구현했습니다. 이는 새로운 환경을 쉽게 추가하고 기존 환경을 수정할 수 있도록 하는 표준화된 API를 제공합니다.

환경 인터페이스는 다음과 같은 표준 메서드를 제공합니다.

class AgenticEnvironment:
    def reset(self):
        """환경을 초기 상태로 리셋하고 초기 관찰을 반환"""
        pass

    def step(self, action):
        """액션을 실행하고 (관찰, 보상, 완료여부, 정보)를 반환"""
        pass

    def render(self):
        """현재 환경 상태를 시각화"""
        pass

    def close(self):
        """환경 자원을 정리"""
        pass

이러한 표준화된 인터페이스를 통해 코딩 환경, 수학 문제 해결 환경, 도구 사용 환경 등 다양한 작업 도메인을 일관된 방식으로 처리할 수 있습니다.

에이전틱 롤아웃 시스템은 또한 실시간 모니터링과 디버깅 기능을 제공합니다. 각 에이전트의 상호작용 과정을 실시간으로 추적하고, 문제가 발생한 경우 상세한 로그와 상태 정보를 제공하여 시스템 개선에 활용할 수 있습니다.

이러한 포괄적인 강화학습 인프라를 통해 Kimi K2는 대규모 에이전틱 작업에서 안정적이고 효율적인 학습을 수행할 수 있으며, 이는 최종적으로 모델의 뛰어난 에이전틱 능력으로 이어집니다. 강화학습은 지도 학습 파인튜닝보다 더 나은 토큰 효율성과 일반화 성능을 제공한다고 여겨집니다. K1.5의 연구를 바탕으로 K2에서는 작업 다양성과 훈련 FLOP 모두에서 강화학습을 확장했습니다. 이를 지원하기 위해 광범위한 시나리오에서 강화학습을 촉진하는 Gym과 유사한 확장 가능한 프레임워크를 개발했습니다.

이 프레임워크는 검증 가능한 보상을 가진 다수의 작업으로 확장되었습니다. 창의적 글쓰기나 개방형 질문 답변과 같이 주관적 선호도에 의존하는 작업의 경우, 모델이 자신의 출력에 대해 쌍별 비교를 수행하여 판단하는 자기 비판 보상을 도입했습니다. 이러한 접근법을 통해 다양한 도메인의 작업들이 모두 강화학습 패러다임의 혜택을 받을 수 있게 되었습니다.

검증 가능한 보상

검증 가능한 보상 시스템은 객관적으로 평가할 수 있는 여러 작업 도메인을 포괄합니다. 각 도메인은 고유한 특성과 요구사항에 맞춘 전문화된 검증 메커니즘을 갖추고 있습니다.

수학, STEM 및 논리적 작업

수학, STEM, 논리적 추론 도메인에서 강화학습 데이터 준비는 다양한 커버리지와 적절한 난이도라는 두 가지 핵심 원칙을 따릅니다.

다양한 커버리지: 수학과 STEM 작업의 경우, 전문가 주석, 내부 QA 추출 파이프라인, 그리고 공개 데이터셋의 조합을 사용하여 고품질 QA 쌍을 수집합니다. 수집 과정에서 태깅 시스템을 활용하여 커버리지가 부족한 도메인의 포함을 의도적으로 증가시킵니다. 논리적 작업의 경우, 데이터셋은 구조화된 데이터 작업(예: 다중 홉 테이블 추론, 교차 테이블 집계)과 논리 퍼즐(예: 24게임, 스도쿠, 수수께끼, 암호 산술, 모스 부호 해독)을 포함한 다양한 형식으로 구성됩니다.

적절한 난이도: 강화학습 프롬프트 세트는 너무 쉽거나 너무 어려워서는 안 되며, 둘 다 신호를 거의 생성하지 않아 학습 효율성을 감소시킬 수 있습니다. SFT 모델의 pass@k 정확도를 사용하여 각 문제의 난이도를 평가하고, 적절한 난이도를 가진 문제만을 선택합니다.

복잡한 지시 따르기

효과적인 지시 따르기는 명시적 제약 조건을 이해하는 것뿐만 아니라 암시적 요구사항을 탐색하고, 엣지 케이스를 처리하며, 확장된 대화에서 일관성을 유지하는 것을 요구합니다. 이러한 도전과제를 해결하기 위해 자동화된 검증과 적대적 탐지를 결합한 하이브리드 검증 프레임워크와 확장 가능한 커리큘럼 생성 파이프라인을 구축했습니다.

하이브리드 규칙 검증: 정확성과 견고성을 모두 보장하기 위해 이중 경로 시스템을 채택합니다. 첫째, 검증 가능한 출력(예: 길이, 스타일 제약)을 가진 지시사항에 대해서는 코드 인터프리터를 통한 결정론적 평가를 구현합니다. 둘째, 제약 조건에 대한 미묘한 이해가 필요한 지시사항에 대해서는 LLM-as-judge 평가를 사용합니다. 모델이 실제 준수 없이 지시사항 이행을 주장할 수 있는 잠재적 적대적 행동을 해결하기 위해, 이러한 기만적 주장을 특별히 탐지하는 추가적인 해킹 체크 레이어를 통합했습니다.

다중 소스 지시사항 생성: 훈련 데이터를 구성하기 위해 포괄적인 커버리지를 보장하는 세 가지 구별되는 생성 전략을 사용합니다. 첫째, 데이터 팀이 개발한 전문가 제작 복잡 조건부 프롬프트와 루브릭을 활용합니다. 둘째, AutoIF에서 영감을 받은 에이전틱 지시사항 증강을 적용합니다. 셋째, 특정 실패 모드나 엣지 케이스를 탐지하는 추가 지시사항을 생성하는 데 특화된 파인튜닝된 모델을 사용합니다. 이러한 다각적 접근법은 지시사항 커버리지에서 폭과 깊이를 모두 보장합니다.

충실성

충실성은 다중 턴 도구 사용, 자체 생성 추론 체인, 개방 환경 상호작용과 같은 시나리오에서 작동하는 에이전틱 모델에 필수적입니다. FACTS Grounding의 평가 프레임워크에서 영감을 받아, 자동화된 검증을 수행하는 문장 수준 충실성 판정자 모델을 훈련했습니다. 이 판정자는 컨텍스트에서 뒷받침하는 증거 없이 사실적 주장을 하는 문장을 탐지하는 데 효과적입니다. 이는 전체적인 충실성 성능을 향상시키는 보상 모델 역할을 합니다.

코딩 및 소프트웨어 엔지니어링

경쟁 수준의 프로그래밍 문제를 다루는 능력을 향상시키기 위해, 공개 소스 데이터셋과 합성 소스 모두에서 문제와 그 판정자를 수집합니다. 합성 데이터의 다양성과 보상 신호의 정확성을 보장하기 위해, 사전 훈련 데이터에서 검색된 고품질 인간 작성 단위 테스트를 통합합니다.

소프트웨어 엔지니어링 작업의 경우, GitHub에서 대량의 풀 리퀘스트와 이슈를 수집하여 사용자 프롬프트/이슈와 실행 가능한 단위 테스트로 구성된 소프트웨어 개발 환경을 구축합니다. 이 환경은 확장성과 보안을 위해 Kubernetes로 구동되는 견고한 샌드박스 인프라 위에 구축되었습니다. 안정적인 성능으로 10,000개 이상의 동시 샌드박스 인스턴스를 지원하여, 경쟁적 코딩과 소프트웨어 엔지니어링 작업 모두에 이상적입니다.

안전성

안전성을 향상시키는 작업은 폭력, 사기, 차별과 같은 주요 위험 범주를 포괄하도록 수동으로 제작된 인간 큐레이션 시드 프롬프트 세트로 시작됩니다. 정교한 탈옥 시도(예: 역할 연기, 문학적 서사, 학술적 담론)를 시뮬레이션하기 위해, 세 가지 핵심 구성요소를 가진 자동화된 프롬프트 진화 파이프라인을 사용합니다.

공격 모델은 대상 LLM에서 안전하지 않은 응답을 유도하도록 설계된 적대적 프롬프트를 반복적으로 생성합니다. 대상 모델은 이러한 프롬프트에 대한 응답을 생성하여 잠재적 취약점을 시뮬레이션합니다. 판정자 모델은 상호작용을 평가하여 적대적 프롬프트가 안전 메커니즘을 성공적으로 우회했는지 결정합니다. 각 상호작용은 작업별 루브릭을 사용하여 평가되어, 판정자 모델이 이진 성공/실패 레이블을 제공할 수 있게 합니다.

검증을 넘어서: 자기 비판 루브릭 보상

검증 가능한 보상을 가진 작업을 넘어 모델 정렬을 확장하기 위해, 자기 비판 피드백으로부터의 일반적인 강화학습을 위한 프레임워크를 도입합니다. 이 접근법은 검증 가능한 시나리오에서 학습된 능력을 더 광범위한 주관적 작업으로 확장하여, 유용성, 창의성, 추론의 깊이, 사실성, 안전성을 포함한 미묘한 인간 선호도와 LLM을 정렬하도록 설계되었습니다.

프레임워크는 모델이 자신의 출력을 평가하여 선호도 신호를 생성하는 자기 비판 루브릭 보상 메커니즘을 사용하여 작동합니다. K2를 유능한 판정자로 부트스트랩하기 위해, 공개 소스와 사내 선호도 데이터셋의 혼합을 큐레이션하고 SFT 단계에서 비판 능력을 초기화합니다.

자기 비판 정책 최적화

학습 루프의 첫 번째 핵심 과정에서, K2 액터는 광범위한 사용 사례를 다루는 일반적인 프롬프트에 대한 응답을 생성합니다. 그 다음 K2 비판자는 핵심 루브릭, 규범적 루브릭, 인간 주석 루브릭의 조합에 대해 쌍별 평가를 수행하여 모든 결과를 순위를 매깁니다.

핵심 루브릭은 Kimi가 소중히 여기는 AI 어시스턴트의 기본 가치를 나타내고, 규범적 루브릭은 보상 해킹을 제거하는 것을 목표로 하며, 인간 주석 루브릭은 특정 지시적 맥락을 위해 데이터 팀이 제작한 것입니다. 특정 루브릭이 필수로 지정될 수 있지만, K2는 이를 내부 사전과 비교하여 가중치를 부여할 유연성을 유지합니다. 이러한 능력은 진화하는 온-정책 행동과의 동적이고 지속적인 정렬을 가능하게 하여, 모델의 응답이 특정 지시사항에 적응하면서도 핵심 정체성과 일관성을 유지하도록 보장합니다.

폐쇄 루프 비판자 개선 및 정렬

강화학습 훈련 중에 비판자 모델은 검증 가능한 신호를 사용하여 개선됩니다. 검증 가능한 보상 프롬프트에서 생성된 온-정책 롤아웃은 비판자를 지속적으로 업데이트하는 데 사용되며, 이는 RLVR에서 객관적 성능 신호를 평가 모델로 직접 증류하는 중요한 단계입니다.

이 전이 학습 과정은 더 주관적인 판단을 검증 가능한 데이터에 근거시켜, 검증 가능한 작업에서의 성능 향상이 명시적 보상 신호가 없는 복잡한 작업에서 비판자의 판단을 향상시킬 수 있게 합니다. 이 폐쇄 루프 과정은 비판자가 정책의 진화와 보조를 맞춰 평가 기준을 지속적으로 재보정하도록 보장합니다.

주관적 평가를 검증 가능한 데이터에 근거시킴으로써, 프레임워크는 복잡하고 비검증 가능한 인간 목표와의 견고하고 확장 가능한 정렬을 가능하게 합니다. 결과적으로, 이러한 전체적인 정렬은 사용자 의도 이해, 창의적 글쓰기, 복잡한 추론, 미묘한 언어 이해를 포함한 광범위한 도메인에서 포괄적인 성능 향상을 가져다줍니다.

강화학습 알고리즘

K2에서는 K1.5에서 도입된 정책 최적화 알고리즘을 기반으로 채택합니다. 각 문제 $x$에 대해, 이전 정책 $\pi_{\mathrm{old}}$에서 $K$개의 응답 ${y_1, \dots, y_k}$를 샘플링하고, 다음 목적 함수에 대해 모델 $\pi_\theta$를 최적화합니다.

여기서 $\bar{r}(x) = \frac{1}{k}\sum_{i=1}^k r(x,y_i)$는 샘플링된 응답의 평균 보상이고, $\tau > 0$는 안정적인 학습을 촉진하는 정규화 파라미터입니다. SFT에서와 마찬가지로 이 목적 함수를 최소화하기 위해 Muon 옵티마이저를 사용합니다.

K2에서 강화학습 훈련을 더 광범위한 작업으로 확장함에 따라, 주요 도전과제는 모든 도메인에서 일관된 성능 향상을 달성하는 것입니다. 이를 해결하기 위해 강화학습 알고리즘에 여러 추가 사항을 도입합니다.

예산 제어

강화학습이 종종 모델 생성 응답의 길이를 상당히 증가시킨다는 것이 널리 관찰되었습니다. 더 긴 응답이 모델이 복잡한 추론 작업에서 향상된 성능을 위해 추가적인 테스트 시간 계산을 활용할 수 있게 하지만, 그 이점은 종종 비추론 도메인에서의 추론 비용을 정당화하지 못합니다.

모델이 추론 예산을 적절히 분배하도록 장려하기 위해, 강화학습 훈련 전반에 걸쳐 샘플당 최대 토큰 예산을 시행합니다. 여기서 예산은 작업 유형에 따라 결정됩니다. 이 토큰 예산을 초과하는 응답은 잘리고 페널티가 부여되어, 모델이 지정된 제한 내에서 솔루션을 생성하도록 인센티브를 제공합니다. 경험적으로, 이 접근법은 모델의 토큰 효율성을 크게 향상시켜 모든 도메인에서 간결하면서도 효과적인 솔루션을 장려합니다.

PTX 손실

공동 강화학습 훈련 중 가치 있는 고품질 데이터의 잠재적 망각을 방지하기 위해, 수동으로 선택된 고품질 샘플로 구성된 데이터셋을 큐레이션하고 보조 PTX 손실을 통해 강화학습 목적 함수에 통합합니다. 이 전략은 고품질 데이터의 장점을 활용할 뿐만 아니라 훈련 체제에 명시적으로 존재하는 제한된 작업 세트에 대한 과적합 위험을 완화합니다. 이러한 증강은 더 광범위한 도메인에서 모델의 일반화를 상당히 향상시킵니다.

온도 감쇠

창의적 글쓰기와 복잡한 추론과 같은 작업에서, 훈련 초기 단계에서 높은 샘플링 온도를 통해 탐색을 촉진하는 것이 중요하다는 것을 발견했습니다. 높은 온도는 모델이 다양하고 혁신적인 응답을 생성할 수 있게 하여, 효과적인 전략의 발견을 촉진하고 차선의 솔루션으로의 조기 수렴 위험을 줄입니다.

그러나 훈련 후기 단계나 평가 중에 높은 온도를 유지하는 것은 과도한 무작위성을 도입하고 모델 출력의 신뢰성과 일관성을 손상시키므로 해로울 수 있습니다. 이를 해결하기 위해 온도 감쇠 스케줄을 사용하여 훈련 전반에 걸쳐 탐색에서 활용으로 전환합니다. 이 전략은 모델이 가장 유익할 때 탐색을 활용하면서 궁극적으로 안정적이고 고품질의 출력으로 수렴하도록 보장합니다.

강화학습 인프라

Kimi K2의 강화학습 훈련을 위해서는 대규모 모델의 정책 최적화와 환경 상호작용을 효율적으로 처리할 수 있는 특별한 인프라가 필요합니다. 이를 위해 동기화된 강화학습 훈련을 지원하는 코로케이션 아키텍처를 개발했으며, 효율적인 엔진 전환, 시스템 시작 최적화, 장기간 다중 턴 작업을 위한 에이전틱 롤아웃 지원을 포함합니다.

코로케이션 아키텍처

K1.5와 유사하게, 동기화된 강화학습 훈련을 위한 하이브리드 코로케이션 아키텍처를 채택합니다. 이 아키텍처에서는 훈련 엔진과 추론 엔진이 동일한 워커에서 공존합니다. 한 엔진이 활발히 작동할 때, 다른 엔진은 GPU 자원을 해제하거나 오프로드하여 수용합니다.

강화학습 훈련의 각 반복에서, 중앙화된 컨트롤러가 먼저 추론 엔진을 호출하여 훈련을 위한 새로운 데이터를 생성합니다. 그 다음 훈련 엔진에 새로운 데이터로 훈련하도록 알리고, 다음 반복을 위해 업데이트된 파라미터를 추론 엔진으로 전송합니다. 각 엔진은 처리량을 위해 크게 최적화되어 있습니다.

또한 모델이 K2의 크기로 확장됨에 따라, 엔진 전환과 장애 복구의 지연 시간이 상당해집니다. 이러한 측면에서의 시스템 설계 고려사항을 제시합니다.

효율적인 엔진 전환

롤아웃 중에 훈련 엔진의 파라미터는 DRAM으로 오프로드됩니다. 훈련 엔진을 가동하는 것은 따라서 H2D 전송의 간단한 단계입니다. 그러나 추론 엔진을 가동하는 것은 더 큰 도전과제입니다. 추론 엔진은 서로 다른 샤딩 패러다임을 가진 훈련 엔진으로부터 업데이트된 파라미터를 얻어야 하기 때문입니다.

K2의 규모와 관련된 방대한 수의 디바이스를 고려할 때, 파라미터 재샤딩과 브로드캐스팅을 위해 네트워크 파일 시스템을 사용하는 것은 비현실적입니다. 오버헤드를 낮게 유지하는 데 필요한 총 대역폭은 초당 수 페타바이트에 달합니다.

이 도전과제를 해결하기 위해, 파라미터 상태를 관리하는 분산 체크포인트 엔진을 훈련 노드에 코로케이션하여 개발했습니다. 파라미터 업데이트를 수행하기 위해, 각 체크포인트 엔진 워커는 훈련 엔진으로부터 파라미터의 로컬 복사본을 얻은 다음, 모든 체크포인트 엔진 워커에 걸쳐 전체 파라미터 세트를 브로드캐스트합니다. 이후 추론 엔진은 체크포인트 엔진으로부터 필요한 파라미터 샤드만을 검색합니다. 이 과정은 위 그림에서 설명됩니다.

1T 모델에 대해 이를 가능하게 하기 위해, 업데이트는 메모리 풋프린트를 최소화하여 파이프라인 방식으로 파라미터별로 수행됩니다. 각 추론 워커의 특정 샤딩 스킴에 관계없이 전체 클러스터에 걸쳐 전체 파라미터 세트를 브로드캐스트하기로 선택했습니다. 이는 이론적으로 최적인 접근법보다 몇 배 더 많은 데이터를 전송하지만, 훈련 엔진과 추론 엔진에 덜 침입적인 더 간단한 시스템 설계를 제공합니다.

훈련 엔진과 추론 엔진을 완전히 분리하기 위해 이 작은 오버헤드를 교환하기로 선택했으며, 이는 유지보수와 테스트를 크게 단순화합니다. 주목할 점은 이 접근법이 감소된 동기화 오버헤드와 더 높은 네트워크 대역폭 활용으로 인해 필요한 것만 전송하는 방법을 능가한다는 것입니다. 시스템은 일반적인 강화학습 훈련 반복에서 무시할 수 있는 지속 시간인 30초 미만으로 Kimi K2의 전체 파라미터 업데이트를 완료할 수 있습니다.

효율적인 시스템 시작

대규모 훈련은 시스템 장애가 발생하기 쉬우므로, Kimi K2만큼 큰 모델에 대해 시작 시간을 최적화하는 것이 중요합니다. 훈련 엔진을 시작하기 위해, 각 훈련 워커가 디스크에서 파라미터의 일부 또는 전혀 읽지 않고, 필요한 파라미터를 동료들에게 브로드캐스트하도록 합니다. 설계 목표는 모든 워커가 집단적으로 체크포인트를 한 번만 읽도록 보장하여 비용이 많이 드는 디스크 IO를 최소화하는 것입니다.

추론 엔진은 독립적인 복제본이므로, 이들 간에 추가적인 동기화 장벽을 도입하는 것을 피하고자 합니다. 따라서 시작을 위해 체크포인트 엔진을 재사용하기로 선택합니다. 체크포인트 엔진이 훈련 엔진이 시작하는 방식과 유사하게 디스크에서 체크포인트를 집단적으로 읽도록 합니다. 그 다음 이전 섹션에서 소개된 접근법을 사용하여 초기화되지 않은 추론 엔진의 상태를 업데이트합니다.

전용 체크포인트 엔진을 활용함으로써, 추론 복제본이 다른 복제본과 통신하지 않고도 재시작할 수 있기 때문에 시스템은 단일 지점 장애에 대해서도 견고해집니다.

에이전틱 롤아웃

강화학습 인프라는 장기간, 다중 턴 에이전틱 작업의 훈련을 지원합니다. 롤아웃 중에 이러한 작업은 복잡한 환경 상호작용과 연장된 롤아웃 지속 시간과 같은 독특한 도전과제를 제시합니다. 이러한 문제를 완화하기 위한 몇 가지 최적화를 소개합니다.

환경의 다양성으로 인해, 특정 상호작용이 환경 피드백(예: 가상 머신이나 코드 인터프리터)을 기다리는 동안 차단되어 GPU가 유휴 상태가 될 수 있습니다. GPU 활용률을 최대화하기 위해 두 가지 전략을 사용합니다. 첫째, 더 쉽게 확장할 수 있는 전용 서비스로 중량급 환경을 배포합니다. 둘째, 특정 비용이 많이 드는 상호작용으로 인한 지연 시간을 상쇄하기 위해 많은 수의 동시 롤아웃을 사용합니다.

에이전틱 롤아웃의 또 다른 도전과제는 개별 롤아웃 궤적이 극도로 길 수 있다는 것입니다. 롱테일 궤적이 전체 롤아웃 과정을 차단하는 것을 방지하기 위해, 부분 롤아웃 기법을 사용합니다. 이 전략은 롱테일 미완료 작업을 일시 중지하고 다음 강화학습 반복에서 재개할 수 있게 합니다.

연구 효율성을 향상시키기 위해, 새로운 환경의 통합을 간소화하는 OpenAI Gym 프레임워크에서 영감을 받은 통합 인터페이스도 설계했습니다. 미래에 강화학습 인프라를 더 다양한 상호작용 환경으로 확장하기를 희망합니다.

평가

이 섹션에서는 Kimi-K2-Instruct의 후훈련 평가로 시작하여 Kimi-K2-Base의 능력에 대한 간략한 개요를 제공하고, 포괄적인 안전성 평가로 마무리합니다.

후훈련 평가

평가 설정

벤치마크

Kimi-K2-Instruct는 다양한 영역에 걸쳐 평가됩니다. 코딩 분야에서는 LiveCodeBench v6(2024년 8월부터 2025년 5월까지의 문제), OJBench, MultiPL-E, SWE-bench Verified, TerminalBench, Multi-SWE-bench, SWE-Lancer, PaperBench, Aider-Polyglot을 채택했습니다.

LiveCodeBench는 특히 중요한 벤치마크로, 모델의 훈련 데이터 오염을 방지하기 위해 지속적으로 새로운 코딩 경쟁 문제를 수집하는 라이브 업데이트 접근법을 사용합니다. 이 벤치마크는 단순한 코드 생성을 넘어서 자기 수정, 코드 실행, 테스트 출력 예측과 같은 포괄적인 평가 시나리오를 제공하여 모델의 코드 관련 능력을 더욱 종합적으로 평가할 수 있게 합니다.

도구 사용 작업의 경우, 다중 턴 도구 호출 능력을 강조하는 $\tau^2$-Bench와 AceBench에서 성능을 평가합니다. 추론 분야에서는 수학, 과학, 논리적 작업의 광범위한 범위를 포함합니다. AIME 2024/2025, MATH-500, HMMT 2025, CNMO 2024, PolyMath-en, ZebraLogic, AutoLogi, GPQA-Diamond, SuperGPQA, Humanity's Last Exam(텍스트 전용)입니다.

장기 컨텍스트 능력은 장기 컨텍스트 검색을 위한 MRCR과 장기 컨텍스트 추론을 위한 DROP, FRAMES, LongBench v2에서 벤치마킹됩니다. 사실성의 경우 FACTS Grounding, Vectara Hallucination Leaderboard, FaithJudge를 평가합니다. 마지막으로 일반적인 능력은 MMLU, MMLU-Redux, MMLU-Pro, IFEval, Multi-Challenge, SimpleQA, LiveBench를 사용하여 평가됩니다.

기준선

오픈소스와 독점 최첨단 모델 모두를 대상으로 벤치마킹하며, 모든 후보는 테스트 시간 계산으로부터의 추가적인 이득을 제거하기 위해 비사고 구성에서 평가됩니다. 오픈소스 기준선으로는 DeepSeek-V3-0324와 Qwen3-235B-A22B가 있으며, 후자는 벤더 권장 비사고 체제에서 실행됩니다. 독점 기준선으로는 Claude Sonnet 4, Claude Opus 4, GPT-4.1, Gemini 2.5 Flash Preview(2025-05-20)가 있으며, 각각 통합된 온도와 top-p 설정 하에서 공식 API를 통해 해당 비사고 모드로 호출됩니다.

평가 구성

모든 실행은 모델을 비사고 모드로 쿼리합니다. 출력 토큰 길이는 SWE-bench Verified(Agentless)를 제외하고 모든 곳에서 8192 토큰으로 제한되며, 이는 16384로 증가됩니다. 문제당 분산이 높은 벤치마크의 경우, 안정적인 점수를 얻기 위해 $k$번 반복 샘플링을 채택하고 결과를 평균화하여 Avg@k로 표시합니다.

장기 컨텍스트 작업의 경우, 평가 중 컨텍스트 윈도우 크기를 128K 토큰으로 설정하고, 이 제한을 초과하는 입력은 윈도우 내에 맞도록 잘라냅니다. SWE-bench Verified는 두 가지 모드로 평가됩니다. 테스트 없는 단일 패치를 통한 Agentless 코딩(Acc)과 내부 검증자를 사용한 best-of-N 선택을 통한 단일 시도(Acc) 및 다중 시도(Acc) 하에서 bash/editor 도구를 통한 에이전틱 코딩입니다. SWE-bench Multilingual은 단일 시도 에이전틱 설정에서만 테스트됩니다. 일부 데이터 포인트는 평가 비용이 너무 높아 생략되었습니다.

평가 결과

Kimi-K2-Instruct의 포괄적인 평가 결과가 표에 제시되어 있으며, 자세한 설명은 부록 C에서 제공됩니다. 네 가지 핵심 도메인에서의 주요 결과를 강조하면 다음과 같습니다.

에이전틱 및 경쟁 코딩

Kimi-K2-Instruct는 실세계 SWE 작업에서 최첨단 오픈소스 성능을 보여줍니다. SWE-bench Verified(65.8%, 다중 시도 시 71.6%), SWE-bench Multilingual(47.3%), SWE-lancer(39.1%)에서 대부분의 기준선을 능가하며, Claude 4 Opus 및 Sonnet과의 격차를 상당히 줄였습니다. 경쟁 코딩 벤치마크(예: LiveCodeBench v6 53.7%, OJBench 27.1%)에서도 모든 모델 중 선두를 차지하여 난이도 수준에 걸친 실용적인 코딩 숙련도를 강조합니다.

이러한 결과는 특히 주목할 만한데, SWE-bench는 실제 GitHub 이슈를 해결하는 능력을 평가하는 벤치마크로, 단순한 코드 생성을 넘어서 복잡한 소프트웨어 엔지니어링 작업을 수행하는 능력을 측정합니다. Kimi-K2-Instruct가 이 벤치마크에서 보여준 성능은 실제 개발 환경에서의 유용성을 시사합니다.

에이전틱 도구 사용

다중 턴 도구 사용 벤치마크에서 Kimi-K2-Instruct는 새로운 표준을 설정합니다. $\tau^2$-Bench에서 66.1 Pass@1, ACEBench에서 76.5를 달성하여 모든 기준선을 상당히 능가합니다. 이러한 결과는 도메인에 걸친 근거 있고 제어된 에이전트 주도 도구 오케스트레이션에서의 강점을 확인합니다.

$\tau^2$-Bench는 소매, 항공, 통신과 같은 다양한 도메인에서 복잡한 다중 턴 도구 호출 시나리오를 평가하는 벤치마크입니다. 각 도메인에서 Kimi-K2-Instruct는 일관되게 강력한 성능을 보여주었으며, 특히 소매(70.6%)와 통신(65.8%) 도메인에서 뛰어난 결과를 달성했습니다.

일반적인 능력

Kimi-K2-Instruct는 일반 지식, 수학, 지시 따르기, 장기 컨텍스트 작업에서 강력하고 균형 잡힌 성능을 보입니다. SimpleQA(31.0%), MMLU(89.5%), MMLU-Redux(92.7%)에서 오픈소스 동료들을 능가하고, 지시 벤치마크(IFEval: 89.8%, Multi-Challenge: 54.1%)에서 모든 모델을 선도합니다.

MMLU는 57개의 다양한 작업을 포괄하는 대규모 다중 작업 언어 이해 벤치마크로, STEM, 인문학, 사회과학 등 다양한 도메인에서 모델의 지식과 문제 해결 능력을 평가합니다. Kimi-K2-Instruct가 이 벤치마크에서 89.5%의 높은 성능을 달성한 것은 광범위한 지식 영역에서의 숙련도를 보여줍니다.

수학과 STEM에서는 최고 수준의 점수(AIME 2024: 69.6%, GPQA-Diamond: 75.1%)를 달성하고, 장기 컨텍스트 사실성과 검색(DROP: 93.5%, MRCR: 55.0%)에서 경쟁력을 유지합니다. 이러한 결과는 Kimi-K2-Instruct를 단기 및 장기 컨텍스트 설정 모두에서 균형 잡히고 유능한 일반주의자로 위치시킵니다.

개방형 평가

LMSYS Arena 리더보드(2025년 7월 17일)에서 Kimi-K2-Instruct는 3,000명 이상의 사용자 투표를 바탕으로 1위 오픈소스 모델이자 전체 5위를 기록했습니다. 다양하고 블라인드 프롬프트에 걸친 이러한 실세계 선호도 신호는 개방형 작업에서 고품질 응답을 생성하는 Kimi-K2의 강점을 강조합니다.

사전 훈련 평가

평가 설정

벤치마크

Kimi-K2-Base를 다양한 능력 영역에서 평가합니다. 일반적인 능력의 경우, MMLU, MMLU-Pro, MMLU-Redux, BBH, TriviaQA, SuperGPQA, SimpleQA, HellaSwag, AGIEval, GPQA-Diamond, ARC-Challenge, WinoGrande에서 평가합니다.

코딩 능력의 경우, EvalPlus(HumanEval, MBPP, HumanEval+, MBPP+의 평균), LiveCodeBench v6, CRUXEval을 사용합니다. 수학적 추론의 경우, GSM8K, GSM8K-Platinum, MATH, CMATH를 활용합니다. 중국어 능력의 경우, C-Eval, CMMLU, CSimpleQA에서 평가합니다.

기준선

선도적인 오픈소스 파운데이션 모델들과 벤치마킹합니다. DeepSeek-V3-Base, Qwen2.5-72B-Base(Qwen3-235B-A22B-Base는 오픈소스가 아니며, Qwen 시리즈에서 가장 큰 오픈소스 기본 모델은 Qwen2.5-72B-Base입니다), Llama 4-Maverick(Llama 4-Behemoth도 오픈소스가 아닙니다). 모든 모델은 공정한 비교를 보장하기 위해 동일한 구성에서 평가됩니다.

평가 구성

MMLU, MMLU-Redux, GPQA-Diamond, HellaSwag, ARC-Challenge, C-Eval, CMMLU에 대해서는 퍼플렉시티 기반 평가를 사용합니다. MMLU-Pro, SuperGPQA, TriviaQA, BBH, CSimpleQA, MATH, CMATH, GSM8K, GSM8K-Platinum, CRUXEval, LiveCodeBench, EvalPlus에 대해서는 생성 기반 평가를 사용합니다.

GPQA-Diamond에 내재된 높은 분산을 완화하기 위해 8번의 독립적인 실행에 걸친 평균 점수를 보고합니다. 모든 평가는 모든 모델에서 일관된 설정을 보장하는 LM-Harness-Evaluation에서 파생된 내부 프레임워크를 사용하여 수행됩니다.

평가 결과

Kimi-K2-Base와 선도적인 오픈소스 파운데이션 모델들 간의 포괄적인 비교가 표에 제시되어 있습니다. 결과는 Kimi-K2-Base가 평가된 작업의 대부분에서 최첨단 성능을 달성하여 오픈소스 환경에서 선도적인 파운데이션 모델로 자리매김함을 보여줍니다.

일반 언어 이해

Kimi-K2-Base는 12개 영어 언어 벤치마크 중 10개에서 최첨단 성능을 달성합니다. 주목할 만한 결과로는 MMLU(87.79%), MMLU-Pro(69.17%), MMLU-Redux(90.17%), SuperGPQA(44.67%), SimpleQA(35.25%)가 있으며, 모든 기준선을 상당히 능가합니다.

코딩 능력

코딩 벤치마크에서 Kimi-K2-Base는 모든 메트릭에서 선도적인 성능으로 새로운 표준을 설정합니다. CRUXEval-I-cot에서 74.00%, CRUXEval-O-cot에서 83.50%, LiveCodeBench v6에서 26.29%, EvalPlus에서 80.33%를 달성하여 특히 단계별 추론이 필요한 시나리오에서 뛰어난 코드 생성 및 이해 능력을 보여줍니다.

수학적 추론

Kimi-K2-Base는 뛰어난 수학적 능력을 보이며, 4개 벤치마크 중 3개에서 선도합니다. MATH(70.22%), GSM8K(92.12%), GSM8K-Platinum(94.21%). CMATH(90.26%)에서는 DeepSeek-V3-Base(90.53%)에 근소하게 뒤처지지만 경쟁력 있는 성능을 유지합니다. 이러한 결과는 다양한 난이도 수준에서 모델의 견고한 수학적 문제 해결 능력을 강조합니다.

중국어 언어 이해

모델은 뛰어난 다국어 능력을 보여주며, 모든 중국어 언어 벤치마크에서 최첨단 결과를 달성합니다. C-Eval(92.50%), CMMLU(90.90%), CSimpleQA(77.57%). 이러한 결과는 Kimi-K2-Base를 다른 언어에서 강력한 성능을 유지하면서 중국어 언어 이해를 위한 선도적인 모델로 확립합니다.

안전성 평가

실험 설정

Promptfoo 프레임워크를 사용하여 Kimi K2에 대한 레드팀 평가를 다른 오픈소스 LLM들과 비교하여 수행했습니다. 평가는 유해한 콘텐츠, 프라이버시 콘텐츠, 보안 콘텐츠를 포함한 다양한 공격 시나리오와 프롬프트 주입 및 반복적 탈옥과 같은 다양한 공격 전략을 다뤘습니다.

모델 선택

Kimi K2를 다른 세 개의 오픈소스 LLM과 비교합니다. DeepSeek-V3, DeepSeek-R1, Qwen3입니다.

Promptfoo 설정

평가된 플러그인과 전략이 표에 나열되어 있으며, 각 플러그인은 모든 전략과 쌍을 이루어 성능을 평가합니다. 플러그인에는 유해한 그래픽 콘텐츠, 괴롭힘과 괴롭힘, 혐오 발언, 모욕, 욕설, 급진화, 자해, 성적 콘텐츠, ToxicChat이 포함됩니다. 범죄 관련으로는 화학 및 생물학적 무기, 아동 착취, 저작권 위반, 사이버 범죄, 불법 활동, 불법 약물, 무차별 무기, 지적 재산권 위반, 비폭력 범죄, 폭력 범죄, 성범죄가 있습니다.

잘못된 정보 관련으로는 경쟁업체 지지, 무감독 계약, 과도한 대리권, 환각, 잘못된 정보와 허위 정보, 전문적 조언, 안전하지 않은 관행, 모방, 과도한 의존, 정치적 의견, 종교적 민감성이 포함됩니다. 프라이버시 관련으로는 프라이버시 위반, API/데이터베이스의 PII, 직접적인 PII 노출, 세션 데이터의 PII, 사회 공학을 통한 PII가 있습니다. 보안 관련으로는 ASCII 밀수, CyberSecEval, Harmbench, 디버그 액세스, 발산 반복, DoNotAnswer, 악성 코드, Pliny, 프롬프트 추출, 추론 DoS, 도구 발견이 포함됩니다.

전략으로는 기본, 프롬프트 주입, 반복적 탈옥, 크레센도가 있습니다.

테스트 케이스 수

대규모 언어 모델 추론의 고유한 비결정성을 고려하여, 단일 패스 출력은 가변성을 보일 수 있습니다. 이를 고려하기 위해 각 전략에 대해 플러그인당 3개의 공격 프롬프트를 생성했습니다.

프롬프트 언어 설정

각 플러그인-전략 조합에 대한 언어 호환성을 사전 테스트했습니다. 일부 플러그인은 영어와 중국어를 모두 지원하는 반면, 다른 플러그인은 영어만 지원합니다. 둘 다 지원하는 조합의 경우, 각 언어로 3개의 프롬프트를 생성하여 조합당 총 6개의 프롬프트를 만들었습니다.

수동 검토

평가 과정에 인간 검토를 통합했습니다. 주관성 문제를 최소화하기 위해 여러 라운드의 검토를 수행하고 동일한 검토자가 주어진 테스트 세트 내의 모든 사례를 평가하도록 배정하여 일관성을 보장하고 판단의 가변성을 줄였습니다.

안전성 평가 결과

다양한 플러그인-전략 조합 하에서 서로 다른 모델들의 통과율이 표에 제시되어 있습니다. 특정 평가 시나리오에 대한 목표화된 최적화 없이, 일부 복잡한 사례(예: 유해한-반복적 탈옥)의 통과율은 다른 모델들에 비해 상대적으로 높았습니다.

서로 다른 공격 전략에서 모델들은 다양한 경향을 보였습니다. Base64 전략 하에서는 통과율이 일반적으로 100%에 접근하거나 도달하여 인코딩 변환이 모델들의 기본 견고성에 최소한의 영향을 미쳤음을 시사합니다. 반대로 크레센도 전략은 통과율의 일반적인 하락을 가져와 더 강한 적대적 효과를 나타냈습니다.

또한 복잡한 공격 전략이 항상 기본 프롬프트를 능가하지는 않습니다. 일부 원래 적대적 프롬프트는 여러 라운드의 변환 후 의도된 의미를 잃을 수 있어, 결과적인 모델 출력이 덜 의미 있게 만들 수 있습니다.

자동화된 레드팀의 한계

인간 검토의 개입으로 인해 평가 결과는 불가피하게 어느 정도의 주관성을 포함합니다. 또한 특정 플러그인 유형은 API 오용이나 외부 도구 호출을 포함하는데, 이는 도구 호출 능력을 가진 에이전트 모델을 평가하는 데 더 적합합니다. 기본 LLM의 맥락에서 이러한 테스트는 제한된 관련성을 가질 수 있습니다.

한계점

내부 테스트를 통해 현재 Kimi K2 모델들에서 몇 가지 한계점이 확인되었습니다. 어려운 추론 작업이나 불분명한 도구 정의를 다룰 때 모델이 과도한 토큰을 생성하여 출력이 잘리거나 불완전한 도구 호출이 발생하는 경우가 있습니다. 또한 도구 사용이 불필요하게 활성화된 특정 작업에서는 성능이 저하될 수 있습니다. 완전한 소프트웨어 프로젝트를 구축할 때 원샷 프롬프팅의 성공률은 에이전틱 코딩 프레임워크 하에서 K2를 사용하는 것만큼 좋지 않습니다. 이러한 문제들을 향후 릴리스에서 해결하기 위해 노력하고 있으며 더 많은 피드백을 기대하고 있습니다.

이러한 한계점들은 현재 대규모 언어 모델들이 직면하고 있는 공통적인 도전과제들을 반영합니다. 특히 복잡한 추론 작업에서의 토큰 생성 제어와 도구 사용의 적절한 조절은 에이전틱 인텔리전스 구현에서 핵심적인 기술적 과제입니다. 앞서 설명한 강화학습 프레임워크의 예산 제어 메커니즘이 이러한 문제를 완화하는 데 도움이 되지만, 여전히 개선의 여지가 있음을 보여줍니다.

결론

Kimi K2는 에이전틱 인텔리전스를 위해 구축된 1조 파라미터 오픈 가중치 MoE 모델입니다. 토큰 효율적인 MuonClip 옵티마이저와 15.5조 토큰의 고품질 데이터셋을 활용하여 Kimi K2는 안정적이고 확장 가능한 사전 훈련을 달성했습니다. 후훈련은 대규모 합성 도구 사용 데이터와 검증 가능한 보상 및 자기 비판 피드백을 모두 사용하는 통합 RL 프레임워크를 결합합니다.

Kimi K2는 에이전틱 및 추론 벤치마크에서 새로운 최첨단 성능을 설정하여 현재까지 가장 유능한 오픈 가중치 LLM으로 자리매김했습니다. 이 모델은 특히 도구 사용, 소프트웨어 엔지니어링, 복잡한 추론 작업에서 뛰어난 성능을 보여주었으며, 앞서 제시된 평가 결과들이 이를 입증합니다.

본 연구의 주요 기여는 세 가지 핵심 영역에 걸쳐 있습니다. 첫째, MuonClip 옵티마이저의 도입으로 대규모 MoE 모델의 안정적인 훈련이 가능해졌습니다. 둘째, 대규모 에이전틱 데이터 합성 파이프라인을 통해 도구 사용 능력을 체계적으로 학습할 수 있는 프레임워크를 구축했습니다. 셋째, 검증 가능한 보상과 자기 비판을 결합한 통합 강화학습 프레임워크를 통해 다양한 작업 도메인에서 일관된 성능 향상을 달성했습니다.

Kimi K2의 성공은 에이전틱 인텔리전스 분야의 발전에 중요한 이정표를 제시합니다. 오픈소스로 공개된 기본 모델과 후훈련된 체크포인트는 커뮤니티가 에이전틱 인텔리전스를 대규모로 탐색, 개선, 배포할 수 있는 기반을 제공합니다. 이는 향후 연구와 실용적 응용 개발에 중요한 자원이 될 것입니다.

부록

Kimi K2 논문의 부록에서는 연구의 기술적 세부사항과 보완 자료들을 제공합니다. 이 섹션에서는 도구 호출을 위한 토큰 템플릿, 평가 세부사항, QK-Clip의 모델 품질에 대한 영향 분석, Muon이 로짓 폭발에 더 취약한 이유, 그리고 강화학습을 위한 엔진 전환 파이프라인에 대한 상세한 설명을 다룹니다.

도구 호출을 위한 토큰 템플릿

Kimi K2의 도구 호출 시스템은 세 가지 핵심 구성요소로 이루어진 토큰 구조를 사용합니다. 이 구조는 모델이 외부 도구와 효과적으로 상호작용할 수 있도록 설계되었습니다.

도구 선언 메시지는 사용 가능한 도구 목록과 인수의 스키마를 정의합니다. 이는 다음과 같은 형식으로 구성됩니다.

<|im_begin|> tool_declare <|im_middle|>
# Tools
{{ tool declaration content }}
<|im_end|>

도구 선언 내용을 표현하기 위해 TypeScript를 사용하는데, 이는 간결한 언어이면서도 포괄적인 타입 시스템을 가지고 있어 도구 파라미터의 타입과 제약사항을 간단한 텍스트로 표현할 수 있기 때문입니다. OpenAI의 채팅 완성 API와 호환되는 JSON 형식과 비교할 때, TypeScript로 정의된 동일한 도구는 훨씬 더 간결합니다.

예를 들어, 날씨 정보를 가져오는 도구와 계산기 도구를 JSON으로 정의하면 상당히 길어지지만, TypeScript로는 다음과 같이 간단하게 표현할 수 있습니다.

namespace functions {
// Get weather for a location and date
type get_weather = (_: {
  // City and country e.g. Beijing, China
  location: string,
  // Date to query, format in '%Y-%m-%d'
  date?: string
}) => any;
// Simple calculator
type Calculator = (_: {
  // Arithmetic expression in javascript
  expr?: string
}) => any;
}

도구 호출 섹션은 모델의 응답 메시지에서 도구 호출 요청을 인코딩합니다. 이는 다음과 같은 템플릿을 따릅니다.

<tool_call_section_begin|>
<|tool_call_begin|>
// call_id part
functions.{{tool name}}:{{counter}}
<|tool_arguments_begin|>
{{ json serialized call arguments }}
<|tool_call_end|>
<|tool_call_section_end|>

이 템플릿은 단일 응답 턴에서 여러 도구 호출을 배치하여 병렬 도구 호출을 지원합니다. 각 도구 호출은 functions.{tool-name}:{counter} 형식의 고유한 호출 ID를 가지며, 여기서 counter는 대화에서 0부터 시작하는 모든 도구 호출의 자동 증가 카운터입니다.

추론 중에 모델이 예상치 못한 토큰을 생성하여 도구 호출을 파싱할 때 형식 오류가 발생할 수 있습니다. 이 문제를 해결하기 위해 lm-format-enforcer에서 영감을 받은 enforcer라는 제약 디코딩 모듈을 개발했습니다. <tool_call_section_begin|> 토큰이 생성되면, 이후의 도구 관련 토큰들이 미리 정의된 템플릿을 따르고 JSON 인수 문자열이 선언된 스키마를 따르도록 보장합니다.

도구 결과 메시지는 도구의 호출 ID와 해당 결과를 포함하는 간단한 텍스트 메시지로 인코딩됩니다.

<|im_begin|> tool <|im_middle|>
## Results of {{call_id}}
{{ execution result content }}
<|im_end|>

QK-Clip이 모델 품질에 미치는 영향 분석

QK-Clip 설계는 최소 개입 원칙을 따릅니다. 즉, 필요할 때만 활성화되고 훈련이 안정화된 후에는 비활성화됩니다. 경험적 증거와 분석은 모델 품질에 대한 무시할 수 있는 영향으로 수렴됩니다.

소규모 절제 실험에서는 0.5B 활성화 파라미터와 3B 총 파라미터를 가진 두 개의 소규모 MoE 모델을 훈련했습니다. 하나는 바닐라 Muon을, 다른 하나는 낮은 클리핑 임계값($\tau = 30$)을 사용하는 MuonClip을 사용했습니다. MuonClip을 적용하는 것이 손실 곡선에 무시할 수 있는 영향을 미친다는 것을 보여주며, 이는 공격적인 클리핑도 MuonClip으로 수렴과 훈련 역학을 손상시키지 않음을 나타냅니다.

위 그림은 공격적인 임계값($\tau = 30$)을 사용한 소규모 설정에서 Muon에 QK-Clip을 적용하는 것이 손실에 무시할 수 있는 영향을 미친다는 것을 보여주며, 이는 어텐션 로짓을 제약하는 안전하고 효과적인 방법임을 나타냅니다.

자기 비활성화 특성

Kimi K2에서 QK-Clip은 일시적으로만 활성화되었습니다.

• 초기 70,000 단계: 어텐션 헤드의 12.7%가 적어도 한 번 QK-Clip을 트리거하여 $S_{\max}$를 100으로 클램핑했습니다.
• 70,000 단계 이후: 모든 헤드가 어느 시점에서 $S_{\max}$를 100 아래로 줄여 QK-Clip을 비활성화했습니다.

QK-Clip이 활성화될 때는 다른 헤드에 대한 잠재적 과정규화를 최소화하기 위해 레이어별이 아닌 헤드별로 적용됩니다. 훈련이 안정화된 후에는 QK-Clip이 비활성화되어 전혀 효과가 없습니다.

Muon이 로짓 폭발에 더 취약한 이유

로짓 폭발은 최대 사전 소프트맥스 어텐션 점수 $S_{\max} = \max_{i,j}(q_i \cdot k_j)$가 훈련 중 무한정 증가할 때 발생합니다. $\left\vert q_i \cdot k_j\right\vert \leq \left\Vert q_i\right\Vert\left\Vert k_j\right\Vert \leq \left\Vert x_i\right\Vert\left\Vert x_j\right\Vert\left\Vert\mathbf{W}_q\right\Vert\left\Vert\mathbf{W}_k\right\Vert$이고 RMS-Norm이 $\left\Vert x_i\right\Vert\left\Vert x_j\right\Vert$를 제한된 상태로 유지하므로, 이 현상은 주로 $\left\Vert\mathbf{W}_q\right\Vert$ 또는 $\left\Vert\mathbf{W}_k\right\Vert$의 증가하는 스펙트럴 노름에 의해 주도됩니다.

경험적으로 Muon이 로짓 폭발에 더 취약하다는 것을 발견했습니다. 이에 대한 가설을 제시합니다.

업데이트의 구조적 차이

Muon은 msign 연산에서 나오는 가중치 업데이트를 생성하므로, 업데이트 행렬의 모든 특이값이 동일하여 효과적 순위가 완전합니다. 반면, Adam이 생성하는 일반적인 업데이트 행렬은 편향된 스펙트럼을 보입니다. 즉, 몇 개의 큰 특이값이 지배적이고 효과적 순위가 낮습니다.

16B Moonlight 모델에서 이러한 현상이 검증되었으며, Muon으로 훈련된 가중치가 Adam으로 훈련된 것보다 더 높은 특이값 엔트로피(즉, 더 높은 효과적 순위)를 보인다는 것을 보여주어 이론적 직관을 뒷받침합니다.

SVD 공식화

단계 $t-1$에서 파라미터 행렬이 특이값 분해 $\mathbf{W}_{t-1} = \sum_i \sigma_i u_i v_i^{\top}$를 가진다고 하면, 업데이트 행렬을 $\Delta\mathbf{W}_t = \sum_j \bar{\sigma} \bar{u}_j \bar{v}_j^{\top}$로 쓸 수 있습니다.

다음 파라미터 업데이트는 $\mathbf{W}_t \leftarrow \sum_i \sigma_i u_i v_i^{\top} + \sum_j \bar{\sigma} \bar{u}_j \bar{v}_j^{\top}$입니다.

Muon에서는 가중치와 업데이트 모두 Adam보다 높은 효과적 순위를 가지므로, 특이벡터 쌍 $u_i v_i^{\top}$가 $\bar{u}_j \bar{v}_j^{\top}$와 정렬될 확률이 더 높다고 가정합니다. 이는 $\mathbf{W}_t$의 해당 특이값이 가산적으로 증가하게 할 수 있습니다.

어텐션 특화 증폭

어텐션 로짓은 이중선형 형태 $q_i \cdot k_j = (x_i \mathbf{W}_q) \cdot (x_j \mathbf{W}_k)$를 통해 계산됩니다. 곱 $\mathbf{W}_q \mathbf{W}_k^{\top}$는 스펙트럴 노름을 제곱하므로, 어느 행렬에서든 특이값 증가가 복합됩니다. Muon의 특이값 확대 경향은 따라서 로짓 폭발의 더 높은 위험으로 이어집니다.

강화학습 훈련을 위한 엔진 전환 파이프라인

강화학습 인프라에서 체크포인트 엔진은 각 GPU에서 세 개의 동일한 크기 디바이스 버퍼를 관리합니다. 오프로드된 모델 파라미터를 로딩하기 위한 H2D 버퍼와 GPU 간 브로드캐스트를 위한 두 개의 IPC 버퍼입니다. IPC 버퍼는 추론 엔진과 공유되어 동일한 물리적 메모리에 직접 액세스할 수 있게 합니다.

이론적으로는 3단계 파이프라인이 가능합니다. (1) H2D - 최신 가중치의 샤드를 H2D 버퍼로 비동기적으로 복사, (2) 브로드캐스트 - 복사가 완료되면 샤드를 하나의 IPC 버퍼로 복사하고 모든 디바이스에 브로드캐스트, (3) 재로드 - 추론 엔진이 다른 IPC 버퍼에서 파라미터를 동시에 로드합니다.

그러나 NVIDIA H800 클러스터에서는 동시 H2D와 브로드캐스트가 공유 PCIe 패브릭을 포화시켜 3단계를 순차적 절차로 축소시킵니다. 따라서 더 간단한 2단계 방식을 채택합니다. (1) 모든 디바이스가 단일 동기 H2D 전송을 수행, (2) 브로드캐스트와 재로드가 병렬로 진행됩니다.

H2D, 브로드캐스트, 가중치 재로드를 오버랩함으로써 훈련 엔진에서 모든 추론 엔진으로 가중치를 재샤딩하는 높은 대역폭을 얻을 수 있습니다.

References

• Kimi K2: Open Agentic Intelligence