PaperlayerView exact revision on arXiv

Source-linked AI summary

A.X K2 Technical Report

Cheolseung Baek, Dhammiko Arya, Eunki Kim, Gun Song, Gyoungeun Han, Hyunho Yang, Hyunjun Eun, Jin Kim, Junyoung Park, Juyun Wee, Minki Hong, Minkyung Park, Minsang Kim, Minsoo Kang, SaeRom Kim, Sangjin Kim, Sangyeol Lee, Seojin Lee, Seokhwan Jo, Seokyoung Hong, Seongho Choi, Seonghye Cho, Seongmin Ok, Sereimony Sek, Seungmo Cho, Seungsik Kim, Singon Kim, Sohee Park, Sooyeon Park, Subin Yi, Sungbin Yoon, Sungeun Lee, Sung Jun Cheon, Sungwan Kim, Sunwoo Lee, Tae Yoon Kim, Wonbeom Jang, Yohan Ra, Yong-jin Han, Youngjin Kim, Youngrang Kim, Yujin Kang, Yujin Lee

arXiv:2608.30181v1cs.AIcs.CL

TL;DR

A.X K2 addresses the challenge of building an agentic foundation model that also supports efficient serving, controllable inference compute, and realistic training constraints. It combines a large MoE design with token-efficient training, sparse gated long-context attention, gated normalization, and Think-Fusion. The model remains competitive with strong open-weight baselines, matching or exceeding them on mathematics and Korean-language benchmarks while enabling thinking-mode control.

  • Problem

    Practical agentic foundation models must combine broad knowledge and reasoning with efficient serving, controllable inference-time compute, and trainability under realistic data and hardware budgets.

  • Method

    A.X K2 combines a 688B-parameter MoE architecture, token-efficient training, SGA and GN for long-context and stable low-precision computation, and Think-Fusion for unified thinking-mode control.

  • Results

    A.X K2 performs competitively across English and Korean benchmarks, matching or exceeding strong open-weight baselines on mathematics and Korean-language tasks.

  • Takeaways & Limitations

    A.X K2 provides a practical, controllable foundation model that supports user-selected reasoning depth within a single unified model.

  • Takeaways & Limitations

    Scaling was constrained to expert parallelism within eight-GPU intra-node NVLink domains, limiting exploration of higher-degree multi-node expert parallelism.

Abstract

from arXiv · show

We introduce A.X K2, a 688B-parameter Mixture-of-Experts (MoE) language model trained from scratch as a high-performance foundation for \emph{agentic} applications. Trained on approximately 8.5T tokens---fewer than its predecessor, A.X K1---on a smaller but higher-quality mixture with substantially expanded agentic and software-engineering data, it nonetheless improves over A.X K1 across the board, by over 30 percentage points on some benchmarks, reflecting large gains in token efficiency. To support long contexts efficiently, we introduce Sparse Gated Attention (SGA), which combines sparse attention with gated attention, and adopt Gated Norm (GN) to stabilize large-scale training. SGA is trained natively at 128K through a \emph{sparse} indexer warmup that optimizes the indexer against its own sparse top-$k$ selection rather than the dense attention distribution, making adaptation markedly cheaper: each query reads only 2,048 positions, yet long-context quality is unchanged and A.X K2 scores 94.6 on RULER out to 256K. The outlier suppression of GN in turn keeps 4-bit NVFP4 serving within one point of FP8 accuracy. A simple yet effective Think-Fusion recipe further lets users switch between thinking and non-thinking modes within a single unified model. Extensive evaluations show that A.X K2 performs competitively against strong open-weight baselines, matching or exceeding them on math and Korean-language benchmarks.

1 Introduction

A.X K2 is a frontier-scale MoE foundation model designed to combine agentic capability with efficient, controllable deployment. Its architecture, training strategy, and evaluation emphasize token efficiency, long-context inference, flexible reasoning modes, and Korean-language performance.

  • Model and training: A.X K2 uses 688B total parameters with 33B active parameters and approximately 8.5T training tokens as an agentic foundation model.It is trained from scratch as A.X K1’s successor.
  • Evaluation: A.X K2 leads representative open-weight comparisons in mathematics, Korean-language understanding, and agentic tool use.Figure 1 identifies AIME26, Apex, KMMLU-Pro, CLIcK, and τ2-Bench Telecom as the relevant benchmarks.
  • Inference control: Think-Fusion enables one unified model to switch explicitly between thinking and non-thinking modes, trading response quality for inference cost per request.Thinking mode supports complex reasoning, whereas non-thinking mode targets concise, low-latency responses.
  • Model and training: Over 30 percentage points of improvement on some benchmarks demonstrates substantial gains over A.X K1 despite using fewer tokens.The report attributes the gains to improved token efficiency and an enhanced multi-stage data-processing pipeline.
  • Architecture: SGA combines sparse attention, a lightweight top-k indexer, and gated attention, while GN suppresses outliers for stable low-precision computation.The associated transformer block selects key–value candidates before latent attention and applies a head-specific output gate.

2 Architecture

A.X K2 is a 688B-parameter MoE model with 33B active parameters, designed for efficient inference and long-context operation. Its architecture combines gated attention, Gated Norm, and Sparse Gated Attention to improve stability, serving efficiency, and context scaling.

  • Model Configuration: A.X K2 uses 688B total parameters and 33B active parameters in an MoE architecture, increasing routed experts from 192 to 256 while keeping activation fixed.The 256-expert design aligns with expert-parallel sharding and kernel tiling under a fixed compute budget.
  • Model Configuration: The backbone uses MLA, gated attention, QK normalization, one dense feed-forward layer, and 60 MoE layers across 61 transformer layers.Gated attention introduces a head-specific output gate, while QK normalization is applied before attention scores are computed.
  • Gated Transformer Blocks: Gated Norm applies input-dependent gating after RMSNorm to suppress hidden-state outliers and smooth loss convergence.GN replaces the earlier dual-normalization design and incurs approximately a 5% training-throughput reduction.
  • Tokenizer and Context: The tokenizer is an unchanged byte-level BPE vocabulary of 163,840 tokens covering English, Korean, Chinese, Japanese, and Spanish.The model reaches 128K context natively and 256K at inference through YaRN scaling.
  • Sparse Gated Attention: Sparse Gated Attention selects k = 2048 tokens per query, reducing attention to 1.6% of positions at 128K and 0.8% at 256K.The sparse indexer is adapted after native 128K training, while gated attention remains part of pre-training.

3 Pre-Training

A.X K2 uses a refreshed, aggressively filtered corpus and a multi-stage curriculum that moves from general knowledge to reasoning and long-context adaptation. Warmup–Stable–Decay scheduling and native context extension organize the training process under a constrained compute budget.

  • Data Mixture: A.X K2 is pre-trained on approximately 8.2T tokens spanning web text, code, STEM, reasoning, books, and synthetic data.The corpus refresh adds newer English and Korean sources, broader multilingual coverage, and stronger synthetic data for agentic and long-horizon objectives.
  • Quality Filtering and Curation: The filtering regime selects a training-efficient subset from a roughly 16.2T-token candidate pool and places higher-value, higher-difficulty data later in training.Quality and domain criteria become progressively stricter across stages, with difficulty scoring introduced from Stage 2 onward.
  • Pre-Training Process: The WSD schedule organizes pre-training into general knowledge, advanced reasoning, and long-context extension stages.Stage 3 progressively extends context to 32K and 128K before sparse-attention adaptation.
  • Pre-Training Stages: Stage 1 trains on 6.4T tokens for general knowledge, while Stage 2 uses 1.4T curated tokens for advanced knowledge and reasoning.Both stages use sequence length 4,096; Stage 2 begins learning-rate decay over higher-complexity data.
  • Optimization: Checkpoint averaging over a narrow window of roughly 8B tokens consistently yields a stronger final model than any individual checkpoint at negligible merge cost.The weighting scheme emphasizes more recent checkpoints while retaining benefits from averaging.

4 Post-Training

Post-training combines indexer and main supervised fine-tuning with multi-stage reinforcement learning for instruction following, alignment, tool use, and safety. The pipeline also addresses precision consistency and rollout throughput for large-scale MoE training.

  • Supervised Fine-Tuning: The SFT pipeline begins with indexer SFT and then fixes the indexer for the main SFT stage.Both stages use packed 128K-token sequences, with 0.4M examples sampled for indexer SFT.
  • Mode Control: Think-Fusion enables one model to support user-specified thinking and non-thinking modes at inference time.Thinking mode produces explicit reasoning chains, while non-thinking mode targets concise, low-latency responses.
  • Data: SFT pairs thinking examples for mathematics, science, code, and tool use with concise instruction examples for chat, safety, and non-thinking responses.This paired construction teaches explicit switching between reasoning and non-thinking modes.
  • Precision Consistency: Trainer–rollout precision mismatch destabilizes reward learning, whereas matched blockwise FP8 execution produces stable training.Divergent quantization schemes cause sampling and update distributions to drift apart, increasing policy KL.
  • Asynchronous Rollout: A TP1/DP8 rollout layout with fixed EP8 increases throughput by 1.6× relative to TP8/DP1 under the same setup.Data-parallel attention replication removes per-layer tensor-parallel all-reduces while preserving trainer-compatible expert sharding.
  • Reinforcement Learning: Multi-stage RL jointly optimizes instruction following, human-preference alignment, and agentic tool use, with later stages adding safety training.All stages use a shared reward framework rather than dedicating each stage to a single capability.

5 Evaluation

Across public and Korean-focused evaluations, A.X K2 improves broadly over A.X K1 and is competitive with strong open-weight models, with particular strength in mathematics, Korean tasks, instruction following, and long-context evaluation.

  • Public benchmark results: A.X K2 leads or matches strong open-weight models in mathematics and Korean benchmarks, while improving over A.X K1 across every evaluated category.It achieves the best compared scores on AIME26, Apex, Apex-shortlist, KMMLU-Pro, and CLIcK, while remaining competitive on KoBALT.
  • Long-context evaluation: 94.6 overall on RULER across lengths up to 256K, while each query reads only 2,048 positions.The model retains 95.1% of its 4K score at 128K and 88.8% at 256K, using 0.8% of the KV cache at 256K.
  • Long-context evaluation: 62.99 sparse versus 62.80 dense on LongBench v1, indicating no measured long-context quality drop after pruning attention to 1.6% of the 128K context.The reported difference is +0.19 in favor of the sparse model and is described as within measurement noise.
  • Korean evaluation: 27.11 pp improvement over A.X K1 on Korean long-context KS-Eval, corresponding to a 68.2% relative gain.KS-Eval covers 8 long-context tasks and 490 items designed around Korean linguistic and cultural context.
  • KS-Eval: 19.47 pp, 27.98 pp, 27.11 pp, and 13.46 pp gains over A.X K1 in Common Sense, Instruction Following, Long Context, and STEM, respectively.A.X K2 scores higher than A.X K1 on all 26 individual KS-Eval tasks; it ranks first in Common Sense and Instruction Following among the compared models.
  • Robustness evaluation: 69.0% overall F1, driven by 61.1% abstention performance, while answerable F1 is 76.9% and ranks fourth of six.Under adversarial conditions, A.X K2 retains 67.8% against incorrect user assertions but reaches 41.4% when a supplied document supports the wrong answer.

6 Inference Optimization

A.X K2 combines low-precision robustness with SGA-based long-context serving and selective EAGLE3 or NVFP4 optimization. These techniques preserve quality while improving throughput and extending the efficient serving regime relative to A.X K1.

  • 6.1 Low-Precision Robustness: 36% of A.X K1’s BF16 serving memory is required by NVFP4, while FP8 requires about 62% and NVFP4 reduces that footprint by roughly 57% relative to FP8.The 276 GB freed relative to FP8 can be redirected to KV cache, increasing concurrency at 128K–256K inputs.
  • 6.1 Low-Precision Robustness: 99.0% of FP8 average accuracy is retained under NVFP4, with a 0.76-point average drop across eleven Korean and English tasks.NVFP4 preserves perfect NIAH scores at 128K, 256K, and 512K; the largest regressions are GSM8K (−2.50) and MATH (−2.42).
  • 6 Inference Optimization: Speculative decoding and quantization can be applied selectively on top of SGA, depending on workload priorities.The evaluation measures serving under controlled B200 single-node conditions, concurrency 32, and 1K output sequence length.
  • 6.2 Long-Context Serving Throughput: 64K inputs mark the FP8 crossover where A.X K2 overtakes A.X K1, reaching +35.4% at 120K with 12.3K versus 9.1K total tok/s.A.X K1 peaks at 64K and declines thereafter, whereas A.X K2 continues rising through 120K.
  • 6.2 Long-Context Serving Throughput: 23–30% throughput gains from EAGLE3 across input lengths move the A.X K1 crossover from 64K to 32K and reach +65.9% at 120K.The 5.6 GiB drafter is under 1% of the 688B target model and requires no additional GPUs or serving infrastructure.
  • 6.2 Long-Context Serving Throughput: NVFP4 exceeds A.X K1 from 1K through 120K and reaches 20.3K total tok/s, or +123.7%, at 120K.In the 32K+ regime, NVFP4 delivers roughly 1.4×–2.2× A.X K1 throughput while retaining 99.0% of FP8 accuracy.
  • 6.2 Long-Context Serving Throughput: EAGLE3 gains depend on acceptance rate, so results measured with vLLM bench serve on a random dataset may differ on real code or dialogue traffic.The reported speculative-decoding figures are therefore workload-dependent rather than universal.

7 Limitations

The authors identify infrastructure, compute-budget, and modality constraints that bound A.X K2’s current scaling and application scope. They present these limitations as directions for future iterations.

  • Infrastructure Constraints on Scaling: Expert parallelism was confined to EP = 8 within the eight-GPU intra-node NVLink domain, limiting exploration of higher-degree multi-node parallelism.Efficient higher-degree scaling would require tighter integration between expert-parallel communication and inter-node interconnects.
  • Optimization under Resource Constraints: Fixed time and GPU budgets caused modest performance shortfalls relative to similarly sized models trained with larger compute budgets.The model configurations were selected using empirical scaling laws under resource constraints.
  • Scope of Modalities: A.X K2 is currently text-only and lacks native multimodal understanding.The authors plan to extend the model with native multimodal capabilities in future work.

8 Conclusion

The conclusion presents A.X K2 as a high-capacity MoE model designed for practical inference efficiency under fixed, time-bounded resources. It highlights controllable reasoning modes and competitive English and Korean benchmark performance, while identifying system-level engineering as central to deployment.

  • 8 Conclusion: A.X K2 uses MoE scaling principles and compute-budget planning to trade off model scale and training tokens under fixed, time-bounded resources.The conclusion frames this system-aware design as yielding globally competitive capability.
  • 8 Conclusion: Think-Fusion trains paired thinking and non-thinking responses, followed by multi-stage on-policy reinforcement learning, for user-controlled mode switching.The unified model supports explicit allocation of inference-time compute according to task needs.
  • 8 Conclusion: A.X K2 is competitive with strong open-weight baselines across mathematics, coding, and general-knowledge benchmarks in English and Korean.The conclusion emphasizes competitive evaluation across both languages and multiple capability areas.
  • 8 Conclusion: System-level engineering addressed gaps between theoretical and realized throughput caused by communication overhead, load imbalance, and execution inefficiencies.The project used infrastructure and optimization insights to make large-scale MoE training effective under real-world constraints.
  • 8 Conclusion: Future work prioritizes native multimodal capabilities and continued scaling toward the trillion-parameter regime.The authors aim to advance both model capability and large-scale training and inference efficiency.

A Contributors

The contributor section lists the paper’s authors alphabetically by first name and marks authors whose affiliations differ from those at the time of the work.

  • A Contributors: Authors are listed alphabetically by first name.Names marked with an asterisk indicate affiliations differing from those at the time of the work.
  • A Contributors: The contributor list includes Cheolseung Baek, Eunki Kim, Hyunho Yang, Hyunjun Eun, and additional authors.The supplied list continues through Tae Yoon Kim.

A.2 Data Engineering

This section lists the contributors to the report. The contributor list spans multiple lines and includes the named researchers below.

  • Dhammiko Arya, Gun Song, and Gyoungeun Han are listed among the report’s contributors.
  • Seungsik Kim and Youngrang Kim are listed as contributors.
  • Seungmo Cho, Sooyeon Park, and Youngjin Kim are also listed as contributors.

B.1 Pre-train Data Category Mixture

The appendix describes how A.X K2’s pre-training and SFT mixtures are composed, constructed, validated, and evaluated under controlled serving and metric procedures.

  • Pre-train mixture: The designed pre-training mixture is reported by language and content category across training stages.Designed percentages may differ from tokens actually consumed; Stage 1 used the first 6.4T tokens of its approximately 8T schedule.
  • SFT mixture: The SFT appendix reports sample statistics and token composition across the SFT stages.Tables 18 and 19 report token ratios and sample distributions by token-length range.
  • Safety data: A.X K2 extends the ESSA-based safety-data recipe with an additional specification-evolution and data-synthesis pass tailored to its training format.
  • Safety data: The safety-data pipeline labels prompts by domain and task to select among 50 evolved safety specifications, without treating labels as supervision targets.
  • Safety data: Converted safety samples undergo filtering for unsafe residue, privacy leakage, language mismatch, malformed boundaries, and unnecessary refusal.
  • Evaluation setup: Open-weight baselines are served through publisher-fixed OpenRouter settings with xhigh reasoning effort and controlled maximum output lengths.
  • Metrics: Table 6 reports pass@1 scores averaged across multiple generations, with eight generations per sample for mathematics benchmarks.AA-Omniscience uses the AAII v4.1 score, defined as (2 × accuracy + (1 − hallucination rate))/3.
  • Evaluation setup: BrowseComp uses only Brave Search’s LLM Context API, with at most 10 searches per problem.

B.5 Details of Evaluation Benchmarks

The evaluation suite covers mathematics, Korean capability, code, science and knowledge, general instruction following, and agentic tasks using specified benchmarks and protocols.

  • Evaluation protocol: Responses are generated until the model reaches its context-length limit, and outputs that hit the limit are graded incorrect.
  • Math: Mathematics evaluation uses AIME26 and Apex, including Apex-shortlist as a companion set of less difficult problems.
  • Math: AIME26 responses are evaluated by comparing the ANSWER field with the ground-truth answer.
  • Korean: Korean evaluation uses KMMLU-Pro, KoBALT, and CLIcK to assess professional, linguistic, cultural, and linguistic knowledge.CLIcK is evaluated in a standard zero-shot setting.
  • Code: Code evaluation uses LiveCodeBench v6, SciCode, and Terminal Bench v2.1; LiveCodeBench correctness is determined by matching execution results with expected outputs.Only the February–May 2025 LiveCodeBench subset is evaluated.
  • Science & Knowledge: Science and knowledge evaluation uses GPQA Diamond, Humanity’s Last Exam, and AA-Omniscience.AA-Omniscience jointly measures factual knowledge and hallucination propensity.
  • General: General evaluation uses IFBench for explicit-constraint compliance and AA-LCR for coherence and key-information retention over long inputs.
  • Agentic: Agentic evaluation uses GDPval, τ2-Bench, τ3-Bench, and BrowseComp for professional tasks, conversational agents, and search-based problem solving.

B.6 Multilingual Evaluation

A.X K2 is evaluated for multilingual capability on Global-MMLU-Lite using four non-Korean languages, with performance reported in thinking mode.

  • Dataset: Global-MMLU-Lite uses human-translated and post-edited parallel items across languages to evaluate multilingual capability.The evaluation uses four non-Korean languages because Korean is covered by separate Korean benchmarks.
  • Reporting: Table 20 reports multilingual performance in thinking mode using exact-match percentages.The highest model score in each column is shown in bold, including ties.
  • Results: 87.56 places A.X K2 ahead of DeepSeek-V4 Flash at 86.81 and MiniMax-M2.7 at 87.31, while trailing the strongest baselines.
  • Results: The cross-language score range falls from 21.50 points for A.X K1 to 4.00 points for A.X K2.A.X K2’s reported English and Japanese scores are 89.75 and 85.75, respectively.

B.7 LongBench v2: Full-Set and Category-Level Results

LongBench v2 reports A.X K2 on both the full and Han-character-filtered sets, showing that filtering changes overall, slice-level, and category-level comparisons.

  • 503 items form the full LongBench v2 set, while 416 items remain after removing the 87 items containing Han characters.
  • A.X K2’s overall score rises from 62.2 to 63.9 (+1.7) after filtering, making it second rather than third among compared models.Qwen3.5 remains ahead by 0.3 points on the filtered set.
  • A.X K2’s Long score decreases by 2.1 points after filtering, with the largest drops in Code Repository Understanding (−14.0) and Long Structured Data Understanding (−10.6).The filter removes 40% and 33% of items in those two categories, respectively.
  • Long-dialogue History Understanding is unchanged because the filter removes none of its items.
  • A.X K2 leads Multi-Document QA on both sets at 60.8 and 63.9, ties for filtered Single-Document QA at 65.8, and leads filtered Long In-context Learning at 75.0.It trails leading models in Code Repository Understanding and Long Structured Data Understanding on both sets.
Loading 2608.30181v1…